From 41f9e449456b6db4f5add42f1bdce0eba7f08743 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 7 Apr 2026 23:19:37 +0000 Subject: [PATCH 1/3] fix: return native prompt queue metadata --- api/job_routes.py | 4 +++- api/queue_orchestration.py | 28 +++++++++++++++++++++++----- tests/api/test_distributed_queue.py | 8 +++++--- utils/async_helpers.py | 15 ++++++++++++++- 4 files changed, 45 insertions(+), 10 deletions(-) diff --git a/api/job_routes.py b/api/job_routes.py index 3843101..d962adf 100644 --- a/api/job_routes.py +++ b/api/job_routes.py @@ -217,7 +217,7 @@ async def distributed_queue_endpoint(request): return await handle_api_error(request, exc, 400) try: - prompt_id, worker_count = await orchestrate_distributed_execution( + prompt_id, prompt_number, worker_count, node_errors = await orchestrate_distributed_execution( payload.prompt, payload.workflow_meta, payload.client_id, @@ -227,6 +227,8 @@ async def distributed_queue_endpoint(request): ) return web.json_response({ "prompt_id": prompt_id, + "number": prompt_number, + "node_errors": node_errors, "worker_count": worker_count, "auto_prepare_supported": True, }) diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 226c645..ab18e7f 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -208,7 +208,7 @@ async def orchestrate_distributed_execution( """Core orchestration logic for the /distributed/queue endpoint. Returns: - tuple[str, int]: (prompt_id, worker_count) + tuple[str, int, int, dict]: (prompt_id, number, worker_count, node_errors) """ ensure_distributed_state() execution_trace_id = trace_execution_id or _generate_execution_trace_id() @@ -317,8 +317,18 @@ async def orchestrate_distributed_execution( if not job_id_map: trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.") - prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id) - return prompt_id, 0 + queue_result = await queue_prompt_payload( + prompt_obj, + workflow_meta, + client_id, + include_queue_metadata=True, + ) + return ( + queue_result["prompt_id"], + queue_result["number"], + 0, + queue_result.get("node_errors", {}), + ) for job_id in job_id_map.values(): await _ensure_distributed_queue(job_id) @@ -392,9 +402,17 @@ async def orchestrate_distributed_execution( ] ) - prompt_id = await queue_prompt_payload(master_prompt, workflow_meta, client_id) + queue_result = await queue_prompt_payload( + master_prompt, + workflow_meta, + client_id, + include_queue_metadata=True, + ) + prompt_id = queue_result["prompt_id"] + prompt_number = queue_result["number"] + node_errors = queue_result.get("node_errors", {}) trace_debug( execution_trace_id, f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}", ) - return prompt_id, len(worker_payloads) + return prompt_id, prompt_number, len(worker_payloads), node_errors diff --git a/tests/api/test_distributed_queue.py b/tests/api/test_distributed_queue.py index b553624..55694a3 100644 --- a/tests/api/test_distributed_queue.py +++ b/tests/api/test_distributed_queue.py @@ -162,7 +162,7 @@ def _load_job_routes_module(): sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration") - queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 1)) + queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 7, 1, {})) sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module @dataclass(frozen=True) @@ -221,7 +221,7 @@ job_routes = _load_job_routes_module() class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase): - async def test_distributed_queue_happy_path_returns_prompt_id(self): + async def test_distributed_queue_happy_path_returns_prompt_metadata(self): request = _FakeRequest( { "prompt": {"1": {"class_type": "Node"}}, @@ -233,12 +233,14 @@ class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase): with patch.object( job_routes, "orchestrate_distributed_execution", - new=AsyncMock(return_value=("prompt_123", 2)), + new=AsyncMock(return_value=("prompt_123", 42, 2, {})), ): response = await job_routes.distributed_queue_endpoint(request) self.assertEqual(response.status, 200) self.assertEqual(response.payload.get("prompt_id"), "prompt_123") + self.assertEqual(response.payload.get("number"), 42) + self.assertEqual(response.payload.get("node_errors"), {}) self.assertTrue(response.payload.get("auto_prepare_supported")) async def test_distributed_queue_missing_prompt_returns_400(self): diff --git a/utils/async_helpers.py b/utils/async_helpers.py index d39abab..7ea3c00 100644 --- a/utils/async_helpers.py +++ b/utils/async_helpers.py @@ -104,7 +104,12 @@ class PromptValidationError(RuntimeError): super().__init__(f"Invalid prompt: {merged}") -async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None): +async def queue_prompt_payload( + prompt_obj, + workflow_meta=None, + client_id=None, + include_queue_metadata=False, +): """Validate and queue a prompt via ComfyUI's prompt queue.""" payload = {"prompt": prompt_obj} payload = prompt_server.trigger_on_prompt(payload) @@ -132,4 +137,12 @@ async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None): prompt_server.number = number + 1 prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive) prompt_server.prompt_queue.put(prompt_queue_item) + + if include_queue_metadata: + return { + "prompt_id": prompt_id, + "number": number, + "node_errors": {}, + } + return prompt_id From 5386de10e370092df23be229984b19098b5ef31b Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 7 Apr 2026 23:34:15 +0000 Subject: [PATCH 2/3] fix: include queue metadata timestamps --- tests/test_async_helpers.py | 91 +++++++++++++++++++++++++++++++++++++ utils/async_helpers.py | 3 +- 2 files changed, 93 insertions(+), 1 deletion(-) create mode 100644 tests/test_async_helpers.py diff --git a/tests/test_async_helpers.py b/tests/test_async_helpers.py new file mode 100644 index 0000000..e5b7759 --- /dev/null +++ b/tests/test_async_helpers.py @@ -0,0 +1,91 @@ +import importlib.util +import sys +import types +import unittest +from pathlib import Path + + +class _PromptQueue: + def __init__(self): + self.items = [] + + def put(self, item): + self.items.append(item) + + +def _load_async_helpers_module(): + module_path = Path("/home/robert/ComfyUI_dev/custom_nodes/ComfyUI-Distributed/utils/async_helpers.py") + package_name = "dist_async_helpers_testpkg" + + for mod_name in list(sys.modules): + if mod_name == package_name or mod_name.startswith(f"{package_name}."): + del sys.modules[mod_name] + + root_pkg = types.ModuleType(package_name) + root_pkg.__path__ = [] + sys.modules[package_name] = root_pkg + + utils_pkg = types.ModuleType(f"{package_name}.utils") + utils_pkg.__path__ = [] + sys.modules[f"{package_name}.utils"] = utils_pkg + + execution_module = types.ModuleType("execution") + + async def _validate_prompt(prompt_id, prompt, partial_execution_targets): + return (True, None, ["9"], {}) + + execution_module.validate_prompt = _validate_prompt + execution_module.SENSITIVE_EXTRA_DATA_KEYS = [] + sys.modules["execution"] = execution_module + + prompt_server = types.SimpleNamespace( + trigger_on_prompt=lambda payload: payload, + number=12, + prompt_queue=_PromptQueue(), + ) + server_module = types.ModuleType("server") + server_module.PromptServer = types.SimpleNamespace(instance=prompt_server) + sys.modules["server"] = server_module + + network_module = types.ModuleType(f"{package_name}.utils.network") + network_module.get_server_loop = lambda: None + sys.modules[f"{package_name}.utils.network"] = network_module + + spec = importlib.util.spec_from_file_location(f"{package_name}.utils.async_helpers", module_path) + module = importlib.util.module_from_spec(spec) + assert spec is not None and spec.loader is not None + spec.loader.exec_module(module) + return module, prompt_server + + +async_helpers, prompt_server = _load_async_helpers_module() + + +class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase): + async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self): + result = await async_helpers.queue_prompt_payload( + {"1": {"class_type": "Node"}}, + workflow_meta={"id": "workflow-1"}, + client_id="client-1", + include_queue_metadata=True, + ) + + self.assertIsInstance(result["prompt_id"], str) + self.assertTrue(result["prompt_id"]) + self.assertEqual(result["number"], 12) + self.assertEqual(result["node_errors"], {}) + + self.assertEqual(prompt_server.number, 13) + self.assertEqual(len(prompt_server.prompt_queue.items), 1) + queued_item = prompt_server.prompt_queue.items[0] + self.assertEqual(queued_item[0], 12) + extra_data = queued_item[3] + self.assertEqual(extra_data["client_id"], "client-1") + self.assertIn("create_time", extra_data) + self.assertIsInstance(extra_data["create_time"], int) + self.assertGreater(extra_data["create_time"], 0) + self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"}) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/async_helpers.py b/utils/async_helpers.py index 7ea3c00..eb34d15 100644 --- a/utils/async_helpers.py +++ b/utils/async_helpers.py @@ -3,6 +3,7 @@ Async helper utilities for ComfyUI-Distributed. """ import asyncio import threading +import time import uuid import execution import server @@ -122,7 +123,7 @@ async def queue_prompt_payload( node_errors = valid[3] if len(valid) > 3 else {} raise PromptValidationError(error_payload, node_errors) - extra_data = {} + extra_data = {"create_time": int(time.time() * 1000)} if workflow_meta: extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta if client_id: From 0e9ee7fbac798bf3a1636a45f32429aba066b1b4 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 7 Apr 2026 23:39:18 +0000 Subject: [PATCH 3/3] test: make async helper regression portable --- tests/test_async_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_async_helpers.py b/tests/test_async_helpers.py index e5b7759..72b61cc 100644 --- a/tests/test_async_helpers.py +++ b/tests/test_async_helpers.py @@ -14,7 +14,7 @@ class _PromptQueue: def _load_async_helpers_module(): - module_path = Path("/home/robert/ComfyUI_dev/custom_nodes/ComfyUI-Distributed/utils/async_helpers.py") + module_path = Path(__file__).resolve().parents[1] / "utils" / "async_helpers.py" package_name = "dist_async_helpers_testpkg" for mod_name in list(sys.modules):