1. 从零手搓AI工程:为什么我不建议你直接调包
很多人一听到“AI工程”这四个字,第一反应就是打开某个云平台,拖几个组件,调一下API,然后跑通一个Demo,就觉得自己已经掌握了。我刚开始也是这么想的,直到有一次线上推理服务在高峰期直接雪崩,日志里全是显存溢出和请求超时,我才意识到——那些被封装好的高级接口,在关键时刻根本救不了你。ai-engineering-from-scratch这个标题,说的就是从最底层开始,把AI工程里那些被隐藏起来的环节,一个一个亲手搭起来。
这篇文章适合谁看?如果你已经会用Python写点脚本,知道什么是张量,但每次遇到模型部署、显存优化、推理加速这些词就心里发虚,那这篇内容就是为你准备的。我不会只给你一堆代码让你复制粘贴,而是把每个设计决策背后的“为什么”讲清楚。比如为什么推理时要固定batch size,为什么量化不是万能药,为什么你的GPU利用率永远上不去。这些问题的答案,在官方文档里往往一笔带过,但在实际工程中,每一个都能让你熬夜到凌晨三点。
我打算从最基础的推理服务搭建开始,一步步走到性能调优和线上排障。整个过程不依赖任何重型框架,核心逻辑用Python和少量C++扩展实现,目的是让你看清AI工程的全貌。你不需要有分布式系统的经验,但最好对Linux命令和Python多进程有点概念。如果你准备好了,我们就从第一个坑开始。
2. 推理服务的最小可行骨架:从HTTP请求到张量计算
2.1 为什么不用现成的模型服务框架
市面上有很多模型服务框架,开箱即用,一行命令就能启动一个推理端点。但我坚持从零手写一个最小服务,原因有三个。第一,那些框架为了通用性,抽象层数太多,一个请求进来要经过路由、预处理、批处理调度、模型执行、后处理五六个模块,每个模块都有配置项,一旦出问题,排查链路极长。第二,框架的默认参数往往针对通用场景,比如动态批处理窗口默认10毫秒,在高并发低延迟场景下这个值就是灾难。第三,也是最关键的,当你亲手写过一遍请求解析、内存分配、线程调度之后,再看那些框架的源码,你会有一种“原来如此”的顿悟感。
我选用的技术栈很简单:Python标准库的http.server做HTTP层,numpy做数据搬运,模型本身用ONNX Runtime或者自己导出的TorchScript。为什么不直接上FastAPI?因为FastAPI的异步模型和GPU推理的同步阻塞特性之间存在微妙的冲突,新手很容易写出看似异步实则串行的代码。用最原始的http.server虽然性能差,但能让你清楚看到每个请求的生命周期。
2.2 请求解析与张量构造的隐藏成本
一个典型的推理请求进来,JSON里包含输入数据,比如一张图片的base64编码或者一个文本序列的token ID列表。很多人直接json.loads然后np.array就完事了,但这里有两个隐藏成本。第一,base64解码是CPU密集操作,一张1080p的JPEG图片解码成RGB数组大约需要15到30毫秒,如果并发上来,CPU会先于GPU成为瓶颈。第二,np.array从Python列表构造时,如果数据类型不匹配,会发生隐式类型转换,比如你的模型期望float32,但JSON里的数字被解析成float64,这个转换在数据量大时能吃掉几毫秒。
我的做法是在服务启动时就预分配好输入缓冲区。对于固定尺寸的输入,直接创建一个np.empty数组,请求到来时用np.frombuffer或者切片赋值填充,避免反复分配内存。对于变长输入,维护一个内存池,按2的幂次方分级管理。这些技巧在常规教程里很少提,但它们是推理服务稳定性的基石。
import numpy as np import base64 from io import BytesIO from PIL import Image class InputBuffer: def __init__(self, shape, dtype=np.float32): self.buffer = np.empty(shape, dtype=dtype) self.shape = shape def fill_from_base64(self, b64_str): img_bytes = base64.b64decode(b64_str) img = Image.open(BytesIO(img_bytes)).convert('RGB') img = img.resize((self.shape[2], self.shape[3])) arr = np.asarray(img, dtype=np.float32) / 255.0 arr = np.transpose(arr, (2, 0, 1)) np.copyto(self.buffer[0], arr) return self.buffer上面这段代码里,np.copyto比直接赋值更安全,因为它会检查形状和类型是否匹配。Image.resize默认使用双线性插值,如果你对精度有要求,可以换成Image.BICUBIC,但速度会慢一倍。这些细节在批量推理时影响巨大。
2.3 模型加载与执行提供者的选择
模型加载阶段,我强烈建议使用ONNX Runtime而不是直接加载PyTorch模型。原因在于ONNX Runtime的执行提供者(Execution Provider)机制允许你灵活切换后端。比如在NVIDIA GPU上,你可以用CUDA EP,在AMD GPU上用ROCM EP,在CPU上用OpenVINO EP。更重要的是,ONNX Runtime内置了图优化,比如算子融合、常量折叠、内存复用,这些优化在PyTorch eager模式下是没有的。
但这里有个坑:ONNX Runtime的默认优化级别是ORT_ENABLE_ALL,它会尝试把所有能融合的算子都融合掉。在某些模型上,过度融合会导致精度下降,尤其是涉及LayerNorm和Softmax的模型。我的经验是,先用ORT_ENABLE_BASIC跑一遍,对比输出差异,如果差异在可接受范围内(比如1e-5),再逐步提高优化级别。另外,intra_op_num_threads这个参数控制算子内部的并行线程数,默认是CPU核心数,但在GPU推理场景下,这个值设得太高反而会因为线程切换开销导致性能下降。我通常设为物理核心数的一半。
import onnxruntime as ort options = ort.SessionOptions() options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC options.intra_op_num_threads = 4 options.inter_op_num_threads = 1 session = ort.InferenceSession( "model.onnx", sess_options=options, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] )注意providers列表的顺序,ONNX Runtime会按顺序尝试,如果CUDA不可用就回退到CPU。这个回退机制在生产环境中非常有用,但你要确保CPU路径也能正常工作,否则回退后直接报错。
3. 批处理与调度:让GPU不再空转的实战策略
3.1 动态批处理的触发条件与超时权衡
GPU最怕的就是空转。一个请求进来,做一次前向传播,然后等下一个请求,GPU利用率可能只有5%。动态批处理的核心思想是:攒几个请求一起算。但攒多久?这就是超时窗口的权衡。窗口设得太短,比如1毫秒,那和单请求没区别;设得太长,比如50毫秒,延迟就上去了,用户体验变差。
我的经验值是这样的:对于在线服务,超时窗口设为5到10毫秒;对于离线批处理,可以设为100毫秒甚至更长。但这不是固定的,你要根据请求到达率动态调整。如果每秒请求数(QPS)很高,比如1000,那5毫秒内能攒到5个请求,批大小为5,GPU利用率就能到60%以上。如果QPS只有10,那5毫秒内可能一个请求都没有,这时候超时窗口应该适当延长,比如20毫秒,保证至少能凑到2个请求。
实现上,我用一个后台线程维护请求队列,主线程收到请求后放入队列并等待条件变量。后台线程每隔一个超时周期检查队列,如果非空就取出所有请求组成一个批次,调用模型推理,然后把结果分发给对应的等待者。这里的关键是条件变量的通知机制,不能用time.sleep轮询,那样CPU占用率会很高。
import threading import time from collections import deque class BatchScheduler: def __init__(self, max_batch_size=8, timeout_ms=10): self.queue = deque() self.max_batch_size = max_batch_size self.timeout = timeout_ms / 1000.0 self.lock = threading.Lock() self.cond = threading.Condition(self.lock) self.results = {} self.running = True self.worker = threading.Thread(target=self._loop, daemon=True) self.worker.start() def submit(self, request_id, input_data): with self.cond: self.queue.append((request_id, input_data)) self.cond.notify() while request_id not in self.results: self.cond.wait() return self.results.pop(request_id) def _loop(self): while self.running: with self.cond: if not self.queue: self.cond.wait(timeout=self.timeout) if not self.queue: continue batch = [] while self.queue and len(batch) < self.max_batch_size: batch.append(self.queue.popleft()) # 在锁外执行推理,避免阻塞提交 inputs = [item[1] for item in batch] outputs = self._infer(inputs) with self.cond: for (req_id, _), out in zip(batch, outputs): self.results[req_id] = out self.cond.notify_all()这段代码里,_infer是实际调用模型的方法,它应该在锁外执行,否则提交请求的线程会被阻塞。另外,self.results字典在请求量极大时会成为内存瓶颈,生产环境应该用更高效的数据结构,比如concurrent.futures.Future。
3.2 批大小对显存和延迟的非线性影响
很多人以为批大小翻倍,显存占用也翻倍,延迟也翻倍。实际上不是这样的。显存占用大致是线性的,因为中间激活值随批大小线性增长。但延迟不是线性的,因为GPU的并行计算单元在批大小较小时利用率不足,批大小增大到一定程度后,计算时间增长缓慢。我实测过一个ResNet-50模型,批大小1时延迟8毫秒,批大小8时延迟12毫秒,批大小16时延迟18毫秒。也就是说,批大小从1到8,吞吐量提升了5倍多,但延迟只增加了50%。
但这里有个临界点。当批大小超过某个值后,显存不够了,或者计算单元饱和了,延迟会急剧上升。这个临界点取决于模型大小和GPU型号。我的做法是写一个简单的压测脚本,从批大小1开始,每次翻倍,记录延迟和显存占用,画出曲线,找到拐点。通常拐点在GPU显存的70%到80%利用率处。
| 批大小 | 延迟(ms) | 显存占用(MB) | 吞吐量(请求/秒) |
|---|---|---|---|
| 1 | 8 | 1200 | 125 |
| 2 | 9 | 1350 | 222 |
| 4 | 10 | 1650 | 400 |
| 8 | 12 | 2250 | 666 |
| 16 | 18 | 3450 | 888 |
| 32 | 35 | 5850 | 914 |
从表里可以看出,批大小16到32,吞吐量几乎没提升,但延迟翻倍。所以最优批大小在16左右。这个数据因模型和硬件而异,但方法论是通用的。
3.3 请求优先级与超时丢弃机制
生产环境里,请求不是平等的。有些是实时交互请求,用户等着看结果,延迟要求高;有些是后台分析请求,可以等几秒钟。如果一视同仁,实时请求会被后台请求拖累。我的做法是在请求头里加一个优先级字段,调度器维护两个队列:高优先级和低优先级。每次组批时,先从高优先级队列取,取完了再用低优先级填充剩余位置。
超时丢弃也很重要。如果一个请求在队列里等了超过阈值,比如200毫秒,还没被处理,就应该直接返回超时错误,而不是继续等。因为用户可能已经重试了,你再算出来也没用,反而浪费GPU资源。实现上,每个请求入队时记录时间戳,调度器组批前检查时间戳,过期的直接丢弃并通知等待者。
注意:超时丢弃的阈值要大于批处理超时窗口,否则请求还没等到组批就被丢了。一般设为批处理窗口的5到10倍。
4. 显存管理与量化:在有限资源下榨取性能
4.1 显存碎片化与内存池设计
GPU显存和CPU内存一样,反复分配释放会产生碎片。尤其是变长输入场景,每次请求的输入尺寸不同,如果每次都cudaMalloc和cudaFree,很快显存就碎片化了,最后明明有足够的总空闲显存,却分配不出一块连续的大内存。解决方案是内存池:启动时一次性分配一大块显存,然后自己管理分配和释放。
我的内存池实现很简单:按2的幂次方分级,比如1KB、2KB、4KB直到256MB。每个级别维护一个空闲链表。分配时向上取整到最近的级别,从对应链表取一块;释放时归还到对应链表。如果某个级别空了,从更高级别切分一块下来。这个策略在请求尺寸分布比较集中的场景下非常高效,碎片率可以控制在5%以内。
但内存池有个缺点:启动时就要确定总大小。设得太小,高峰期不够用;设得太大,浪费显存。我的经验是,先跑一轮压测,记录峰值显存占用,然后内存池大小设为峰值的1.2倍。另外,要留出至少20%的显存给模型权重和CUDA上下文,否则会OOM。
4.2 量化不是银弹:INT8精度的实际损失评估
量化是显存优化的利器,FP32转INT8,显存直接降到四分之一,推理速度也能提升2到3倍。但量化会损失精度,而且损失程度因模型而异。我见过一个文本分类模型,INT8量化后准确率只掉了0.1%,完全可用;也见过一个目标检测模型,量化后小目标的召回率掉了15%,直接不可用。
所以量化之前,一定要做精度评估。步骤是:准备一个验证集,至少1000个样本;用FP32模型跑一遍,记录每个样本的输出;用INT8模型跑一遍,对比输出差异。对于分类模型,看Top-1准确率变化;对于检测模型,看mAP变化;对于生成模型,看BLEU或者ROUGE变化。如果掉点超过可接受范围,就要考虑混合量化:只量化对精度不敏感的层,比如卷积层,保留全连接层为FP32。
ONNX Runtime提供了动态量化和静态量化两种模式。动态量化不需要校准数据,直接对权重做量化,激活值在推理时动态计算量化参数。静态量化需要校准数据,提前计算好激活值的量化参数,精度通常更好。我的建议是,如果手头有代表性数据,优先用静态量化;如果没有,动态量化也能用,但要做好精度下降的心理准备。
from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( model_input="model_fp32.onnx", model_output="model_int8.onnx", weight_type=QuantType.QUInt8, optimize_model=True )QuantType.QUInt8是无符号8位整数量化,适合ReLU激活的模型。如果是LeakyReLU或者GELU,用QInt8有符号量化更合适。optimize_model=True会在量化前先做图优化,通常能提升量化后的精度。
4.3 显存不足时的降级策略
即使做了量化和内存池,高峰期仍然可能显存不足。这时候不能直接崩掉,要有降级策略。我的降级顺序是这样的:第一级,降低批大小,从16降到8,显存占用减半,吞吐量下降但服务不中断;第二级,切换到CPU推理,延迟大幅上升但至少能返回结果;第三级,返回503错误,让上游重试或降级。
实现上,用一个显存监控线程,每隔100毫秒查询一次torch.cuda.memory_allocated()或者ONNX Runtime的显存统计。如果超过阈值,触发降级信号。降级信号通过一个全局状态变量传递给调度器,调度器动态调整批大小上限。这个机制在Kubernetes环境下尤其重要,因为容器有显存限制,超了直接被杀。
提示:降级策略要提前演练,不能等线上出问题了才第一次跑。我一般会在预发环境用压力测试工具模拟显存打满,验证降级链路是否通畅。
5. 线上排障实录:一次推理延迟毛刺的完整排查链路
5.1 问题现象与初步定位
那是一个周五晚上,监控系统突然报警:推理服务P99延迟从15毫秒飙升到200毫秒,但QPS没有明显变化。更奇怪的是,P50延迟还是正常的15毫秒,只有P99异常。这说明大部分请求正常,少数请求特别慢。我第一反应是某个请求的输入尺寸特别大,导致计算时间过长。但查了日志,输入尺寸分布和平时一样。
接着我怀疑是GPU降频。用nvidia-smi -q -d PERFORMANCE查看,发现GPU时钟确实从1.8GHz降到了1.2GHz。但为什么降频?温度正常,功耗也正常。继续查,发现nvidia-smi显示有一个进程占用了大量显存,但不是我们的推理进程。原来是一台共享GPU的机器上,另一个团队跑了一个训练任务,把显存吃掉了大半,导致我们的推理进程频繁触发显存回收,进而引发延迟毛刺。
5.2 排查过程中的三个关键命令
第一个命令是nvidia-smi --query-compute-apps=pid,used_memory --format=csv,列出所有占用GPU的进程和显存。这个命令比直接nvidia-smi更清晰,因为它只显示计算进程,不显示图形进程。
第二个命令是py-spy dump --pid <推理进程PID>,这是Python的采样分析工具,能直接打印出当前所有线程的调用栈。我用它发现推理主线程卡在了cudaMemcpy上,说明数据在CPU和GPU之间搬运时被阻塞了。
第三个命令是nsys profile --stats=true -o profile_result python infer.py,这是NVIDIA的系统级性能分析工具,能生成时间线,显示每个CUDA核函数的执行时间和内存拷贝时间。从时间线上看,正常情况下cudaMemcpy只占5%的时间,出问题时占到了60%。
5.3 根因分析与修复方案
根因很明确:共享GPU环境下,另一个进程的显存占用导致我们的进程在分配显存时触发了同步等待。CUDA的显存分配是同步操作,当显存不足时,驱动会尝试回收其他进程的缓存,这个回收过程会阻塞当前进程。修复方案有三个:第一,申请独占GPU,不让其他进程共享;第二,设置CUDA_VISIBLE_DEVICES隔离;第三,在代码里预分配所有需要的显存,避免运行时分配。
我选了第三个方案,因为最可控。具体做法是在服务启动时,用torch.cuda.memory_reserved()或者ONNX Runtime的arena_extend_strategy参数,一次性预留足够显存。ONNX Runtime的arena_extend_strategy设为kSameAsRequested,表示按需扩展,但扩展后不释放,这样后续请求就不会触发新的分配。
options = ort.SessionOptions() options.enable_cpu_mem_arena = False options.add_session_config_entry("session.use_env_allocators", "1") options.add_session_config_entry("session.arena_extend_strategy", "kSameAsRequested")修复后,P99延迟回落到18毫秒,虽然比P50的15毫秒略高,但已经可以接受。这次排查让我深刻体会到,AI工程不只是模型和算法,更是对系统资源的精细管理。
5.4 从这次故障中提炼的检查清单
后来我把这次排查经验整理成了一个检查清单,每次上线新服务前过一遍:
- GPU是否独占?用
nvidia-smi确认没有其他计算进程。 - 显存是否预分配?启动后观察
nvidia-smi的显存占用是否稳定。 - 是否有降级策略?模拟显存打满,验证降级链路。
- 监控是否覆盖P99?只看平均值会漏掉毛刺。
- 日志是否记录输入尺寸?方便定位大请求。
- 是否有性能基线?每次变更后对比基线,发现退化。
这个清单帮我避免了好几次潜在故障。比如有一次,新来的同事在服务里加了一个日志打印,把每个请求的完整输入都打出来了,导致磁盘IO飙升,推理延迟跟着上涨。用清单里的“性能基线”一对比,立刻发现了问题。
6. 从手搓到生产:还需要补上的几块拼图
6.1 健康检查与优雅退出
手搓的服务往往忽略健康检查和优雅退出。健康检查不是简单的返回200,而是要检查模型是否加载成功、GPU是否可用、显存是否充足。我的做法是暴露一个/health端点,里面依次检查:模型session是否非空、cudaGetDeviceCount是否大于0、显存空闲是否大于阈值。只有全部通过才返回200,否则返回503。
优雅退出是指收到SIGTERM信号后,不再接受新请求,等待正在处理的请求完成,然后释放资源再退出。这个在Kubernetes滚动更新时特别重要,否则正在处理的请求会被直接杀掉,用户看到502错误。实现上,用一个全局的shutting_down标志,主循环检查这个标志,如果为真就关闭监听套接字,然后等待所有工作线程结束。
6.2 日志与指标暴露
日志要结构化,用JSON格式,包含时间戳、请求ID、输入尺寸、批大小、推理耗时、显存占用。这些字段在排查问题时缺一不可。指标暴露用Prometheus格式,暴露inference_latency_seconds直方图、batch_size直方图、gpu_memory_used_bytes仪表盘。有了这些指标,你就能在Grafana上画出延迟分布和显存趋势,提前发现异常。
我见过很多团队只打日志不暴露指标,结果每次排查都要grep日志,效率极低。指标是聚合的,能一眼看出趋势;日志是离散的,用于定位具体请求。两者互补,缺一不可。
6.3 版本管理与回滚
模型文件要版本化,每次更新模型都要记录版本号、训练数据哈希、评估指标。服务启动时加载指定版本的模型,并在日志里打印版本信息。如果新模型上线后指标恶化,要能一键回滚到旧版本。我的做法是把模型文件放在对象存储里,用版本号做路径,服务启动时从配置中心读取当前版本号,然后下载对应模型。回滚就是改配置中心的版本号,重启服务。
这个机制看起来简单,但很多团队直到出了事故才想起来做。我经历过一次模型更新后准确率暴跌,因为没有版本管理,花了两个小时才找到旧模型文件。从那以后,版本管理成了我的必选项。
6.4 压测与容量规划
上线前必须压测。压测不是简单地用ab或者wrk打流量,而是要模拟真实请求分布:输入尺寸有长有短,请求到达有突发有平稳。我用Locust写压测脚本,定义多个任务,每个任务有不同的输入尺寸和权重。压测时逐步增加并发用户数,观察延迟和显存的变化,找到服务能承受的最大QPS。
容量规划就是根据压测结果,计算需要多少台机器。比如单机最大QPS是100,预计峰值QPS是800,那至少需要8台机器,再留20%余量,就是10台。这个计算要定期回顾,因为模型更新后性能会变化。
注意:压测环境要和生产环境硬件一致,否则数据没有参考价值。我见过在CPU机器上压测,然后上线到GPU机器,结果完全对不上。
7. 一些让我少走弯路的个人习惯
我刚开始做AI工程的时候,总想一步到位,把服务写得尽善尽美。结果往往是过度设计,代码复杂到自己也看不懂。后来我养成了一个习惯:先写一个最笨的版本,能跑通就行,然后在此基础上逐步优化。比如批处理,第一版就是批大小1,跑通了再加动态批处理;量化,第一版就是FP32,跑通了再试INT8。这样每一步都有基线,出了问题也知道是哪个改动引入的。
另一个习惯是记录“失败日志”。每次遇到坑,解决之后花五分钟写下来:现象是什么、排查过程、根因、修复方案。这些记录后来成了我自己的知识库,下次遇到类似问题,直接搜关键词就能找到答案。我建议你也这么做,不用写得多正式,几句话就行,关键是坚持。
还有一个反直觉的经验:不要过早优化。我见过一个团队,服务还没上线就开始搞模型蒸馏、算子融合、多流并行,结果上线后发现QPS只有个位数,根本用不上这些优化。先让服务跑起来,有了真实流量,再根据瓶颈优化。瓶颈可能在CPU预处理,可能在网络IO,也可能在GPU计算,不跑起来你永远不知道。
最后,保持对底层的好奇心。当你用nvidia-smi看到GPU利用率只有30%的时候,不要满足于“能用就行”,去查查为什么只有30%,是批大小不够,还是数据加载拖了后腿,还是核函数本身效率低。每一次深挖,都会让你对AI工程的理解更深一层。这个领域变化很快,但底层原理变化很慢,把底层吃透了,上层的新框架新工具,你都能快速上手。