🔧 模型部署与服务化¶

import torch
import torch.nn as nn
import numpy as np

# 定义模型
class TextEncoder(nn.Module):
    def __init__(self, vocab_size=30000, embed_dim=256, hidden_dim=512, num_classes=10):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True, bidirectional=True)
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim, num_classes)
        )

    def forward(self, input_ids):
        embeds = self.embedding(input_ids)
        lstm_out, (hidden, _) = self.lstm(embeds)
        # 拼接前向和后向最终隐状态
        hidden = torch.cat([hidden[-2], hidden[-1]], dim=1)  # 负索引：从末尾倒数访问元素
        logits = self.classifier(hidden)
        return logits

model = TextEncoder()
model.eval()

# ===== 导出ONNX =====
dummy_input = torch.randint(0, 30000, (1, 128))  # [batch_size, seq_len]

torch.onnx.export(
    model,
    dummy_input,
    "text_encoder.onnx",
    export_params=True,
    opset_version=17,
    do_constant_folding=True,           # 常量折叠优化
    input_names=["input_ids"],
    output_names=["logits"],
    dynamic_axes={                       # 动态维度
        "input_ids": {0: "batch_size", 1: "seq_len"},
        "logits": {0: "batch_size"}
    }
)

print("ONNX model exported successfully!")

# ===== 验证ONNX模型 =====
import onnx
onnx_model = onnx.load("text_encoder.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX model validation passed!")

# ===== ONNX Runtime推理 =====
import onnxruntime as ort

session = ort.InferenceSession(
    "text_encoder.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
)

# 推理
input_data = np.random.randint(0, 30000, (2, 128)).astype(np.int64)
outputs = session.run(None, {"input_ids": input_data})
print(f"Output shape: {outputs[0].shape}")  # [2, 10]

import onnxruntime as ort
from onnxruntime.quantization import quantize_dynamic, QuantType
from onnxruntime.transformers import optimizer

# ===== 1. 动态量化 =====
quantize_dynamic(
    model_input="text_encoder.onnx",
    model_output="text_encoder_int8.onnx",
    weight_type=QuantType.QInt8
)

# ===== 2. 图优化（Transformer模型专用）=====
# optimized_model = optimizer.optimize_model(
#     "bert_model.onnx",
#     model_type="bert",
#     num_heads=12,
#     hidden_size=768
# )
# optimized_model.save_model_to_file("bert_optimized.onnx")

# ===== 3. 性能基准测试 =====
import time

def benchmark_onnx(model_path: str, input_data: dict, num_runs: int = 100):
    """ONNX模型性能基准测试"""
    session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])

    # 预热
    for _ in range(10):
        session.run(None, input_data)

    # 基准测试
    latencies = []
    for _ in range(num_runs):
        start = time.perf_counter()
        session.run(None, input_data)
        latencies.append((time.perf_counter() - start) * 1000)

    print(f"Model: {model_path}")
    print(f"  Mean latency: {np.mean(latencies):.2f} ms")
    print(f"  P50 latency:  {np.percentile(latencies, 50):.2f} ms")
    print(f"  P95 latency:  {np.percentile(latencies, 95):.2f} ms")
    print(f"  P99 latency:  {np.percentile(latencies, 99):.2f} ms")

    return latencies

# test_input = {"input_ids": np.random.randint(0, 30000, (1, 128)).astype(np.int64)}
# benchmark_onnx("text_encoder.onnx", test_input)
# benchmark_onnx("text_encoder_int8.onnx", test_input)

import torch

model = TextEncoder()
model.eval()

# ===== 方式一：Tracing（推荐简单模型）=====
dummy_input = torch.randint(0, 30000, (1, 128))
traced_model = torch.jit.trace(model, dummy_input)
traced_model.save("text_encoder_traced.pt")

# ===== 方式二：Scripting（支持动态控制流）=====
class DynamicModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        # 包含条件分支，必须用script
        if x.shape[0] > 1:
            return self.linear(x).mean(dim=0)
        else:
            return self.linear(x).squeeze(0)

