diff --git a/api/job_routes.py b/api/job_routes.py index b35655b..b160bad 100644 --- a/api/job_routes.py +++ b/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) diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index ab18e7f..8c23c57 100644 --- a/api/queue_orchestration.py +++ b/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() diff --git a/nodes/collector.py b/nodes/collector.py index d577b0f..d1e581f 100644 --- a/nodes/collector.py +++ b/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,10 +515,6 @@ 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}") @@ -490,11 +522,22 @@ def mark_worker_done(done_worker_id): 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: diff --git a/tests/api/test_distributed_queue.py b/tests/api/test_distributed_queue.py index ac6a333..1e548a8 100644 --- a/tests/api/test_distributed_queue.py +++ b/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( { diff --git a/tests/test_async_helpers.py b/tests/test_async_helpers.py index 72b61cc..c42b9cb 100644 --- a/tests/test_async_helpers.py +++ b/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() diff --git a/tests/test_collector_list_inputs.py b/tests/test_collector_list_inputs.py index d62312d..fd50176 100644 --- a/tests/test_collector_list_inputs.py +++ b/tests/test_collector_list_inputs.py @@ -46,7 +46,7 @@ def get(self, _path): 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 update(self, value): 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 @@ async def _fake_get_client_session(): 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() diff --git a/utils/async_helpers.py b/utils/async_helpers.py index eb34d15..bc74670 100644 --- a/utils/async_helpers.py +++ b/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 diff --git a/utils/constants.py b/utils/constants.py index caca4f9..a419845 100644 --- a/utils/constants.py +++ b/utils/constants.py @@ -45,6 +45,16 @@ # 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