凌晨两点,屏幕上的日志还在疯狂滚动。第37次尝试启动TorchServe服务,依然卡在“Loading model…”然后超时。同事发来的微信还在闪烁:“客户明天要看演示,模型部署必须搞定。”这场景太熟悉了——模型在本地跑得好好的,一到生产环境就各种水土不服。

从TorchServe的坑说起

先看那段折磨我半宿的代码:

# 错误示范:直接拿训练脚本里的模型定义来部署
model = MyComplexModel()
model.load_state_dict(torch.load('best_model.pth'))
model.eval()

# 然后试图直接扔给TorchServe... 等着报错吧

问题在哪?训练时的模型类定义和推理时的环境根本是两回事。TorchServe要求你把整个推理逻辑打包成一个自包含的handler,而很多从研究转工程的同学会忽略这个关键点。

正确的打开方式应该是这样:

# my_handler.py
class MyModelHandler(BaseHandler):
    def initialize(self, context):
        # 这里踩过大坑:路径要用绝对路径
        model_path = context.system_properties["model_dir"]
        self.model = torch.jit.load(f"{model_path}/model.pt")
        
    def preprocess(self, data):
        # 客户端传过来的可能是各种格式
        # 别假设输入一定是numpy数组,实际可能是base64字符串
        inputs = data[0].get("body")
        if isinstance(inputs, str):
            inputs = base64.b64decode(inputs)
        return torch.tensor(inputs)
    
    def inference(self, inputs):
        with torch.no_grad():
            outputs = self.model(inputs)
        return outputs
    
    def postprocess(self, inference_output):
        # 实际部署时这里经常要转成JSON友好格式
        # 直接返回tensor客户端会骂人的
        return [inference_output.cpu().numpy().tolist()]

打包模型时更要注意:

# 新手常忘的一步:先转TorchScript
traced_model = torch.jit.trace(model, example_input)
torch.jit.save(traced_model, "model.pt")

# 打包mar文件(这里版本号管理很重要)
torch-model-archiver \
  --model-name mymodel \
  --version 1.0 \
  --serialized-file model.pt \
  --handler my_handler.py \
  --export-path model_store \
  --extra-files config.json  # 配置文件别漏了!

启动服务时遇到的典型问题:

# 错误:内存不足直接OOM
torchserve --start --model-store ./model_store --models mymodel.mar

# 正确:生产环境一定要限制资源
torchserve --start \
  --model-store ./model_store \
  --models mymodel.mar \
  --ncs \
  --ts-config config.properties  # 这里面要配好内存和线程数

Triton:当TorchServe不够用时

上周遇到个场景:客户要求同时支持PyTorch、TensorFlow和ONNX模型,还要做模型流水线。TorchServe这时候就有点力不从心了,这时候NVIDIA Triton Inference Server上场了。

Triton的配置哲学完全不同,它是声明式的:

# config.pbtxt
name: "ensemble_model"
platform: "ensemble"
max_batch_size: 32

input [
  {
    name: "input_data"
    data_type: TYPE_FP32
    dims: [224, 224, 3]
  }
]

output [
  {
    name: "classification_result"
    data_type: TYPE_FP32
    dims: [1000]
  }
]

ensemble_scheduling {
  step [
    {
      model_name: "preprocess_model"
      model_version: -1
      input_map {
        key: "raw_input"
        value: "input_data"
      }
      output_map {
        key: "processed_output"
        value: "preprocessed_data"
      }
    },
    {
      model_name: "inference_model"
      model_version: -1
      input_map {
        key: "model_input"
        value: "preprocessed_data"
      }
      output_map {
        key: "model_output"
        value: "classification_result"
      }
    }
  ]
}

Triton最香的地方是它的客户端库,支持多种语言:

# 客户端调用比TorchServe简洁多了
import tritonclient.http as httpclient

client = httpclient.InferenceServerClient(url="localhost:8000")
inputs = httpclient.InferInput("input_data", [1, 224, 224, 3], "FP32")
inputs.set_data_from_numpy(image_array)

