From 1e06196ff2fed5365f385d746bf6f903494f94f4 Mon Sep 17 00:00:00 2001 From: TianyeDong Date: Sat, 30 May 2026 00:14:36 -0400 Subject: [PATCH] refactor(miles): share router URL acquisition logic --- miles/router/router.py | 14 ++++++++------ tests/fast/router/test_router.py | 8 ++++++++ 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/miles/router/router.py b/miles/router/router.py index 63b95a5fa7..24f04d6161 100644 --- a/miles/router/router.py +++ b/miles/router/router.py @@ -538,6 +538,12 @@ def _candidate_set(self) -> set[str]: return self.enabled_workers - self.dead_workers return set(self.worker_request_counts) - self.dead_workers + def _pick_least_loaded(self, candidates: set[str]) -> str: + """Select the least-loaded candidate and increment its in-flight count.""" + url = min(candidates, key=lambda u: self.worker_request_counts.get(u, 0)) + self.worker_request_counts[url] += 1 + return url + def _use_url(self): """Synchronous, raise-on-empty selector. @@ -553,9 +559,7 @@ def _use_url(self): candidates = self._candidate_set() if not candidates: raise RuntimeError("No enabled live workers available in the pool") - url = min(candidates, key=lambda u: self.worker_request_counts.get(u, 0)) - self.worker_request_counts[url] += 1 - return url + return self._pick_least_loaded(candidates) async def _use_url_async(self): """C20 0-active suspend selector (production dispatch path). @@ -574,9 +578,7 @@ async def _use_url_async(self): async with self._workers_changed: await self._workers_changed.wait_for(lambda: bool(self._candidate_set())) candidates = self._candidate_set() - url = min(candidates, key=lambda u: self.worker_request_counts.get(u, 0)) - self.worker_request_counts[url] += 1 - return url + return self._pick_least_loaded(candidates) async def _notify_workers_changed(self) -> None: """Wake every dispatcher suspended in :meth:`_use_url`. diff --git a/tests/fast/router/test_router.py b/tests/fast/router/test_router.py index 39383f4135..ceca2634ac 100644 --- a/tests/fast/router/test_router.py +++ b/tests/fast/router/test_router.py @@ -142,6 +142,14 @@ def test_use_url_selects_min_load(self, router_factory): assert selected == "http://w2:8000" assert router.worker_request_counts["http://w2:8000"] == 3 + def test_use_url_async_selects_min_load(self, router_factory): + router = router_factory() + router.worker_request_counts = {"http://w1:8000": 5, "http://w2:8000": 2, "http://w3:8000": 8} + + selected = asyncio.run(router._use_url_async()) + assert selected == "http://w2:8000" + assert router.worker_request_counts["http://w2:8000"] == 3 + def test_use_url_excludes_dead_workers(self, router_factory): router = router_factory() router.worker_request_counts = {"http://w1:8000": 5, "http://w2:8000": 1, "http://w3:8000": 3}