027、模型服务化:深夜调不通的TorchServe和那个救场的Triton
凌晨两点,屏幕上的日志还在疯狂滚动。第37次尝试启动TorchServe服务,依然卡在“Loading model…”然后超时。同事发来的微信还在闪烁:“客户明天要看演示,模型部署必须搞定。”这场景太熟悉了——模型在本地跑得好好的,一到生产环境就各种水土不服。
凌晨两点,屏幕上的日志还在疯狂滚动。第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])
性能调优那些事儿
模型服务化不是能跑就行,性能差照样过不了验收。几个关键指标:
- 延迟:从收到请求到返回结果的时间
- 吞吐:每秒能处理多少请求
- 资源利用率: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%的性能提升花两周优化,不如加台机器来得实在。知道什么时候该优化代码,什么时候该加硬件,这是工程师和研究员的重要区别。
凌晨四点了,服务终于稳定跑起来。看着监控面板上平稳的曲线,我知道明天能给客户交差了。模型部署这条路,从来不是一蹴而就的,每个深夜的调试,都是第二天稳定运行的基石。
更多推荐
所有评论(0)