fix: prevent stuck distributed job completion

api/job_routes.py

@@ -14,7 +14,7 @@ from PIL import Image
 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 comfy.utils import ProgressBar
 
 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 @@ class DistributedCollectorNode:
     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 @@ class DistributedCollectorNode:
                     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 @@ class DistributedCollectorNode:
 
             # 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 @@ class DistributedCollectorNode:
             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 @@ class DistributedCollectorNode:
                         # 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,10 +515,6 @@ class DistributedCollectorNode:
                                     )
                                     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}")
@@ -490,11 +522,22 @@ class DistributedCollectorNode:
                             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 @@ class DistributedCollectorNode:
                 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 @@ class DistributedCollectorNode:
             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 @@ def _load_job_routes_module():
     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 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
         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 @@ class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
         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()

tests/test_collector_list_inputs.py

@@ -46,7 +46,7 @@ def _load_collector_module():
     comfy_module = types.ModuleType("comfy")
     model_management = types.ModuleType("comfy.model_management")
 
-    class InterruptProcessingException(Exception):
+    class InterruptProcessingException(BaseException):
         pass
 
     model_management.InterruptProcessingException = InterruptProcessingException
@@ -85,7 +85,9 @@ def _load_collector_module():
     sys.modules[f"{package_name}.utils.config"] = config_module
 
     constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.CLOSED_JOB_TTL_SECONDS = 3600.0
     constants_module.HEARTBEAT_INTERVAL = 1.0
+    constants_module.MAX_COLLECTOR_BUSY_GRACE_PERIODS = 10
     sys.modules[f"{package_name}.utils.constants"] = constants_module
 
     image_module = types.ModuleType(f"{package_name}.utils.image")
@@ -248,6 +250,90 @@ def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final
     assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}
 
 
+def test_worker_treats_closed_master_job_as_terminal_acknowledgement():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    images = torch.zeros(2, 2, 2, 3)
+    posted_payloads = []
+
+    class _FakeImage:
+        def save(self, fp, format=None, compress_level=None):
+            fp.write(b"png-bytes")
+
+    class _FakeResponse:
+        status = 410
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc_type, exc, tb):
+            return False
+
+        async def json(self, content_type=None):
+            return {"code": "job_closed", "error": "job no longer accepts results"}
+
+        def raise_for_status(self):
+            raise AssertionError("known closed jobs must not raise")
+
+    class _FakeSession:
+        def post(self, url, json, timeout):
+            posted_payloads.append(json)
+            return _FakeResponse()
+
+    async def _fake_get_client_session():
+        return _FakeSession()
+
+    module.tensor_to_pil = lambda *_args, **_kwargs: _FakeImage()
+    module.get_client_session = _fake_get_client_session
+    module.encode_audio_payload = lambda _audio: None
+
+    asyncio.run(
+        collector.send_batch_to_master(
+            images,
+            None,
+            "closed-job",
+            "http://master",
+            "worker-a",
+        )
+    )
+
+    assert len(posted_payloads) == 1
+
+
+def test_master_busy_probe_has_a_finite_grace_limit():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    module.prompt_server.distributed_pending_jobs = {}
+    module.MAX_COLLECTOR_BUSY_GRACE_PERIODS = 1
+    module.get_worker_timeout_seconds = lambda: 0.05
+    module.load_config = lambda: {
+        "workers": [{"id": "worker-a", "host": "127.0.0.1", "port": 8189}]
+    }
+
+    async def _busy_probe(*_args, **_kwargs):
+        return {"exec_info": {"queue_remaining": 1}}
+
+    module.probe_worker = _busy_probe
+    images = torch.zeros(1, 2, 2, 3)
+
+    combined, _audio = asyncio.run(
+        asyncio.wait_for(
+            collector.execute(
+                images=images,
+                audio=None,
+                multi_job_id="bounded-job",
+                enabled_worker_ids='["worker-a"]',
+            ),
+            timeout=0.5,
+        )
+    )
+
+    assert torch.equal(combined, images)
+    assert "bounded-job" not in module.prompt_server.distributed_pending_jobs
+    assert "bounded-job" in module.prompt_server.distributed_closed_jobs
+
+
 def test_audio_only_worker_sends_one_completion_without_image():
     module = _load_collector_module()
     collector = module.DistributedCollectorNode()

utils/async_helpers.py

@@ -2,7 +2,6 @@
 Async helper utilities for ComfyUI-Distributed.
 """
 import asyncio
-import threading
 import time
 import uuid
 import execution
@@ -26,32 +25,19 @@ def run_async_in_server_loop(coro: Coroutine, timeout: Optional[float] = None) -
         
     Raises:
         TimeoutError: If the operation times out
-        Exception: Any exception raised by the coroutine
+        BaseException: Any error or processing interrupt raised by the coroutine
     """
-    event = threading.Event()
-    result = None
-    error = None
-    
-    async def wrapper():
-        nonlocal result, error
-        try:
-            result = await coro
-        except Exception as e:
-            error = e
-        finally:
-            event.set()
-    
-    # Schedule on server's event loop
     loop = get_server_loop()
-    asyncio.run_coroutine_threadsafe(wrapper(), loop)
-    
-    # Wait for completion
-    if not event.wait(timeout):
-        raise TimeoutError(f"Async operation timed out after {timeout} seconds")
-    
-    if error:
-        raise error
-    return result
+    future = asyncio.run_coroutine_threadsafe(coro, loop)
+
+    try:
+        return future.result(timeout)
+    except TimeoutError:
+        # Do not misclassify a TimeoutError raised by the coroutine itself.
+        if future.done():
+            raise
+        future.cancel()
+        raise TimeoutError(f"Async operation timed out after {timeout} seconds") from None
 
 
 prompt_server = server.PromptServer.instance

utils/constants.py

@@ -45,6 +45,16 @@ MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20'))  # Maximum items per
 # Heartbeat monitoring
 HEARTBEAT_INTERVAL = float(os.environ.get('COMFYUI_HEARTBEAT_INTERVAL', '10'))  # Heartbeat/check interval in seconds
 HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60'))  # Worker heartbeat timeout in seconds (default 60s)
+# Cap consecutive worker-timeout extensions when probes keep reporting "busy".
+MAX_COLLECTOR_BUSY_GRACE_PERIODS = max(
+    0,
+    int(os.environ.get('COMFYUI_MAX_COLLECTOR_BUSY_GRACE_PERIODS', '10')),
+)
+# Retain terminal job IDs long enough to classify delayed worker callbacks.
+CLOSED_JOB_TTL_SECONDS = max(
+    0.0,
+    float(os.environ.get('COMFYUI_CLOSED_JOB_TTL_SECONDS', '3600')),
+)
 
 # USDU result collection
 DYNAMIC_MODE_MAX_POLL_TIMEOUT = 10.0