dynamic_model = DynamicModel()
scripted_model = torch.jit.script(dynamic_model)
scripted_model.save("dynamic_model_scripted.pt")

# 加载并推理
loaded_model = torch.jit.load("text_encoder_traced.pt")
with torch.no_grad():
    output = loaded_model(dummy_input)
    print(f"TorchScript output shape: {output.shape}")

import tensorflow as tf

# 定义模型
class TFClassifier(tf.keras.Model):
    def __init__(self, num_classes=10):
        super().__init__()
        self.dense1 = tf.keras.layers.Dense(256, activation='relu')
        self.dropout = tf.keras.layers.Dropout(0.3)
        self.dense2 = tf.keras.layers.Dense(num_classes)

    def call(self, inputs, training=False):
        x = self.dense1(inputs)
        x = self.dropout(x, training=training)
        return self.dense2(x)

# 构建并保存
model = TFClassifier()
model(tf.random.normal((1, 768)))  # 初始化权重

# SavedModel格式
tf.saved_model.save(model, "saved_model/classifier")

# 带签名保存（推荐）
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 768], dtype=tf.float32)])
def serve(inputs):
    return {"predictions": model(inputs, training=False)}

tf.saved_model.save(model, "saved_model/classifier_v2", signatures={"serving_default": serve})

# 加载推理
loaded = tf.saved_model.load("saved_model/classifier_v2")
result = loaded.signatures["serving_default"](tf.random.normal((2, 768)))
print(f"TF output shape: {result['predictions'].shape}")

"""
简单的模型推理服务
运行：uvicorn serve:app --host 0.0.0.0 --port 8000
"""
import os
import time
import numpy as np
import onnxruntime as ort
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel, Field

MODEL_PATH = os.getenv("MODEL_PATH", "/app/model/model.onnx")
MODEL_PROVIDERS = os.getenv("MODEL_PROVIDERS", "CPUExecutionProvider").split(",")

