1. 为什么要在 Spark on YARN 上跑 GPU 推理先把场景说清楚。很多团队早期做 AI 推理套路都很简单模型训练完导出权重写个 Flask 或者 FastAPI 服务扔到一台带显卡的机器上前面挂个 Nginx完事。这个方案在请求量小、模型单一的时候确实够用但只要业务稍微长起来问题就全冒出来了。模型从 1 个变成 8 个每个模型对显存、算力、批大小的要求都不一样离线批量推理任务和在线服务抢同一张卡某台机器挂了整个服务就不可用资源利用率忽高忽低白天在线服务吃满晚上离线任务又闲着。这时候你自然会想能不能把推理任务也纳入到已有的资源调度体系里如果公司已经有一套 Spark on YARN 的集群上面跑着大量的离线数据处理任务那把 GPU 推理也塞进去复用同一套资源管理、同一套监控告警、同一套提交作业的流程听起来就非常顺理成章。这就是“基于 Spark on YARN with GPU 的 AI 模型推理”这个方向的核心动机——不是再造一个推理框架而是把推理当成一类特殊的 Spark 作业交给 YARN 去调度让 GPU 资源像 CPU、内存一样被统一管理。它解决的问题可以归纳成三条。第一是资源统一GPU 卡不再是被某个服务独占的黑盒而是 YARN 里一种可申请、可回收、可配额管理的资源。第二是弹性与隔离不同业务线的推理任务通过队列隔离谁也别抢谁的卡任务结束资源立刻归还。第三是批流一体同一套代码既能跑离线的大批量推理比如每天凌晨把上亿条数据过一遍模型也能跑准实时的微批推理不用维护两套技术栈。适合谁来参考如果你已经在用 Spark 做数据处理手里有 GPU 机器又不想引入一套全新的推理服务平台那这套方案非常契合。如果你是从零开始、只做在线低延迟推理那说实话专门的推理服务框架可能更直接Spark 的调度开销和启动延迟不一定划算。所以先想清楚自己的场景再决定要不要往下走。2. 整体架构设计与关键选型考量2.1 三层结构YARN 管资源Spark 管计算GPU 管加速整套方案的骨架其实很清晰我把它拆成三层来看。最底层是YARN 资源管理层它负责整个集群的 CPU、内存、GPU 的分配。中间层是Spark 计算层它把用户的推理逻辑翻译成分布式任务向 YARN 申请容器Container在容器里启动 Executor。最上层是GPU 加速层Executor 里的推理代码通过 CUDA 或者推理运行时把计算卸载到显卡上。关键在于YARN 本身对 GPU 的支持是通过Resource Types机制实现的。早期版本里 GPU 是作为一种扩展资源类似yarn.io/gpu来管理的后来引入了更通用的ResourceTypes可以自定义gpu这种资源类型。NodeManager 启动时会探测本机显卡数量上报给 ResourceManager调度器在分配容器时就会把 GPU 数量作为约束条件之一。这意味着一个 Executor 申请 1 张卡YARN 就会确保这个容器所在的节点真的有可用的卡并且在这张卡被占用期间不会分配给别的容器。2.2 为什么是 Spark on YARN 而不是 standalone 或 Kubernetes这里要解释一下选型逻辑因为很多人会问既然都上 GPU 了为什么不直接上 Kubernetes答案取决于你现有的基础设施。如果你的团队已经有成熟的 YARN 集群运维体系、监控、权限、队列配额都在 YARN 上那强行迁到 K8s 意味着重建一整套运维能力成本很高。Spark on YARN 的优势在于复用存量已有的 Hadoop 生态、HDFS 数据本地性、YARN 队列的 ACL 控制全都能直接用。另一个原因是数据本地性。推理任务往往要读取 HDFS 或者 Hive 表里的大量数据Spark on YARN 天然能利用数据本地性调度把计算推到数据所在的节点减少网络传输。而独立的推理服务通常要把数据拉过来这在批量场景下是很大的开销。至于 Spark standalone 模式它虽然也能管 GPU通过spark.executor.resource.gpu.amount之类的配置但资源隔离和队列管理能力远不如 YARN多租户场景下容易出问题。所以如果集群是多业务共享的YARN 是更稳妥的选择。2.3 推理任务的两种形态批式与微批在 Spark 上跑推理任务形态主要分两种。一种是离线批式推理典型场景是每天定时把一张大表里的样本过一遍模型结果写回 Hive 或者 HDFS。这种任务对延迟不敏感但对吞吐要求高适合用大 batch、多 Executor 并行。另一种是微批推理用 Spark Structured Streaming 消费消息队列的数据每个微批做一次推理延迟在秒级到分钟级。这种场景下 GPU 利用率可能不高因为微批之间的间隙 GPU 是空闲的需要仔细调 batch 大小和触发间隔。我个人的经验是批式推理是 Spark on YARN with GPU 最舒服的场景因为它的资源申请和释放是粗粒度的和 YARN 的调度模型天然匹配。微批推理虽然能做但要接受一定的延迟抖动而且 GPU 利用率往往上不去除非你能把微批间隔压得很短、batch 攒得足够大。3. 环境准备与 GPU 资源接入实操3.1 集群侧的前置条件在动手写推理代码之前集群本身要满足几个条件。首先是Hadoop 和 YARN 版本要支持 GPU 资源调度。较新的版本对自定义资源类型支持更完善配置也更清晰。其次是NodeManager 所在节点必须装好显卡驱动并且驱动版本要和后续推理运行时比如 CUDA、TensorRT匹配。这一点经常被忽略——YARN 只负责数卡不负责保证驱动能用驱动装错了任务照样跑不起来。然后是资源类型配置。需要在yarn-site.xml里启用自定义资源类型让 YARN 认识gpu这个资源。配置的核心是告诉 YARN集群里有gpu这种资源NodeManager 要探测并上报。配置完成后重启 YARN用yarn node -list或者资源管理界面确认每个节点的 GPU 数量被正确识别。如果显示为 0那多半是探测脚本没配好或者驱动路径不对。注意GPU 资源的探测依赖 NodeManager 的硬件探测能力不同发行版和版本的配置项名称可能不同务必对照你所用版本的官方文档不要照搬网上旧版本的配置。3.2 Spark 侧的资源申请配置Spark 要能申请 GPU需要在提交作业时显式声明。核心配置项包括Executor 数量、每个 Executor 的 GPU 数量、以及对应的发现脚本Discovery Script。发现脚本的作用是告诉 Spark这个 Executor 实际能用哪几张卡。因为 YARN 分配了 GPU 数量但具体是哪几张卡的编号需要脚本去查通常通过环境变量或者nvidia-smi的输出解析。一个典型的提交命令大致长这样spark-submit \ --master yarn \ --deploy-mode cluster \ --num-executors 8 \ --executor-cores 4 \ --executor-memory 16g \ --conf spark.executor.resource.gpu.amount1 \ --conf spark.task.resource.gpu.amount1 \ --conf spark.executor.resource.gpu.discoveryScript./getGpusResources.sh \ --conf spark.yarn.am.resource.gpu.amount0 \ inference_job.py这里有几个点值得展开。spark.executor.resource.gpu.amount1表示每个 Executor 申请 1 张卡。spark.task.resource.gpu.amount1表示每个 task 占用 1 张卡这意味着同一个 Executor 内的 task 是串行使用 GPU 的不会出现多个 task 抢同一张卡的情况。如果你希望一个 Executor 内多个 task 共享一张卡比如小模型可以把 task 的 GPU 需求设成小数但这需要运行时支持且容易踩坑新手不建议。spark.yarn.am.resource.gpu.amount0是让 ApplicationMaster 不占卡因为 AM 只做调度不需要 GPU。这个细节很多人会漏结果 AM 白占一张卡浪费资源。3.3 发现脚本怎么写发现脚本是连接 YARN 分配和 Spark 实际使用的桥梁。它的职责很简单输出一个 JSON告诉 Spark 当前 Executor 能用哪些 GPU。YARN 在分配容器时会把分配给这个容器的 GPU 编号通过环境变量传进来具体变量名依版本而定脚本读取这个变量转换成 Spark 期望的格式。一个简化版的脚本逻辑是读取环境变量里的 GPU 设备列表如果为空则回退到用nvidia-smi查询然后输出类似{name: gpu, addresses: [0, 1]}的 JSON。脚本必须可执行并且路径要在提交时正确指定。我踩过的坑是脚本没有执行权限或者路径是相对路径导致在集群节点上找不到任务直接失败。所以用绝对路径并且提前在所有 NodeManager 节点上放好脚本是最稳的做法。4. 推理代码的编写与分布式改造4.1 把单机推理改造成 Spark 作业单机推理的代码通常是加载模型读一批数据逐条或者按 batch 推理输出结果。改造成 Spark 作业核心是把“读数据”和“推理”这两步映射到 RDD 或者 DataFrame 的算子上去。最常见的方式是用mapPartitions因为模型加载是重操作不应该每条数据都加载一次。mapPartitions保证每个分区只初始化一次模型然后在这个分区内批量推理。结构大致是这样Driver 端负责读取数据源、定义 schema、广播一些配置每个 Executor 在mapPartitions里初始化模型从 HDFS 或者本地缓存加载权重然后遍历分区内的数据攒够一个 batch 就送进 GPU 推理最后 yield 出结果。这里的关键是模型加载要放在分区函数内部而不是外部否则模型对象会被序列化分发既慢又容易出错。def infer_partition(iterator): model load_model(/path/to/model) # 每个分区加载一次 batch [] for row in iterator: batch.append(row) if len(batch) BATCH_SIZE: yield from run_inference(model, batch) batch [] if batch: yield from run_inference(model, batch) result spark.read.parquet(input_path).rdd.mapPartitions(infer_partition)4.2 批大小的选择与显存管理批大小batch size是影响 GPU 利用率和吞吐的核心参数。太小GPU 算力吃不饱大量时间花在数据搬运和 kernel 启动上太大显存溢出OOM任务直接挂掉。选择批大小要综合考虑模型大小、输入维度、显卡显存。一个实用的做法是先用单机脚本测出在目标显卡上不 OOM 的最大批大小然后在 Spark 里取它的 70% 到 80% 作为默认值留出余量应对数据波动。显存管理还有个容易被忽视的点PyTorch 和 TensorFlow 默认会缓存显存即使推理结束也不一定立刻归还。在 Spark 这种长时间运行多个 task 的环境里如果每个 task 都占着显存不放很快就会 OOM。解决办法是在推理结束后显式释放或者设置显存分配策略为按需增长。另外如果同一个 Executor 要跑多个 task务必确保前一个 task 的显存已经释放再启动下一个。提示在 Executor 里跑推理建议把spark.task.resource.gpu.amount设为 1让 task 串行使用 GPU。这样虽然牺牲了一点并发但避免了多 task 抢卡导致的显存冲突和性能抖动稳定性优先。4.3 数据读取与结果写回的优化数据读取这块尽量用列式存储Parquet、ORC只读需要的列减少 IO。如果输入数据在 Hive 表里用 Spark SQL 读进来再转 RDD能利用 Spark 的谓词下推和分区裁剪。结果写回时如果数据量大注意控制输出文件数量避免产生大量小文件拖垮 HDFS。可以在写之前用repartition或者coalesce调整分区数。还有一个实战技巧把模型文件提前分发到各个节点。如果每次任务启动都从 HDFS 拉几百 MB 甚至几个 GB 的模型权重启动时间会很长。可以预先放到节点的本地磁盘或者用 Spark 的--files参数分发让 Executor 直接从本地读。这个优化在频繁提交任务的场景下效果非常明显。5. 常见问题与排查技巧实录5.1 GPU 申请不到或者数量不对最常见的现象是任务一直处于等待状态或者报错说资源不足。排查思路分几步。先看 YARN 的资源界面确认集群里 GPU 总量和可用量是否正确。如果总量是 0说明 NodeManager 没探测到卡检查驱动和探测配置。如果总量对但可用量一直是 0说明卡被别的任务占着或者有僵尸容器没释放。这时候可以看 YARN 的容器列表找到占用 GPU 的容器确认是不是自己的历史任务没清理干净。还有一种情况是申请了 GPU 但 Spark 说找不到。这通常是发现脚本的问题。检查脚本是否有执行权限、路径是否正确、输出的 JSON 格式是否符合要求。可以在脚本里加日志把环境变量和输出打印出来方便定位。5.2 显存溢出OOM的几种典型原因OOM 是 GPU 推理最常见的故障。原因通常有三类。第一类是批大小设得太大这个最直接调小即可。第二类是显存泄漏比如模型对象或者中间张量没有被释放跑着跑着显存就满了。这种要用显存监控工具观察显存随时间的变化曲线如果是持续上升那就是泄漏。第三类是多 task 共享 GPU 导致的冲突前面提到过把 task 的 GPU 需求设为 1 就能避免。下面这张表是我整理的常见问题速查方便对照排查现象可能原因排查方向解决手段任务一直等待GPU 资源不足看 YARN 资源界面释放僵尸容器或增加配额找不到 GPU发现脚本异常检查脚本权限和输出修正脚本并分发到所有节点显存 OOM批大小过大单机测试最大批大小调小批大小留余量显存持续上涨显存泄漏监控显存曲线显式释放张量和模型推理速度慢GPU 利用率低看 GPU 利用率和批大小增大批大小或减少数据搬运任务启动慢模型加载耗时看启动阶段日志预分发模型到本地磁盘5.3 性能不达预期的调优思路如果任务能跑通但速度不理想先别急着改代码先用监控工具看清楚瓶颈在哪。GPU 利用率低说明数据供给跟不上或者批太小GPU 利用率高但吞吐还是低说明模型本身计算量大或者数据搬运开销大。CPU 利用率高而 GPU 低通常是数据预处理比如图像解码、文本分词在 CPU 上成了瓶颈这时候可以考虑把预处理也放到 GPU 上或者增加预处理并行度。另一个容易被忽视的点是Executor 的并发度。如果每个 Executor 只申请 1 张卡、跑 1 个 task那集群里有多少张卡就只能并行多少个 task。如果单个 task 的 GPU 利用率不高整体吞吐就上不去。这时候可以考虑一个 Executor 申请多张卡或者用更小的批但更多的并发。具体怎么选要实测。6. 我踩过的坑和几条实用建议先说一个最坑的驱动版本和推理运行时版本不匹配。有一次集群升级了显卡驱动但推理用的运行时还是旧版本结果任务在单机上跑得好好的一到集群就报奇怪的 CUDA 错误。排查了大半天才发现是版本问题。所以集群节点的驱动版本、CUDA 版本、推理框架版本三者必须对齐并且要在所有节点上保持一致。建议做一个版本清单升级时统一操作。第二个坑是模型加载路径用了本地相对路径。单机测试时没问题因为工作目录就是脚本所在目录。但到了集群Executor 的工作目录是 YARN 分配的临时目录相对路径直接失效。所以模型路径、配置文件路径一律用绝对路径或者通过参数传入别偷懒。第三个建议是给推理任务单独建一个 YARN 队列。推理任务和离线数据处理任务混在一个队列里容易互相影响。推理任务往往对启动时间敏感如果被大数据任务堵住体验很差。单独队列加上 GPU 资源配额能保证推理任务的资源供给也方便做成本核算。最后分享一个提高稳定性的小技巧在推理代码里加异常捕获和重试。GPU 推理偶尔会因为驱动、显存、数据异常等原因失败如果整个任务因为一条数据挂掉代价太大。可以在分区函数里对单条或者单批推理做 try-catch失败的数据记录下来跳过保证整体任务能跑完。跑完之后再单独分析失败数据比整个任务重跑高效得多。这套方案我在几个批量推理场景里用过稳定性确实比裸服务好资源利用率也上来了。但它不是银弹微批场景的延迟和 GPU 利用率问题依然存在需要根据业务容忍度去权衡。如果你的场景是每天定时跑大批量数据那这套组合拳值得一试。
