Stop running iterators when js client disconnects (#7835)

* Add code

* add changeset

* move

* add changeset

* Add code

* Format notebook

* Address comments

---------

Co-authored-by: gradio-pr-bot <gradio-pr-bot@users.noreply.github.com>

.changeset/eighty-ends-invent.md

@@ -0,0 +1,5 @@
+---
+"gradio": patch
+---
+
+feat:Stop running iterators when js client disconnects

.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/*

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}
+{"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}

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()

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(

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");
+});