Compare commits

...
3 changed files with 158 additions and 12 deletions
+6 -6
View File
@@ -347,12 +347,12 @@ async def orchestrate_distributed_execution(
requires_values = any(end_override_names(node.get("inputs", {}))
for node in prompt_obj.values() if node.get("class_type") == END_CLASS)
value_preflight = {"require_value_results": True} if requires_values else {}
idle_ids = (await preflight_image_list_workers(active_workers, require_idle=delegate_master,
**value_preflight)
if active_workers else {})
active_workers = [worker for worker in active_workers if worker["id"] in idle_ids]
if not active_workers and delegate_master:
raise RuntimeError("No feature-compatible direct worker is idle in orchestrator-only mode")
# Persist candidates, not this submission-time idle snapshot: a preceding
# master prompt may still be using them. The image-list runtime rechecks
# capability and availability before assigning the first item.
if active_workers:
await preflight_image_list_workers(active_workers, require_idle=False,
**value_preflight)
else:
delegate_master = selected_delegate_master
+2 -1
View File
@@ -89,7 +89,8 @@ File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations an
## Scope and failure behavior
- This version supports the master and direct workers, not Runpod serverless pools. Unsupported pools are excluded without contacting the provider for this job.
- Worker eligibility is determined before execution and checked again before dispatch. Busy or transiently unreachable workers are excluded; healthy candidates can continue. Malformed capability/queue responses and incompatible versions fail explicitly. Workers joining later are not added to an already-running region.
- Queue submission checks selected direct workers for reachability and feature compatibility, but does not freeze the worker list to that moment's idle snapshot. A selected worker busy with an earlier prompt remains a candidate for the queued run. Immediately before a region starts assigning images, its saved candidates are checked again: workers still busy or transiently unreachable are excluded; healthy idle candidates can continue. Malformed capability/queue responses and incompatible versions fail explicitly. Workers excluded by the initial reachability probe are not retained, and newly selected workers are not added to an already-queued run. Workers becoming free after the region starts are not added mid-region.
- **Orchestrator Only** requires a selected, reachable, feature-compatible direct worker at submission, but can queue while all selected workers are busy. At region start it requires at least one idle compatible direct worker and fails if none are available; it never falls back to master processing. The queue response's `worker_count` counts saved remote candidates, not guaranteed execution-time participants.
- Omitted worker IDs select enabled direct workers. Explicit IDs select the named configured direct workers, including disabled ones; an empty list stays empty, unknown IDs fail, and serverless workers are excluded. Local workers use the master's loopback callback; remote workers use its configured public address.
- Only the enclosed processing is distributed. List creation and downstream saving run on the master in both modes. With participation enabled, the master is a normal scheduler participant and can run alone when no direct workers are available. **Orchestrator-only mode never silently falls back to master processing.** Failed items are not reassigned automatically.
- Local items execute while the enclosing master prompt waits; they are not queued behind that prompt on the master's `/prompt` queue. The enclosed branch is not also evaluated as a full list by the outer prompt.
+150 -5
View File
@@ -1,8 +1,9 @@
"""Image-list participation uses the existing master orchestrator setting."""
import json
import sys
import types
import unittest
from unittest.mock import AsyncMock, patch
from unittest.mock import AsyncMock, Mock, patch
from test_serverless_integration import PKG, load_modules
@@ -76,7 +77,7 @@ class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
self.config["settings"]["master_delegate_only"] = True
await self.queue()
self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
self.assertTrue(self.preflight.call_args.kwargs["require_idle"])
self.assertIs(self.preflight.call_args.kwargs["require_idle"], False)
# Upstream source stays on master even in orchestrator-only mode.
self.assertIn("1", self.q.queue_prompt_payload.call_args.args[0])
@@ -107,10 +108,154 @@ class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
async def test_busy_workers_do_not_prevent_enabled_master(self):
self.preflight.return_value = {}
def use_real_preflight(self, queues):
"""Exercise capability/queue parsing with read-only fake HTTP responses."""
from test_dispatch_selection import dispatch, _FakeSession, _valid_capability
responses = {(worker_id, "capabilities"): (200, _valid_capability())
for worker_id in queues}
responses.update({(worker_id, "queue"): (200, queue)
for worker_id, queue in queues.items()})
session = _FakeSession(responses)
for target, value in (
("get_client_session", AsyncMock(return_value=session)),
("build_worker_url", lambda worker, path: f"http://{worker['id']}{path}"),
):
patched = patch.object(dispatch, target, value)
patched.start()
self.addCleanup(patched.stop)
self.preflight = AsyncMock(wraps=dispatch.preflight_image_list_workers)
sys.modules[PKG + ".api.orchestration.dispatch"].preflight_image_list_workers = self.preflight
return responses
async def run_prepared_selection(self):
"""Run real runtime selection up to the scheduler, without any item I/O."""
rt = sys.modules[PKG + ".image_list.runtime"]
inputs = self.prepared_inputs()
image = object()
context = rt.create_master_job(
"2", "3", self.q.queue_prompt_payload.call_args.args[0],
[{"image": image, "index": 0}],
json.loads(inputs["image_list_enabled_worker_ids"]),
inputs["image_list_master_callback"],
master_participates=inputs["image_list_master_participates"],
explicit_targets=True,
)
self.addCleanup(rt.remove_master_job, context)
self.scheduler = AsyncMock(return_value=[{"image": image}])
self.local_slot = Mock(return_value=types.SimpleNamespace(close=AsyncMock()))
self.runtime_session = AsyncMock(return_value=object())
config = sys.modules[PKG + ".utils.config"]
network = sys.modules[PKG + ".utils.network"]
with patch.object(config, "load_config", lambda: self.config), \
patch.object(config, "get_worker_timeout_seconds", lambda: 60, create=True), \
patch.object(config, "get_image_list_job_timeout_seconds", lambda: 0, create=True), \
patch.object(network, "get_client_session", self.runtime_session, create=True), \
patch.object(network, "build_worker_url", Mock(), create=True), \
patch.object(rt, "dynamic_schedule", self.scheduler), \
patch.dict(sys.modules, {PKG + ".image_list.local_executor": types.SimpleNamespace(
LocalExecutionSlot=self.local_slot)}):
await rt.run_image_list_job(context)
self.assertIsNone(rt.get_master_job(context))
return self.scheduler.call_args.args[1]
async def test_busy_worker_at_enqueue_is_available_to_scheduler_when_idle_at_start(self):
responses = self.use_real_preflight({
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
})
result = await self.queue()
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
self.assertEqual(result[2], 1) # Candidates, not a promise of execution-time availability.
self.q.dispatch_worker_prompt.assert_not_awaited()
responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
# A newly enabled worker is not implicitly added to this queued selection.
self.config["workers"].append({"id": "later", "enabled": True})
slots = await self.run_prepared_selection()
self.assertEqual([slot for slot in slots if isinstance(slot, str)], ["worker"])
self.assertEqual(len(slots), 2)
self.assertEqual(self.preflight.await_count, 2)
self.assertEqual(self.preflight.await_args_list[0].args, self.preflight.await_args_list[1].args)
async def test_worker_still_busy_at_start_is_excluded_from_scheduler(self):
self.use_real_preflight({
"worker": {"queue_running": [], "queue_pending": [[1, "prompt-a"]]},
})
await self.queue()
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
slots = await self.run_prepared_selection()
self.assertEqual(len(slots), 1)
self.assertNotIn("worker", slots)
self.assertEqual(self.preflight.await_count, 2)
self.runtime_session.assert_not_awaited()
async def test_explicit_empty_selection_stays_master_only_at_start(self):
self.use_real_preflight({"worker": {"queue_running": [], "queue_pending": []}})
await self.queue(enabled_worker_ids=[])
slots = await self.run_prepared_selection()
self.assertEqual(len(slots), 1)
self.assertNotIn("worker", slots)
self.preflight.assert_not_awaited()
self.runtime_session.assert_not_awaited()
async def test_busy_incompatible_worker_fails_before_master_queue(self):
responses = self.use_real_preflight({
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
})
capability = responses[("worker", "capabilities")][1]
for field, invalid, message in (
("protocol_version", 1, "incompatible"),
("execution_targets", False, "execution_targets"),
("value_results", False, "value_results"),
):
with self.subTest(field=field), patch.dict(capability, {field: invalid}):
if field == "value_results":
self.prompt["3"]["inputs"]["value_10"] = ["2", 2]
with self.assertRaisesRegex(RuntimeError, message):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
self.q.dispatch_worker_prompt.assert_not_awaited()
async def test_orchestrator_with_all_workers_busy_queues_for_idle_start(self):
responses = self.use_real_preflight({
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
})
result = await self.queue(delegate_master=True)
self.q.queue_prompt_payload.assert_awaited_once()
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
self.assertEqual(result[2], 1)
self.q.dispatch_worker_prompt.assert_not_awaited()
responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
self.assertEqual(await self.run_prepared_selection(), ["worker"])
self.local_slot.assert_not_called()
self.assertEqual(self.preflight.await_count, 2)
self.assertIs(self.preflight.await_args_list[0].kwargs["require_idle"], False)
self.assertIs(self.preflight.await_args_list[1].kwargs["require_idle"], True)
async def test_orchestrator_rechecks_availability_without_master_fallback(self):
responses = self.use_real_preflight({"worker": {"queue_running": [], "queue_pending": []}})
await self.queue(delegate_master=True)
responses[("worker", "queue")] = (200, {"queue_running": [[1, "other"]], "queue_pending": []})
with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"):
await self.run_prepared_selection()
self.assertEqual(self.preflight.await_count, 2)
self.scheduler.assert_not_awaited()
self.local_slot.assert_not_called()
async def test_orchestrator_retains_busy_candidates_when_another_worker_is_idle(self):
self.config["workers"].append({"id": "idle", "enabled": True})
responses = self.use_real_preflight({
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
"idle": {"queue_running": [], "queue_pending": []},
})
await self.queue(delegate_master=True)
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker", "idle"])
responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
responses[("idle", "queue")] = (200, {"queue_running": [[1, "other"]], "queue_pending": []})
self.assertEqual(await self.run_prepared_selection(), ["worker"])
self.local_slot.assert_not_called()
async def test_orchestrator_without_workers_fails_before_queue(self):
self.config["settings"]["master_delegate_only"] = True