#937 LiteLLM failover 沙箱设计 + 演练脚本

backend/python/tests/test_litellm_failover_sandbox.py

@@ -0,0 +1,141 @@
+"""
+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())