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:
kshitijk4poor 2026-08-03 17:30:28 +05:30 committed by kshitij
parent 9e99a335a7
commit 0e4daade14
2 changed files with 46 additions and 6 deletions

View File

@ -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:

View File

@ -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