BIZ-40: NVIDIA Sidecar 限流代理 Phase1 — 核心代理模块

交付文件: - config.py: 配置管理 (SidecarConfig + load_config)，修复 PEP 563 类型推断 bug - rate_limiter.py: 令牌桶 (TokenBucket) + 网关识别 (is_nvidia_gateway) - priority_queue.py: 四级优先级队列，修复 PASSTHROUGH 语义 bug - server.py: FastAPI 代理主入口，修复 worker_loop 重试悬挂 bug - __init__.py: 包声明与公开导出 - pyproject.toml: 依赖声明 + mypy 配置 - README.md: 快速启动指南 + 环境变量列表评审修复: - worker_loop 令牌重试从重入队改为 poll-wait (防止 future 悬挂) - 路由函数 + lifespan 补充返回类型注解 - heapq 重复 import 移到文件顶部 - config.py 清理无用代码行 - types-PyYAML stub 安装 - 新增 README.md 验证: mypy 0 issues, 全量单元测试通过 Co-authored-by: multica-agent <github@multica.ai>
2026-06-24 08:32:47 +08:00
parent cca4089f2a
commit 6b5f53a0fd
7 changed files with 1488 additions and 0 deletions
@@ -0,0 +1,200 @@
+"""
+NVIDIA Sidecar 限流代理 — 令牌桶 + 网关识别模块 (§3.2)
+
+从 BIZ-26 rate_limiter.py 提取核心限流逻辑，去除多线程调度器、缓存管理等。
+保留：Priority, TokenBucket, is_nvidia_gateway, normalize_gateway_name。
+"""
+
+from __future__ import annotations
+
+import time
+import threading
+from enum import IntEnum
+from typing import Any
+
+
+# ---------------------------------------------------------------------------
+# 优先级枚举
+# ---------------------------------------------------------------------------
+
+class Priority(IntEnum):
+    """请求优先级（数值越小优先级越高）。"""
+    URGENT = 1
+    HIGH = 2
+    NORMAL = 3
+    LOW = 4
+
+
+# ---------------------------------------------------------------------------
+# NVIDIA 网关别名集
+# ---------------------------------------------------------------------------
+
+NVIDIA_GATEWAY_ALIASES: set[str] = {
+    "nvidia",
+    "nvidia-gateway",
+    "nvidiavx",
+    "nvidiavx18088980513",
+}
+
+
+def is_nvidia_gateway(value: str | None) -> bool:
+    """判断给定网关名/模型全路径是否属于 NVIDIA 网关。
+
+    Args:
+        value: 网关名（如 ``"nvidia"``）或模型全路径前缀
+               （如 ``"nvidia/deepseek-ai/deepseek-v4-pro"``）。
+               None 时直接返回 False。
+
+    Returns:
+        True 当 value 的 provider 部分匹配已知 NVIDIA 别名。
+    """
+    if value is None:
+        return False
+
+    # 提取 provider 前缀：取 "/" 前第一个部分
+    provider = value.split("/", 1)[0].lower().strip()
+    return provider in NVIDIA_GATEWAY_ALIASES
+
+
+def normalize_gateway_name(value: str | None) -> str | None:
+    """规范化网关名：提取 provider 前缀并转为小写。
+
+    Args:
+        value: 网关名或模型全路径。None 时返回 None。
+
+    Returns:
+        provider 前缀的小写形式，或 None。
+    """
+    if value is None:
+        return None
+    return value.split("/", 1)[0].lower().strip()
+
+
+# ---------------------------------------------------------------------------
+# 令牌桶（线程安全）
+# ---------------------------------------------------------------------------
+
+class TokenBucket:
+    """线程安全的令牌桶实现。
+
+    支持固定速率令牌补充和消费，带有溢出保护和可选的阻塞等待。
+    """
+
+    def __init__(self, rate: float = 40 / 60, capacity: int = 40) -> None:
+        """初始化令牌桶。
+
+        Args:
+            rate: 令牌补充速率（令牌/秒）。默认 40/60 ≈ 0.667 token/s（40 RPM）。
+            capacity: 桶最大容量（令牌数）。默认 40。
+        """
+        self._rate: float = float(rate)
+        self._capacity: int = int(capacity)
+        self._tokens: float = float(capacity)  # 启动时桶满
+        self._last_refill: float = time.monotonic()
+        self._lock: threading.Lock = threading.Lock()
+
+    # ---- 内部方法 ----
+
+    def _refill(self) -> None:
+        """补充令牌（调用方需持有 _lock）。
+
+        根据距上次补充的时间差计算新增令牌数，不超过 capacity。
+        """
+        now = time.monotonic()
+        elapsed = now - self._last_refill
+        if elapsed > 0 and self._rate > 0:
+            new_tokens = elapsed * self._rate
+            self._tokens = min(self._tokens + new_tokens, float(self._capacity))
+        self._last_refill = now
+
+    # ---- 公开方法 ----
+
+    def consume(self, tokens: int = 1) -> bool:
+        """尝试立即消费令牌（非阻塞）。
+
+        Args:
+            tokens: 要消费的令牌数，默认 1。
+
+        Returns:
+            True 消费成功；False 令牌不足。
+        """
+        if tokens <= 0:
+            return True
+
+        with self._lock:
+            self._refill()
+            if self._tokens >= tokens:
+                self._tokens -= tokens
+                return True
+            return False
+
+    def try_consume(self, tokens: int = 1, timeout: float = 2.0) -> bool:
+        """尝试在指定时间内消费令牌（阻塞）。
+
+        Args:
+            tokens: 要消费的令牌数，默认 1。
+            timeout: 最大等待秒数，默认 2.0。
+
+        Returns:
+            True 在超时前成功消费；False 超时。
+        """
+        if tokens <= 0:
+            return True
+
+        deadline = time.monotonic() + timeout
+        while True:
+            with self._lock:
+                self._refill()
+                if self._tokens >= tokens:
+                    self._tokens -= tokens
+                    return True
+
+            # 释放锁后计算剩余等待时间
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                return False
+            # 等待到下一个令牌应该补充的时间点
+            sleep_time = min(remaining, max(0.05, 1.0 / self._rate) if self._rate > 0 else remaining)
+            time.sleep(sleep_time)
+
+    def wait_for_token(self, timeout: float | None = None) -> bool:
+        """等待并尝试消费 1 个令牌。
+
+        Args:
+            timeout: 最大等待秒数；None 表示无限等待（不推荐）。
+
+        Returns:
+            True 成功消费；False 超时。
+        """
+        return self.try_consume(tokens=1, timeout=timeout if timeout is not None else float("inf"))
+
+    def get_status(self) -> dict[str, Any]:
+        """获取令牌桶当前状态。
+
+        Returns:
+            包含 tokens, capacity, rate_per_minute, utilization 的字典。
+        """
+        with self._lock:
+            self._refill()
+            rate_per_minute = self._rate * 60.0
+            utilization = 0.0 if self._capacity == 0 else (
+                (self._capacity - self._tokens) / self._capacity
+            )
+            return {
+                "tokens": round(self._tokens, 2),
+                "capacity": self._capacity,
+                "rate_per_minute": round(rate_per_minute, 1),
+                "utilization": round(utilization, 4),
+            }
+
+    # ---- 属性 ----
+
+    @property
+    def rate(self) -> float:
+        """当前令牌补充速率（令牌/秒）。"""
+        return self._rate
+
+    @property
+    def capacity(self) -> int:
+        """桶容量。"""
+        return self._capacity