1 Token Bucket 机制
Token Bucket 机制(令牌桶算法):
TokenBucket 是一个常用的速率限制算法,用于控制请求速率,确保请求以平稳的速率发送。它的工作原理如下:
初始化:
令牌桶以一定的速率被填充(例如,每秒增加 query_per_second 个令牌)。
桶中令牌的最大数量是有限的,一旦达到上限,令牌不再增加。
获取令牌:
当发送请求时,需要从桶中取出一个令牌。
如果桶中有足够的令牌,请求立即发送。
如果桶中没有足够的令牌,get_token 方法将阻塞,直到有可用的令牌为止
优点:
- 精细控制:可以精细地控制请求的速率,适合需要平稳控制请求频率的情况。
- 适合多线程:适合多线程或高并发环境,因为令牌桶机制可以在各个线程之间共享速率限制。
缺点:
- 实现复杂度:实现起来相对复杂,需要手动管理令牌生成和消耗的逻辑。
- 处理瞬时错误:对于瞬时网络错误或 API 服务暂时不可用的情况,令牌桶机制不如重试机制直接有效。
适用场景:
- 高并发请求控制:需要对高并发请求进行速率限制。
- 持续负载下的速率控制:需要在持续负载下平稳地控制请求速率,例如防止服务器过载。
import threading
import time
from queue import Queue
from logging import getLogger
class TokenBucket:
"""A token bucket for rate limiting.
Args:
query_per_second (float): The rate of the token bucket.
"""
def __init__(self, rate, verbose=False):
self._rate = rate
self._tokens = threading.Semaphore(0) # Semaphore to control the number of tokens
self.started = False # Flag to start the token adding thread only once
self._request_queue = Queue() # Queue to track request timestamps for verbose logging
self.logger = getLogger() # Logger for verbose output
self.verbose = verbose # Flag to enable verbose logging
def _add_tokens(self):
"""Add tokens to the bucket periodically."""
while True:
if self._tokens._value < self._rate:
self._tokens.release() # Release a token, increasing the semaphore count
time.sleep(1 / self._rate) # Sleep for the interval based on rate
def get_token(self):
"""Get a token from the bucket."""
if not self.started:
self.started = True
threading.Thread(target=self._add_tokens, daemon=True).start() # Start token adding thread
self._tokens.acquire() # Acquire a token (decrement semaphore)
if self.verbose:
cur_time = time.time()
# Clean up old entries in the queue
while not self._request_queue.empty():
if cur_time - self._request_queue.queue[0] > 60:
self._request_queue.get()
else:
break
# Add the current request timestamp
self._request_queue.put(cur_time)
self.logger.info(f'Current RPM {self._request_queue.qsize()}.')
import requests
import time
class RateLimitedAPI:
def __init__(self, query_per_second, verbose=False):
self.token_bucket = TokenBucket(query_per_second, verbose)
def call_external_api(self, text: str) -> str:
self.token_bucket.get_token() # Acquire a token to proceed
try:
# Simulate API call
response = requests.post(
"https://example.com/api",
json={'question': text},
timeout=120
)
response.raise_for_status()
return response.json()
except requests.RequestException as e:
# Handle request exception
print(f"Request failed: {e}")
return ''
finally:
self.token_bucket.get_token() # Release the token (this might be redundant)
# 示例用法
if __name__ == "__main__":
rate_limited_api = RateLimitedAPI(query_per_second=3, verbose=True)
for i in range(5):
response = rate_limited_api.call_external_api(f"Test query {i+1}")
print(response)
time.sleep(1) # Observe the rate limit effect
init 构造函数:
rate:每秒允许的请求数。
verbose:如果为 True,启用详细日志。
初始化 Semaphore(信号量)为 0,表示没有令牌可用。
started 用于启动令牌添加线程,只在第一次调用 get_token 时启动。
_add_tokens 方法:
负责向令牌桶中添加令牌。
每隔一段时间(1 / rate 秒)释放一个令牌,增加信号量。
这是一个无限循环的守护线程。
get_token 方法:
获取一个令牌,调用 acquire 减少信号量。
如果 verbose 为 True,将请求的时间戳放入队列,并清理超过 60 秒的旧条目以记录当前的每分钟请求数(RPM)。
TokenBucket 控制 API 请求速率,通过信号量和线程添加令牌,确保请求按速率进行。
verbose 模式用于详细日志记录,帮助监控请求的实际速率。
RateLimitedAPI 类展示了如何将 TokenBucket 集成到 API 请求流程中,控制请求的速率。
3 重试机制
重试机制(结合 requests
库和指数退避策略):
- 使用
requests
库的Session
对象发送 HTTP 请求,并结合指数退避策略进行重试。 - 在遇到网络错误或请求失败时,通过逐渐增加等待时间来重试请求。
- 主要优点是处理网络故障和暂时的 API 限制,在重试失败后会有较长的等待时间,有助于避免持续的请求失败。
优点:
- 简单易用:易于实现,尤其是结合
requests
库的Session
对象,可以轻松处理 HTTP 请求。 - 自动处理错误:能够自动处理瞬时错误和暂时的 API 服务不可用的情况,通过指数退避策略避免持续失败。
缺点:
- 缺乏精细控制:不能像令牌桶机制那样精细地控制请求速率,适合处理短期的网络问题或 API 限制。
- 等待时间增加:在高失败率情况下,指数退避策略可能导致较长时间的等待。
适用场景:
- 处理不稳定的网络连接:适合在网络连接不稳定时自动重试请求。
- API 服务不稳定:适合在 API 服务偶尔不可用或负载过高时自动重试请求。
import requests
import time
import logging
logger = logging.getLogger(__name__)
# 定义自定义异常类
class InferenceServiceBusyException(Exception):
pass
def make_request_with_retry(apiurl, query, _headers, max_retries=8):
with requests.Session() as session:
backoff_time = 1 # 初始等待时间(秒)
retry_count = 0
last_exception = None
while retry_count < max_retries:
try:
response = session.post(apiurl,
json=query,
headers=_headers,
timeout=120)
logger.info('response {}'.format(response.text))
response.raise_for_status() # 检查 HTTP 错误
result = response.json()["bo"]["result"]
if '模型推理服务繁忙' in result:
raise InferenceServiceBusyException("模型推理服务繁忙,尝试重试...")
return result
except InferenceServiceBusyException as busy_err:
last_exception = busy_err
logger.error(f"{busy_err} Waiting {backoff_time} seconds before retrying...")
except requests.RequestException as req_err:
last_exception = req_err
logger.error(f"Request error: {req_err} Waiting {backoff_time} seconds before retrying...")
except Exception as err:
last_exception = err
logger.error(f"Unexpected error: {err} Waiting {backoff_time} seconds before retrying...")
time.sleep(backoff_time)
retry_count += 1
backoff_time *= 2 # 指数退避
# 如果达到最大重试次数,抛出最后一次的异常
if last_exception:
raise last_exception
return '' # 达到最大重试次数后,返回空字符串
apiurl = 'your_api_url'
query = { 'your': 'query_data' }
headers = { 'Authorization': 'Bearer your_token' }
try:
result = make_request_with_retry(apiurl, query, headers)
print(result)
except Exception as e:
print(f"Request failed: {e}")
自定义异常类 InferenceServiceBusyException:
用于明确区分 ‘模型推理服务繁忙’ 的错误。这使得错误处理更清晰,也能提供更具体的错误信息。
变量 last_exception:
用于存储最后一次捕获的异常。如果达到最大重试次数,将抛出这个异常。
异常处理:
InferenceServiceBusyException:如果 result 中包含 ‘模型推理服务繁忙’,引发自定义异常并记录日志。
requests.RequestException:捕获 requests 相关的所有 HTTP 错误(如超时、连接错误等),记录日志并更新 last_exception。
Exception:捕获所有其他异常,记录日志并更新 last_exception。
指数退避:
每次重试前,等待时间成倍增加,从 1 秒开始。
最大重试次数后处理:
如果达到最大重试次数(max_retries),抛出最后一次的异常。这确保了调用者可以获得最准确的失败原因。
3 选择的建议
- 如果你的系统需要精细控制请求频率,尤其是高并发的环境,选择 Token Bucket 机制。
- 如果你主要处理的是网络不稳定或 API 服务偶尔不可用的情况,使用 重试机制 可能更简单直接。
在一些实际场景中,你可能会结合这两种机制。例如,使用 Token Bucket 来控制总体的请求速率,同时使用重试机制来处理瞬时的网络错误。这种组合方式能够提供更全面的鲁棒性和控制能力。