AI 推理服务架构：从单机到分布式弹性扩展

阶段1: 单机单GPU           阶段2: 单机多GPU         阶段3: 分布式集群
┌──────────────┐         ┌──────────────┐       ┌──────────────────┐
│  FastAPI     │         │  FastAPI     │       │  API Gateway     │
│  + HF Model  │    →    │  + vLLM      │   →   │  + vLLM Cluster  │
│  + 1 GPU     │         │  + 4 GPUs    │       │  + K8s Autoscale │
└──────────────┘         └──────────────┘       └──────────────────┘
延迟: 500ms               延迟: 200ms             延迟: 150ms
QPS: 5                    QPS: 80                 QPS: 1000+

# vLLM 服务启动配置
VLLM_CONFIG = {
    "model": "meta-llama/Llama-3-70B",
    "tensor_parallel_size": 4,          # 4 GPU 张量并行
    "gpu_memory_utilization": 0.90,     # GPU 内存利用率
    "max_model_len": 8192,              # 最大上下文长度
    "max_num_seqs": 256,                # 最大并发序列
    "swap_space": 16,                   # CPU 交换空间 (GB)
    "enable_prefix_caching": True,      # 前缀缓存
    "enable_chunked_prefill": True,     # 分块预填充
    "enforce_eager": False,             # 使用 CUDA Graph
    "dtype": "float16",                 # 精度
    "quantization": "awq",              # 量化
    "served_model_name": "llama-3-70b", # 对外模型名
}

# 启动命令
# python -m vllm.entrypoints.openai.api_server --config config.json

import tensorrt as trt
import tensorrt_llm
from tensorrt_llm.runtime import ModelRunner, GenerationSession

class TensorRTEngine:
    """
    TensorRT-LLM 推理引擎
    构建 → 序列化 → 反序列化 → 推理
    """
    def __init__(self, model_dir: str, engine_dir: str):
        self.model_dir = model_dir
        self.engine_dir = engine_dir
        self.runner: ModelRunner = None

    def build_engine(self, config: dict):
        """构建 TensorRT 引擎（一次性操作）"""
        builder = tensorrt_llm.Builder()
        builder_config = builder.create_builder_config(
            name="llama",
            precision=config.get("precision", "fp16"),
            tensor_parallel_size=config.get("tp_size", 1),
            pipeline_parallel_size=config.get("pp_size", 1),
            max_batch_size=config.get("max_batch_size", 256),
            max_input_len=config.get("max_input_len", 2048),
            max_output_len=config.get("max_output_len", 512),
            max_beam_width=config.get("max_beam_width", 1),
            use_gpt_attention=config.get("use_gpt_attention", True),
            use_plugin=True,
        )
        # 构建并序列化引擎
        engine = builder.build_model(builder_config, self.model_dir)
        engine.serialize(self.engine_dir)

    def load_engine(self):
        """加载已构建的引擎"""
        self.runner = ModelRunner.from_dir(
            self.engine_dir,
            rank=0,
            is_pp_first=True,
            is_pp_last=True
        )

    async def generate(self, prompt_ids: list[int], max_new_tokens: int = 256) -> list[int]:
        session = GenerationSession(
            model_config=self.runner.model_config,
            engine=self.runner.engine,
            input_ids=prompt_ids
        )
        session.decode(max_new_tokens=max_new_tokens)
        return session.output_ids

import asyncio
from dataclasses import dataclass, field
from collections import deque
import time

@dataclass
class InferenceRequest:
    request_id: str
    input_ids: list[int]
    max_tokens: int
    temperature: float = 1.0
    top_p: float = 0.9
    stream: bool = False
    arrival_time: float = field(default_factory=time.time)
    future: asyncio.Future = field(default_factory=asyncio.get_event_loop().create_future)

class DynamicBatcher:
    """
    动态批处理器：收集请求到批次中，达到阈值或超时后提交
    核心：在延迟和吞吐之间权衡
    """
    def __init__(self, engine, max_batch_size: int = 32,
                 max_wait_ms: float = 50, max_seq_len: int = 2048):
        self.engine = engine
        self.max_batch_size = max_batch_size
        self.max_wait_ms = max_wait_ms
        self.max_seq_len = max_seq_len
        self.queue: deque[InferenceRequest] = deque()
        self.running = False

    async def submit(self, request: InferenceRequest) -> str:
        self.queue.append(request)
        return await request.future

    async def start(self):
        self.running = True
        while self.running:
            batch = await self._collect_batch()
            if batch:
                await self._process_batch(batch)

    async def _collect_batch(self) -> list[InferenceRequest]:
        batch = []
        deadline = time.time() + self.max_wait_ms / 1000

        while len(batch) < self.max_batch_size and time.time() < deadline:
            if self.queue:
                req = self.queue.popleft()
                batch.append(req)
            else:
                await asyncio.sleep(0.001)

        return batch

    async def _process_batch(self, batch: list[InferenceRequest]):
        # Padding 到相同长度
        max_len = min(max(len(r.input_ids) for r in batch), self.max_seq_len)
        padded_inputs = []
        attention_masks = []

        for req in batch:
            padding_len = max_len - len(req.input_ids[:max_len])
            padded = req.input_ids[:max_len] + [0] * padding_len
            mask = [1] * min(len(req.input_ids), max_len) + [0] * padding_len
            padded_inputs.append(padded)
            attention_masks.append(mask)

        # 执行批量推理
        outputs = await self.engine.generate_batch(
            input_ids=padded_inputs,
            attention_mask=attention_masks,
            max_new_tokens=max(r.max_tokens for r in batch),
            temperature=max(r.temperature for r in batch),
        )

        # 分发结果
        for req, output in zip(batch, outputs):
            if not req.future.done():
                req.future.set_result(output)

