从零手搓AI工程链路：推理服务、批处理调度与性能优化实战

1. 为什么我要从零手搓一套AI工程链路第一次看到ai-engineering-from-scratch这个项目名的时候我脑子里蹦出来的不是“又一个教程仓库”而是过去两年带团队踩过的那些坑。市面上讲AI工程的内容要么停留在调API的层面要么直接甩给你一个封装好的框架你照着跑能出结果但一旦线上出问题连从哪查都不知道。这个项目标题里的from-scratch才是真正的价值点——它逼着你把每一层都掀开看一遍。所谓AI工程说白了就是把一个训练好的模型变成能稳定对外提供服务的东西。这中间涉及数据管道、特征处理、模型推理、服务编排、监控告警、版本管理这一整条链路。很多人以为模型训完就完事了实际上训完只是开始后面80%的工作量都在工程化上。ai-engineering-from-scratch这个项目要解决的就是让你在没有现成大框架兜底的情况下用最基础的组件把这条链路搭起来从而真正理解每个环节在干什么。这套内容适合谁如果你已经会用 PyTorch 或 TensorFlow 跑通一个demo但不知道模型怎么上线、怎么扛并发、怎么处理脏数据那这个项目就是为你准备的。如果你是从后端转过来做AI的工程师想搞清楚推理服务和普通Web服务的区别同样适用。甚至你是做算法的想补上工程这块短板从零搭一遍也是最快的路径。我个人的判断是凡是涉及“从零”的项目最大的收益不是最后跑通的那个结果而是过程中被迫做的每一个决策。2. 整体架构设计与技术选型思路2.1 为什么不用现成框架而选择手写这是很多人第一个会问的问题有FastAPI、有TorchServe、有Triton为什么还要自己写我的答案很直接——因为你得先知道框架帮你做了什么才有资格用框架。ai-engineering-from-scratch的核心思路就是先把轮子拆开看再决定用不用。从工程角度讲手写一遍推理服务能让你直面几个关键问题模型加载放在进程启动时还是首次请求时批处理batching在哪个层级做GPU显存怎么管理请求队列满了怎么降级这些问题在框架里往往被一行配置掩盖了但线上出事的时候恰恰是这些配置在决定生死。我见过太多团队用现成框架上线结果QPS一高就雪崩排查半天发现是默认的worker数量不对。选型上这个项目走的是极简路线Python做服务层NumPy和PyTorch做计算Redis做队列和缓存Docker做环境隔离。为什么不用Kubernetes因为从零阶段引入K8s只会增加认知负担单机Docker Compose足够把链路跑通。等你理解了单机怎么调度再上K8s才有意义。这个取舍很关键很多教程一上来就堆一整套云原生结果学习者连模型怎么加载都没搞明白。2.2 分层设计把推理链路切成四段我把整条链路拆成四层每一层职责单一方便单独测试和替换。第一层是接入层负责HTTP请求的接收、参数校验、限流。这一层不碰任何模型逻辑只做协议转换。第二层是调度层负责把请求攒成批次、管理队列、分配计算资源。第三层是推理层真正跑模型前向计算的地方。第四层是数据层处理输入预处理和输出后处理包括tokenize、归一化、反归一化这些。这么切的好处是每一层都可以独立压测。比如你想知道瓶颈在调度还是推理把调度层换成直通模式一测便知。我实际做下来发现大部分性能问题不在模型本身而在调度层的批处理策略上。批处理窗口设成10毫秒还是50毫秒吞吐量能差三倍。2.3 关键技术点预判在动手之前有几个技术点必须提前想清楚否则写到一半会推翻重来。并发模型的选择。Python有GIL纯多线程跑推理没意义。所以推理层要么用多进程要么把计算释放到C扩展里。这个项目采用的是多进程worker加进程间队列的方案主进程只负责调度worker进程各自持有模型副本。代价是显存占用翻倍但换来的是真正的并行。批处理的实现位置。批处理放在调度层做而不是推理层。原因是调度层能看到全局请求可以根据队列长度动态调整批次大小。如果放在推理层每个worker各自为战批次永远攒不大。模型版本管理。从零阶段不搞复杂的模型仓库但至少要支持热加载。方案是用一个版本号加文件路径的映射表新版本加载成功后原子切换指针。这个细节后面会展开讲。3. 核心模块的实操拆解3.1 推理服务的骨架搭建先看服务骨架。我用的是Python标准库的http.server加socketserver的ThreadingMixIn而不是FastAPI。为什么因为从零的意义在于暴露底层。FastAPI帮你做了太多事你连请求是怎么被解析的都不清楚。用标准库写一遍你会真正理解WSGI和ASGI的区别理解为什么异步框架在高并发下更有优势。骨架大概长这样一个RequestHandler类继承BaseHTTPRequestHandler重写do_POST方法。在方法里读取Content-Length从rfile读bodyjson解析然后丢给调度层。返回的时候设置状态码和Content-Type把结果写进wfile。就这么简单但每一步都有坑。比如Content-Length没读对会导致body截断json解析失败要返回400而不是500这些细节在框架里是默认行为手写就得自己兜。class InferenceHandler(BaseHTTPRequestHandler): def do_POST(self): length int(self.headers.get(Content-Length, 0)) body self.rfile.read(length) try: payload json.loads(body) except json.JSONDecodeError: self.send_error(400, invalid json) return result scheduler.submit(payload) resp json.dumps(result).encode() self.send_response(200) self.send_header(Content-Type, application/json) self.send_header(Content-Length, str(len(resp))) self.end_headers() self.wfile.write(resp)这段代码看着朴素但它把HTTP服务的本质暴露得很清楚读请求、处理、写响应。理解了这三步再看任何Web框架都是降维打击。3.2 批处理调度器的实现细节调度器是整个项目最核心也最难的部分。它的任务是把零散到达的请求攒成批次凑够一定数量或者等够一定时间就发给推理层。这里有个经典的权衡批次越大吞吐越高但延迟也越高。ai-engineering-from-scratch采用的是动态窗口策略。具体做法是维护一个请求队列主循环每隔一个极短的时间片比如2毫秒检查一次队列。如果队列长度超过阈值比如8立刻组批发送如果没超过但队列非空且距离上次发送超过最大等待时间比如20毫秒也发送。这样在低负载时延迟可控高负载时自动攒大批次。def schedule_loop(self): last_flush time.time() while self.running: time.sleep(0.002) if not self.queue: continue now time.time() if len(self.queue) self.batch_size or (now - last_flush) self.max_wait: batch self.drain(self.batch_size) self.dispatch(batch) last_flush now参数怎么定batch_size取决于模型大小和显存。我实测下来一个中等规模的Transformer批次从1加到16吞吐能提升近10倍但延迟只增加几毫秒。超过16之后收益递减因为GPU算力打满了。max_wait一般设在10到50毫秒之间超过50毫秒用户就能感知到卡顿。这两个参数没有万能值必须根据你的模型和硬件实测。注意批处理的前提是同一批次内所有请求的输入形状一致。如果输入长度不一要么padding到统一长度要么按长度分桶。padding会浪费算力分桶实现复杂但效率高。从零阶段建议先用padding跑通再优化。3.3 模型加载与热更新机制模型加载最忌讳的是每次请求都重新加载。正确做法是进程启动时加载一次常驻内存。但这里有个多进程的坑如果你用fork启动worker模型会在每个子进程里各加载一份显存直接爆掉。解决方案有两种一是用spawn启动方式让子进程自己加载二是主进程加载后用共享内存传递。前者简单但启动慢后者快但实现复杂。项目里用的是spawn因为简单可靠。热更新的实现思路是维护一个全局的模型引用新模型在后台线程加载加载完成后用一个锁保护地替换引用。正在处理的请求持有旧引用不受影响新请求拿到新引用。这样做到零停机更新。class ModelManager: def __init__(self): self._model None self._lock threading.Lock() def load(self, path): new_model torch.load(path, map_locationcpu) new_model.eval() with self._lock: self._model new_model def get(self): with self._lock: return self._model这个模式看着简单但有个隐患旧模型什么时候释放如果一直有请求持有旧引用显存就回收不了。实际做法是加引用计数计数归零时手动del并调用torch.cuda.empty_cache()。这个细节很多教程不讲但线上跑久了必然遇到显存泄漏。3.4 输入预处理与输出后处理的工程化预处理和后处理是最容易被低估的部分。模型本身可能只占代码量的20%剩下80%都在处理数据。ai-engineering-from-scratch把这块单独抽成模块就是为了强调它的重要性。预处理要处理的问题包括文本长度截断、特殊字符清洗、编码转换、数值归一化。以文本为例tokenize之后要检查长度超长的截断不足的padding。这里有个坑是padding的位置左padding还是右padding会影响模型输出必须和训练时保持一致。我见过有人训练用右padding推理用左padding结果准确率掉了一大截排查了两天才发现。后处理要处理的是模型原始输出的解读。比如分类任务要取argmax检测任务要做NMS生成任务要处理停止符。这些逻辑如果写在推理层里会导致推理层越来越臃肿。抽出来之后推理层只负责算后处理层负责解释职责清晰。def preprocess(texts, max_len128): encoded tokenizer(texts, paddingmax_length, truncationTrue, max_lengthmax_len) return {k: torch.tensor(v) for k, v in encoded.items()} def postprocess(logits): probs torch.softmax(logits, dim-1) scores, indices torch.topk(probs, k5, dim-1) return [{label: int(i), score: float(s)} for i, s in zip(indices[0], scores[0])]4. 性能压测与线上问题排查4.1 压测方案与关键指标压测不是随便跑个ab或者wrk就完事。AI服务的压测要关注几个特殊指标首token延迟、吞吐量tokens per second、显存占用、批次利用率。普通Web服务看QPS和P99就够了AI服务还得看这些。我用的是自己写的压测脚本模拟不同并发下的请求分布。关键是要模拟真实场景比如请求长度符合某种分布而不是全部等长。全部等长的压测结果会严重偏乐观因为批处理效率被高估了。并发数QPSP50延迟P99延迟显存占用批次利用率11282ms95ms2.1GB8%878102ms180ms2.3GB52%32210152ms420ms2.8GB87%64235270ms980ms3.2GB91%从这张表能看出几个规律。并发从1到8QPS涨了6倍多因为批处理开始起作用。到32之后QPS增长放缓因为GPU算力接近饱和。到64的时候QPS几乎不涨但P99延迟翻倍说明系统已经过载。这个拐点就是你的服务容量上限超过它就该扩容或者限流了。4.2 常见问题速查表实际跑下来遇到的问题五花八门我整理成一张表方便对照排查。现象可能原因排查方法解决方案首次请求特别慢模型懒加载看日志加载时间启动时预热显存持续增长引用未释放nvidia-smi观察加引用计数和empty_cache批次利用率低请求间隔大打印批次大小分布调大max_waitP99延迟毛刺GC停顿开GC日志调优GC阈值或换语言输出结果不稳定padding方向不一致对比训练推理代码统一padding策略吞吐上不去worker数不足看CPU利用率增加worker进程这张表里的每一条都是我实际踩过的。比如显存泄漏那条当时服务跑了一天显存从2G涨到8G最后发现是日志里保存了tensor引用。这种问题不看代码根本想不到。4.3 几个反直觉的实操心得第一个心得批处理不是越大越好。我一开始把batch_size设成64想着吞吐最大化结果P99延迟飙到2秒。后来降到16吞吐只掉了15%但延迟降了一半。用户体验的提升远比那点吞吐重要。第二个心得预热比优化更重要。服务刚启动时CUDA上下文初始化、cuDNN算法选择都要时间第一次推理可能比稳态慢10倍。上线前一定要用真实数据预热几十次否则第一个用户会以为服务挂了。第三个心得日志要打对地方。AI服务的日志不能只打请求和响应还要打批次大小、队列长度、显存占用。这些指标是排查性能问题的关键。我习惯每处理100个批次打一条汇总日志既不刷屏又能看出趋势。提示压测时一定要监控GPU利用率和显存光看QPS会误导。有时候QPS上不去不是代码问题是GPU已经打满了这时候优化代码没用得换卡或者加卡。5. 从单机到可扩展的演进路径5.1 单机方案的边界在哪里单机方案能扛多少以我实测的数据一张中端GPU加8核CPU跑中等规模模型稳定QPS在200左右。这个量级对很多内部系统够用了但对外服务肯定不够。单机的瓶颈通常出现在三个地方GPU算力、CPU预处理、网络IO。判断瓶颈在哪看资源利用率就行。GPU利用率90%以上说明算力是瓶颈加卡或者换卡。CPU利用率高但GPU低说明预处理拖后腿可以把预处理也放到GPU上或者用更快的tokenizer。网络IO瓶颈一般出现在返回大结果的时候比如生成任务返回长文本这时候要考虑流式返回。5.2 水平扩展的关键改造点从单机到多机需要改造的地方比想象中多。首先是调度层要独立出来不能和推理进程绑在一起。调度层变成无状态服务推理层变成可水平扩展的worker池。请求进来先到调度层调度层根据worker负载分发。其次是模型分发。每个worker都要有模型怎么保证版本一致方案是用共享存储加版本号worker启动时拉取指定版本。更新时先更新共享存储再滚动重启worker。这个过程要保证新旧版本兼容否则会出现请求打到旧worker返回旧格式的情况。最后是会话保持。如果模型是有状态的比如带缓存的生成模型请求必须打到同一个worker。解决方案是在调度层做一致性哈希按session id路由。这个改造不小但如果不做缓存命中率会暴跌。5.3 监控与告警的最小可用配置从零阶段不需要上Prometheus加Grafana那一套但至少要有个地方能看到关键指标。我的做法是用一个简单的HTTP端点暴露metrics格式用Prometheus的文本格式这样以后想接监控系统直接就能接。def metrics_handler(self): lines [ fai_batch_size {self.stats[batch_size]}, fai_queue_length {self.stats[queue_length]}, fai_gpu_memory_mb {self.stats[gpu_memory]}, fai_requests_total {self.stats[requests]}, ] body \n.join(lines).encode() self.send_response(200) self.send_header(Content-Type, text/plain) self.end_headers() self.wfile.write(body)告警规则设三条就够队列长度持续超过阈值、P99延迟超过阈值、显存占用超过阈值。这三条覆盖了绝大多数线上故障。告警渠道用邮件或者内部IM都行关键是有人看。6. 我在这套链路里踩过的坑6.1 那些文档不会告诉你的细节第一个坑是CUDA上下文和fork的冲突。如果你在主进程初始化了CUDA再用fork启动子进程子进程里的CUDA会处于未定义状态一跑就崩。这个问题的表现是子进程直接段错误没有任何Python异常。解决方案就是前面说的用spawn或者干脆在子进程里初始化CUDA。第二个坑是PyTorch的eval模式。模型加载后一定要调model.eval()否则BatchNorm和Dropout会按训练模式跑输出完全不对。这个坑隐蔽在于它不会报错只是结果不准你可能调半天参数才发现是模式没切。第三个坑是多进程下的随机种子。如果预处理里有随机操作每个worker的种子必须不同否则所有worker产生相同的随机序列批处理就失去意义了。解决方案是在worker启动时用进程ID做种子。6.2 性能优化的优先级排序优化要有优先级不能瞎调。我的经验是先看批次利用率再看预处理耗时最后看模型本身。批次利用率低的话优化模型没用因为GPU根本没吃满。预处理耗时占比高的话把预处理移到GPU上收益最大。模型本身的优化量化、剪枝放在最后因为改动大、风险高。具体数据一个典型请求预处理占30%推理占60%后处理占10%。如果把预处理优化掉一半整体提升15%。如果把推理量化到INT8推理时间减半整体提升30%。但量化的代价是精度可能下降需要评估。6.3 代码组织的一点个人习惯我习惯把配置全部抽到一个config.py里用dataclass定义。这样调参的时候不用满代码找数字改一个地方就行。而且dataclass支持类型提示写错了编辑器会提醒。from dataclasses import dataclass dataclass class Config: batch_size: int 16 max_wait_ms: int 20 num_workers: int 4 max_queue_size: int 256 model_path: str models/v1.pt max_seq_len: int 128另一个习惯是给每个模块写一个__main__入口能单独跑。这样调试的时候不用启动整个服务直接跑单个模块就行。比如预处理模块单独跑喂几条数据看输出对不对。这个习惯能省大量调试时间。这套从零搭建的链路我前后迭代了三版才稳定。第一版跑通但性能差第二版性能上来了但代码乱第三版才兼顾了性能和可维护性。如果你也在做类似的事我的建议是先把链路跑通别一上来就追求完美。跑通之后你才知道瓶颈在哪才知道该优化什么。上来就设计一套完美架构大概率是过度设计最后发现优化的地方根本不是你以为的那个。