diff --git a/backends/comfyui-media/bridge/jobs.py b/backends/comfyui-media/bridge/jobs.py index aaf4d86..2b1fd46 100644 --- a/backends/comfyui-media/bridge/jobs.py +++ b/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, diff --git a/backends/comfyui-media/bridge/server.py b/backends/comfyui-media/bridge/server.py index dba7a93..51bf0fc 100644 --- a/backends/comfyui-media/bridge/server.py +++ b/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): diff --git a/backends/comfyui-media/bridge/tests/fake_comfy.py b/backends/comfyui-media/bridge/tests/fake_comfy.py index e2b511f..7866d3b 100644 --- a/backends/comfyui-media/bridge/tests/fake_comfy.py +++ b/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() diff --git a/backends/comfyui-media/bridge/tests/test_cancellation.py b/backends/comfyui-media/bridge/tests/test_cancellation.py index dc6bc7c..a3f60ab 100644 --- a/backends/comfyui-media/bridge/tests/test_cancellation.py +++ b/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() diff --git a/recipes/linux-nvidia-comfyui-ace-step-1-5-xl-sft.json b/recipes/linux-nvidia-comfyui-ace-step-1-5-xl-sft.json index bfe5ace..5e2d97c 100644 --- a/recipes/linux-nvidia-comfyui-ace-step-1-5-xl-sft.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-ace-step-1-5-xl-turbo.json b/recipes/linux-nvidia-comfyui-ace-step-1-5-xl-turbo.json index b4491fe..6c7cad5 100644 --- a/recipes/linux-nvidia-comfyui-ace-step-1-5-xl-turbo.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-flux-2-klein-4b.json b/recipes/linux-nvidia-comfyui-flux-2-klein-4b.json index 3cd4634..c81f8d4 100644 --- a/recipes/linux-nvidia-comfyui-flux-2-klein-4b.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-ideogram-4.json b/recipes/linux-nvidia-comfyui-ideogram-4.json index c2e171d..a0b6a49 100644 --- a/recipes/linux-nvidia-comfyui-ideogram-4.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-krea-2-turbo.json b/recipes/linux-nvidia-comfyui-krea-2-turbo.json index f6b5625..33a4085 100644 --- a/recipes/linux-nvidia-comfyui-krea-2-turbo.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-ltx-2-5-full.json b/recipes/linux-nvidia-comfyui-ltx-2-5-full.json index a34bde7..01474ad 100644 --- a/recipes/linux-nvidia-comfyui-ltx-2-5-full.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-ltx-2-5.json b/recipes/linux-nvidia-comfyui-ltx-2-5.json index 9cdfb16..241eecf 100644 --- a/recipes/linux-nvidia-comfyui-ltx-2-5.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-minimax-h3-turbo.json b/recipes/linux-nvidia-comfyui-minimax-h3-turbo.json index 1f7ca70..962447c 100644 --- a/recipes/linux-nvidia-comfyui-minimax-h3-turbo.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-minimax-h3.json b/recipes/linux-nvidia-comfyui-minimax-h3.json index e351992..ff6659d 100644 --- a/recipes/linux-nvidia-comfyui-minimax-h3.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-minimax-music3.json b/recipes/linux-nvidia-comfyui-minimax-music3.json index 0199e30..1284eb3 100644 --- a/recipes/linux-nvidia-comfyui-minimax-music3.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-qwen-image-2-1-nvfp4.json b/recipes/linux-nvidia-comfyui-qwen-image-2-1-nvfp4.json index 9ab7cb9..ca39d3d 100644 --- a/recipes/linux-nvidia-comfyui-qwen-image-2-1-nvfp4.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-qwen-image-2-1.json b/recipes/linux-nvidia-comfyui-qwen-image-2-1.json index b507636..80e3b22 100644 --- a/recipes/linux-nvidia-comfyui-qwen-image-2-1.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-qwen-image-2512-lightning.json b/recipes/linux-nvidia-comfyui-qwen-image-2512-lightning.json index f5c8217..d9cdaf0 100644 --- a/recipes/linux-nvidia-comfyui-qwen-image-2512-lightning.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-qwen-image-2512.json b/recipes/linux-nvidia-comfyui-qwen-image-2512.json index 11c02a5..63ac28a 100644 --- a/recipes/linux-nvidia-comfyui-qwen-image-2512.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-qwen-image-edit-2511.json b/recipes/linux-nvidia-comfyui-qwen-image-edit-2511.json index 3bf8332..7d2175b 100644 --- a/recipes/linux-nvidia-comfyui-qwen-image-edit-2511.json +++ b/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", diff --git a/recipes/linux-nvidia-comfyui-yue2-3b.json b/recipes/linux-nvidia-comfyui-yue2-3b.json index 8b48b63..283902f 100644 --- a/recipes/linux-nvidia-comfyui-yue2-3b.json +++ b/recipes/linux-nvidia-comfyui-yue2-3b.json @@ -135,7 +135,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",