4. 常见限流算法（上）：令牌桶算法原理、实现与适用场景
在模型风险管理系统中，API 网关需要应对突发流量、保护下游模型推理服务不被击穿。令牌桶算法（Token Bucket）是业界应用最广泛的限流算法之一，它允许一定程度的突发流量，同时保证长期平均速率可控，非常适合模型推理 API 的流量整形场景。
4.1 核心原理
令牌桶算法的基本思想是：系统以恒定速率向一个桶中放入令牌，每个请求必须从桶中获取一个令牌才能被处理。如果桶中令牌不足，请求被拒绝或排队等待。
关键参数：
- 桶容量（bucket_capacity）：桶最多能存放的令牌数，决定了突发流量的上限。
- 令牌填充速率（refill_rate）：每秒向桶中添加的令牌数量，决定了长期平均请求速率。
- 当前令牌数（current_tokens）：桶中实时可用的令牌数量。
4.2 算法实现
令牌桶的实现通常有两种方式：定时填充和延迟计算。在高并发场景下，延迟计算方式性能更优，因为它避免了定时器的开销。
4.2.1 延迟计算实现（推荐）
延迟计算的核心思想是：在每次请求到达时，根据时间差计算应补充的令牌数，而不是依赖后台定时任务。
import time
class TokenBucket:
def __init__(self, capacity, refill_rate):
"""
:param capacity: 桶容量（最大突发令牌数）
:param refill_rate: 令牌填充速率（个/秒）
"""
self.capacity = capacity
self.refill_rate = refill_rate
self.current_tokens = capacity # 初始满桶
self.last_refill_time = time.time()
def allow_request(self, tokens=1):
"""
判断是否允许请求通过
:param tokens: 本次请求消耗的令牌数
:return: True 表示放行，False 表示拒绝
"""
now = time.time()
# 计算从上一次填充到现在应补充的令牌数
elapsed = now - self.last_refill_time
refill_tokens = elapsed * self.refill_rate
self.current_tokens = min(self.capacity, self.current_tokens + refill_tokens)
self.last_refill_time = now
# 判断是否足够令牌
if self.current_tokens >= tokens:
self.current_tokens -= tokens
return True
else:
return False
实现要点：
- 原子性：在多线程/协程环境下，
current_tokens和last_refill_time的读写需要加锁或使用原子操作。 - 精度控制：
time.time()精度可能不够，生产环境建议使用time.monotonic()或高性能时钟。 - 令牌消耗：模型推理请求的“成本”可能不同（如大模型请求消耗 2 个令牌，小模型消耗 1 个），可通过
tokens参数实现差异化限流。
4.2.2 参数配置对比
| 参数 | 含义 | 模型风险场景示例 |
|---|---|---|
| capacity | 桶容量，决定最大突发请求数 | 模型上线初期，允许 100 个请求同时涌入进行压测 |
| refill_rate | 令牌填充速率，决定稳态 QPS | 模型推理服务每秒最多处理 50 个请求 |
| tokens | 单次请求消耗令牌数 | 高风险模型（如信用评分）消耗 2 个令牌，低风险模型消耗 1 个 |
4.3 适用场景分析
令牌桶算法在模型风险管理中有三个典型应用场景：
- 模型推理 API 突发保护：允许客户端在短时间内突发请求（如批量数据回测），但长期平均速率受控，避免模型服务过载。
- 多模型共享网关：不同模型分配不同的桶容量和填充速率，实现租户级别的流量隔离。
- 模型版本切换流量控制：新版本上线时，通过调整令牌桶参数逐步放量，从 10% 流量平滑过渡到 100%。
避坑提示：
- 桶容量设置过大：如果 capacity 远大于 refill_rate，突发流量可能瞬间打垮后端模型服务。建议 capacity = refill_rate × 允许突发秒数。
- 时钟回拨问题：服务器时钟回拨会导致令牌桶瞬间被填满，引发流量洪峰。生产环境应使用单调递增时钟。
- 分布式一致性：在分布式网关中，每个节点独立维护令牌桶会导致整体限流不准确。建议使用 Redis 等中心化存储实现分布式令牌桶。
4.4 与漏桶算法的对比
令牌桶和漏桶是两种最基础的限流算法，理解它们的差异有助于在模型风险管理中选择合适的方案。
| 对比维度 | 令牌桶算法 | 漏桶算法 |
|---|---|---|
| 突发处理 | 允许突发，桶容量越大突发能力越强 | 严格平滑，不允许突发 |
| 实现复杂度 | 中等，需维护时间戳和令牌计数 | 简单，队列 + 固定速率消费 |
| 适用场景 | 需要容忍突发流量的模型推理 API | 流量必须严格平滑的场景（如数据库写入） |
| 内存开销 | 低，仅需几个变量 | 中等，需要维护请求队列 |
4.5 模型风险场景实战建议
- 初始配置：建议 capacity = refill_rate × 2，允许 2 秒的突发流量，同时保证稳态速率。
- 动态调整：结合模型监控指标（如推理延迟、内存使用率），当模型服务压力增大时自动降低 refill_rate。
- 预热策略：模型刚启动时，令牌桶从空桶开始，避免冷启动阶段被限流。可设置初始令牌数为 capacity × 0.5。