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:
parent
2b0d58e88b
commit
9891f4b63f
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue