perf(zai): parallelize endpoint detection probes

Z.AI has separate billing for general vs coding plans and global vs
China endpoints. On startup, detect_zai_endpoint() probes up to 4
endpoints sequentially with 8s timeout each, taking 8-9 seconds when
the first endpoints return non-200 (rate limited) before a working one
is found.

Replace the sequential loop with concurrent.futures.ThreadPoolExecutor
to probe all 4 endpoints in parallel. Results are returned in
ZAI_ENDPOINTS priority order so the preference chain is preserved.

Benchmark on macOS M4 Max, Python 3.11, Hermes v0.8.0:
  Before: 8.8s (sequential: global=0.9s/429, cn=1.6s/429,
           coding-global=4.3s/200, coding-cn=2.0s/200)
  After:  ~4.5s (single round-trip, bounded by slowest endpoint)

Signed-off-by: Merlin <merlin@merlin.me>
This commit is contained in:
light-merlin-dark 2026-04-11 11:12:55 -05:00 committed by kshitij
parent 2b0d58e88b
commit 9891f4b63f
1 changed files with 64 additions and 30 deletions

View File

@ -648,41 +648,75 @@ ZAI_ENDPOINTS = [
]
def _probe_single_zai_endpoint(
api_key: str, endpoint: tuple, timeout: float,
) -> Optional[Dict[str, str]]:
"""Probe a single Z.AI endpoint. Returns endpoint info dict or None.
Preserves the per-endpoint candidate-model loop: endpoints carry a
``probe_models`` LIST and each model is tried in order until one
succeeds (some plans only accept newer/older GLM slugs).
"""
ep_id, base_url, probe_models, label = endpoint
for model in probe_models:
try:
resp = httpx.post(
f"{base_url}/chat/completions",
headers={
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
},
json={
"model": model,
"stream": False,
"max_tokens": 1,
"messages": [{"role": "user", "content": "ping"}],
},
timeout=timeout,
)
if resp.status_code == 200:
logger.debug("Z.AI endpoint probe: %s (%s) model=%s OK", ep_id, base_url, model)
return {
"id": ep_id,
"base_url": base_url,
"model": model,
"label": label,
}
logger.debug("Z.AI endpoint probe: %s model=%s returned %s", ep_id, model, resp.status_code)
except Exception as exc:
logger.debug("Z.AI endpoint probe: %s model=%s failed: %s", ep_id, model, exc)
return None
def detect_zai_endpoint(api_key: str, timeout: float = 8.0) -> Optional[Dict[str, str]]:
"""Probe z.ai endpoints to find one that accepts this API key.
"""Probe z.ai endpoints in parallel to find one that accepts this API key.
Returns {"id": ..., "base_url": ..., "model": ..., "label": ...} for the
first working endpoint, or None if all fail. For endpoints with multiple
candidate models, tries each in order and returns the first that succeeds.
first working endpoint (in ZAI_ENDPOINTS priority order), or None if all
fail. For endpoints with multiple candidate models, each worker tries
its endpoint's models in order and returns the first that succeeds.
"""
for ep_id, base_url, probe_models, label in ZAI_ENDPOINTS:
for model in probe_models:
from concurrent.futures import ThreadPoolExecutor, as_completed
with ThreadPoolExecutor(max_workers=len(ZAI_ENDPOINTS)) as pool:
futures = {
pool.submit(_probe_single_zai_endpoint, api_key, ep, timeout): ep[0]
for ep in ZAI_ENDPOINTS
}
results: Dict[str, Dict[str, str]] = {}
for future in as_completed(futures):
ep_id = futures[future]
try:
resp = httpx.post(
f"{base_url}/chat/completions",
headers={
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
},
json={
"model": model,
"stream": False,
"max_tokens": 1,
"messages": [{"role": "user", "content": "ping"}],
},
timeout=timeout,
)
if resp.status_code == 200:
logger.debug("Z.AI endpoint probe: %s (%s) model=%s OK", ep_id, base_url, model)
return {
"id": ep_id,
"base_url": base_url,
"model": model,
"label": label,
}
logger.debug("Z.AI endpoint probe: %s model=%s returned %s", ep_id, model, resp.status_code)
except Exception as exc:
logger.debug("Z.AI endpoint probe: %s model=%s failed: %s", ep_id, model, exc)
result = future.result()
if result is not None:
results[ep_id] = result
except Exception:
pass
# 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