fix: prevent stuck distributed job completion

api/job_routes.py

@@ -14,7 +14,7 @@
 from ..utils.logging import debug_log
 from ..utils.image import pil_to_tensor, ensure_contiguous
 from ..utils.network import handle_api_error
-from ..utils.constants import JOB_INIT_GRACE_PERIOD, MEMORY_CLEAR_DELAY
+from ..utils.constants import CLOSED_JOB_TTL_SECONDS, JOB_INIT_GRACE_PERIOD, MEMORY_CLEAR_DELAY
 try:
     from .queue_orchestration import ensure_distributed_state, orchestrate_distributed_execution
 except ImportError:
@@ -149,6 +149,9 @@ async def prepare_job_endpoint(request):
 
         ensure_distributed_state()
         async with prompt_server.distributed_jobs_lock:
+            closed_jobs = getattr(prompt_server, "distributed_closed_jobs", None)
+            if closed_jobs is not None:
+                closed_jobs.pop(multi_job_id, None)
             if multi_job_id not in prompt_server.distributed_pending_jobs:
                 prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
 
@@ -335,6 +338,19 @@ async def job_complete_endpoint(request):
                     queue_size = pending.qsize()
                     break
 
+                closed_jobs = getattr(prompt_server, "distributed_closed_jobs", {})
+                closed_at = closed_jobs.get(multi_job_id)
+                if closed_at is not None:
+                    if time.monotonic() - closed_at <= CLOSED_JOB_TTL_SECONDS:
+                        return web.json_response(
+                            {
+                                "code": "job_closed",
+                                "error": "job no longer accepts results",
+                            },
+                            status=410,
+                        )
+                    closed_jobs.pop(multi_job_id, None)
+
             if time.monotonic() > deadline:
                 return await handle_api_error(request, "job not initialized", 404)
             await asyncio.sleep(0.05)

api/queue_orchestration.py

@@ -44,6 +44,8 @@ def ensure_distributed_state(server_instance=None):
     ps = server_instance or prompt_server
     if not hasattr(ps, "distributed_pending_jobs"):
         ps.distributed_pending_jobs = {}
+    if not hasattr(ps, "distributed_closed_jobs"):
+        ps.distributed_closed_jobs = {}
     if not hasattr(ps, "distributed_jobs_lock"):
         ps.distributed_jobs_lock = asyncio.Lock()
 
@@ -56,6 +58,7 @@ async def _ensure_distributed_queue(job_id):
     """Ensure a queue exists for the given distributed job ID."""
     ensure_distributed_state()
     async with prompt_server.distributed_jobs_lock:
+        prompt_server.distributed_closed_jobs.pop(job_id, None)
         if job_id not in prompt_server.distributed_pending_jobs:
             prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()
 

nodes/collector.py

@@ -12,7 +12,11 @@
 
 from ..utils.logging import debug_log, log
 from ..utils.config import get_worker_timeout_seconds, load_config, is_master_delegate_only
-from ..utils.constants import HEARTBEAT_INTERVAL
+from ..utils.constants import (
+    CLOSED_JOB_TTL_SECONDS,
+    HEARTBEAT_INTERVAL,
+    MAX_COLLECTOR_BUSY_GRACE_PERIODS,
+)
 from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous
 from ..utils.network import build_worker_url, get_client_session, probe_worker
 from ..utils.audio_payload import encode_audio_payload
@@ -57,6 +61,22 @@ def INPUT_TYPES(s):
     RETURN_NAMES = ("images", "audio")
     FUNCTION = "run"
     CATEGORY = "image"
+
+    @staticmethod
+    def _mark_job_closed(multi_job_id):
+        closed_jobs = getattr(prompt_server, "distributed_closed_jobs", None)
+        if closed_jobs is None:
+            closed_jobs = {}
+            prompt_server.distributed_closed_jobs = closed_jobs
+
+        now = time.monotonic()
+        cutoff = now - CLOSED_JOB_TTL_SECONDS
+        stale_job_ids = [
+            job_id for job_id, closed_at in closed_jobs.items() if closed_at < cutoff
+        ]
+        for job_id in stale_job_ids:
+            closed_jobs.pop(job_id, None)
+        closed_jobs[str(multi_job_id)] = now
 
     @staticmethod
     def _unwrap_list_input(value):
@@ -196,6 +216,17 @@ async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_ur
                     json=payload,
                     timeout=aiohttp.ClientTimeout(total=timeout_seconds),
                 ) as response:
