feat: implement startup audit to block over-limit models on restart

- Add startup_audit() to check Redis usage against limits at startup
- Call startup_audit() in active_monitor_loop
- Add AGENTS.md for agent guidelines
This commit is contained in:
kim
2026-02-02 13:39:34 +08:00
parent 6430c2e828
commit bdd51e8f9c
2 changed files with 176 additions and 0 deletions
+103
View File
@@ -0,0 +1,103 @@
# OpenCode Agent Protocol - LiteLLM VolcSync
此文件定义了在该代码库中工作的 AI Agent 的行为准则、代码风格和工作流。
## 1. 项目上下文
此仓库包含用于 **LiteLLM** 的自定义回调脚本,核心功能是限制模型用量(特别是火山引擎/VolcEngine)。
- **核心逻辑**:通过 `quota_interceptor` 拦截请求,并通过 `DailyLimitHandler` 记录用量到 Redis。
- **关键文件**:
- `volc_limit_sync.py`: 主同步逻辑,从火山引擎 API 拉取配额并更新 Redis。
- `daily_limit.py`: 轻量级回调,用于 LiteLLM 运行时检查。
- `config.yaml`: 包含模型定义的 LiteLLM 配置文件。
## 2. 环境与运行 (Build & Run)
本项目主要由 Python 脚本组成,没有复杂的构建过程。
### 依赖安装
```bash
pip install litellm redis requests
```
### 运行与调试
由于是脚本项目,通常直接运行 Python 文件进行测试。
- **启动同步监控**:
```bash
# 必须配置必要的环境变量
export VOLC_ACCESS_KEY="your_key"
export VOLC_SECRET_KEY="your_secret"
export REDIS_HOST="localhost"
python volc_limit_sync.py
```
- **调试模式**:
设置环境变量 `DAILY_LIMIT_LOG_LEVEL=DEBUG` 可启用详细日志。
## 3. 代码风格指南 (Code Style)
### 3.1 基础规范
- **Python 版本**:兼容 Python 3.8+。
- **命名规范**:
- 变量/函数:`snake_case` (e.g., `clean_model_name`, `get_daily_limit`)
- 常量:`UPPER_CASE` (e.g., `LOG_LEVEL`, `REDIS_HOST`)
- 类名:`PascalCase` (e.g., `DailyLimitHandler`)
- **类型提示 (Typing)**:
- **不要** 强制添加类型提示(Type Hints),保持与现有代码一致的动态类型风格,除非涉及复杂的数据结构。
- **导入排序**:标准库 -> 第三方库 (`litellm`, `redis`, `requests`) -> 本地模块。
### 3.2 日志与输出 (Logging)
**重要**:本项目使用自定义的日志风格,**不要**使用 Python 标准 `logging` 模块。
- **必须**使用 `log(msg, level="INFO")` 或 `dprint(msg)` 函数。
- **Emoji 前缀**:日志消息通常以 Emoji 开头以增强可读性:
- ℹ️ `INFO`: 一般信息
- 🔍 `DEBUG`: 调试信息
- ⚠️ `WARNING`: 警告/阻断
- ❌ `ERROR`: 错误/异常
- 📊 `Usage`: 用量统计
- ♻️ `RESTORE`: 恢复模型
- 🚫 `BLOCK`: 阻断模型
### 3.3 错误处理 (Error Handling)
由于这些脚本运行在 LiteLLM 主进程中,**稳定性是第一优先级**。
- **防御性编程**:关键逻辑(特别是涉及 Redis 或外部 API 的部分)**必须**包裹在 `try-except` 块中。
- **禁止崩溃**:回调函数内的异常不应导致 LiteLLM 崩溃。捕获异常并记录错误日志 (`traceback.format_exc()`),然后安全返回(`pass` 或返回默认值)。
- **示例**:
```python
try:
# 危险操作
redis_client.incr(key)
except Exception as e:
# 记录但不抛出,以免影响主业务
log(f"❌ Redis failed: {e}", level="ERROR")
```
### 3.4 特定模式
- **LiteLLM 集成**:
- 检查 `LITELLM_AVAILABLE` 标志,确保脚本可以在没有 litellm 环境下独立运行(用于调试)。
- 处理 `litellm.exceptions` 中的特定异常(如 `ServiceUnavailableError`)。
- **时间处理**:
- 明确处理时区问题。本项目强制使用 **北京时间 (UTC+8)** 进行每日限额重置(00:00 CST)。
## 4. 测试 (Testing)
目前没有集成的测试框架(如 pytest)。
- **验证方式**:
- 编写或运行 `debug_*.py` 脚本进行手动验证。
- 使用 `if __name__ == "__main__":` 块在脚本底部包含自测逻辑。
## 5. Git 提交规范
- 使用约定式提交 (Conventional Commits):
- `feat`: 新功能 (e.g., `feat: add support for deepseek-v3`)
- `fix`: 修复 (e.g., `fix: resolve redis connection timeout`)
- `docs`: 文档更新
- `style`: 代码格式调整
- `refactor`: 重构
## 6. Cursor/Copilot 规则
- 优先保持现有文件的代码结构,不要为了"整洁"而破坏原有的 `try-except` 结构。
- 生成代码时,优先复用已有的 `clean_model_name`, `generate_api_base_hash` 等工具函数。
+73
View File
@@ -470,8 +470,73 @@ def restore_blocked_models():
log(f"❌ Restore check failed: {e}", level="ERROR") log(f"❌ Restore check failed: {e}", level="ERROR")
def startup_audit():
"""
🚀 Startup Audit:
System initialization check. Iterate through all models in the router.
If a model is found to be already over the limit in Redis, kick it out immediately.
系统启动自检:遍历 Router 中所有模型,如果发现 Redis 中记录已超额,直接踢出。
"""
if not redis_client or not LITELLM_AVAILABLE:
return
try:
from litellm.proxy.proxy_server import llm_router
if not llm_router or not llm_router.model_list:
dprint("⚠️ Startup audit skipped: llm_router not ready.")
return
dprint("🚀 Starting Startup Audit (Pre-emptive Block)...")
blocked_count = 0
# Create a copy to iterate safely
# We only need to read configs to find targets
for model_entry in llm_router.model_list:
# Skip if already blocked (shouldn't happen on fresh start, but good practice)
if "_daily_limit_backup" in model_entry:
continue
params = model_entry.get("litellm_params", {})
raw_model = params.get("model", "")
api_base = params.get("api_base", "")
# 1. Clean Name
clean_model = clean_model_name(raw_model)
if not clean_model:
continue
# 2. Check Limit & Usage
# We use effective limit logic which checks both static and dynamic
effective_limit = get_effective_limit(clean_model, api_base)
if not effective_limit:
continue
usage_key = get_usage_key(clean_model, api_base)
current_val = redis_client.get(usage_key)
if current_val and float(current_val) >= float(effective_limit):
dprint(
f"🚫 [STARTUP BLOCK] Found {raw_model} over limit ({int(float(current_val))}/{int(effective_limit)}). Evicting now."
)
soft_evict_from_router(raw_model, api_base)
blocked_count += 1
if blocked_count > 0:
dprint(f"✅ Startup Audit Complete: Blocked {blocked_count} models.")
else:
dprint("✅ Startup Audit Complete: All models healthy.")
except ImportError:
dprint("⚠️ Startup audit skipped: litellm not found.")
except Exception as e:
log(f"❌ Startup audit failed: {e}", level="ERROR")
def quota_interceptor(kwargs, user_api_key_dict=None): def quota_interceptor(kwargs, user_api_key_dict=None):
try: try:
# 1. 尝试恢复(定期)
restore_blocked_models() restore_blocked_models()
except (ServiceUnavailableError, ContextWindowExceededError): except (ServiceUnavailableError, ContextWindowExceededError):
raise raise
@@ -832,7 +897,15 @@ def active_monitor_loop():
f"✅ Redis connected at / Redis 已连接于 {REDIS_HOST}:{REDIS_PORT}", flush=True f"✅ Redis connected at / Redis 已连接于 {REDIS_HOST}:{REDIS_PORT}", flush=True
) )
# ==========================================
# 🚀 Run Startup Audit Immediately
# ==========================================
# Give LiteLLM a moment to initialize router if running concurrently
time.sleep(2)
startup_audit()
last_usage_map = {} last_usage_map = {}
last_sync_time = 0 last_sync_time = 0
current_sync_interval = 0 current_sync_interval = 0
global_min_remaining = 0 global_min_remaining = 0