255 lines
9.0 KiB
Python
255 lines
9.0 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 _ClientConnectorError(Exception):
|
|
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
|
|
aiohttp_module.ClientConnectorError = _ClientConnectorError
|
|
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()
|
|
|
|
|
|
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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|