From bac16afaade8151fe505327c5f53d9e76b581e59 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sun, 11 Oct 2026 20:57:20 +0000 Subject: [PATCH] fix(image-list): isolate unavailable workers during preflight --- api/orchestration/dispatch.py | 76 +++++--- api/orchestration/prompt_transform.py | 24 ++- tests/test_dispatch_selection.py | 267 +++++++++++++++++++++++++- 3 files changed, 335 insertions(+), 32 deletions(-) diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index 3ef6da4..e1e5bbb 100644 --- a/api/orchestration/dispatch.py +++ b/api/orchestration/dispatch.py @@ -173,9 +173,10 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0): async def inspect(worker): worker_id = str(worker.get("id") or "unknown") - async with semaphore: - cap_url = build_worker_url(worker, "/distributed/image-list/capabilities") - try: + check = "capability" + try: + async with semaphore: + cap_url = build_worker_url(worker, "/distributed/image-list/capabilities") async with session.get(cap_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response: if response.status != 200: raise RuntimeError( @@ -183,37 +184,58 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0): f"(HTTP {response.status}); upgrade that worker before running this feature" ) capability = await response.json() - except asyncio.TimeoutError as exc: - raise RuntimeError(f"Image-list worker {worker_id} capability check timed out") from exc - except aiohttp.ClientError as exc: - raise RuntimeError(f"Image-list worker {worker_id} capability check failed: {exc}") from exc - nodes = set(capability.get("node_ids", [])) if isinstance(capability, dict) else set() - if (not isinstance(capability, dict) or capability.get("supported") is not True - or capability.get("protocol_version") != 1 or not required_nodes.issubset(nodes)): - raise RuntimeError( - f"Image-list worker {worker_id} is incompatible: requires protocol v1 and nodes " - f"{sorted(required_nodes)}; no work was submitted" - ) - queue_url = build_worker_url(worker, "/queue") - try: + if not isinstance(capability, dict): + raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data") + supported = capability.get("supported") + protocol_version = capability.get("protocol_version") + nodes = capability.get("node_ids") + if ( + not isinstance(supported, bool) + or type(protocol_version) is not int + or not isinstance(nodes, list) + or any(not isinstance(node_id, str) for node_id in nodes) + ): + raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data") + if ( + supported is not True + or protocol_version != 1 + or not required_nodes.issubset(nodes) + ): + raise RuntimeError( + f"Image-list worker {worker_id} is incompatible: requires protocol v1 and nodes " + f"{sorted(required_nodes)}; no work was submitted" + ) + + check = "queue" + queue_url = build_worker_url(worker, "/queue") async with session.get(queue_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response: if response.status != 200: raise RuntimeError(f"Image-list worker {worker_id} queue preflight returned HTTP {response.status}") queue = await response.json() - except asyncio.TimeoutError as exc: - raise RuntimeError(f"Image-list worker {worker_id} queue preflight timed out") from exc - except aiohttp.ClientError as exc: - raise RuntimeError(f"Image-list worker {worker_id} queue preflight failed: {exc}") from exc - if not isinstance(queue, dict) or not isinstance(queue.get("queue_running"), list) or not isinstance(queue.get("queue_pending"), list): - raise RuntimeError(f"Image-list worker {worker_id} returned malformed /queue state") - idle = not queue["queue_running"] and not queue["queue_pending"] - return worker_id, idle + if ( + not isinstance(queue, dict) + or not isinstance(queue.get("queue_running"), list) + or not isinstance(queue.get("queue_pending"), list) + ): + raise RuntimeError(f"Image-list worker {worker_id} returned malformed /queue state") + idle = not queue["queue_running"] and not queue["queue_pending"] + return worker_id, idle + except (asyncio.TimeoutError, aiohttp.ClientConnectionError) as exc: + log( + f"[Distributed] Image-list worker {worker_id} unavailable during {check} check; " + f"skipping: {exc}" + ) + return worker_id, None statuses = await asyncio.gather(*(inspect(worker) for worker in workers)) - result = {worker_id: idle for worker_id, idle in statuses} - if require_idle and not any(result.values()): + idle_workers = { + worker_id: True + for worker_id, idle in statuses + if idle is True + } + if require_idle and not idle_workers: raise RuntimeError("No selected feature-compatible direct worker is idle; unrelated worker prompts were left untouched") - return {worker_id: True for worker_id, idle in result.items() if idle} + return idle_workers async def select_active_workers( diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index 2bc79db..c03a039 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -442,14 +442,30 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids): def _resolve_direct_workers_for_image_list(config, requested_ids): - """Resolve only enabled direct workers already selected by orchestration.""" + """Resolve enabled workers by default, or exact explicitly selected direct workers.""" requested = None if requested_ids is None else {str(worker_id) for worker_id in requested_ids} + configured_workers = (config or {}).get("workers", []) + configured_ids = { + str(worker.get("id") or "").strip() + for worker in configured_workers + if str(worker.get("id") or "").strip() + } + if requested is not None: + unknown_ids = requested - configured_ids + if unknown_ids: + raise RuntimeError( + "Unknown explicitly selected image-list worker ID(s): " + + ", ".join(sorted(unknown_ids)) + ) direct = [] - for worker in (config or {}).get("workers", []): + for worker in configured_workers: worker_id = str(worker.get("id") or "").strip() - if not worker_id or not worker.get("enabled", False): + if not worker_id: continue - if requested is not None and worker_id not in requested: + if requested is None: + if not worker.get("enabled", False): + continue + elif worker_id not in requested: continue if str(worker.get("type") or "local").strip().lower() == "serverless": continue diff --git a/tests/test_dispatch_selection.py b/tests/test_dispatch_selection.py index 8792e72..fedf562 100644 --- a/tests/test_dispatch_selection.py +++ b/tests/test_dispatch_selection.py @@ -55,7 +55,13 @@ def _load_dispatch_module(): def __init__(self, total=None): self.total = total - class _ClientConnectorError(Exception): + class _ClientError(Exception): + pass + + class _ClientConnectorError(_ClientError): + pass + + class _ClientConnectionError(_ClientError): pass class _WSMsgType: @@ -76,7 +82,9 @@ def _load_dispatch_module(): 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 @@ -100,6 +108,75 @@ def _load_dispatch_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": 1, + "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 = [ @@ -249,6 +326,194 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase): 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": 1, "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"] = 2 + 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()