comfyui-media: a dropped request interrupts its own GPU job

backends/comfyui-media/bridge/jobs.py

@@ -101,10 +101,10 @@ async def run(self, build, *, kind: str, upload_ref: list[str] | None = None) ->
         """Submit one job and return (artifact bytes, mime type).
 
         Cancellation semantics: if this coroutine is cancelled (client
-        disconnect) while the backend job is still running, the job is *not*
-        abandoned and the admission slot is *not* released here. A reaper task
-        is kept alive until ComfyUI reports a terminal state; only then is the
-        slot freed. On an unrecoverable reap the runner is marked unhealthy and
+        disconnect) while the backend job is still running, the job is
+        interrupted but the admission slot is *not* released here. A reaper
+        task is kept alive until ComfyUI reports a terminal state; only then is
+        the slot freed. On an unrecoverable reap the runner is marked unhealthy and
         the slot is retained forever rather than risking a second GPU job.
 
         Every failure path is classified by what we can *prove*:
@@ -200,15 +200,9 @@ async def run(self, build, *, kind: str, upload_ref: list[str] | None = None) ->
             if upload_ref:
                 upload_names = list(upload_ref)
             if submitted:
-                # The prompt is on the GPU; hold the slot until it is provably
-                # terminal.
-                self._schedule_reap(
-                    owner,
-                    prompt_id,
-                    output_node=output_node,
-                    upload_names=upload_names,
-                    stop_after_seconds=REAP_BUDGET_SECONDS,
-                )
+                # The prompt is on the GPU. Nobody is waiting for it any more,
+                # so stop it, and hold the slot until it is provably terminal.
+                self._schedule_interrupt(owner, prompt_id, output_node, upload_names)
             elif submit_started:
                 # The shielded submit task may have reached Comfy, so let it
                 # finish and reap based on the outcome rather than guessing.
@@ -304,6 +298,45 @@ def _schedule_reap(
         self._reaps.add(task)
         task.add_done_callback(self._reaps.discard)
 
+    def _schedule_interrupt(
+        self,
+        owner: str,
+        prompt_id: str,
+        output_node: str | None,
+        upload_names: list[str] | None,
+    ) -> None:
+        task = asyncio.get_running_loop().create_task(
+            self._interrupt_and_reap(owner, prompt_id, output_node, upload_names)
+        )
+        self._reaps.add(task)
+        task.add_done_callback(self._reaps.discard)
+
+    async def _interrupt_and_reap(
+        self,
+        owner: str,
+        prompt_id: str,
+        output_node: str | None,
+        upload_names: list[str] | None,
+    ) -> None:
+        """Stop an abandoned job, then release the slot once it is terminal.
+
+        The targeted delete/interrupt only touches our own prompt id. Comfy
+        notices an interrupt at its next sampler step, so the slot stays held
+        (busy, not quarantined) by the ordinary reaper until history proves
+        the job ended, exactly as for a job that finishes on its own.
+        """
+        try:
+            await self.comfy.cancel_own_job(prompt_id)
+        except Exception:
+            log.warning("interrupting an abandoned job failed; waiting for it to finish")
+        await self._reap(
+            owner,
+            prompt_id,
+            output_node=output_node,
+            upload_names=upload_names,
+            stop_after_seconds=REAP_BUDGET_SECONDS,
+        )
+
     def _schedule_reap_after_submit(
         self,
         owner: str,

backends/comfyui-media/bridge/server.py

@@ -68,6 +68,8 @@
 # Startup readiness probe against the fixed ComfyUI endpoint.
 READY_ATTEMPTS = 30
 READY_DELAY = 1.0
+# How often a long generation checks whether its caller is still connected.
+DISCONNECT_POLL_SECONDS = 1.0
 
 
 def _load_graphs() -> tuple[dict, object]:
@@ -274,14 +276,17 @@ async def video_generations(request: Request) -> Response:
                     except Exception:
                         pass
                     raise
-        data, mime = await _run_generation(
-            "video",
-            model,
-            payload,
-            image=image,
-            last_frame=last_frame,
-            audio=audio,
-            video=video,
+        data, mime = await _until_disconnect(
+            request,
+            _run_generation(
+                "video",
+                model,
+                payload,
+                image=image,
+                last_frame=last_frame,
+                audio=audio,
+                video=video,
+            ),
         )
         if mime != "video/mp4":
             raise backend_error("The media backend produced an unexpected output type.")
@@ -462,6 +467,32 @@ async def build(comfy: ComfyClient):
     return app
 
 
+async def _until_disconnect(request: Request, work):
+    """Await ``work``, cancelling it if the client hangs up first.
+
+    Starlette does not cancel a plain handler when its client disconnects, so
+    without this a dropped video request kept the GPU busy until the clip was
+    finished and then threw it away. Cancelling the generation lets the runner
+    interrupt its own ComfyUI job.
+    """
+    task = asyncio.ensure_future(work)
+    try:
+        while True:
+            done, _ = await asyncio.wait({task}, timeout=DISCONNECT_POLL_SECONDS)
+            if done:
+                return task.result()
+            if await request.is_disconnected():
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+                raise BridgeError(499, "client_closed", "The client closed the request.", "client_closed")
+    finally:
+        if not task.done():
+            task.cancel()
+
+
 def _validate_prompt(payload: dict) -> None:
     prompt = payload.get("prompt")
     if prompt is not None and not isinstance(prompt, str):

backends/comfyui-media/bridge/tests/fake_comfy.py

@@ -28,6 +28,9 @@ def __init__(
         self.running: list[str] = []
         self.deleted: list[str] = []
         self.interrupts = 0
+        # Real ComfyUI notices an interrupt at its next sampler step; a test can
+        # hold the job running past the interrupt to model a slow step.
+        self.interrupt_finishes = True
         self.uploaded: list[tuple[str, bytes, str]] = []
         self.view_requests: list[dict] = []
         self.job_state = "pending"  # pending -> running -> done
@@ -104,6 +107,8 @@ async def handler(self, request: httpx.Request) -> httpx.Response:
             return httpx.Response(200)
         if path == "/interrupt":
             self.interrupts += 1
+            if not self.interrupt_finishes:
+                return httpx.Response(200, json={})
             self.running.clear()
             self.job_state = "done"
             self.terminal.set()

backends/comfyui-media/bridge/tests/test_cancellation.py

@@ -17,6 +17,8 @@ def video_builder(model_id, payload, image_filename=None, prefix="lloom"):
 
 async def test_cancelled_request_keeps_slot_until_terminal():
     fake = FakeComfy()
+    # The interrupt has not landed yet: the sampler is mid-step.
+    fake.interrupt_finishes = False
     app, comfy = make_bridge(fake, video_builder)
     runner = app.state.bridge["runner"]
     async with asgi_client(app) as client:
@@ -226,3 +228,88 @@ async def test_cancelled_reaper_cleans_uploaded_and_output_files(tmp_path, monke
         assert not upload_path.exists()
         assert not output_path.exists()
         await comfy.aclose()
+
+
+async def _wait_until_free(runner, budget: float = 4.0) -> None:
+    for _ in range(int(budget / 0.02)):
+        if not runner.busy:
+            return
+        await asyncio.sleep(0.02)
+    raise AssertionError("slot was never released")
+
+
+async def test_cancelled_request_interrupts_its_job():
+    """A dropped request stops its own GPU job instead of letting it finish."""
+    fake = FakeComfy()
+    app, comfy = make_bridge(fake, video_builder)
+    runner = app.state.bridge["runner"]
+    async with asgi_client(app) as client:
+        await comfy.start()
+        task = asyncio.ensure_future(
+            client.post("/v1/videos/generations", json={"model": "video-model", "prompt": "x"})
+        )
+        await asyncio.wait_for(fake.submitted.wait(), timeout=3)
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+        # No fake.finish(): only the interrupt can end this job.
+        await _wait_until_free(runner)
+        assert fake.interrupts == 1
+        [prompt_id] = fake.prompts
+        assert prompt_id in fake.deleted, "a still-pending job is dropped from the queue too"
+        assert runner.unhealthy is False
+        await comfy.aclose()
+
+
+async def test_client_disconnect_cancels_the_generation():
+    """Starlette does not cancel a plain handler when its client goes away.
+
+    Drive the app over raw ASGI so the client can hang up the way uvicorn
+    reports it: an ``http.disconnect`` message after the body was read.
+    """
+    import json
+
+    fake = FakeComfy()
+    app, comfy = make_bridge(fake, video_builder)
+    runner = app.state.bridge["runner"]
+    await comfy.start()
+    body = json.dumps({"model": "video-model", "prompt": "x"}).encode()
+    hung_up = asyncio.Event()
+    sent_body = False
+
+    async def receive():
+        nonlocal sent_body
+        if not sent_body:
+            sent_body = True
+            return {"type": "http.request", "body": body, "more_body": False}
+        await hung_up.wait()
+        return {"type": "http.disconnect"}
+
+    messages = []
+
+    async def send(message):
+        messages.append(message)
+
+    scope = {
+        "type": "http",
+        "asgi": {"version": "3.0"},
+        "http_version": "1.1",
+        "method": "POST",
+        "scheme": "http",
+        "path": "/v1/videos/generations",
+        "raw_path": b"/v1/videos/generations",
+        "query_string": b"",
+        "root_path": "",
+        "headers": [(b"content-type", b"application/json"), (b"content-length", str(len(body)).encode())],
+        "client": ("127.0.0.1", 1),
+        "server": ("bridge.test", 80),
+    }
+    request = asyncio.ensure_future(app(scope, receive, send))
+    await asyncio.wait_for(fake.submitted.wait(), timeout=3)
+    hung_up.set()
+    await asyncio.wait_for(request, timeout=5)
+    await _wait_until_free(runner)
+    assert fake.interrupts == 1
+    assert len(fake.prompts) == 1
+    await comfy.aclose()

recipes/linux-nvidia-comfyui-ace-step-1-5-xl-sft.json

@@ -152,7 +152,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-ace-step-1-5-xl-turbo.json

@@ -152,7 +152,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-flux-2-klein-4b.json

@@ -189,7 +189,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-ideogram-4.json

@@ -198,7 +198,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-krea-2-turbo.json

@@ -165,7 +165,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-ltx-2-5-full.json

@@ -176,7 +176,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-ltx-2-5.json

@@ -166,7 +166,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-minimax-h3-turbo.json

@@ -198,7 +198,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-minimax-h3.json

@@ -160,7 +160,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-minimax-music3.json

@@ -150,7 +150,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-qwen-image-2-1-nvfp4.json

@@ -163,7 +163,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-qwen-image-2-1.json

@@ -161,7 +161,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-qwen-image-2512-lightning.json

@@ -169,7 +169,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-qwen-image-2512.json

@@ -148,7 +148,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",

recipes/linux-nvidia-comfyui-qwen-image-edit-2511.json

@@ -170,7 +170,7 @@
           "healthUrl": "http://127.0.0.1:${port}/health",
           "bootstrap": {
             "adapter": "docker",
-            "image": "lloom/comfyui-media:source-f478dab442b7ce625ac98db9968c5c30cc7574c259f2b2515c01530779312ddd",
+            "image": "lloom/comfyui-media:source-5f3794c99fdb3c3ccaac2b4f28b11e895b93ee0da0ec662be47d40ea08251cf6",
             "pull": false,
             "createArgs": [
               "--restart",