diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index e1e5bbb..0aa8251 100644 --- a/api/orchestration/dispatch.py +++ b/api/orchestration/dispatch.py @@ -77,27 +77,44 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0): ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://") ws_url = f"{ws_url}/distributed/worker_ws" session = await get_client_session() + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + ws = None - async with session.ws_connect(ws_url, heartbeat=20, timeout=timeout) as ws: - await ws.send_json(ws_payload) - async for msg in ws: - if msg.type == aiohttp.WSMsgType.TEXT: - data = json.loads(msg.data or "{}") - if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id: - if data.get("ok"): - return data.get("prompt_id") - error_text = data.get("error") or "Worker rejected websocket dispatch." - validation_error = data.get("validation_error") - node_errors = data.get("node_errors") - if validation_error: - error_text = f"{error_text} | validation_error={validation_error}" - if node_errors: - error_text = f"{error_text} | node_errors={node_errors}" - raise RuntimeError(error_text) - elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED): - raise RuntimeError(f"Worker websocket closed unexpectedly: {msg.type}") - - raise RuntimeError("Worker websocket closed before dispatch_ack was received.") + try: + async with asyncio.timeout_at(deadline): + ws = await session.ws_connect(ws_url, heartbeat=20) + await ws.send_json(ws_payload) + while True: + msg = await ws.receive() + if msg.type == aiohttp.WSMsgType.TEXT: + data = json.loads(msg.data or "{}") + if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id: + if data.get("ok"): + return data.get("prompt_id") + error_text = data.get("error") or "Worker rejected websocket dispatch." + validation_error = data.get("validation_error") + node_errors = data.get("node_errors") + if validation_error: + error_text = f"{error_text} | validation_error={validation_error}" + if node_errors: + error_text = f"{error_text} | node_errors={node_errors}" + raise RuntimeError(error_text) + elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED): + raise RuntimeError(f"Worker websocket closed before acknowledgement: {msg.type}; submission status is ambiguous") + except asyncio.TimeoutError as exc: + raise RuntimeError( + "Worker websocket dispatch timed out before dispatch_ack; submission status is ambiguous." + ) from exc + finally: + if ws is not None: + # Closing must not discard an already acknowledged prompt ID or + # replace a definite rejection with an ambiguous timeout. + try: + async with asyncio.timeout_at(min(deadline, loop.time() + 1.0)): + await ws.close() + except (asyncio.TimeoutError, aiohttp.ClientError): + pass async def dispatch_worker_prompt( diff --git a/tests/test_image_list_timeout_config.py b/tests/test_image_list_timeout_config.py new file mode 100644 index 0000000..f4a1739 --- /dev/null +++ b/tests/test_image_list_timeout_config.py @@ -0,0 +1,106 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + + +def _load_config_module(): + module_path = Path(__file__).resolve().parents[1] / "utils" / "config.py" + package_name = "image_list_timeout_config_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__ = [str(module_path.parent)] + sys.modules[package_name] = root_pkg + + logging_module = types.ModuleType(f"{package_name}.logging") + logging_module.log = lambda *_args, **_kwargs: None + logging_module.debug_log = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.logging"] = logging_module + + constants_module = types.ModuleType(f"{package_name}.constants") + constants_module.HEARTBEAT_TIMEOUT = 30 + sys.modules[f"{package_name}.constants"] = constants_module + + spec = importlib.util.spec_from_file_location(f"{package_name}.config", 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 + + +config_module = _load_config_module() + + +SETTING_KEY = "image_list_job_timeout_seconds" + + +def _config_with_timeout(value): + return {"settings": {SETTING_KEY: value}} + + +def test_default_config_sets_image_list_job_timeout_to_zero(): + assert config_module.get_default_config()["settings"][SETTING_KEY] == 0.0 + + +def test_missing_value_defaults_to_disabled_timeout(): + assert config_module.get_image_list_job_timeout_seconds({}) == 0.0 + assert config_module.get_image_list_job_timeout_seconds({"settings": {}}) == 0.0 + + +def test_none_config_loads_current_config(monkeypatch): + cfg = _config_with_timeout("12.5") + monkeypatch.setattr(config_module, "load_config", lambda: cfg) + + assert config_module.get_image_list_job_timeout_seconds() == 12.5 + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + (0, 0.0), + (0.0, 0.0), + ("0", 0.0), + (" 0.0 ", 0.0), + (1, 1.0), + (2.5, 2.5), + ("3", 3.0), + (" 4.25 ", 4.25), + ], +) +def test_accepts_zero_and_positive_numeric_values(value, expected): + result = config_module.get_image_list_job_timeout_seconds(_config_with_timeout(value)) + + assert result == expected + assert isinstance(result, float) + + +@pytest.mark.parametrize( + "value", + [ + True, + False, + -1, + -0.01, + float("nan"), + float("inf"), + float("-inf"), + "not-a-number", + "", + " ", + "-1", + "nan", + "inf", + None, + [], + {}, + ], +) +def test_rejects_invalid_values_with_setting_name(value): + with pytest.raises(ValueError, match=SETTING_KEY): + config_module.get_image_list_job_timeout_seconds(_config_with_timeout(value)) diff --git a/tests/test_image_list_websocket_deadline.py b/tests/test_image_list_websocket_deadline.py new file mode 100644 index 0000000..0bcadbc --- /dev/null +++ b/tests/test_image_list_websocket_deadline.py @@ -0,0 +1,212 @@ +import asyncio +import importlib.util +import sys +import types +import unittest +import uuid +from pathlib import Path +from typing import Any, cast +from unittest.mock import patch + +import aiohttp +from aiohttp import web + + +def _load_dispatch_module(): + module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py" + package_name = f"image_list_ws_deadline_testpkg_{uuid.uuid4().hex}" + package_modules = {} + + for suffix in ("", ".api", ".api.orchestration", ".utils"): + package = types.ModuleType(package_name + suffix) + package.__path__ = [] + package_modules[package_name + suffix] = package + + logging_module = types.ModuleType(f"{package_name}.utils.logging") + setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None) + setattr(logging_module, "log", lambda *_args, **_kwargs: None) + package_modules[logging_module.__name__] = logging_module + + network_module = types.ModuleType(f"{package_name}.utils.network") + setattr(network_module, "build_worker_url", lambda *_args, **_kwargs: "http://example.invalid") + + async def _unused_probe_worker(*_args, **_kwargs): + return None + + async def _unused_get_client_session(*_args, **_kwargs): + raise RuntimeError("get_client_session must be patched by this test") + + setattr(network_module, "probe_worker", _unused_probe_worker) + setattr(network_module, "get_client_session", _unused_get_client_session) + package_modules[network_module.__name__] = network_module + + module_name = f"{package_name}.api.orchestration.dispatch" + spec = importlib.util.spec_from_file_location(module_name, module_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + package_modules[module_name] = module + + previous_modules = {name: sys.modules.get(name) for name in package_modules} + try: + sys.modules.update(package_modules) + spec.loader.exec_module(module) + finally: + for name, previous in previous_modules.items(): + if previous is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous + return module + + +dispatch = _load_dispatch_module() + + +class WebsocketDispatchDeadlineTests(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.mode = None + self.received_payload = None + self.dispatch_received = asyncio.Event() + self.peer_finished = asyncio.Event() + self.irrelevant_messages_sent = 0 + + app = web.Application() + app.router.add_get("/distributed/worker_ws", self._handle_websocket) + self.runner = web.AppRunner(app) + await self.runner.setup() + self.site = web.TCPSite(self.runner, "127.0.0.1", 0) + await self.site.start() + server = cast(Any, self.site._server) + port = server.sockets[0].getsockname()[1] + self.worker_url = f"http://127.0.0.1:{port}" + self.session = aiohttp.ClientSession() + + async def asyncTearDown(self): + await self.session.close() + await self.runner.cleanup() + + async def _get_client_session(self): + return self.session + + async def _handle_websocket(self, request): + ws = web.WebSocketResponse() + await ws.prepare(request) + try: + payload = await ws.receive_json() + self.received_payload = payload + self.dispatch_received.set() + + if self.mode in ("success", "slow-close"): + await ws.send_json({ + "type": "dispatch_ack", + "request_id": "unrelated-request", + "ok": True, + "prompt_id": "wrong-prompt", + }) + await ws.send_json({"type": "heartbeat"}) + await ws.send_json({ + "type": "dispatch_ack", + "request_id": payload["request_id"], + "ok": True, + "prompt_id": "prompt-123", + }) + if self.mode == "slow-close": + await asyncio.sleep(0.15) + elif self.mode == "rejection": + await ws.send_json({ + "type": "dispatch_ack", + "request_id": payload["request_id"], + "ok": False, + "error": "Invalid prompt", + "validation_error": {"node": "bad"}, + "node_errors": {"bad": "missing input"}, + }) + elif self.mode in ("no-ack", "wait-for-cancel"): + async def send_irrelevant_messages(): + while True: + try: + await ws.send_json({"type": "heartbeat"}) + self.irrelevant_messages_sent += 1 + await ws.send_json({ + "type": "dispatch_ack", + "request_id": "unrelated-request", + "ok": True, + "prompt_id": "not-this-request", + }) + self.irrelevant_messages_sent += 1 + except (ConnectionError, RuntimeError): + return + await asyncio.sleep(0.01) + + sender = asyncio.create_task(send_irrelevant_messages()) + try: + async for _message in ws: + pass + finally: + sender.cancel() + try: + await sender + except asyncio.CancelledError: + pass + finally: + self.peer_finished.set() + return ws + + async def _dispatch(self, timeout): + with patch.object(dispatch, "get_client_session", self._get_client_session): + return await dispatch._dispatch_via_websocket( + self.worker_url, + {"prompt": {"node": {"class_type": "Example"}}, "workflow": {"name": "test"}}, + "client-456", + timeout=timeout, + ) + + async def test_returns_matching_acknowledgement_prompt_id(self): + self.mode = "success" + + prompt_id = await self._dispatch(timeout=0.5) + + self.assertEqual(prompt_id, "prompt-123") + received_payload = self.received_payload + assert received_payload is not None + self.assertEqual(received_payload["type"], "dispatch_prompt") + self.assertTrue(received_payload["request_id"]) + self.assertEqual(received_payload["client_id"], "client-456") + self.assertEqual(received_payload["prompt"], {"node": {"class_type": "Example"}}) + self.assertEqual(received_payload["workflow"], {"name": "test"}) + + async def test_acknowledged_prompt_id_survives_slow_close(self): + self.mode = "slow-close" + self.assertEqual(await self._dispatch(timeout=0.03), "prompt-123") + + async def test_matching_rejection_preserves_worker_error_details(self): + self.mode = "rejection" + + with self.assertRaisesRegex(RuntimeError, "Invalid prompt.*validation_error=.*node_errors="): + await self._dispatch(timeout=0.5) + + async def test_repeated_irrelevant_messages_do_not_reset_ack_deadline(self): + self.mode = "no-ack" + started = asyncio.get_running_loop().time() + + with self.assertRaisesRegex(RuntimeError, "submission status is ambiguous"): + await asyncio.wait_for(self._dispatch(timeout=0.12), timeout=0.5) + + elapsed = asyncio.get_running_loop().time() - started + self.assertLess(elapsed, 0.5) + self.assertGreaterEqual(self.irrelevant_messages_sent, 4) + await asyncio.wait_for(self.peer_finished.wait(), timeout=0.5) + + async def test_outer_cancellation_propagates_and_closes_peer(self): + self.mode = "wait-for-cancel" + task = asyncio.create_task(self._dispatch(timeout=2.0)) + await asyncio.wait_for(self.dispatch_received.wait(), timeout=0.5) + + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=0.5) + await asyncio.wait_for(self.peer_finished.wait(), timeout=0.5) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/config.py b/utils/config.py index 324762e..388b5c3 100644 --- a/utils/config.py +++ b/utils/config.py @@ -5,6 +5,7 @@ import asyncio import copy import os import json +import math from contextlib import asynccontextmanager from .logging import log from .secret_fields import is_secret_field_name as _is_secret_field_name @@ -81,7 +82,8 @@ def get_default_config(): "worker_probe_concurrency": 8, "worker_prep_concurrency": 4, "media_sync_concurrency": 2, - "media_sync_timeout_seconds": 120 + "media_sync_timeout_seconds": 120, + "image_list_job_timeout_seconds": 0.0 }, "tunnel": { "status": "stopped", @@ -215,6 +217,37 @@ def get_worker_timeout_seconds(default: int = HEARTBEAT_TIMEOUT) -> int: return max(1, int(default)) +def get_image_list_job_timeout_seconds(config=None) -> float: + """Return the validated image-list total-job timeout in seconds. + + A value of zero disables the total-job deadline. Missing settings default + to zero; configured values must be finite, non-negative numbers. + """ + setting_key = "settings.image_list_job_timeout_seconds" + cfg = load_config() if config is None else config + if not isinstance(cfg, dict): + raise ValueError(f"{setting_key} must be a finite non-negative number") + + settings = cfg.get("settings", {}) + if not isinstance(settings, dict): + raise ValueError(f"{setting_key} must be a finite non-negative number") + + value = settings.get("image_list_job_timeout_seconds", 0.0) + if isinstance(value, bool) or not isinstance(value, (int, float, str)): + raise ValueError(f"{setting_key} must be a finite non-negative number") + if isinstance(value, str) and not value.strip(): + raise ValueError(f"{setting_key} must be a finite non-negative number") + + try: + timeout = float(value) + except (OverflowError, ValueError): + raise ValueError(f"{setting_key} must be a finite non-negative number") from None + + if not math.isfinite(timeout) or timeout < 0: + raise ValueError(f"{setting_key} must be a finite non-negative number") + return timeout + + def is_master_delegate_only() -> bool: """Returns True when master should skip local workload and act as orchestrator only.""" try: