终章：部署上线+性能调优——从Dev到Prod的AI协作闭环

import torch
import onnx
import onnxruntime as ort
from torch.testing import assert_close

# 1. 构建示例模型与输入
model = torch.hub.load('pytorch/vision:v0.15.0', 'resnet18', pretrained=True).eval()
x = torch.randn(1, 3, 224, 224)

# 2. 导出 ONNX（关键：指定 dynamic_axes 支持变长 batch）
onnx_path = "resnet18.onnx"
torch.onnx.export(
    model, x,
    onnx_path,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
    opset_version=17
)

# 3. 加载并推理 ONNX
ort_session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
ort_out = ort_session.run(None, {"input": x.numpy()})[0]

# 4. PyTorch 原生推理
with torch.no_grad():
    pt_out = model(x).numpy()

# 5. 断言严格一致性（容忍 1e-5 数值误差）
assert_close(
    torch.from_numpy(ort_out),
    torch.from_numpy(pt_out),
    atol=1e-5, rtol=1e-5,
    msg="ONNX output deviates from PyTorch beyond tolerance!"
)

# 构建阶段：编译依赖、打包模型
FROM nvidia/cuda:12.1.1-runtime-ubuntu22.04 AS builder
RUN apt-get update && apt-get install -y python3-pip && rm -rf /var/lib/apt/lists/*
COPY requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt

# 运行阶段：极简镜像
FROM nvidia/cuda:12.1.1-runtime-ubuntu22.04
RUN apt-get update && apt-get install -y curl && rm -rf /var/lib/apt/lists/*
COPY --from=builder /usr/local/lib/python3.10/site-packages /usr/local/lib/python3.10/site-packages
COPY model-store/ /home/model-server/model-store/
COPY config.properties /home/model-server/config.properties
EXPOSE 8080 8081
CMD ["torchserve", "--start", "--model-store", "/home/model-server/model-store", "--ts-config", "/home/model-server/config.properties"]

http {
  limit_req_zone $binary_remote_addr zone=api:10m rate=50r/s;

  upstream torchserve {
    server 127.0.0.1:8080;
  }

  server {
    location /predictions/ {
      limit_req zone=api burst=100 nodelay;
      proxy_pass http://torchserve;
      proxy_next_upstream error timeout http_500;
      client_max_body_size 100M;  # 同步调整！解决 413 错误
    }
  }
}

# 注册服务
curl -X POST http://kong:8001/services \
  --data name=resnet-service \
  --data url="http://torchserve:8080"

# 绑定路由
curl -X POST http://kong:8001/services/resnet-service/routes \
  --data paths[]="/v1/predict"

# 启用插件链
curl -X POST http://kong:8001/plugins \
  --data name=rate-limiting \
  --data config.minute=3000 \
  --data service.id=$(curl -s http://kong:8001/services/resnet-service | jq -r .id)

curl -X POST http://kong:8001/plugins \
  --data name=request-transformer \
  --data config.add.headers="X-Trace-ID:${uuid()}" \
  --data service.id=$(curl -s http://kong:8001/services/resnet-service | jq -r .id)

# locustfile.py
from locust import HttpUser, task, between
import json

class TorchServeUser(HttpUser):
    wait_time = between(0.1, 0.5)
    
    @task
    def predict(self):
        with open("test.jpg", "rb") as f:
            self.client.post(
                "/predictions/resnet18",
                files={"data": f},
                timeout=30
            )

from evidently.report import Report
from evidently.metrics import DataDriftMetrics, ClassificationPerformanceMetrics
from datetime import datetime

report = Report(metrics=[
    DataDriftMetrics(),
    ClassificationPerformanceMetrics()
])
report.run(reference_data=ref_df, current_data=prod_df)
report.save_html(f"drift_report_{datetime.now().strftime('%Y%m%d')}.html")

- name: Check drift score
  id: drift
  run: |
    SCORE=$(python -c "import json; print(json.load(open('drift.json'))['metrics'][0]['result']['dataset_drift'])")
    echo "DRIFT_SCORE=$SCORE" >> $GITHUB_ENV
    echo "score=${SCORE}" >> $GITHUB_OUTPUT

- name: Trigger retrain if drift > 0.3
  if: ${{ env.DRIFT_SCORE > 0.3 }}
  run: |
    python train.py
    docker build -t acme/model:latest .
    docker push acme/model:latest
    kubectl rollout restart deployment/torchserve