处理 API 请求和速率限制的2种方法

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 来控制总体的请求速率，同时使用重试机制来处理瞬时的网络错误。这种组合方式能够提供更全面的鲁棒性和控制能力。