@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理：启动时加载模型"""
    app.state.model_session = ort.InferenceSession(
        MODEL_PATH,
        providers=MODEL_PROVIDERS,
    )
    print("Model loaded successfully!")
    yield
    print("Shutting down...")

app = FastAPI(title="Text Classification API", version="1.0", lifespan=lifespan)

class PredictRequest(BaseModel):
    input_ids: list[list[int]] = Field(..., description="[batch_size, seq_len]")

class PredictResponse(BaseModel):
    predictions: list[int]
    probabilities: list[list[float]]
    latency_ms: float

def softmax(logits: np.ndarray) -> np.ndarray:
    shifted = logits - np.max(logits, axis=-1, keepdims=True)
    exp_logits = np.exp(shifted)
    return exp_logits / exp_logits.sum(axis=-1, keepdims=True)

@app.post("/predict", response_model=PredictResponse)
async def predict(payload: PredictRequest, request: Request):
    """模型推理接口"""
    try:
        session: ort.InferenceSession = request.app.state.model_session
        start = time.perf_counter()

        input_array = np.array(payload.input_ids, dtype=np.int64)
        input_name = session.get_inputs()[0].name
        outputs = session.run(None, {input_name: input_array})
        logits = outputs[0]

        probs = softmax(logits)
        predictions = np.argmax(logits, axis=-1).tolist()
        latency = (time.perf_counter() - start) * 1000

        return PredictResponse(
            predictions=predictions,
            probabilities=probs.tolist(),
            latency_ms=round(latency, 2)
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health():
    """健康检查"""
    return {
        "status": "healthy",
        "model_path": MODEL_PATH,
        "providers": MODEL_PROVIDERS,
    }

# Triton客户端调用
import tritonclient.http as httpclient
import numpy as np

def triton_inference(input_ids: np.ndarray):
    """调用Triton服务"""
    client = httpclient.InferenceServerClient(url="localhost:8000")

    # 检查服务状态
    if not client.is_server_ready():
        raise RuntimeError("Triton server not ready")

    # 构建输入
    inputs = [
        httpclient.InferInput("input_ids", input_ids.shape, "INT64")
    ]
    inputs[0].set_data_from_numpy(input_ids)

    # 构建输出
    outputs = [
        httpclient.InferRequestedOutput("logits")
    ]

    # 推理
    result = client.infer(
        model_name="text_encoder",
        model_version="1",
        inputs=inputs,
        outputs=outputs
    )

    logits = result.as_numpy("logits")
    return logits

# 使用
# input_data = np.random.randint(0, 30000, (4, 128)).astype(np.int64)
# result = triton_inference(input_data)
# print(f"Triton result shape: {result.shape}")

# 1. 创建模型处理器（handler.py）
"""
自定义TorchServe Handler
"""
import json
import torch
import numpy as np
from ts.torch_handler.base_handler import BaseHandler

class TextClassificationHandler(BaseHandler):

    def __init__(self):
        super().__init__()
        self.initialized = False

    def initialize(self, context):
        """模型初始化"""
        properties = context.system_properties
        model_dir = properties.get("model_dir")

        # 加载模型
        self.model = torch.jit.load(f"{model_dir}/model.pt")
        self.model.eval()

        # 加载标签映射
        with open(f"{model_dir}/labels.json", "r") as f:  # with自动管理资源，确保文件正确关闭
            self.labels = json.load(f)  # json.loads将JSON字符串转为Python对象

        self.initialized = True

    def preprocess(self, data):
        """预处理输入"""
        inputs = []
        for row in data:
            input_data = row.get("data") or row.get("body")
            if isinstance(input_data, (bytes, bytearray)):  # isinstance检查对象类型
                input_data = input_data.decode("utf-8")
            if isinstance(input_data, str):
                input_data = json.loads(input_data)
            inputs.append(input_data["input_ids"])
        return torch.tensor(inputs, dtype=torch.long)

    def inference(self, input_batch):
        """模型推理"""
        with torch.no_grad():
            logits = self.model(input_batch)
            probs = torch.softmax(logits, dim=-1)
            predictions = torch.argmax(probs, dim=-1)
        return predictions, probs

    def postprocess(self, inference_output):
        """后处理输出"""
        predictions, probs = inference_output
        results = []
        for pred, prob in zip(predictions, probs):  # zip并行遍历多个可迭代对象
            results.append({
                "label": self.labels[pred.item()],
                "confidence": prob[pred].item(),
                "all_probabilities": {
                    self.labels[i]: p.item()
                    for i, p in enumerate(prob)  # enumerate同时获取索引和值
                }
            })
        return results

# KServe自定义Predictor
from kserve import Model, ModelServer
import numpy as np
import onnxruntime as ort

class TextClassifierModel(Model):
    def __init__(self, name: str):
        super().__init__(name)  # super()调用父类方法
        self.name = name
        self.model = None
        self.ready = False

    def load(self):
        """加载模型"""
        self.model = ort.InferenceSession(
            "/mnt/models/model.onnx",
            providers=["CPUExecutionProvider"]
        )
        self.ready = True

    def predict(self, payload: dict, headers: dict = None) -> dict:
        """推理"""
        input_ids = np.array(payload["instances"], dtype=np.int64)
        outputs = self.model.run(None, {"input_ids": input_ids})
        predictions = np.argmax(outputs[0], axis=-1).tolist()
        return {"predictions": predictions}

if __name__ == "__main__":
    model = TextClassifierModel("text-classifier")
    model.load()
    ModelServer().start([model])

"""
金丝雀发布管理器
逐步将流量从旧版本迁移到新版本
"""
import time
from dataclasses import dataclass

@dataclass
class CanaryConfig:
    """金丝雀发布配置"""
    initial_weight: int = 5       # 初始流量比例 %
    step_weight: int = 10         # 每步增加 %
    max_weight: int = 100         # 最大流量
    step_interval: int = 300      # 每步间隔（秒）
    error_threshold: float = 0.01 # 错误率阈值
    latency_threshold_ms: float = 200  # 延迟阈值
    rollback_on_failure: bool = True

class CanaryDeployment:
    """金丝雀部署管理器"""

    def __init__(self, config: CanaryConfig):
        self.config = config
        self.current_weight = 0
        self.is_active = False

    def start(self, old_version: str, new_version: str):
        """开始金丝雀发布"""
        self.old_version = old_version
        self.new_version = new_version
        self.current_weight = self.config.initial_weight
        self.is_active = True

        print(f"Starting canary: {old_version} -> {new_version}")
        print(f"Initial traffic split: {self.current_weight}% to new version")

        self._update_traffic_split()

    def check_health(self) -> dict:
        """检查新版本的健康状况"""
        # 实际实现中从Prometheus/监控系统获取指标
        metrics = {
            "error_rate": 0.005,       # 模拟：0.5%错误率
            "p99_latency_ms": 150,     # 模拟：150ms p99延迟
            "accuracy": 0.92,          # 模拟：92%准确率
        }
        return metrics

    def should_proceed(self, metrics: dict) -> bool:
        """判断是否可以继续推进"""
        if metrics["error_rate"] > self.config.error_threshold:
            print(f"ERROR: Error rate {metrics['error_rate']:.3f} exceeds threshold")
            return False
        if metrics["p99_latency_ms"] > self.config.latency_threshold_ms:
            print(f"ERROR: Latency {metrics['p99_latency_ms']:.0f}ms exceeds threshold")
            return False
        return True

    def advance(self):
        """推进金丝雀流量"""
        metrics = self.check_health()

        if not self.should_proceed(metrics):
            if self.config.rollback_on_failure:
                self.rollback()
            return False

        self.current_weight = min(
            self.current_weight + self.config.step_weight,
            self.config.max_weight
        )

        print(f"Advancing canary: {self.current_weight}% traffic to {self.new_version}")
        self._update_traffic_split()

        if self.current_weight >= self.config.max_weight:
            self.promote()
            return True

        return True

    def rollback(self):
        """回滚到旧版本"""
        print(f"ROLLBACK: Reverting to {self.old_version}")
        self.current_weight = 0
        self._update_traffic_split()
        self.is_active = False

    def promote(self):
        """完全切换到新版本"""
        print(f"PROMOTED: {self.new_version} is now serving 100% traffic")
        self.is_active = False

    def _update_traffic_split(self):
        """更新流量分配（实际中调用K8s API）"""
        old_weight = 100 - self.current_weight
        print(f"  Traffic: {self.old_version}={old_weight}%, {self.new_version}={self.current_weight}%")

# 使用示例
config = CanaryConfig(initial_weight=5, step_weight=15, step_interval=60)
canary = CanaryDeployment(config)
canary.start("v1.0", "v2.0")

# 模拟逐步推进
for step in range(8):
    print(f"\n--- Step {step + 1} ---")
    if not canary.advance():
        break
    if not canary.is_active:
        break

import asyncio
import time
import numpy as np
from collections import deque
from dataclasses import dataclass, field
from threading import Lock  # 线程池/多线程：并发执行任务

@dataclass  # 自动生成__init__等方法
class InferenceRequest:
    """推理请求"""
    input_ids: np.ndarray
    future: asyncio.Future = field(default=None)
    timestamp: float = field(default_factory=time.time)

class DynamicBatcher:
    """
    动态批处理：将多个请求合并为批次，提高GPU利用率
    """

    def __init__(
        self,
        model_session,
        max_batch_size: int = 32,
        max_wait_time_ms: float = 50,
    ):
        self.model = model_session
        self.max_batch_size = max_batch_size
        self.max_wait_time = max_wait_time_ms / 1000
        self.queue: deque = deque()
        self.lock = Lock()
        self._running = True

    async def predict(self, input_ids: np.ndarray) -> np.ndarray:
        """提交推理请求，等待批处理结果"""
        loop = asyncio.get_running_loop()
        future = loop.create_future()

        request = InferenceRequest(
            input_ids=input_ids,
            future=future
        )

        with self.lock:
            self.queue.append(request)

        # 如果达到批大小，立即处理
        if len(self.queue) >= self.max_batch_size:
            await self._process_batch()

        return await future

    async def _process_batch(self):
        """处理一个批次"""
        with self.lock:
            batch_size = min(len(self.queue), self.max_batch_size)
            if batch_size == 0:
                return

            requests = [self.queue.popleft() for _ in range(batch_size)]

        # 合并输入
        batch_input = np.concatenate([r.input_ids for r in requests], axis=0)

        # 批量推理
        outputs = self.model.run(None, {"input_ids": batch_input})
        results = outputs[0]

        # 分发结果
        offset = 0
        for req in requests:
            size = req.input_ids.shape[0]
            result = results[offset:offset + size]
            if not req.future.done():
                req.future.set_result(result)
            offset += size

    async def batch_processor(self):
        """后台批处理循环"""
        while self._running:
            if len(self.queue) > 0:
                oldest = self.queue[0].timestamp
                if (time.time() - oldest) >= self.max_wait_time or \
                   len(self.queue) >= self.max_batch_size:
                    await self._process_batch()
            await asyncio.sleep(0.001)  # 1ms检查间隔

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import asyncio
import json

app = FastAPI()

async def async_inference(input_data: dict):  # async def定义异步函数；用await调用
    """异步推理"""
    # 模拟耗时推理
    await asyncio.sleep(0.1)  # await等待异步操作完成
    return {"prediction": 1, "confidence": 0.95}

@app.post("/predict/async")
async def predict_async(request: dict):
    """异步推理接口 - 返回任务ID"""
    task_id = f"task_{id(request)}"
    # 提交后台任务
    asyncio.create_task(process_task(task_id, request))
    return {"task_id": task_id, "status": "processing"}

@app.post("/predict/stream")
async def predict_stream(request: dict):
    """流式推理接口 - 适用于LLM生成"""
    async def generate():
        tokens = ["Hello", " world", "!", " This", " is", " streaming", "."]
        for token in tokens:
            yield f"data: {json.dumps({'token': token})}\n\n"  # yield生成器：惰性产出值，节省内存
            await asyncio.sleep(0.1)
        yield "data: [DONE]\n\n"

    return StreamingResponse(
        generate(),
        media_type="text/event-stream"
    )

async def process_task(task_id: str, request: dict):
    """后台处理任务"""
    result = await async_inference(request)
    # 存储结果到Redis/数据库
    print(f"Task {task_id} completed: {result}")

"""
推理性能优化策略汇总
"""

class InferenceOptimizer:
    """推理优化策略集合"""

    @staticmethod  # @staticmethod静态方法，不需要实例
    def model_optimization():
        """模型层面优化"""
        strategies = {
            "量化": "FP32 → FP16/INT8，减少计算量和内存",
            "剪枝": "移除冗余参数，减少模型大小",
            "蒸馏": "大模型指导小模型，保持性能降低成本",
            "ONNX优化": "图优化、常量折叠、算子融合",
            "TensorRT": "NVIDIA GPU专用优化，FP16/INT8加速",
        }
        return strategies

    @staticmethod
    def serving_optimization():
        """服务层面优化"""
        strategies = {
            "动态批处理": "合并请求，提高GPU利用率",
            "模型预热": "启动时预加载模型和运行warmup推理",
            "连接池": "复用HTTP/gRPC连接",
            "异步处理": "非阻塞IO，提高并发能力",
            "缓存": "相同输入缓存结果，减少重复计算",
            "多Worker": "Gunicorn多进程，利用多核CPU",
        }
        return strategies

    @staticmethod
    def infrastructure_optimization():
        """基础设施优化"""
        strategies = {
            "弹性伸缩": "HPA根据负载自动扩缩容",
            "GPU共享": "MIG/MPS多模型共享GPU",
            "负载均衡": "Nginx/Istio分配请求",
            "CDN缓存": "静态资产缓存，减少带宽",
            "就近部署": "多区域部署，降低网络延迟",
        }
        return strategies

# 模型预热示例
def warmup_model(session, input_shape, num_warmup: int = 10):
    """模型预热，消除首次推理的冷启动延迟"""
    import numpy as np

    dummy_input = np.random.randint(0, 1000, input_shape).astype(np.int64)

    for i in range(num_warmup):
        session.run(None, {"input_ids": dummy_input})

    print(f"Model warmup completed ({num_warmup} runs)")