3. 实现动态批处理推理优化

任务目标

实现一个支持动态批处理（Dynamic Batching）的推理服务核心逻辑：当多个推理请求在短时间内到达时，将它们聚合为同一批次送入模型执行，从而提升GPU利用率并降低单请求的有效延迟。本案例聚焦于动态批处理调度器本身，不涉及具体模型实现，使用模拟耗时函数代替真实推理。

核心设计

动态批处理的核心是等待窗口机制：请求到达后不立即执行，而是进入队列等待最多 max_wait_ms 毫秒；窗口关闭时，将队列中所有请求合并为一个批次执行。关键参数：

• max_batch_size：单批次最大请求数，超过则立即触发执行
• max_wait_ms：最大等待时间，到期强制触发执行

完整实现

import time
import threading
from collections import deque
from dataclasses import dataclass
from typing import List, Callable, Optional

@dataclass
class Request:
    req_id: int
    payload: str

class DynamicBatcher:
    def __init__(self, model_fn: Callable[[List[str]], List[float]],
                 max_batch_size: int = 4, max_wait_ms: int = 50):
        self.model_fn = model_fn
        self.max_batch_size = max_batch_size
        self.max_wait_s = max_wait_ms / 1000.0
        self.queue: deque = deque()
        self.lock = threading.Lock()
        self.batch_counter = 0
        self._shutdown = False
        self._timer: Optional[threading.Timer] = None
        self._pending_results: dict = {}

    def submit(self, req: Request) -> float:
        """提交请求，返回推理结果（阻塞等待）"""
        with self.lock:
            self.queue.append(req)
            if len(self.queue) >= self.max_batch_size:
                self._trigger_batch_locked()
            elif self._timer is None:
                self._start_timer_locked()
        # 等待结果
        while req.req_id not in self._pending_results:
            time.sleep(0.001)
        return self._pending_results.pop(req.req_id)

    def _start_timer_locked(self):
        self._timer = threading.Timer(self.max_wait_s, self._on_timer)
        self._timer.daemon = True
        self._timer.start()

    def _on_timer(self):
        with self.lock:
            self._timer = None
            if self.queue:
                self._trigger_batch_locked()

    def _trigger_batch_locked(self):
        if not self.queue:
            return
        batch = []
        while self.queue and len(batch) < self.max_batch_size:
            batch.append(self.queue.popleft())
        if self._timer is not None:
            self._timer.cancel()
            self._timer = None
        self.batch_counter += 1
        batch_id = self.batch_counter
        # 在锁外执行推理，避免阻塞提交
        threading.Thread(target=self._execute_batch, args=(batch_id, batch)).start()

    def _execute_batch(self, batch_id: int, batch: List[Request]):
        payloads = [r.payload for r in batch]
        # 模拟批量推理：单条耗时约20ms，批量后总耗时约30ms（GPU并行收益）
        time.sleep(0.03)
        results = self.model_fn(payloads)
        with self.lock:
            for req, res in zip(batch, results):
                self._pending_results[req.req_id] = res

    def shutdown(self):
        self._shutdown = True
        with self.lock:
            if self._timer:
                self._timer.cancel()

# ---------- 模拟模型：批量推理耗时低于逐个推理 ----------
def mock_model(payloads: List[str]) -> List[float]:
    """模拟GPU批量推理：返回每个请求的处理耗时（ms）"""
    # 批量推理总耗时固定为30ms，均摊到每个请求
    per_request_ms = 30.0 / len(payloads)
    return [per_request_ms] * len(payloads)

# ---------- 验证入口 ----------
if __name__ == "__main__":
    batcher = DynamicBatcher(model_fn=mock_model, max_batch_size=4, max_wait_ms=50)

    # 模拟并发提交8个请求（间隔5ms到达）
    results = {}
    def send(req):
        t0 = time.time()
        res = batcher.submit(req)
        elapsed_ms = (time.time() - t0) * 1000
        results[req.req_id] = (res, elapsed_ms)

    threads = []
    for i in range(8):
        t = threading.Thread(target=send, args=(Request(i, f"data-{i}"),))
        threads.append(t)
        t.start()
        time.sleep(0.005)  # 5ms间隔到达

    for t in threads:
        t.join()

    batcher.shutdown()

    # 输出统计
    total_batches = batcher.batch_counter
    avg_latency = sum(v[1] for v in results.values()) / len(results)
    print(f"总批次数: {total_batches}")
    print(f"平均端到端延迟: {avg_latency:.1f} ms")
    print(f"各请求结果: { {k: round(v[0], 2) for k, v in results.items()} }")
    print(f"各请求等待时间(ms): { {k: round(v[1], 1) for k, v in results.items()} }")

运行方式

python dynamic_batcher.py

验证方法与通过标准

运行后观察输出，满足以下条件即验证通过：

• 批次数 < 请求数：8个请求应被聚合为2~3个批次（取决于到达时序），而非8次独立推理
• 平均延迟显著低于串行推理：串行推理8个请求每个20ms共需160ms；动态批处理后，每个请求的端到端延迟应明显低于160ms（预期约30~80ms区间，因等待窗口和批次拆分而异）
• 结果正确性：每个请求均获得返回值，且返回值为该批次均摊耗时（如批次大小为4时返回7.5ms）

若批次数为8，说明动态批处理未生效；若平均延迟接近160ms，说明等待窗口或触发逻辑有误。

关键说明

• 本实现使用threading.Timer实现等待窗口，达到max_wait_ms或队列满max_batch_size时触发批处理
• 推理在独立线程中执行，避免阻塞请求提交；submit方法通过轮询_pending_results等待结果返回
• 真实场景中，将mock_model替换为Triton/TensorRT-LLM等推理引擎的批量执行接口即可