3. 令牌桶算法原理与实现

令牌桶算法是模型风险管理 API 限流中最常用的算法之一。它通过一个“桶”来平滑突发流量，同时保证长期平均速率不超过预设阈值。在模型推理、特征计算、批量评分等高频调用场景中，令牌桶能有效防止下游模型服务被瞬时流量冲垮。

3.1 工作机制

令牌桶的核心思想是：系统以固定速率向桶中放入令牌，每个请求必须从桶中获取一个（或多个）令牌才能被处理。如果桶中令牌不足，请求被拒绝或排队等待。

3.2 参数配置

令牌桶的两个核心参数直接影响限流效果。在模型风险管理场景中，需要根据模型推理耗时、并发数、SLA 要求进行调优。

3.3 Python 代码实现

以下实现基于时间戳计算令牌增量，避免使用定时器，适合在模型推理中间件中嵌入。代码包含完整的参数校验、令牌获取和状态查询。

import time
import threading

class TokenBucket:
    """令牌桶限流器，用于模型风险管理 API 限流"""

    def __init__(self, rate: float, capacity: int):
        """
        初始化令牌桶
        :param rate: 令牌填充速率（个/秒）
        :param capacity: 桶容量（最大令牌数）
        """
        if rate <= 0:
            raise ValueError("rate 必须大于 0")
        if capacity <= 0:
            raise ValueError("capacity 必须大于 0")

        self.rate = rate
        self.capacity = capacity
        self.tokens = capacity          # 初始装满令牌
        self.last_refill_time = time.time()
        self.lock = threading.Lock()    # 线程安全

    def _refill(self):
        """计算并补充自上次填充以来新增的令牌"""
        now = time.time()
        elapsed = now - self.last_refill_time
        # 新增令牌数 = 速率 × 经过时间
        new_tokens = elapsed * self.rate
        self.tokens = min(self.capacity, self.tokens + new_tokens)
        self.last_refill_time = now

    def try_acquire(self, tokens: int = 1) -> bool:
        """
        尝试获取指定数量的令牌
        :param tokens: 需要的令牌数（默认 1）
        :return: True 表示获取成功，False 表示令牌不足
        """
        if tokens <= 0:
            raise ValueError("tokens 必须大于 0")

        with self.lock:
            self._refill()
            if self.tokens >= tokens:
                self.tokens -= tokens
                return True
            return False

    def get_available_tokens(self) -> float:
        """获取当前可用令牌数（近似值）"""
        with self.lock:
            self._refill()
            return self.tokens

    def reset(self):
        """重置令牌桶为满状态"""
        with self.lock:
            self.tokens = self.capacity
            self.last_refill_time = time.time()

# ========== 使用示例 ==========
if __name__ == "__main__":
    # 配置：速率 100 QPS，容量 200（允许 2 秒突发）
    bucket = TokenBucket(rate=100, capacity=200)

    # 模拟请求
    for i in range(10):
        if bucket.try_acquire():
            print(f"请求 {i+1}: 通过 ✅ (剩余令牌: {bucket.get_available_tokens():.1f})")
        else:
            print(f"请求 {i+1}: 限流 ❌")
        time.sleep(0.01)  # 模拟请求间隔 10ms

3.4 代码关键设计说明

• 惰性填充：只在获取令牌时计算时间差并补充令牌，避免后台线程开销，适合模型推理这种请求间隔不固定的场景。
• 线程安全：使用 threading.Lock 保证多线程环境下令牌计数的准确性，防止超发。
• 浮点精度：elapsed * self.rate 在长时间运行时可能产生微小误差，但模型 API 限流场景下可忽略。若需高精度，可使用 decimal.Decimal。
• 容量限制：min(self.capacity, self.tokens + new_tokens) 确保令牌数不超过桶容量，避免累积过多令牌导致突发流量失控。