模型部署：REST API 与微服务

# 批处理推理：最大化 GPU 利用率
import torch
from torch.utils.data import DataLoader

def batch_inference(model, dataset, batch_size=256, device="cuda"):
    model.eval()
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
    predictions = []

    with torch.no_grad():
        for batch in loader:
            batch = batch.to(device)
            preds = model(batch)
            predictions.append(preds.cpu().numpy())

    return np.concatenate(predictions, axis=0)

# 在线推理：低延迟单条处理
class OnlineInferenceService:
    def __init__(self, model_path):
        self.model = torch.jit.load(model_path)
        self.model.eval()
        self.model.to("cuda")
        self._warmup()

    def _warmup(self):
        dummy = torch.randn(1, 768, device="cuda")
        for _ in range(10):
            _ = self.model(dummy)

    def predict(self, input_tensor):
        with torch.no_grad():
            return self.model(input_tensor.to("cuda"))

# PyTorch 导出 ONNX
import torch

model = MyModel()
model.load_state_dict(torch.load("model.pth"))
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"},
    },
    opset_version=17,
)

# PyTorch 导出 TorchScript
import torch

# 方式一：trace（适合无控制流的模型）
model = MyModel().eval()
traced = torch.jit.trace(model, torch.randn(1, 768))
traced.save("model_traced.pt")

# 方式二：script（支持控制流）
class ScriptableModel(torch.nn.Module):
    def forward(self, x):
        if x.dim() > 2:
            return self.encoder(x)
        return self.decoder(x)

scripted = torch.jit.script(ScriptableModel())
scripted.save("model_scripted.pt")

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from contextlib import asynccontextmanager
import onnxruntime as ort
import numpy as np

class InferenceRequest(BaseModel):
    features: list
    request_id: str | None = None

class InferenceResponse(BaseModel):
    prediction: list
    request_id: str | None = None
    latency_ms: float

@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.session = ort.InferenceSession("model.onnx")
    yield
    app.state.session = None

app = FastAPI(lifespan=lifespan)

