diff --git a/.changeset/eighty-ends-invent.md b/.changeset/eighty-ends-invent.md new file mode 100644 index 000000000000..36ddad77259a --- /dev/null +++ b/.changeset/eighty-ends-invent.md @@ -0,0 +1,5 @@ +--- +"gradio": patch +--- + +feat:Stop running iterators when js client disconnects diff --git a/.gitignore b/.gitignore index 51efba9ddea6..5dbd7349b153 100644 --- a/.gitignore +++ b/.gitignore @@ -44,6 +44,7 @@ demo/all_demos/requirements.txt demo/*/config.json demo/annotatedimage_component/*.png demo/fake_diffusion_with_gif/*.gif +demo/cancel_events/cancel_events_output_log.txt # Etc .idea/* diff --git a/demo/cancel_events/run.ipynb b/demo/cancel_events/run.ipynb index d821e0034d66..d2b61c048a4b 100644 --- a/demo/cancel_events/run.ipynb +++ b/demo/cancel_events/run.ipynb @@ -1 +1 @@ -{"cells": [{"cell_type": "markdown", "id": "302934307671667531413257853548643485645", "metadata": {}, "source": ["# Gradio Demo: cancel_events"]}, {"cell_type": "code", "execution_count": null, "id": "272996653310673477252411125948039410165", "metadata": {}, "outputs": [], "source": ["!pip install -q gradio "]}, {"cell_type": "code", "execution_count": null, "id": "288918539441861185822528903084949547379", "metadata": {}, "outputs": [], "source": ["import time\n", "import gradio as gr\n", "\n", "\n", "def fake_diffusion(steps):\n", " for i in range(steps):\n", " print(f\"Current step: {i}\")\n", " time.sleep(0.5)\n", " yield str(i)\n", "\n", "\n", "def long_prediction(*args, **kwargs):\n", " time.sleep(10)\n", " return 42\n", "\n", "\n", "with gr.Blocks() as demo:\n", " with gr.Row():\n", " with gr.Column():\n", " n = gr.Slider(1, 10, value=9, step=1, label=\"Number Steps\")\n", " run = gr.Button(value=\"Start Iterating\")\n", " output = gr.Textbox(label=\"Iterative Output\")\n", " stop = gr.Button(value=\"Stop Iterating\")\n", " with gr.Column():\n", " textbox = gr.Textbox(label=\"Prompt\")\n", " prediction = gr.Number(label=\"Expensive Calculation\")\n", " run_pred = gr.Button(value=\"Run Expensive Calculation\")\n", " with gr.Column():\n", " cancel_on_change = gr.Textbox(label=\"Cancel Iteration and Expensive Calculation on Change\")\n", " cancel_on_submit = gr.Textbox(label=\"Cancel Iteration and Expensive Calculation on Submit\")\n", " echo = gr.Textbox(label=\"Echo\")\n", " with gr.Row():\n", " with gr.Column():\n", " image = gr.Image(sources=[\"webcam\"], label=\"Cancel on clear\", interactive=True)\n", " with gr.Column():\n", " video = gr.Video(sources=[\"webcam\"], label=\"Cancel on start recording\", interactive=True)\n", "\n", " click_event = run.click(fake_diffusion, n, output)\n", " stop.click(fn=None, inputs=None, outputs=None, cancels=[click_event])\n", " pred_event = run_pred.click(fn=long_prediction, inputs=[textbox], outputs=prediction)\n", "\n", " cancel_on_change.change(None, None, None, cancels=[click_event, pred_event])\n", " cancel_on_submit.submit(lambda s: s, cancel_on_submit, echo, cancels=[click_event, pred_event])\n", " image.clear(None, None, None, cancels=[click_event, pred_event])\n", " video.start_recording(None, None, None, cancels=[click_event, pred_event])\n", "\n", " demo.queue(max_size=20)\n", "\n", "if __name__ == \"__main__\":\n", " demo.launch()\n"]}], "metadata": {}, "nbformat": 4, "nbformat_minor": 5} \ No newline at end of file +{"cells": [{"cell_type": "markdown", "id": "302934307671667531413257853548643485645", "metadata": {}, "source": ["# Gradio Demo: cancel_events"]}, {"cell_type": "code", "execution_count": null, "id": "272996653310673477252411125948039410165", "metadata": {}, "outputs": [], "source": ["!pip install -q gradio "]}, {"cell_type": "code", "execution_count": null, "id": "288918539441861185822528903084949547379", "metadata": {}, "outputs": [], "source": ["import time\n", "import gradio as gr\n", "import atexit\n", "import pathlib\n", "\n", "log_file = (pathlib.Path(__file__).parent / \"cancel_events_output_log.txt\").resolve()\n", "\n", "def fake_diffusion(steps):\n", " log_file.write_text(\"\")\n", " for i in range(steps):\n", " print(f\"Current step: {i}\")\n", " with log_file.open(\"a\") as f:\n", " f.write(f\"Current step: {i}\\n\")\n", " time.sleep(0.2)\n", " yield str(i)\n", "\n", "\n", "def long_prediction(*args, **kwargs):\n", " time.sleep(10)\n", " return 42\n", "\n", "\n", "with gr.Blocks() as demo:\n", " with gr.Row():\n", " with gr.Column():\n", " n = gr.Slider(1, 10, value=9, step=1, label=\"Number Steps\")\n", " run = gr.Button(value=\"Start Iterating\")\n", " output = gr.Textbox(label=\"Iterative Output\")\n", " stop = gr.Button(value=\"Stop Iterating\")\n", " with gr.Column():\n", " textbox = gr.Textbox(label=\"Prompt\")\n", " prediction = gr.Number(label=\"Expensive Calculation\")\n", " run_pred = gr.Button(value=\"Run Expensive Calculation\")\n", " with gr.Column():\n", " cancel_on_change = gr.Textbox(label=\"Cancel Iteration and Expensive Calculation on Change\")\n", " cancel_on_submit = gr.Textbox(label=\"Cancel Iteration and Expensive Calculation on Submit\")\n", " echo = gr.Textbox(label=\"Echo\")\n", " with gr.Row():\n", " with gr.Column():\n", " image = gr.Image(sources=[\"webcam\"], label=\"Cancel on clear\", interactive=True)\n", " with gr.Column():\n", " video = gr.Video(sources=[\"webcam\"], label=\"Cancel on start recording\", interactive=True)\n", "\n", " click_event = run.click(fake_diffusion, n, output)\n", " stop.click(fn=None, inputs=None, outputs=None, cancels=[click_event])\n", " pred_event = run_pred.click(fn=long_prediction, inputs=[textbox], outputs=prediction)\n", "\n", " cancel_on_change.change(None, None, None, cancels=[click_event, pred_event])\n", " cancel_on_submit.submit(lambda s: s, cancel_on_submit, echo, cancels=[click_event, pred_event])\n", " image.clear(None, None, None, cancels=[click_event, pred_event])\n", " video.start_recording(None, None, None, cancels=[click_event, pred_event])\n", "\n", " demo.queue(max_size=20)\n", " atexit.register(lambda: log_file.unlink())\n", "\n", "if __name__ == \"__main__\":\n", " demo.launch()\n"]}], "metadata": {}, "nbformat": 4, "nbformat_minor": 5} \ No newline at end of file diff --git a/demo/cancel_events/run.py b/demo/cancel_events/run.py index cbc43e36cf3f..409858c5748a 100644 --- a/demo/cancel_events/run.py +++ b/demo/cancel_events/run.py @@ -1,11 +1,17 @@ import time import gradio as gr +import atexit +import pathlib +log_file = (pathlib.Path(__file__).parent / "cancel_events_output_log.txt").resolve() def fake_diffusion(steps): + log_file.write_text("") for i in range(steps): print(f"Current step: {i}") - time.sleep(0.5) + with log_file.open("a") as f: + f.write(f"Current step: {i}\n") + time.sleep(0.2) yield str(i) @@ -45,6 +51,7 @@ def long_prediction(*args, **kwargs): video.start_recording(None, None, None, cancels=[click_event, pred_event]) demo.queue(max_size=20) + atexit.register(lambda: log_file.unlink()) if __name__ == "__main__": demo.launch() diff --git a/gradio/routes.py b/gradio/routes.py index ce105b906075..77b6c374adce 100644 --- a/gradio/routes.py +++ b/gradio/routes.py @@ -803,11 +803,11 @@ async def sse_stream(request: fastapi.Request): message=str(e), ) response = process_msg(message) - if response is not None: - yield response if isinstance(e, asyncio.CancelledError): del blocks._queue.pending_messages_per_session[session_hash] await blocks._queue.clean_events(session_hash=session_hash) + if response is not None: + yield response raise e return StreamingResponse( diff --git a/js/app/test/cancel_events.spec.ts b/js/app/test/cancel_events.spec.ts index ed56470c65c9..b7677fb3c5fe 100644 --- a/js/app/test/cancel_events.spec.ts +++ b/js/app/test/cancel_events.spec.ts @@ -1,4 +1,5 @@ import { test, expect } from "@gradio/tootils"; +import { readFileSync } from "fs"; test("when using an iterative function the UI should update over time as iteration results are received", async ({ page @@ -44,3 +45,27 @@ test("when using an iterative function it should be possible to cancel the funct await page.waitForTimeout(1000); await expect(textbox).toHaveValue("0"); }); + +test("when using an iterative function and the user closes the page, the python function should stop running", async ({ + page +}) => { + const start_button = await page.locator("button", { + hasText: /Start Iterating/ + }); + + await start_button.click(); + await page.waitForTimeout(300); + await page.close(); + + // wait for the duration of the entire iteration + // check that the final value did not get written + // to the log file. That's our proof python stopped + // running + await new Promise((resolve) => setTimeout(resolve, 2000)); + const data = readFileSync( + "../../demo/cancel_events/cancel_events_output_log.txt", + "utf-8" + ); + expect(data).toContain("Current step: 0"); + expect(data).not.toContain("Current step: 8"); +});