class ContinuousBatcher:
    """
    连续批处理（Iteration-Level Scheduling）
    与传统批处理不同：不需要等所有请求完成才接受新请求
    每个解码步后都可以加入新请求、移除已完成请求
    """
    def __init__(self, engine, max_batch_size: int = 256):
        self.engine = engine
        self.max_batch_size = max_batch_size
        self.active_sequences: list[Sequence] = []
        self.waiting_queue: deque[Sequence] = deque()

    @dataclass
    class Sequence:
        request_id: str
        input_ids: list[int]
        output_ids: list[int] = field(default_factory=list)
        max_tokens: int = 256
        is_finished: bool = False
        future: asyncio.Future = None

    async def submit(self, request: InferenceRequest) -> str:
        seq = self.Sequence(
            request_id=request.request_id,
            input_ids=request.input_ids,
            max_tokens=request.max_tokens,
            future=asyncio.get_event_loop().create_future()
        )
        self.waiting_queue.append(seq)
        return await seq.future

    async def run(self):
        while True:
            # 1. 从等待队列添加新序列到活跃批次
            while (self.waiting_queue and
                   len(self.active_sequences) < self.max_batch_size):
                self.active_sequences.append(self.waiting_queue.popleft())

            if not self.active_sequences:
                await asyncio.sleep(0.001)
                continue

            # 2. 执行单步解码
            step_outputs = await self.engine.step_decode(
                [s.input_ids + s.output_ids for s in self.active_sequences]
            )

            # 3. 处理输出
            finished = []
            for seq, output in zip(self.active_sequences, step_outputs):
                seq.output_ids.append(output.token_id)
                if output.is_eos or len(seq.output_ids) >= seq.max_tokens:
                    seq.is_finished = True
                    seq.future.set_result(seq.output_ids)
                    finished.append(seq)

            # 4. 移除已完成序列
            self.active_sequences = [s for s in self.active_sequences if not s.is_finished]

class GPUScheduler:
    """GPU 资源调度器：管理多模型在同一 GPU 上的调度"""
    def __init__(self, gpu_memory_total: int = 80_000):  # MB
        self.total_memory = gpu_memory_total
        self.allocated: dict[str, int] = {}  # model_id -> memory_mb
        self.model_streams: dict[str, int] = {}  # model_id -> cuda_stream

    def can_load(self, model_memory: int) -> bool:
        used = sum(self.allocated.values())
        return (self.total_memory - used) >= model_memory

    def allocate(self, model_id: str, memory_mb: int) -> int:
        if not self.can_load(memory_mb):
            raise RuntimeError(f"Insufficient GPU memory: need {memory_mb}MB, "
                             f"available {self.total_memory - sum(self.allocated.values())}MB")
        self.allocated[model_id] = memory_mb
        return len(self.allocated)  # stream id

    def deallocate(self, model_id: str):
        self.allocated.pop(model_id, None)
        self.model_streams.pop(model_id, None)

    def get_utilization(self) -> float:
        return sum(self.allocated.values()) / self.total_memory

class ModelWarmPool:
    """模型热池：常用模型常驻 GPU，不常用模型按需加载"""
    def __init__(self, scheduler: GPUScheduler, storage_path: str):
        self.scheduler = scheduler
        self.storage = storage_path
        self.loaded: dict[str, Any] = {}  # model_id -> model instance
        self.loading_locks: dict[str, asyncio.Lock] = {}
        self.access_times: dict[str, float] = {}
        self.idle_threshold = 600  # 10分钟未使用则卸载

    async def get_model(self, model_id: str, model_config: dict):
        if model_id in self.loaded:
            self.access_times[model_id] = time.time()
            return self.loaded[model_id]

        # 防止并发加载同一模型
        if model_id not in self.loading_locks:
            self.loading_locks[model_id] = asyncio.Lock()

        async with self.loading_locks[model_id]:
            if model_id in self.loaded:
                return self.loaded[model_id]

            # 检查 GPU 内存是否足够
            required_mem = model_config["memory_mb"]
            if not self.scheduler.can_load(required_mem):
                await self._evict_idle_models(required_mem)

            # 加载模型
            model = await self._load_model(model_id, model_config)
            self.scheduler.allocate(model_id, required_mem)
            self.loaded[model_id] = model
            self.access_times[model_id] = time.time()
            return model

    async def _evict_idle_models(self, required_mem: int):
        """LRU 淘汰空闲模型"""
        sorted_models = sorted(self.access_times.items(), key=lambda x: x[1])
        for model_id, last_access in sorted_models:
            if time.time() - last_access > self.idle_threshold:
                await self._unload_model(model_id)
                if self.scheduler.can_load(required_mem):
                    break

    async def _unload_model(self, model_id: str):
        if model_id in self.loaded:
            del self.loaded[model_id]
            self.scheduler.deallocate(model_id)
            self.access_times.pop(model_id, None)
            import torch
            torch.cuda.empty_cache()

