资讯详情 从零手搓AI工程:推理服务性能优化与稳定性实战
📅 2026/10/4 20:59:33
1. 从零手搓AI工程为什么我不建议你直接调包很多人一听到“AI工程”这四个字第一反应就是打开某个云平台调几个API把模型接进去跑通就完事了。我刚开始接触这块的时候也是这么想的直到有一次线上服务在高峰期直接雪崩排查了整整两天才发现问题出在一个我从来没关注过的环节——推理请求的批处理队列积压。那一刻我才意识到只会调包的人永远只能停留在“能用”的层面一旦出了问题连从哪下手都不知道。ai-engineering-from-scratch这个标题核心讲的其实就是一件事把AI工程当成一门真正的工程学科来对待从最底层的数据管道、模型加载、推理优化、服务部署一层一层自己搭起来。它适合那些已经会用现成框架跑demo但想真正理解“为什么这样设计”“出了问题怎么定位”的开发者。你不需要是算法专家但你得愿意动手写代码、看日志、调参数。我写这篇东西的出发点很简单市面上讲AI应用的文章百分之九十都在教你调API剩下百分之十在讲模型原理但中间那层——工程实现——几乎没人系统讲。而这层恰恰是决定一个AI服务能不能扛住真实流量、能不能稳定运行的关键。接下来我会按照我自己从零搭建一套AI推理服务的完整路径把每个环节的选型逻辑、踩过的坑、以及那些文档里不会写的经验全部摊开来讲。2. 环境搭建别急着装框架先把依赖关系理清楚2.1 为什么我选择从Python虚拟环境开始而不是直接全局安装我见过太多人拿到一台新机器上来就是pip install torch transformers然后跑着跑着发现版本冲突整个环境废掉。AI工程涉及的东西太多了——深度学习框架、推理加速库、服务框架、监控工具这些库之间的依赖关系极其复杂。PyTorch 2.x 和某些版本的CUDA之间有严格的对应关系transformers库的某个版本可能依赖特定版本的tokenizers而tokenizers又对Rust编译器版本有要求。我的做法是每个项目独立虚拟环境用venv或者conda都行但一定要隔离。创建环境之后第一件事不是装框架而是先确定CUDA版本。你可以用nvidia-smi查看驱动支持的CUDA最高版本然后去PyTorch官网查对应的安装命令。这一步看起来简单但我敢说至少一半的初学者在这里翻过车。# 查看显卡驱动和CUDA版本 nvidia-smi # 创建独立虚拟环境 python -m venv ai-eng-env source ai-eng-env/bin/activate # 根据CUDA版本安装对应PyTorch # 例如CUDA 12.1 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121注意不要盲目复制网上的安装命令一定要去PyTorch官网根据你的CUDA版本生成对应的命令。版本不匹配的报错信息往往非常隐晦可能表现为推理结果异常而不是直接报错。2.2 推理加速库的选择ONNX Runtime还是TensorRT模型训练和模型推理是两回事。训练的时候你关心的是收敛速度和显存占用推理的时候你关心的是延迟和吞吐。这两个目标对应的优化手段完全不同。我一开始也是直接用PyTorch做推理后来发现单次请求延迟高得离谱才开始研究推理加速。ONNX Runtime和TensorRT是两条主流路线。ONNX Runtime的优势是跨平台、支持多种硬件后端、部署简单适合快速上线。TensorRT的优势是在NVIDIA显卡上性能极致但只支持NVIDIA硬件而且模型转换过程比较折腾。我的建议是如果你的服务需要跑在多种硬件上或者团队没有专门的推理优化工程师先用ONNX Runtime。等业务量上来了再针对核心模型做TensorRT优化。# ONNX Runtime推理示例 import onnxruntime as ort import numpy as np # 加载模型 session ort.InferenceSession(model.onnx) # 查看输入输出信息 for inp in session.get_inputs(): print(f输入名: {inp.name}, 形状: {inp.shape}, 类型: {inp.type}) # 构造输入并推理 input_data np.random.randn(1, 3, 224, 224).astype(np.float32) outputs session.run(None, {input: input_data})这里有个细节ONNX模型转换的时候动态维度dynamic axes的设置非常关键。如果你把batch维度固定死了那服务就只能处理固定batch size的请求灵活性大打折扣。我一般会把batch维度和序列长度维度都设为动态虽然这样会损失一点点性能但换来的是服务层面的灵活性。2.3 服务框架选型FastAPI够用但别忽略并发模型服务框架这块FastAPI是目前最主流的选择异步支持好、自动生成文档、类型校验完善。但很多人用FastAPI写AI服务的时候犯了一个致命错误在异步函数里直接调用同步的推理代码。这会导致整个事件循环被阻塞并发能力直接归零。正确的做法有两种一是用run_in_executor把推理放到线程池里执行二是用专门的推理服务框架比如Triton Inference Server。前者适合小规模服务后者适合生产环境。我一开始用的是第一种方案后来请求量上来之后切换到了Triton性能提升非常明显。from fastapi import FastAPI import asyncio from concurrent.futures import ThreadPoolExecutor app FastAPI() executor ThreadPoolExecutor(max_workers4) def sync_inference(data): # 同步推理逻辑 return model.predict(data) app.post(/predict) async def predict(data: dict): loop asyncio.get_event_loop() result await loop.run_in_executor(executor, sync_inference, data) return {result: result}提示线程池的worker数量不是越多越好。对于GPU推理来说过多的并发线程反而会导致显存竞争和上下文切换开销。一般建议设置为GPU数量的2到4倍具体需要压测确定。3. 数据管道AI工程里最容易被低估的环节3.1 输入预处理为什么不能放在请求线程里做我见过很多AI服务的代码结构是这样的请求进来在接口函数里做图像解码、缩放、归一化然后送进模型推理。这种写法在低并发下没问题但一旦请求量上来预处理就会成为瓶颈。因为图像解码和缩放是CPU密集型操作而推理是GPU密集型操作两者串行执行意味着GPU有大量时间在等CPU。我的优化方案是把预处理做成独立的流水线阶段。具体来说请求进来之后先丢进一个队列由专门的预处理worker池做解码和缩放处理完再送进推理队列。这样CPU和GPU可以并行工作整体吞吐能提升好几倍。这个思路借鉴了经典的生产者-消费者模型在AI工程里同样适用。import queue import threading preprocess_queue queue.Queue(maxsize100) inference_queue queue.Queue(maxsize50) def preprocess_worker(): while True: raw_data preprocess_queue.get() processed decode_and_resize(raw_data) inference_queue.put(processed) def inference_worker(): while True: data inference_queue.get() result model.predict(data) # 后续处理3.2 批处理策略动态批处理到底怎么实现批处理是提升GPU利用率最有效的手段之一。原理很简单把多个请求合并成一个batch送进模型GPU一次前向传播就能处理多个样本。但实现起来有几个关键问题需要解决batch size怎么定、等待时间怎么控制、超时了怎么办。我的经验是采用动态批处理策略设置一个最大batch size和一个最大等待时间。当队列里的请求数达到最大batch size或者最早的那个请求等待时间超过阈值就立即触发一次推理。这样在低负载时延迟低高负载时吞吐高。Triton Inference Server内置了这个机制如果自己实现的话需要注意线程安全问题。参数建议值说明最大batch size8-32根据显存和模型大小调整最大等待时间10-50ms根据业务延迟要求调整队列最大长度100-500超过则拒绝请求防止雪崩推理超时时间5-30s超时则返回错误避免请求堆积3.3 数据格式的统一别让格式转换吃掉你的性能AI服务里数据格式转换的开销经常被忽略。比如图像数据在传输过程中可能是base64编码的字符串解码成bytes之后需要转成numpy数组再转成tensor最后还要做归一化。每一步转换都有开销而且很容易出错。我的做法是在服务入口处就把数据统一成标准格式后续所有环节都基于这个标准格式操作。对于图像服务标准格式就是numpy数组形状为(H, W, C)数据类型为uint8。预处理阶段再做归一化和维度变换。这样每个环节的输入输出都很明确调试起来也方便。import base64 import numpy as np from PIL import Image import io def decode_image(base64_str: str) - np.ndarray: 将base64图像解码为标准numpy数组 img_bytes base64.b64decode(base64_str) img Image.open(io.BytesIO(img_bytes)) img img.convert(RGB) return np.array(img) # 形状 (H, W, 3), dtype uint8 def preprocess(img: np.ndarray, target_size: tuple) - np.ndarray: 缩放、归一化、维度变换 img Image.fromarray(img).resize(target_size) arr np.array(img).astype(np.float32) / 255.0 mean np.array([0.485, 0.456, 0.406]) std np.array([0.229, 0.224, 0.225]) arr (arr - mean) / std arr arr.transpose(2, 0, 1) # HWC - CHW return np.expand_dims(arr, axis0) # 增加batch维度4. 模型服务化从单机脚本到可扩展服务4.1 模型加载的时机与内存管理模型加载看起来简单但里面有不少门道。首先模型应该在服务启动时加载一次而不是每次请求都加载。这个道理大家都懂但实际写代码的时候很多人会把模型加载放在接口函数里导致每次请求都要重新加载模型性能极差。其次模型加载之后要确认它真的在GPU上。我遇到过好几次模型加载成功但实际还在CPU上的情况推理速度慢了几十倍。排查方法很简单加载完之后打印模型参数的device属性。另外如果服务需要加载多个模型要注意显存分配避免OOM。import torch class ModelService: def __init__(self, model_path: str, device: str cuda): self.device torch.device(device if torch.cuda.is_available() else cpu) self.model torch.load(model_path, map_locationself.device) self.model.eval() self.model.to(self.device) # 确认模型在正确的设备上 param_device next(self.model.parameters()).device print(f模型加载完成设备: {param_device}) # 预热 self._warmup() def _warmup(self): 用假数据预热触发CUDA内核编译 dummy torch.randn(1, 3, 224, 224).to(self.device) with torch.no_grad(): for _ in range(3): self.model(dummy) print(预热完成)注意模型预热非常重要。第一次推理时CUDA需要编译内核耗时可能是后续推理的几十倍。如果不预热第一个真实请求的延迟会非常高用户体验很差。4.2 健康检查与优雅关闭生产环境的服务必须要有健康检查接口。这个接口不能只返回一个固定的200状态码而应该真正检查模型是否可用、GPU是否正常、队列是否积压。我一般会实现一个/health接口检查项包括模型是否加载、GPU显存是否正常、最近一次推理是否成功。优雅关闭同样重要。服务收到终止信号时应该先停止接受新请求等待正在处理的请求完成然后释放GPU资源再退出。如果不做优雅关闭正在处理的请求会直接失败用户会看到错误。import signal import sys from fastapi import FastAPI app FastAPI() shutting_down False app.get(/health) async def health(): if shutting_down: return {status: shutting_down}, 503 # 检查模型状态 if not model_service.is_ready(): return {status: unhealthy}, 503 return {status: healthy} def handle_signal(signum, frame): global shutting_down shutting_down True # 等待正在处理的请求完成 model_service.drain() sys.exit(0) signal.signal(signal.SIGTERM, handle_signal) signal.signal(signal.SIGINT, handle_signal)4.3 日志与监控出了问题怎么快速定位AI服务的日志和普通后端服务不太一样需要额外记录推理相关的信息输入数据的形状和类型、推理耗时、GPU显存占用、batch size等。这些信息在排查性能问题时非常关键。我的做法是用结构化日志每条推理记录都包含请求ID、预处理耗时、推理耗时、后处理耗时、batch size、GPU显存。这样当延迟升高时我可以快速判断是哪个环节出了问题。监控方面除了常规的QPS和延迟还要关注GPU利用率和显存使用率这两个指标能提前预警容量问题。import time import logging import json logger logging.getLogger(inference) def log_inference(request_id, preprocess_time, inference_time, postprocess_time, batch_size): log_entry { request_id: request_id, preprocess_ms: round(preprocess_time * 1000, 2), inference_ms: round(inference_time * 1000, 2), postprocess_ms: round(postprocess_time * 1000, 2), batch_size: batch_size, total_ms: round((preprocess_time inference_time postprocess_time) * 1000, 2) } logger.info(json.dumps(log_entry))5. 性能调优那些压测报告不会告诉你的细节5.1 延迟与吞吐的权衡没有银弹AI服务调优的核心矛盾就是延迟和吞吐的权衡。增大batch size能提升吞吐但会增加单个请求的等待时间。降低batch size能减少延迟但GPU利用率上不去。这个平衡点取决于你的业务场景如果是实时交互场景延迟优先如果是离线批处理场景吞吐优先。我的一般做法是先确定延迟上限然后在这个约束下最大化吞吐。比如业务要求P99延迟不超过200ms那我就从batch size1开始压测逐步增大batch size直到P99延迟接近200ms为止。这个过程中要同时观察GPU利用率如果GPU利用率还很低但延迟已经到上限了说明瓶颈不在GPU可能在预处理或者网络传输。场景类型延迟要求推荐batch策略硬件建议实时交互P99 100ms小batch短等待高端GPU多实例准实时P99 500ms中等batch中端GPU离线批处理无硬性要求大batch任意GPU可CPU5.2 显存优化模型量化与内存复用显存是AI服务最宝贵的资源。除了增大batch size还有很多手段可以优化显存使用。模型量化是最直接的方法FP32转FP16能省一半显存INT8量化能省四分之三。但量化会带来精度损失需要评估对业务的影响。另一个容易被忽略的点是内存复用。PyTorch的CUDA缓存分配器会缓存已分配的内存如果频繁创建和销毁tensor缓存会碎片化。我的做法是尽量复用tensor比如预分配输入输出缓冲区避免在推理循环里反复分配。# FP16量化示例 model.half() # 转为FP16 input_tensor input_tensor.half() # 预分配缓冲区 input_buffer torch.zeros(1, 3, 224, 224, dtypetorch.float16, devicecuda) output_buffer torch.zeros(1, 1000, dtypetorch.float16, devicecuda) def inference_with_buffer(data): input_buffer.copy_(data) with torch.no_grad(): output_buffer.copy_(model(input_buffer)) return output_buffer5.3 多实例部署什么时候该加卡什么时候该加实例当单实例性能到瓶颈时下一步就是多实例部署。但加卡和加实例是两回事。加卡指的是在同一台机器上跑多个GPU实例加实例指的是在多台机器上各跑一个实例。前者通信开销小但受单机资源限制后者扩展性好但需要负载均衡。我的经验是先尝试单机多卡用CUDA_VISIBLE_DEVICES环境变量隔离不同实例。如果单机资源不够了再考虑多机部署。多机部署时要注意模型文件的同步和版本一致性我一般会用共享存储或者对象存储来管理模型文件。# 单机多卡部署示例 # 实例1使用GPU 0 CUDA_VISIBLE_DEVICES0 python serve.py --port 8001 # 实例2使用GPU 1 CUDA_VISIBLE_DEVICES1 python serve.py --port 8002 # Nginx负载均衡配置 upstream inference_backend { server 127.0.0.1:8001; server 127.0.0.1:8002; }提示多实例部署时每个实例的显存占用要留有余量。如果每个实例都把显存用满当某个实例崩溃重启时可能会因为显存不足而启动失败。6. 踩坑实录那些让我熬夜排查的典型问题6.1 推理结果不一致浮点精度还是并发问题有一次线上服务出现了一个诡异的现象同一个输入有时候返回结果A有时候返回结果B两者差异很小但不为零。我一开始怀疑是浮点精度问题但同样的输入在单机测试时结果完全一致。排查了很久才发现问题出在批处理上当请求被分到不同的batch里时由于batch内其他样本的影响浮点运算的顺序不同导致了微小的数值差异。这个问题在大多数场景下可以忽略但如果你的业务对结果一致性有严格要求就需要特别注意。解决方案是固定batch的组成方式或者使用确定性算法。PyTorch提供了torch.use_deterministic_algorithms(True)来强制确定性但会牺牲一些性能。6.2 内存泄漏不是Python的锅服务跑了一段时间之后内存持续增长最终OOM。第一反应是Python内存泄漏但用工具排查后发现Python对象数量正常。真正的原因是CUDA缓存没有释放。PyTorch的CUDA缓存分配器会缓存已释放的显存如果推理过程中频繁创建不同形状的tensor缓存会越来越大。解决方法有两个一是尽量使用固定形状的tensor二是定期调用torch.cuda.empty_cache()。但要注意empty_cache()会强制同步频繁调用会影响性能。我一般是在服务空闲时或者每处理N个请求后调用一次。import torch import gc class MemoryManager: def __init__(self, cleanup_interval1000): self.counter 0 self.cleanup_interval cleanup_interval def maybe_cleanup(self): self.counter 1 if self.counter self.cleanup_interval: gc.collect() torch.cuda.empty_cache() self.counter 0 # 记录清理后的显存 allocated torch.cuda.memory_allocated() / 1024**2 cached torch.cuda.memory_reserved() / 1024**2 print(f显存清理完成 - 已分配: {allocated:.1f}MB, 缓存: {cached:.1f}MB)6.3 队列积压限流比扩容更重要高峰期请求量突增队列迅速积压延迟飙升最终服务不可用。我一开始的思路是扩容加机器加实例。但后来发现扩容的速度永远赶不上流量突增的速度而且扩容之后流量回落资源又浪费了。更有效的做法是限流。在服务入口处设置一个令牌桶或者漏桶超过容量的请求直接拒绝返回429状态码。这样虽然会损失一部分请求但保证了已接受请求的服务质量。配合自动扩缩容可以在流量持续高位时逐步增加容量。import time from threading import Lock class RateLimiter: def __init__(self, rate: float, capacity: int): self.rate rate # 每秒允许的请求数 self.capacity capacity # 桶容量 self.tokens capacity self.last_time time.time() self.lock Lock() def allow(self) - bool: with self.lock: now time.time() elapsed now - self.last_time self.tokens min(self.capacity, self.tokens elapsed * self.rate) self.last_time now if self.tokens 1: self.tokens - 1 return True return False6.4 模型更新热更新还是滚动更新模型迭代是常态但更新模型时如何不影响线上服务是个问题。热更新指的是在不重启服务的情况下替换模型滚动更新指的是逐个替换实例。热更新的优点是快缺点是如果新模型有问题影响面是全部实例。滚动更新更安全但需要额外的编排工具。我的做法是小版本更新用滚动更新大版本更新用蓝绿部署。蓝绿部署就是同时运行两套环境流量先切一小部分到新环境观察一段时间没问题再全量切换。这样即使新模型有问题也能快速回滚。# 模型版本管理示例 class ModelManager: def __init__(self): self.models {} self.active_version None def load_model(self, version: str, path: str): model torch.load(path) model.eval() self.models[version] model def switch_version(self, version: str): if version not in self.models: raise ValueError(f版本 {version} 未加载) self.active_version version print(f已切换到模型版本: {version}) def predict(self, data): model self.models[self.active_version] with torch.no_grad(): return model(data)7. 从能跑到好用一些让服务更稳的工程习惯7.1 配置管理别把参数写死在代码里我见过太多项目把模型路径、batch size、超时时间这些参数直接写在代码里改一个参数就要重新部署。正确的做法是用配置文件或者环境变量管理这些参数。我一般用YAML配置文件配合环境变量覆盖这样不同环境可以用同一份代码不同的配置。# config.yaml model: path: /models/bert-base device: cuda max_batch_size: 16 warmup_iterations: 3 server: host: 0.0.0.0 port: 8000 workers: 4 timeout: 30 inference: max_queue_size: 200 max_wait_ms: 20 rate_limit: 100import yaml import os def load_config(path: str config.yaml) - dict: with open(path, r) as f: config yaml.safe_load(f) # 环境变量覆盖 if os.getenv(MODEL_PATH): config[model][path] os.getenv(MODEL_PATH) if os.getenv(PORT): config[server][port] int(os.getenv(PORT)) return config7.2 单元测试与集成测试AI服务也需要测试很多人觉得AI服务没法做单元测试因为输出是不确定的。但其实可以测的东西很多预处理函数的输入输出形状和范围、后处理函数的逻辑、批处理队列的行为、限流器的正确性。这些都不依赖模型的具体输出。集成测试则是用真实模型跑一遍完整流程验证端到端的正确性。我一般会准备一组固定的测试输入和期望输出允许一定的误差范围每次代码变更后自动运行。这样能及时发现回归问题。import pytest import numpy as np def test_preprocess_output_shape(): img np.random.randint(0, 255, (480, 640, 3), dtypenp.uint8) result preprocess(img, target_size(224, 224)) assert result.shape (1, 3, 224, 224) assert result.dtype np.float32 def test_preprocess_normalization(): img np.full((224, 224, 3), 128, dtypenp.uint8) result preprocess(img, target_size(224, 224)) # 归一化后均值应该在0附近 assert abs(result.mean()) 0.5 def test_rate_limiter(): limiter RateLimiter(rate10, capacity5) # 前5个请求应该通过 for _ in range(5): assert limiter.allow() is True # 第6个应该被拒绝 assert limiter.allow() is False7.3 文档与交接让下一个人能看懂你的服务AI工程项目的文档特别重要因为涉及的东西多且杂。我一般会维护三份文档架构文档说明整体设计和服务依赖关系运维文档说明如何部署、监控、排查问题API文档说明接口的输入输出格式和错误码。架构文档里我会画一张服务拓扑图用文字描述也行标明每个组件的职责和数据流向。运维文档里会记录常见的故障现象和对应的排查步骤比如“延迟升高先看GPU利用率再看队列长度再看预处理耗时”。这些文档在交接和应急时能救命。8. 写在最后一些个人体会从零搭建AI工程这件事最大的收获不是学会了某个框架或者工具而是建立了一套排查问题的思维框架。当服务出问题时我知道从哪一层开始查每一层有哪些可能的故障点每个故障点对应什么现象。这种能力是调包调不出来的。另外一点体会是AI工程和传统后端工程没有本质区别都是围绕稳定性、性能、可维护性做文章。区别在于AI服务多了GPU这个变量多了模型这个不确定因素。把GPU当成一种特殊的资源来管理把模型当成一种特殊的依赖来对待很多问题就迎刃而解了。如果你也在从零搭建AI服务我的建议是先跑通最小闭环然后逐步优化。不要一开始就追求完美架构而是在实践中发现问题、解决问题。每一次踩坑都是对系统理解加深的机会。我到现在还在不断调整自己的服务架构因为业务在变、模型在变、硬件也在变没有一劳永逸的方案。