🔧 模型部署与服务化¶

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
"""
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel  # Pydantic数据验证模型
import numpy as np
import onnxruntime as ort
import time
from contextlib import asynccontextmanager

# 全局模型引用
model_session = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理：启动时加载模型"""
    global model_session
    model_session = ort.InferenceSession(
        "text_encoder.onnx",
        providers=["CPUExecutionProvider"]
    )
    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]]  # [batch_size, seq_len]

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

@app.post("/predict", response_model=PredictResponse)
async def predict(request: PredictRequest):
    """模型推理接口"""
    try:  # try/except捕获异常
        start = time.perf_counter()

        input_array = np.array(request.input_ids, dtype=np.int64)
        outputs = model_session.run(None, {"input_ids": input_array})
        logits = outputs[0]

        # Softmax概率
        exp_logits = np.exp(logits - np.max(logits, axis=-1, keepdims=True))
        probs = exp_logits / exp_logits.sum(axis=-1, keepdims=True)

        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_loaded": model_session is not None}

# 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)")