213 lines
8.2 KiB
Python
213 lines
8.2 KiB
Python
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()
|