perf(zai): early-exit when the highest-priority endpoint wins (simplify finding)
The as_completed drain + `with` join made the parallel version WORSE than sequential main in the common case (first endpoint succeeds fast, others slow/unreachable): main returned at first success, the parallel version waited for every straggler. Now: after each completion, walk endpoints in priority order and return as soon as a success is unbeatable (all higher-priority probes already finished); pool uses shutdown(wait=False) so losers drain in the background. Mutation-checked: removing the early exit makes the new timing test fail (8.2s vs <1.5s).
This commit is contained in:
parent
9e99a335a7
commit
0e4daade14
|
|
@ -698,11 +698,17 @@ def detect_zai_endpoint(api_key: str, timeout: float = 8.0) -> Optional[Dict[str
|
|||
"""
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
with ThreadPoolExecutor(max_workers=len(ZAI_ENDPOINTS)) as pool:
|
||||
# No `with` block: a context manager would join ALL probe threads on
|
||||
# exit, defeating the early return below. shutdown(wait=False) lets the
|
||||
# surviving daemon-style probes drain in the background instead of
|
||||
# blocking the caller on slow/unreachable endpoints.
|
||||
pool = ThreadPoolExecutor(max_workers=len(ZAI_ENDPOINTS))
|
||||
try:
|
||||
futures = {
|
||||
pool.submit(_probe_single_zai_endpoint, api_key, ep, timeout): ep[0]
|
||||
for ep in ZAI_ENDPOINTS
|
||||
}
|
||||
by_id = {ep_id: f for f, ep_id in futures.items()}
|
||||
results: Dict[str, Dict[str, str]] = {}
|
||||
for future in as_completed(futures):
|
||||
ep_id = futures[future]
|
||||
|
|
@ -712,12 +718,24 @@ def detect_zai_endpoint(api_key: str, timeout: float = 8.0) -> Optional[Dict[str
|
|||
results[ep_id] = result
|
||||
except Exception:
|
||||
pass
|
||||
# Early exit in PRIORITY order: walk endpoints highest-priority
|
||||
# first; if one has succeeded and every higher-priority probe
|
||||
# has already finished (without success), no later completion
|
||||
# can win — return now instead of waiting out slow endpoints
|
||||
# (main's sequential loop also stopped at first success).
|
||||
for ep in ZAI_ENDPOINTS:
|
||||
if not by_id[ep[0]].done():
|
||||
break # a higher-priority probe is still in flight
|
||||
if ep[0] in results:
|
||||
return results[ep[0]]
|
||||
|
||||
# Return first match in priority order (ZAI_ENDPOINTS list order)
|
||||
for ep in ZAI_ENDPOINTS:
|
||||
if ep[0] in results:
|
||||
return results[ep[0]]
|
||||
return None
|
||||
# All probes finished: first match in priority order, if any.
|
||||
for ep in ZAI_ENDPOINTS:
|
||||
if ep[0] in results:
|
||||
return results[ep[0]]
|
||||
return None
|
||||
finally:
|
||||
pool.shutdown(wait=False)
|
||||
|
||||
|
||||
def _resolve_zai_base_url(api_key: str, default_url: str, env_override: str) -> str:
|
||||
|
|
|
|||
|
|
@ -712,6 +712,28 @@ class TestZaiParallelProbe:
|
|||
monkeypatch.setattr("hermes_cli.auth.httpx.post", self._mock_post({}))
|
||||
assert detect_zai_endpoint("bad-key", timeout=1.0) is None
|
||||
|
||||
def test_early_exit_does_not_wait_for_slow_losers(self, monkeypatch):
|
||||
"""When the highest-priority endpoint succeeds fast, the caller must
|
||||
return without waiting for slow lower-priority probes to finish."""
|
||||
import time as _time
|
||||
|
||||
from hermes_cli.auth import ZAI_ENDPOINTS, detect_zai_endpoint
|
||||
|
||||
first = ZAI_ENDPOINTS[0]
|
||||
inner = self._mock_post({(first[1], first[2][0]): True})
|
||||
|
||||
def _slow_losers(url, headers=None, json=None, timeout=None):
|
||||
if not url.startswith(first[1]):
|
||||
_time.sleep(2.0) # slow lower-priority endpoints
|
||||
return inner(url, headers=headers, json=json, timeout=timeout)
|
||||
|
||||
monkeypatch.setattr("hermes_cli.auth.httpx.post", _slow_losers)
|
||||
t0 = _time.perf_counter()
|
||||
result = detect_zai_endpoint("test-key", timeout=5.0)
|
||||
elapsed = _time.perf_counter() - t0
|
||||
assert result is not None and result["id"] == first[0]
|
||||
assert elapsed < 1.5, f"early exit failed: waited {elapsed:.2f}s for losers"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Kimi / Moonshot model list isolation tests
|
||||
|
|
|
|||
Loading…
Reference in New Issue