outputs = httpclient.InferRequestedOutput("classification_result")
results = client.infer("ensemble_model", inputs=[inputs], outputs=[outputs])

性能调优那些事儿

模型服务化不是能跑就行,性能差照样过不了验收。几个关键指标:

  1. 延迟:从收到请求到返回结果的时间
  2. 吞吐:每秒能处理多少请求
  3. 资源利用率:GPU/CPU别闲着也别过载

TorchServe的调优主要在config.properties里:

# 工作线程数(不是越多越好!)
default_workers_per_model=2

# 批量处理超时时间(毫秒)
batch_max_delay=100

# 最大批量大小(根据显存调整)
max_batch_size=16

Triton的调优更细致:

# 针对不同硬件配置实例数
instance_group [
  {
    count: 2  # 两个实例
    kind: KIND_GPU
    gpus: [0, 1]  # 分别跑在两个GPU上
  }
]

# 动态批处理配置
dynamic_batching {
  max_queue_delay_microseconds: 500
  preferred_batch_size: [4, 8, 16]
}

监控与日志:别等出事了才看

生产环境最怕的就是“昨天还好好的,今天突然挂了”。一定要提前埋点:

# 在handler里加监控埋点
def inference(self, inputs):
    start_time = time.time()
    try:
        outputs = self.model(inputs)
        # 记录成功率和延迟
        metrics.add_timer("inference_latency", time.time() - start_time)
        metrics.increment_counter("inference_success")
        return outputs
    except Exception as e:
        # 错误类型要分类记录
        metrics.increment_counter(f"inference_error_{type(e).__name__}")
        logger.error(f"Inference failed: {e}", exc_info=True)
        raise

日志配置也有讲究:

# log4j2.xml里这么配
<AsyncLogger name="ACCESS_LOG" level="info" additivity="false">
  <AppenderRef ref="AccessLog"/>
  <!-- 生产环境别用同步日志,性能掉得厉害 -->
</AsyncLogger>

<AsyncLogger name="MODEL_LOG" level="debug" additivity="false">
  <AppenderRef ref="ModelLog"/>
  <!-- 模型推理日志单独存,方便排查 -->
</AsyncLogger>

经验之谈

干了这么多年模型部署,有些教训是实打实踩坑踩出来的:

关于选型:如果团队主要用PyTorch且模型不太复杂,TorchServe上手更快。但如果要搞多框架支持、模型流水线、或者对性能要求极高,直接上Triton,长远来看省心。别想着“先用简单的,以后再迁移”,迁移成本往往比重做还高。

关于版本管理:模型版本一定要和代码版本绑定。我吃过亏——部署了v1.2模型,用的却是v1.1的预处理代码,线上结果和测试环境对不上。现在我们的CI流程强制要求模型mar文件和handler代码一起打包,版本号必须一致。

关于测试:别只测一张图片的成功请求。要模拟生产环境的并发场景,用locust或者jmeter压测。重点观察内存泄漏——有些模型推理几次没问题,跑一晚上内存就炸了。记得测试异常情况:发畸形数据、并发关停服务、网络闪断恢复。

关于配置:所有参数必须可配置化,别写死在代码里。批量大小、超时时间、工作线程数,这些都要能根据部署环境动态调整。我们在K8s里用ConfigMap管理这些配置,不同环境(测试、预发、生产)用不同配置。

最后说个心态问题:模型服务化是个工程活,不是研究活。追求的不是最新最炫的模型结构,而是稳定、可维护、可监控。有时候为了5%的性能提升花两周优化,不如加台机器来得实在。知道什么时候该优化代码,什么时候该加硬件,这是工程师和研究员的重要区别。

凌晨四点了,服务终于稳定跑起来。看着监控面板上平稳的曲线,我知道明天能给客户交差了。模型部署这条路,从来不是一蹴而就的,每个深夜的调试,都是第二天稳定运行的基石。

Logo

腾讯云面向开发者汇聚海量精品云计算使用和开发经验,营造开放的云计算技术生态圈。

更多推荐