class Autoscaler:
    """基于指标的自定扩缩容"""
    def __init__(self, min_replicas: int = 1, max_replicas: int = 10,
                 target_gpu_util: float = 0.7, target_latency_ms: float = 500):
        self.min_replicas = min_replicas
        self.max_replicas = max_replicas
        self.target_gpu_util = target_gpu_util
        self.target_latency = target_latency_ms
        self.current_replicas = min_replicas

    def decide(self, metrics: dict) -> dict:
        gpu_util = metrics.get("gpu_utilization", 0)
        avg_latency = metrics.get("avg_latency_ms", 0)
        queue_depth = metrics.get("queue_depth", 0)
        qps = metrics.get("qps", 0)

        # 扩容信号
        scale_up_score = 0
        if gpu_util > self.target_gpu_util * 1.2:
            scale_up_score += 2
        if avg_latency > self.target_latency * 1.5:
            scale_up_score += 2
        if queue_depth > 50:
            scale_up_score += 1

        # 缩容信号
        scale_down_score = 0
        if gpu_util < self.target_gpu_util * 0.3:
            scale_down_score += 2
        if avg_latency < self.target_latency * 0.3 and queue_depth < 5:
            scale_down_score += 1
        if qps < 1 and self.current_replicas > 1:
            scale_down_score += 1

        if scale_up_score >= 2 and self.current_replicas < self.max_replicas:
            return {"action": "scale_up", "replicas": self.current_replicas + 1}
        elif scale_down_score >= 2 and self.current_replicas > self.min_replicas:
            return {"action": "scale_down", "replicas": self.current_replicas - 1}
        else:
            return {"action": "none", "replicas": self.current_replicas}

# vllm-deployment.yaml
apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-llama-70b
spec:
  replicas: 2
  selector:
    matchLabels:
      app: vllm-llama-70b
  template:
    metadata:
      labels:
        app: vllm-llama-70b
    spec:
      containers:
      - name: vllm
        image: vllm/vllm-openai:latest
        args:
          - --model=meta-llama/Llama-3-70B
          - --tensor-parallel-size=4
          - --gpu-memory-utilization=0.90
          - --max-model-len=8192
        resources:
          limits:
            nvidia.com/gpu: 4
            memory: 128Gi
          requests:
            nvidia.com/gpu: 4
            memory: 64Gi
        readinessProbe:
          httpGet:
            path: /health
            port: 8000
          initialDelaySeconds: 120
          periodSeconds: 10
---
apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: vllm-hpa
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: vllm-llama-70b
  minReplicas: 2
  maxReplicas: 10
  metrics:
  - type: Pods
    pods:
      metric:
        name: gpu_utilization
      target:
        type: AverageValue
        averageValue: "70"
  - type: Resource
    resource:
      name: memory
      target:
        type: Utilization
        averageUtilization: 80

class InferenceLoadBalancer:
    """推理负载均衡器：感知 GPU 负载和模型位置"""
    def __init__(self):
        self.backends: dict[str, BackendInfo] = {}

    @dataclass
    class BackendInfo:
        url: str
        model_id: str
        gpu_util: float = 0.0
        active_requests: int = 0
        queue_depth: int = 0
        avg_latency_ms: float = 0.0
        healthy: bool = True

    def select(self, model_id: str) -> str:
        candidates = [
            (url, info) for url, info in self.backends.items()
            if info.model_id == model_id and info.healthy
        ]
        if not candidates:
            raise RuntimeError(f"No healthy backend for model {model_id}")

        # 最少连接 + 最低延迟加权
        def score(info: BackendInfo) -> float:
            load_score = info.active_requests * 0.5 + info.queue_depth * 0.3
            latency_score = info.avg_latency_ms / 1000 * 0.2
            return load_score + latency_score

        return min(candidates, key=lambda x: score(x[1]))[0]

    def update_health(self, url: str, metrics: dict):
        if url in self.backends:
            info = self.backends[url]
            info.gpu_util = metrics.get("gpu_util", info.gpu_util)
            info.active_requests = metrics.get("active_requests", info.active_requests)
            info.avg_latency_ms = metrics.get("avg_latency_ms", info.avg_latency_ms)
            info.healthy = metrics.get("healthy", True)