fix: return native prompt queue metadata
This commit is contained in:
+3
-1
@@ -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,
|
||||
})
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
+14
-1
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user