+                    if getattr(response, "status", 200) == 410:
+                        response_payload = await response.json(content_type=None)
+                        if (
+                            isinstance(response_payload, dict)
+                            and response_payload.get("code") == "job_closed"
+                        ):
+                            log(
+                                f"Worker - Master closed job {multi_job_id} before the result arrived; "
+                                "discarding the remaining completion payloads."
+                            )
+                            return
                     response.raise_for_status()
             except Exception as e:
                 media_type = "image/audio" if "image" in payload else "audio-only"
@@ -345,6 +376,9 @@ async def execute(self, images, audio, load_balance=False, multi_job_id="", is_w
 
             # Create the queue before any expensive local work to avoid job_complete race.
             async with prompt_server.distributed_jobs_lock:
+                closed_jobs = getattr(prompt_server, "distributed_closed_jobs", None)
+                if closed_jobs is not None:
+                    closed_jobs.pop(multi_job_id, None)
                 if multi_job_id not in prompt_server.distributed_pending_jobs:
                     prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
                     debug_log(f"Master - Initialized queue early for job {multi_job_id}")
@@ -380,6 +414,7 @@ async def execute(self, images, audio, load_balance=False, multi_job_id="", is_w
             base_timeout = float(get_worker_timeout_seconds())
             slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
             last_activity = time.time()
+            busy_grace_periods = 0
 
 
             # Get queue size before starting
@@ -434,6 +469,7 @@ def mark_worker_done(done_worker_id):
                         # Record activity and refresh timeout baseline
                         last_activity = time.time()
                         base_timeout = float(get_worker_timeout_seconds())
+                        busy_grace_periods = 0
 
                         if is_last:
                             mark_worker_done(worker_id)
@@ -479,22 +515,29 @@ def mark_worker_done(done_worker_id):
                                     )
                                     if payload is not None and queue_remaining and queue_remaining > 0:
                                         any_busy = True
-                                        log(
-                                            f"Master - Probe grace: worker {wid} appears busy "
-                                            f"(queue_remaining={queue_remaining}). Continuing to wait."
-                                        )
                                         break
                                 except Exception as e:
                                     debug_log(f"Collector probe failed for worker {wid}: {e}")
                         except Exception as e:
                             debug_log(f"Collector probe setup error: {e}")
 
                         if any_busy:
-                            # Refresh last_activity and continue waiting
-                            last_activity = time.time()
-                            # Refresh base timeout in case the user changed it in UI
-                            base_timeout = float(get_worker_timeout_seconds())
-                            continue
+                            if busy_grace_periods < MAX_COLLECTOR_BUSY_GRACE_PERIODS:
+                                busy_grace_periods += 1
+                                log(
+                                    "Master - Probe grace: a missing worker still appears busy; "
+                                    f"continuing wait period {busy_grace_periods}/"
+                                    f"{MAX_COLLECTOR_BUSY_GRACE_PERIODS}."
+                                )
+                                last_activity = time.time()
+                                # Refresh base timeout in case the user changed it in UI.
+                                base_timeout = float(get_worker_timeout_seconds())
+                                continue
+                            log(
+                                "Master - Busy-worker grace exhausted after "
+                                f"{MAX_COLLECTOR_BUSY_GRACE_PERIODS} additional wait periods; "
+                                "finishing with the results received so far."
+                            )
 
                         # Check queue size again with lock
                         async with prompt_server.distributed_jobs_lock:
@@ -529,6 +572,7 @@ def mark_worker_done(done_worker_id):
                 async with prompt_server.distributed_jobs_lock:
                     if multi_job_id in prompt_server.distributed_pending_jobs:
                         del prompt_server.distributed_pending_jobs[multi_job_id]
+                    self._mark_job_closed(multi_job_id)
                 raise
 
             total_collected = sum(len(imgs) for imgs in worker_images.values())
@@ -537,6 +581,7 @@ def mark_worker_done(done_worker_id):
             async with prompt_server.distributed_jobs_lock:
                 if multi_job_id in prompt_server.distributed_pending_jobs:
                     del prompt_server.distributed_pending_jobs[multi_job_id]
+                self._mark_job_closed(multi_job_id)
 
             combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
             try:

tests/api/test_distributed_queue.py

@@ -1,5 +1,6 @@
 import importlib.util
 import sys
+import time
 import types
 import unittest
 import asyncio
@@ -153,6 +154,7 @@ async def _handle_api_error(_request, error, status=500):
     sys.modules[f"{package_name}.utils.network"] = network_module
 
     constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.CLOSED_JOB_TTL_SECONDS = 3600.0
     constants_module.MEMORY_CLEAR_DELAY = 0.0
     constants_module.JOB_INIT_GRACE_PERIOD = 10.0
     sys.modules[f"{package_name}.utils.constants"] = constants_module
@@ -325,6 +327,53 @@ async def test_job_complete_accepts_audio_without_image(self):
         self.assertIsNone(queued["tensor"])
         self.assertEqual(queued["audio"]["sample_rate"], 44100)
 
+    async def test_job_complete_returns_gone_for_known_closed_job(self):
+        job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
+        job_routes.prompt_server.distributed_pending_jobs = {}
+        job_routes.prompt_server.distributed_closed_jobs = {
+            "closed-job": time.monotonic()
+        }
+        request = _FakeRequest(
+            {
+                "job_id": "closed-job",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "image": "data:image/png;base64,AAAA",
+                "is_last": True,
+            }
+        )
+
+        with (
+            patch.object(job_routes, "JOB_INIT_GRACE_PERIOD", 0.0),
+            patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"),
+        ):
+            response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 410)
+        self.assertEqual(response.payload.get("code"), "job_closed")
+
+    async def test_job_complete_keeps_unknown_job_distinct_from_closed_job(self):
+        job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
+        job_routes.prompt_server.distributed_pending_jobs = {}
+        job_routes.prompt_server.distributed_closed_jobs = {}
+        request = _FakeRequest(
+            {
+                "job_id": "unknown-job",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "image": "data:image/png;base64,AAAA",
+                "is_last": True,
+            }
+        )
+
+        with (
+            patch.object(job_routes, "JOB_INIT_GRACE_PERIOD", 0.0),
+            patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"),
+        ):
+            response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 404)
+
     async def test_job_complete_rejects_payload_without_image_or_audio(self):
         request = _FakeRequest(
             {

tests/test_async_helpers.py

@@ -1,5 +1,7 @@
+import asyncio
 import importlib.util
 import sys
+import threading
 import types
 import unittest
 from pathlib import Path
@@ -87,5 +89,54 @@ async def test_queue_prompt_payload_includes_create_time_and_client_metadata(sel
         self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
 
 
+class RunAsyncInServerLoopTests(unittest.TestCase):
+    def test_propagates_comfy_style_base_exception(self):
+        loop = asyncio.new_event_loop()
+        loop_thread = threading.Thread(target=loop.run_forever, daemon=True)
+        loop_thread.start()
+        original_get_server_loop = async_helpers.get_server_loop
+        async_helpers.get_server_loop = lambda: loop
+
+        class ComfyStyleInterrupt(BaseException):
+            pass
+
+        async def raise_interrupt():
+            raise ComfyStyleInterrupt("cancelled")
+
+        try:
+            with self.assertRaisesRegex(ComfyStyleInterrupt, "cancelled"):
+                async_helpers.run_async_in_server_loop(raise_interrupt(), timeout=1.0)
+        finally:
+            async_helpers.get_server_loop = original_get_server_loop
+            loop.call_soon_threadsafe(loop.stop)
+            loop_thread.join(timeout=1.0)
+            loop.close()
+
+    def test_timeout_cancels_scheduled_coroutine(self):
+        loop = asyncio.new_event_loop()
+        loop_thread = threading.Thread(target=loop.run_forever, daemon=True)
+        loop_thread.start()
+        original_get_server_loop = async_helpers.get_server_loop
+        async_helpers.get_server_loop = lambda: loop
+        cancelled = threading.Event()
+
+        async def wait_forever():
+            try:
+                await asyncio.Future()
+            except asyncio.CancelledError:
+                cancelled.set()
+                raise
+
+        try:
+            with self.assertRaisesRegex(TimeoutError, "timed out after 0.05 seconds"):
+                async_helpers.run_async_in_server_loop(wait_forever(), timeout=0.05)
+            self.assertTrue(cancelled.wait(timeout=1.0))
+        finally:
+            async_helpers.get_server_loop = original_get_server_loop
+            loop.call_soon_threadsafe(loop.stop)
+            loop_thread.join(timeout=1.0)
+            loop.close()
+
+
 if __name__ == "__main__":
     unittest.main()