Files
robertvoy-ComfyUI-Distributed/tests/test_dispatch_selection.py
T
Robert Wojciechowski 7ed60b93ca feat: dynamically distribute image lists across direct workers (#98)
* feat: add Runpod serverless image collector pools

* fix: prevent worker settings clipping and remove preview copy

* fix: initialize worker settings before mounting the form

* fix: validate complete serverless requests against 10 MB limit

* fix: simplify serverless cards and keep status dots circular

* fix: match serverless endpoint box to master card

* fix: preserve serverless capacity and recovery across job failures

Retain durable reservations until provider termination or explicit recovery, protect unresolved pools and credentials, and add scoped reconciliation. Consolidate lifecycle and secret policies and split prompt/media stages with regression coverage.

* feat: add dynamic per-image list queue with Start and End nodes

* fix(image-list): reject upstream data bypasses around Start

* fix(image-list): make optional End text a connection-only input

* fix(image-list): isolate unavailable workers during preflight

* fix(image-list): bound acknowledgements and configure run deadline

* fix(media-sync): preserve output and temp source roots

* fix(image-list): own job lifecycle and preserve worker preparation

* feat(image-list): let master participate unless orchestrator-only

* feat(image-list)!: replace fixed Start extras with growable value_N slots

Distributed Image List Start no longer has reference_image, text,
strength, seed or per-slot mode dropdowns. It now takes images, mask and
up to 10 generic value_N inputs (IMAGE, MASK, STRING, INT, FLOAT,
BOOLEAN) that grow via native V3 Autogrow, each with a matching output.

- One value is shared by every image; one per image pairs in order; any
  other count fails before dispatch. Only images is split into frames.
- Outputs are declared up front (image, mask, item_index, context,
  value_1..value_10). web/imageListStart.js hides unused value outputs,
  mirrors input types, and follows Autogrow compaction so downstream
  links keep their value.
- Image-list protocol bumped to v2; older workers fail preflight.

Breaking: saved workflows using the old Start inputs must recreate Start.

* fix(serverless): disable Clear saved API key while the worker is enabled

The checkbox read "Clear saved API key (disable pool first)", but the
panel never defines "pool" and saving with it ticked on an enabled worker
was rejected by the backend.

- Disable the checkbox while the worker is enabled, with a tooltip saying
  to turn the worker off first. Also disable it when no key is saved.
- Use "worker" instead of "pool" in user-facing serverless errors, the
  serverless guide and the README entry.

* feat(serverless): replace Concurrent jobs with a single Images per run setting

- Remove the separate Concurrent jobs field; rename Jobs per run to
  Images per run (1-16) with a tooltip explaining that Runpod's max
  workers decides how many GPUs run them.
- Backend uses a fixed per-worker limit of 16 outstanding jobs, so all
  images of a run are submitted at once. Unresolved jobs still count
  against it. The limit is constant, so it never has to change while
  jobs are unresolved.
- Legacy 'concurrency' in saved configs and older browser tabs is
  accepted and ignored/dropped; it no longer triggers the busy guard.
- Revert the enable toggle and show the reason when the backend rejects it.
- Trim a pasted api.runpod.ai/v2/<id> URL to the endpoint ID; plainer
  endpoint-ID error otherwise.
- Confirm before a type change deletes a saved Runpod API key.
- Update docs and tests.
2026-10-13 15:20:05 -07:00

520 lines
20 KiB
Python

import asyncio
import importlib.util
import sys
import types
import unittest
from pathlib import Path
from unittest.mock import patch
def _load_dispatch_module():
module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py"
package_name = "dist_dispatch_testpkg"
root_pkg = types.ModuleType(package_name)
root_pkg.__path__ = []
sys.modules[package_name] = root_pkg
api_pkg = types.ModuleType(f"{package_name}.api")
api_pkg.__path__ = []
sys.modules[f"{package_name}.api"] = api_pkg
orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
orch_pkg.__path__ = []
sys.modules[f"{package_name}.api.orchestration"] = orch_pkg
utils_pkg = types.ModuleType(f"{package_name}.utils")
utils_pkg.__path__ = []
sys.modules[f"{package_name}.utils"] = utils_pkg
logging_module = types.ModuleType(f"{package_name}.utils.logging")
logging_module.debug_log = lambda *_args, **_kwargs: None
logging_module.log = lambda *_args, **_kwargs: None
sys.modules[f"{package_name}.utils.logging"] = logging_module
network_module = types.ModuleType(f"{package_name}.utils.network")
network_module.build_worker_url = lambda *_args, **_kwargs: "http://example.invalid"
async def _probe_worker(*_args, **_kwargs):
return None
network_module.probe_worker = _probe_worker
async def _fake_session():
raise RuntimeError("get_client_session should be mocked in these tests")
network_module.get_client_session = _fake_session
sys.modules[f"{package_name}.utils.network"] = network_module
created_aiohttp_stub = False
if "aiohttp" not in sys.modules:
created_aiohttp_stub = True
aiohttp_module = types.ModuleType("aiohttp")
class _ClientTimeout:
def __init__(self, total=None):
self.total = total
class _ClientError(Exception):
pass
class _ClientConnectorError(_ClientError):
pass
class _ClientConnectionError(_ClientError):
pass
class _WSMsgType:
TEXT = "TEXT"
ERROR = "ERROR"
CLOSED = "CLOSED"
class _TCPConnector:
def __init__(self, *args, **kwargs):
self.args = args
self.kwargs = kwargs
class _ClientSession:
def __init__(self, *args, **kwargs):
self.closed = False
async def close(self):
self.closed = True
aiohttp_module.ClientTimeout = _ClientTimeout
setattr(aiohttp_module, "ClientError", _ClientError)
aiohttp_module.ClientConnectorError = _ClientConnectorError
setattr(aiohttp_module, "ClientConnectionError", _ClientConnectionError)
aiohttp_module.WSMsgType = _WSMsgType
aiohttp_module.TCPConnector = _TCPConnector
aiohttp_module.ClientSession = _ClientSession
aiohttp_module.web = types.SimpleNamespace(
json_response=lambda payload, status=200: {"payload": payload, "status": status}
)
sys.modules["aiohttp"] = aiohttp_module
spec = importlib.util.spec_from_file_location(
f"{package_name}.api.orchestration.dispatch",
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)
if created_aiohttp_stub:
sys.modules.pop("aiohttp", None)
return module
dispatch = _load_dispatch_module()
def _load_prompt_transform_module():
module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "prompt_transform.py"
package_name = "prompt_transform_worker_selection_testpkg"
for suffix in ("", ".api", ".api.orchestration", ".utils"):
package = types.ModuleType(package_name + suffix)
package.__path__ = []
sys.modules[package_name + suffix] = package
logging_module = types.ModuleType(f"{package_name}.utils.logging")
setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None)
sys.modules[f"{package_name}.utils.logging"] = logging_module
spec = importlib.util.spec_from_file_location(
f"{package_name}.api.orchestration.prompt_transform",
module_path,
)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
prompt_transform = _load_prompt_transform_module()
class _FakeResponse:
def __init__(self, status=200, payload=None, error=None):
self.status = status
self.payload = payload
self.error = error
async def __aenter__(self):
if self.error is not None:
raise self.error
return self
async def __aexit__(self, *_args):
return None
async def json(self):
return self.payload
class _FakeSession:
def __init__(self, responses):
self.responses = responses
def get(self, url, timeout=None):
del timeout
worker_id = url.removeprefix("http://").split("/", 1)[0]
endpoint = "capabilities" if url.endswith("/capabilities") else "queue"
response = self.responses[(worker_id, endpoint)]
if isinstance(response, BaseException):
return _FakeResponse(error=response)
status, payload = response
return _FakeResponse(status=status, payload=payload)
def _valid_capability():
return {
"supported": True,
"protocol_version": 2,
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
}
def _valid_queue(running=None, pending=None):
return {"queue_running": running or [], "queue_pending": pending or []}
class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
async def test_select_active_workers_filters_offline(self):
workers = [
{"id": "w1", "name": "Worker 1"},
{"id": "w2", "name": "Worker 2"},
{"id": "w3", "name": "Worker 3"},
]
async def fake_probe(worker):
return worker["id"] != "w2"
with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
active_workers, delegate_master = await dispatch.select_active_workers(
workers,
use_websocket=False,
delegate_master=False,
probe_concurrency=3,
)
self.assertEqual([w["id"] for w in active_workers], ["w1", "w3"])
self.assertFalse(delegate_master)
async def test_select_active_workers_disables_delegate_when_all_offline(self):
workers = [{"id": "w1", "name": "Worker 1"}]
async def fake_probe(_worker):
return False
with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
active_workers, delegate_master = await dispatch.select_active_workers(
workers,
use_websocket=False,
delegate_master=True,
probe_concurrency=1,
)
self.assertEqual(active_workers, [])
self.assertFalse(delegate_master)
async def test_select_active_workers_uses_websocket_probe_when_enabled(self):
workers = [{"id": "w1", "name": "Worker 1"}, {"id": "w2", "name": "Worker 2"}]
async def fake_http_probe(_worker):
return False
async def fake_ws_probe(_worker):
return True
with patch.object(dispatch, "worker_is_active", side_effect=fake_http_probe) as http_probe, patch.object(
dispatch,
"worker_ws_is_active",
side_effect=fake_ws_probe,
) as ws_probe:
active_workers, _ = await dispatch.select_active_workers(
workers,
use_websocket=True,
delegate_master=False,
probe_concurrency=4,
)
self.assertEqual([w["id"] for w in active_workers], ["w1", "w2"])
self.assertEqual(ws_probe.call_count, 2)
self.assertEqual(http_probe.call_count, 0)
async def test_probe_concurrency_is_bounded(self):
workers = [{"id": f"w{i}", "name": f"Worker {i}"} for i in range(6)]
state = {"in_flight": 0, "max_in_flight": 0}
async def fake_probe(_worker):
state["in_flight"] += 1
state["max_in_flight"] = max(state["max_in_flight"], state["in_flight"])
await asyncio.sleep(0.01)
state["in_flight"] -= 1
return True
with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
active_workers, _ = await dispatch.select_active_workers(
workers,
use_websocket=False,
delegate_master=False,
probe_concurrency=2,
)
self.assertEqual(len(active_workers), len(workers))
self.assertLessEqual(state["max_in_flight"], 2)
self.assertGreaterEqual(state["max_in_flight"], 2)
async def test_select_least_busy_worker_round_robins_idle_workers(self):
workers = [
{"id": "w1", "name": "Worker 1"},
{"id": "w2", "name": "Worker 2"},
{"id": "w3", "name": "Worker 3"},
]
queue_map = {"w1": 0, "w2": 0, "w3": 2}
async def fake_probe(worker_url, timeout=3.0):
worker_id = worker_url.rsplit("/", 1)[-1]
return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
dispatch,
"probe_worker",
side_effect=fake_probe,
):
dispatch._least_busy_rr_index = 0
selected1 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
selected2 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
selected3 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
self.assertEqual(selected1["id"], "w1")
self.assertEqual(selected2["id"], "w2")
self.assertEqual(selected3["id"], "w1")
async def test_select_least_busy_worker_chooses_smallest_queue_when_all_busy(self):
workers = [
{"id": "w1", "name": "Worker 1"},
{"id": "w2", "name": "Worker 2"},
{"id": "w3", "name": "Worker 3"},
]
queue_map = {"w1": 5, "w2": 2, "w3": 4}
async def fake_probe(worker_url, timeout=3.0):
worker_id = worker_url.rsplit("/", 1)[-1]
return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
dispatch,
"probe_worker",
side_effect=fake_probe,
):
selected = await dispatch.select_least_busy_worker(workers, probe_concurrency=2)
self.assertEqual(selected["id"], "w2")
async def test_select_least_busy_worker_returns_none_when_all_probes_fail(self):
workers = [{"id": "w1", "name": "Worker 1"}]
async def fake_probe(_worker_url, timeout=3.0):
return None
with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
dispatch,
"probe_worker",
side_effect=fake_probe,
):
selected = await dispatch.select_least_busy_worker(workers, probe_concurrency=1)
self.assertIsNone(selected)
async def test_image_list_preflight_skips_unreachable_worker_and_keeps_healthy_worker(self):
workers = [{"id": "healthy"}, {"id": "offline"}]
responses = {
(worker_id, "capabilities"): (200, _valid_capability())
for worker_id in ("healthy", "offline")
}
responses.update({
(worker_id, "queue"): (200, _valid_queue())
for worker_id in ("healthy", "offline")
})
responses[("offline", "capabilities")] = dispatch.aiohttp.ClientConnectionError("connection refused")
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
), patch.object(dispatch, "log") as log_mock:
result = await dispatch.preflight_image_list_workers(workers)
self.assertEqual(result, {"healthy": True})
self.assertTrue(any("offline" in str(call) for call in log_mock.call_args_list))
async def test_image_list_preflight_skips_timeout_and_keeps_healthy_worker(self):
workers = [{"id": "healthy"}, {"id": "slow"}]
responses: dict = {
(worker_id, "capabilities"): (200, _valid_capability())
for worker_id in ("healthy", "slow")
}
responses.update({
(worker_id, "queue"): (200, _valid_queue())
for worker_id in ("healthy", "slow")
})
responses[("slow", "queue")] = asyncio.TimeoutError()
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
), patch.object(dispatch, "log") as log_mock:
result = await dispatch.preflight_image_list_workers(workers)
self.assertEqual(result, {"healthy": True})
self.assertTrue(any("slow" in str(call) for call in log_mock.call_args_list))
async def test_image_list_preflight_excludes_busy_workers(self):
workers = [{"id": "idle"}, {"id": "busy"}]
responses = {
(worker_id, "capabilities"): (200, _valid_capability())
for worker_id in ("idle", "busy")
}
responses.update({
("idle", "queue"): (200, _valid_queue()),
("busy", "queue"): (200, _valid_queue(running=[{"prompt": "existing"}])),
})
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
result = await dispatch.preflight_image_list_workers(workers)
self.assertEqual(result, {"idle": True})
async def test_image_list_preflight_reports_no_idle_workers_when_all_are_unavailable(self):
workers = [{"id": "offline-a"}, {"id": "offline-b"}]
responses = {
(worker_id, "capabilities"): dispatch.aiohttp.ClientConnectionError("unreachable")
for worker_id in ("offline-a", "offline-b")
}
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
self.assertEqual(
await dispatch.preflight_image_list_workers(workers, require_idle=False),
{},
)
with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"):
await dispatch.preflight_image_list_workers(workers, require_idle=True)
async def test_image_list_preflight_fails_closed_for_malformed_capabilities(self):
malformed = {"supported": True, "protocol_version": 2, "node_ids": None}
session = _FakeSession({("worker", "capabilities"): (200, malformed)})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "malformed.*capabilit"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_preflight_fails_closed_for_incompatible_protocol(self):
incompatible = _valid_capability()
incompatible["protocol_version"] = 1
session = _FakeSession({("worker", "capabilities"): (200, incompatible)})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "incompatible"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_preflight_fails_closed_for_missing_required_node(self):
incompatible = _valid_capability()
incompatible["node_ids"] = ["DistributedImageListStart"]
session = _FakeSession({("worker", "capabilities"): (200, incompatible)})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "incompatible"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_preflight_fails_closed_for_malformed_queue(self):
responses = {
("worker", "capabilities"): (200, _valid_capability()),
("worker", "queue"): (200, {"queue_running": [], "queue_pending": None}),
}
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "malformed /queue"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_preflight_does_not_hide_unexpected_errors(self):
session = _FakeSession({("worker", "capabilities"): ValueError("programming error")})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(ValueError, "programming error"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
class ImageListDirectWorkerResolutionTests(unittest.TestCase):
def setUp(self):
self.config = {"workers": [
{"id": "enabled", "enabled": True, "host": "enabled-host", "port": 8188},
{"id": "disabled", "enabled": False, "host": "disabled-host", "port": 8189},
{"id": "serverless", "enabled": True, "type": "serverless"},
{"id": "other", "enabled": True, "host": "other-host", "port": 8190},
]}
def test_implicit_selection_uses_only_enabled_direct_workers(self):
workers = prompt_transform._resolve_direct_workers_for_image_list(self.config, None)
self.assertEqual([worker["id"] for worker in workers], ["enabled", "other"])
def test_explicit_selection_includes_exact_disabled_direct_worker(self):
workers = prompt_transform._resolve_direct_workers_for_image_list(self.config, ["disabled"])
self.assertEqual([worker["id"] for worker in workers], ["disabled"])
def test_explicit_empty_selection_returns_no_workers(self):
self.assertEqual(prompt_transform._resolve_direct_workers_for_image_list(self.config, []), [])
def test_unknown_explicit_id_raises_instead_of_falling_back(self):
with self.assertRaisesRegex(RuntimeError, "missing"):
prompt_transform._resolve_direct_workers_for_image_list(self.config, ["enabled", "missing"])
def test_explicit_serverless_selection_is_excluded(self):
self.assertEqual(prompt_transform._resolve_direct_workers_for_image_list(self.config, ["serverless"]), [])
if __name__ == "__main__":
unittest.main()