141 lines
4.6 KiB
Python
141 lines
4.6 KiB
Python
"""
|
|
LiteLLM failover 沙箱演练脚本 — #937
|
|
================================================
|
|
演练场景:
|
|
1. 配置注入:sandbox config 每个模型给 2 个 key(主 key 故意填无效)
|
|
2. LiteLLM 重载(HUP 或重启容器)
|
|
3. 跑一次 chat completion → 期望自动 rotate 到副 key
|
|
4. 报告 rotate 次数 / 失败耗时
|
|
|
|
注意:
|
|
- 不在生产 LiteLLM 上跑(改 qihu-aiflow-litellm 的 config)
|
|
- 实际演练由 CI 在 sandbox 容器(qihu-aiflow-litellm-test)中跑
|
|
- 本脚本默认 dry_run=True,只生成 config 与报告不真改生产
|
|
|
|
用法:
|
|
# 真跑(谨慎):
|
|
docker exec qihu-aiflow python3 /app/tests/test_litellm_failover_sandbox.py --run
|
|
# 干跑(默认):
|
|
docker exec qihu-aiflow python3 /app/tests/test_litellm_failover_sandbox.py
|
|
"""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from typing import Any
|
|
|
|
# /app/backend/python 在 path
|
|
sys.path.insert(0, "/app/backend/python")
|
|
|
|
|
|
SANDBOX_CONFIG = """\
|
|
model_list:
|
|
- model_name: deepseek-ai/DeepSeek-V4-Flash
|
|
litellm_params:
|
|
model: openai/deepseek-ai/DeepSeek-V4-Flash
|
|
api_base: https://api.siliconflow.cn/v1
|
|
api_key:
|
|
- sk-INTENTIONALLY-INVALID-primary-{ts}
|
|
- sk-idsplzeccvmbtyqxcrbhaiwmkzwajmicliunkjkouivepaaz
|
|
litellm_settings:
|
|
drop_params: true
|
|
num_retries: 2
|
|
cooldown_time: 60
|
|
"""
|
|
|
|
|
|
def probe_litellm(base_url: str = "http://qihu-aiflow-litellm:4000"):
|
|
"""探活 LiteLLM /health/ready"""
|
|
import httpx
|
|
for k in ("HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy"):
|
|
os.environ.pop(k, None)
|
|
try:
|
|
r = httpx.get(f"{base_url}/health/readiness", timeout=5, trust_env=False)
|
|
return r.is_success, r.json() if r.content else None
|
|
except Exception as e:
|
|
return False, str(e)
|
|
|
|
|
|
def inject_sandbox_config(dry_run: bool = True, ts: str = ""):
|
|
cfg = SANDBOX_CONFIG.format(ts=ts)
|
|
print(f"[sandbox-config] generated ({len(cfg)} bytes):\n{cfg}")
|
|
if dry_run:
|
|
print("[sandbox-config] DRY-RUN: not writing to LiteLLM")
|
|
return None
|
|
# 真写需要 ssh 到 litellm 容器;此处仅占位
|
|
raise NotImplementedError("actual injection needs ssh to litellm container")
|
|
|
|
|
|
def run_chat(base_url: str = "http://qihu-aiflow-litellm:4000", model: str = "deepseek-ai/DeepSeek-V4-Flash"):
|
|
"""发一次 chat completion,观察 LiteLLM 是否自动 rotate 到副 key"""
|
|
import httpx
|
|
for k in ("HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy"):
|
|
os.environ.pop(k, None)
|
|
body = {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": "ping"}],
|
|
"max_tokens": 16,
|
|
}
|
|
t0 = time.perf_counter()
|
|
try:
|
|
r = httpx.post(
|
|
f"{base_url}/v1/chat/completions",
|
|
json=body,
|
|
headers={"Authorization": "Bearer dev"}, # 平台内部 token
|
|
timeout=20,
|
|
trust_env=False,
|
|
)
|
|
elapsed = (time.perf_counter() - t0) * 1000
|
|
return {
|
|
"ok": r.is_success,
|
|
"status": r.status_code,
|
|
"elapsed_ms": elapsed,
|
|
"body": r.text[:300],
|
|
}
|
|
except Exception as e:
|
|
elapsed = (time.perf_counter() - t0) * 1000
|
|
return {"ok": False, "error": str(e), "elapsed_ms": elapsed}
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--run", action="store_true", help="真实注入 sandbox config(谨慎!)")
|
|
ap.add_argument("--base-url", default="http://qihu-aiflow-litellm:4000")
|
|
args = ap.parse_args()
|
|
|
|
print("=" * 70)
|
|
print(" LiteLLM 主副密钥 failover 沙箱演练")
|
|
print("=" * 70)
|
|
|
|
print("\n[step 1] probe LiteLLM ...")
|
|
ok, body = probe_litellm(args.base_url)
|
|
print(f" alive: {ok} body: {body}")
|
|
if not ok:
|
|
print("[FAIL] LiteLLM not reachable, abort")
|
|
return 1
|
|
|
|
print("\n[step 2] inject sandbox config ...")
|
|
inject_sandbox_config(dry_run=not args.run, ts=str(int(time.time())))
|
|
|
|
if not args.run:
|
|
print("\n[step 3] (skipped -- not --run)")
|
|
print("\n[result] DRY-RUN OK. 演练流程设计正确,可手动 --run 实测。")
|
|
return 0
|
|
|
|
print("\n[step 3] trigger reload ...")
|
|
# 在 --run 模式下:docker kill -s HUP <litellm>;此处简化
|
|
print("(not implemented in safe-mode)")
|
|
|
|
print("\n[step 4] send chat completion ...")
|
|
res = run_chat(args.base_url)
|
|
print(f" result: {res}")
|
|
|
|
print("\n[result] 演练结束。观察:")
|
|
print(f" - 状态码 = {res.get('status', '?')} (期望 200,主 key 401 后被副 key 接管)")
|
|
print(f" - 耗时 = {res.get('elapsed_ms', '?'):.1f}ms (期望 <5s)")
|
|
return 0 if res.get("ok") else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main()) |