@app.post("/predict", response_model=InferenceResponse)
async def predict(req: InferenceRequest):
    import time
    t0 = time.time()
    try:
        input_data = np.array([req.features], dtype=np.float32)
        result = app.state.session.run(
            None, {"input": input_data}
        )
        return InferenceResponse(
            prediction=result[0].tolist(),
            request_id=req.request_id,
            latency_ms=(time.time() - t0) * 1000,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# 启动服务并配置 Workers
# Uvicorn 多 worker 模式适合 CPU 推理
uvicorn main:app \
    --host 0.0.0.0 \
    --port 8000 \
    --workers 4 \
    --log-level info

# 测试 API
curl -X POST http://localhost:8000/predict \
    -H "Content-Type: application/json" \
    -d '{"features": [0.5, 1.2, -0.3], "request_id": "test-001"}'

# 查看自动生成的文档
# 浏览器打开 http://localhost:8000/docs

# GPU 推理服务的 Dockerfile
FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04

WORKDIR /app

# 安装 Python 和依赖
RUN apt-get update && apt-get install -y \
    python3.11 python3.11-venv && \
    rm -rf /var/lib/apt/lists/*

RUN python3.11 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY app/ ./app/
# 模型文件不打包进镜像，启动时从 S3 下载
COPY scripts/download_model.sh ./

EXPOSE 8000

CMD ["sh", "download_model.sh", "&&", "uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]

# 构建和运行
docker build -t inference-service:v1 .

# CPU 模式运行
docker run -d --name inference \
    -p 8000:8000 \
    -e MODEL_URL=s3://bucket/model.onnx \
    inference-service:v1

# GPU 模式运行（需要 nvidia-container-toolkit）
docker run -d --gpus all --name inference-gpu \
    -p 8000:8000 \
    -e MODEL_URL=s3://bucket/model.onnx \
    -e CUDA_VISIBLE_DEVICES=0 \
    inference-service:v1

# 查看容器日志
docker logs -f inference-gpu

apiVersion: apps/v1
kind: Deployment
metadata:
  name: inference-service
spec:
  replicas: 3
  selector:
    matchLabels:
      app: inference
  template:
    metadata:
      labels:
        app: inference
    spec:
      containers:
        - name: inference
          image: registry.example.com/inference:v1
          ports:
            - containerPort: 8000
          resources:
            requests:
              cpu: "2"
              memory: "4Gi"
              nvidia.com/gpu: 1
            limits:
              cpu: "4"
              memory: "8Gi"
              nvidia.com/gpu: 1
          env:
            - name: MODEL_URL
              value: "s3://bucket/model.onnx"
          readinessProbe:
            httpGet:
              path: /health
              port: 8000
            initialDelaySeconds: 30
            periodSeconds: 10

apiVersion: v1
kind: Service
metadata:
  name: inference-service
spec:
  selector:
    app: inference
  ports:
    - protocol: TCP
      port: 80
      targetPort: 8000
  type: ClusterIP
---
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
  name: inference-ingress
spec:
  rules:
    - host: api.example.com
      http:
        paths:
          - path: /predict
            pathType: Prefix
            backend:
              service:
                name: inference-service
                port:
                  number: 80

# 基于自定义指标的 HPA
apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: inference-hpa
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: inference-service
  minReplicas: 2
  maxReplicas: 20
  metrics:
    - type: Pods
      pods:
        metric:
          name: inference_queue_depth
        target:
          type: AverageValue
          averageValue: "10"
    - type: Resource
      resource:
        name: cpu
        target:
          type: Utilization
          averageUtilization: 70
  behavior:
    scaleUp:
      stabilizationWindowSeconds: 60
      policies:
        - type: Pods
          value: 4
          periodSeconds: 60
    scaleDown:
      stabilizationWindowSeconds: 300
      policies:
        - type: Percent
          value: 10
          periodSeconds: 120

# Prometheus 自定义指标导出器
from prometheus_client import Histogram, Counter, Gauge
import time

# 推理延迟直方图
INFERENCE_LATENCY = Histogram(
    "inference_latency_seconds",
    "Inference latency",
    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0],
)

# 请求队列深度
QUEUE_DEPTH = Gauge(
    "inference_queue_depth",
    "Number of pending requests",
)

# 请求计数器
REQUEST_COUNT = Counter(
    "inference_requests_total",
    "Total inference requests",
    ["status"],
)

@INFERENCE_LATENCY.time()
def run_inference(input_data):
    REQUEST_COUNT.labels(status="success").inc()
    return model.predict(input_data)

# GitHub Actions CI/CD Pipeline
name: ML Deploy Pipeline
on:
  push:
    branches: [main]

jobs:
  test-and-build:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4

      - name: Run model tests
        run: |
          pip install -r requirements.txt
          pytest tests/ -v
          python scripts/validate_model.py

      - name: Build and push
        run: |
          docker build -t registry/inference:${{ github.sha }} .
          docker push registry/inference:${{ github.sha }}

  deploy:
    needs: test-and-build
    runs-on: ubuntu-latest
    steps:
      - name: Update K8s manifest
        run: |
          sed -i "s|image:.*|image: registry/inference:${{ github.sha }}|" k8s/deployment.yaml
          kubectl apply -f k8s/

      - name: Wait and verify
        run: |
          kubectl rollout status deployment/inference-service --timeout=300s
          python scripts/smoke_test.py

# Argo Rollouts 金丝雀发布配置
apiVersion: argoproj.io/v1alpha1
kind: Rollout
metadata:
  name: inference-rollout
spec:
  replicas: 5
  strategy:
    canary:
      steps:
        - setWeight: 5
        - pause: { duration: 5m }
        - setWeight: 25
        - pause: { duration: 5m }
        - setWeight: 50
        - pause: { duration: 10m }
        - setWeight: 100
      analysis:
        templates:
          - templateName: inference-analysis
        startingStep: 1
---
apiVersion: argoproj.io/v1alpha1
kind: AnalysisTemplate
metadata:
  name: inference-analysis
spec:
  metrics:
    - name: error-rate
      interval: 1m
      successCondition: result[0] < 0.01
      provider:
        prometheus:
          query: |
            sum(rate(inference_requests_total{status="error"}[1m]))
            /
            sum(rate(inference_requests_total[1m]))