1. 为什么要在 Spark on YARN 上跑 GPU 推理把 AI 模型推理塞进 Spark on YARN 这套体系里第一次听到的人多半会皱眉推理不是应该用专门的在线服务框架吗为什么要绕这么大一圈我最初接手这个方向时也有同样的疑问直到真正面对每天要跑几十万条离线样本、每条样本都要过一遍大模型这种场景才明白这套组合的价值在哪里。核心矛盾在于批处理的数据规模和推理的算力需求之间的错配。纯 CPU 集群跑大模型推理吞吐低到让人绝望纯 GPU 服务器上用脚本跑又缺乏分布式调度、容错、资源隔离这些能力。Spark on YARN 恰好补上了后半段——它本来就是为大规模分布式批处理设计的YARN 负责资源调度和隔离Spark 负责数据分片和任务编排。把 GPU 作为 YARN 的一种可调度资源接进来就能让数据在哪、算力就在哪这件事变得自然。具体来说这套方案解决的是三类问题。第一类是离线批量推理比如给几千万条用户行为文本打标签、给海量图片做特征提取这类任务对延迟不敏感但对吞吐极其敏感。第二类是模型评测与回归每次模型迭代后要在固定测试集上跑一遍需要可复现、可追溯的分布式执行环境。第三类是数据预处理管线中的模型环节比如在 ETL 流程里嵌入一个 embedding 模型把非结构化数据转成向量再入库。适合读这篇内容的人我大致分三类一是有 Spark 使用经验、想把手里的 GPU 机器利用起来的工程师二是做推荐、搜索、风控等方向、需要大规模离线推理的数据团队三是正在评估自建推理平台 vs 复用现有大数据平台这个选型问题的技术负责人。如果你完全没碰过 Spark建议先补一下 RDD、Executor、Stage 这些基础概念否则后面讲资源分配时会比较吃力。需要提前说明的是这套方案不是银弹。它的优势在吞吐和与现有数据链路复用劣势在启动开销和对交互式低延迟场景不友好。如果你的场景是 QPS 几千的在线服务老老实实用专门的推理服务框架更合适。判断标准很简单任务是不是一次提交、批量处理、分钟级到小时级完成如果是那这套组合就值得考虑。2. YARN 的 GPU 资源调度机制到底怎么工作很多人以为只要在 Spark 配置里写个spark.executor.resource.gpu.amount1就能用上 GPU结果提交上去发现 Executor 根本没拿到卡或者多个 Executor 抢同一张卡。问题出在对 YARN GPU 调度机制的理解上——它不是自动发现而是需要显式配置和插件支持。2.1 从 NodeManager 发现 GPU 到容器分配YARN 本身对 GPU 的支持是通过Resource Types机制实现的。在 Hadoop 3.x 之前YARN 只认 CPU 和内存两种资源GPU 需要靠外部插件比如早期的yarn-gpu补丁来扩展。Hadoop 3.x 之后引入了可配置资源类型GPU 才成为一等公民。整个链路是这样的NodeManager 启动时通过配置的发现脚本discovery script扫描本机的 GPU 设备把数量上报给 ResourceManager。RM 在调度时把 GPU 当作一种和 CPU、内存并列的资源维度按容量分配给 ApplicationMaster 申请的容器。容器启动后NodeManager 会把分配给它的 GPU 设备通过环境变量通常是NVIDIA_VISIBLE_DEVICES暴露给容器内的进程。这里有个关键点容易被忽略发现脚本必须自己写或者用官方提供的。Hadoop 自带了一个nvidia-smi的封装脚本示例但很多发行版没有默认启用。如果发现脚本没配好NodeManager 上报的 GPU 数量就是 0后面所有配置都是白搭。配置上主要涉及resource-types.xml和node-resources.xml两个文件。前者定义资源类型后者在特定节点上覆盖默认值。一个典型的resource-types.xml长这样configuration property nameyarn.resource-types/name valueyarn.io/gpu/value /property property nameyarn.resource-types.yarn.io/gpu.units/name value/value /property /configurationunits留空表示这是计数型资源不是带单位的可压缩资源。这一点很重要GPU 不能像内存那样用 1.5 张卡必须是整数。2.2 发现脚本的坑与验证方法发现脚本的写法直接决定了 YARN 能不能正确识别 GPU。官方示例脚本大致逻辑是调用nvidia-smi --query-gpuindex --formatcsv,noheader | wc -l得到卡数然后按 YARN 要求的 JSON 格式输出。但实际部署时会遇到几个问题。第一个问题是脚本执行用户权限。NodeManager 通常以yarn用户运行如果nvidia-smi对yarn用户不可见比如驱动安装时限制了权限脚本会返回 0。验证方法很简单切到 yarn 用户手动跑一遍脚本看输出对不对。第二个问题是MIG 模式下的卡数统计。如果 GPU 开了 MIGMulti-Instance GPU一张物理卡会被切成多个实例nvidia-smi看到的可能是实例而不是物理卡。这时候发现脚本要按实例数上报否则调度会错乱。第三个问题是脚本路径和超时。发现脚本默认在 NodeManager 的yarn.nodemanager.resource-plugins.gpu.discovery-script.path指定如果脚本执行超过yarn.nodemanager.resource-plugins.gpu.discovery-script.timeout默认 10 秒会被判定失败。我见过因为脚本里加了sleep调试导致超时的案例排查了半天。验证 GPU 是否被正确识别最直接的方法是看 RM 的 Web UI在 Nodes 页面能看到每个节点的资源列表里有没有yarn.io/gpu这一项。如果没有说明发现脚本或配置有问题别急着往下走。2.3 容器内的 GPU 隔离与可见性YARN 分配 GPU 后怎么保证容器 A 用的卡和容器 B 不冲突答案是Docker 或 LinuxContainerExecutor 的设备隔离。在非 Docker 模式下YARN 通过NVIDIA_VISIBLE_DEVICES环境变量告诉进程你只能用这几张卡但这只是软隔离——如果进程硬要访问其他卡理论上是可以的。真正的硬隔离需要配合 cgroups device controller 或者 Docker 的--gpus参数。生产环境我强烈建议用Docker 容器执行器。YARN 的 Docker 支持可以把 GPU 设备通过--gpus或者 device 映射的方式传给容器隔离性更好环境一致性也更好。配置上需要在container-executor.cfg里允许 GPU 设备并在yarn-site.xml里开启 Docker 相关选项。这里有个实操心得Docker 镜像里的 CUDA 版本要和宿主机驱动兼容。宿主机驱动是向下兼容的但镜像里的 CUDA runtime 版本不能超过驱动支持的上限。我踩过一次坑镜像里装了 CUDA 12.x宿主机驱动只支持到 11.8结果容器一启动就报CUDA driver version is insufficient。解决办法要么升级驱动要么降镜像里的 CUDA 版本。3. Spark 侧如何把 GPU 用起来YARN 把 GPU 资源准备好了接下来是 Spark 怎么申请和使用。Spark 从 3.0 开始正式支持 GPU 调度但配置项比较分散容易漏配。3.1 Executor 的 GPU 申请配置最核心的几个配置项如下表配置项作用典型值spark.executor.resource.gpu.amount每个 Executor 申请的 GPU 数1spark.task.resource.gpu.amount每个 Task 占用的 GPU 数1spark.driver.resource.gpu.amountDriver 申请的 GPU 数0spark.executor.resource.gpu.discoveryScriptExecutor 侧发现脚本同 YARN 脚本spark.executor.resource.gpu.vendor厂商标识nvidia.com这里最容易搞错的是executor.amount和task.amount的关系。如果 Executor 申请 1 张卡Task 也申请 1 张卡那这个 Executor 同时只能跑 1 个 Task——因为卡被占满了。这会导致并发度极低一个 Executor 一个 Task资源利用率很差。正确的做法取决于你的推理模型大小。如果是小模型比如 BERT-base 级别一张卡可以同时跑多个推理实例那可以把task.resource.gpu.amount设成 0.5 甚至 0.25让一个 Executor 里并行跑多个 Task 共享一张卡。但要注意共享 GPU 的 Task 之间会有显存竞争如果模型显存占用大很容易 OOM。我的经验是先用task.amount1跑通确认单 Task 推理正常再根据显存余量逐步降低 task.amount 提升并发。不要一上来就追求高并发显存溢出比低并发更麻烦。3.2 推理代码里怎么拿到 GPU配置对了代码里怎么用Spark 的 GPU 调度会把分配到的设备通过CUDA_VISIBLE_DEVICES环境变量传给 Executor 进程。所以你的推理代码里不需要手动指定 devicePyTorch 或 TensorFlow 会自动看到被分配的那张卡。但有个细节要注意在 Executor 里初始化模型时要确保每个 Task 用的是自己的模型实例。如果模型是全局单例多个 Task 共享同一个模型对象在 GPU 上会出问题——因为不同 Task 可能被分配到不同的卡共享模型会导致设备不匹配。正确的模式是在mapPartitions里初始化模型让每个分区对应一个 Task有自己的模型实例def infer_partition(iterator): import torch device torch.device(cuda if torch.cuda.is_available() else cpu) model load_model().to(device) model.eval() for batch in iterator: with torch.no_grad(): inputs tokenizer(batch, return_tensorspt).to(device) outputs model(**inputs) yield process(outputs) rdd.mapPartitions(infer_partition)用mapPartitions而不是map的原因很实际模型加载开销大map是逐条处理每条都要重新加载模型性能会崩。mapPartitions一个分区只加载一次分区内所有数据复用这个模型。3.3 批处理与显存控制的平衡推理吞吐的关键在于batch size。单条推理 GPU 利用率极低因为大部分时间花在 kernel launch 和数据搬运上。把多条数据攒成一个 batch 送进模型GPU 的并行能力才能发挥出来。但 batch size 不能无限大受限于显存。这里有个实操技巧在分区内做动态 batching。不要固定 batch size而是根据当前显存余量动态调整。简单做法是设一个保守的初始值比如 16跑起来后用torch.cuda.memory_allocated()监控如果余量充足就翻倍接近上限就减半。另一个技巧是用torch.cuda.amp做混合精度推理。FP16 推理相比 FP32 能省将近一半显存速度也有提升而且对大多数模型的精度影响很小。开启方式很简单with torch.no_grad(), torch.cuda.amp.autocast(): outputs model(**inputs)注意autocast只影响前向计算模型权重还是 FP32。如果想进一步省显存可以把模型也转成 FP16但要小心数值稳定性问题有些模型比如涉及 softmax 的FP16 下容易溢出。4. 从零搭一套可跑的推理管线前面讲的是原理和配置这一节把整个流程串起来给一个可以照着做的完整方案。我按环境准备 → 数据读取 → 推理执行 → 结果写出的顺序来讲每一步都说明为什么这么做。4.1 环境准备与依赖打包环境这块最大的坑是Python 依赖的分发。Spark 集群上每个 Executor 都需要有 PyTorch、transformers 这些库但集群节点不一定都装了。有三种方案第一种是预装到所有节点用 Ansible 之类的工具批量装。优点是启动快缺点是版本管理麻烦升级要动所有节点。第二种是用--archives打包虚拟环境把 venv 打成 tar.gz 上传Executor 解压后用它。优点是隔离性好缺点是包大、上传慢。第三种是用 Docker 镜像把依赖全打进镜像YARN 拉镜像启动容器。这是生产环境最推荐的方式一致性和隔离性都最好。我个人的选择是开发调试用第二种生产用第三种。开发阶段用 venv 打包灵活改依赖方便生产用 Docker 镜像一次构建到处运行。如果用 venv 方案打包命令大致是python -m venv infer_env source infer_env/bin/activate pip install torch transformers venv-pack -o infer_env.tar.gz提交时用--archives infer_env.tar.gz#env然后在代码里把env/bin加到sys.path。注意venv-pack要装而且打包出来的环境路径是相对的解压后路径会变所以要用#env指定别名代码里按别名找。4.2 数据分片与并行度设计数据读取这块核心问题是分区数怎么定。分区数决定了 Task 数Task 数决定了并行度。并行度太低GPU 喂不饱太高调度开销和显存压力都上来了。一个经验公式是分区数 GPU 总数 × 每卡并发 Task 数 × 1.5。比如你有 8 张卡每卡跑 2 个 Task那分区数设 24 左右比较合适。多出来的 0.5 是为了应对数据倾斜和任务失败重试。数据源如果是 HDFS 上的文件sc.textFile的第二个参数可以直接指定分区数。如果是 Parquet用spark.read.parquet后调repartition(n)。注意repartition会触发 shuffle如果原始分区已经比较均匀用coalesce更省。数据倾斜是另一个要防的问题。如果某些 key 的数据量特别大对应的 Task 会跑很久拖慢整个 Stage。检测方法是看 Spark UI 里各 Task 的耗时分布如果 max 和 median 差好几倍就是倾斜了。解决办法可以是加盐打散或者对大数据量的 key 单独处理。4.3 推理执行的容错与重试分布式推理最怕的是单点失败导致整个任务重跑。Spark 本身有 Task 重试机制默认失败 4 次才放弃。但 GPU 任务的重试有个特殊问题显存泄漏。如果 Task 失败时没有正确释放显存重试的 Task 可能因为显存不足再次失败形成死循环。防范措施是在代码里用try/finally确保显存释放def infer_partition(iterator): model None try: model load_model().cuda() for batch in iterator: yield infer(model, batch) finally: if model is not None: del model torch.cuda.empty_cache()torch.cuda.empty_cache()会释放缓存分配器持有的显存但不释放还被引用的张量。所以del model要在前面确保模型对象被回收。另一个容错点是推理结果的幂等性。如果 Task 重试同一批数据可能被处理两次结果写出时要保证不重复。用DataFrame写出时可以按主键去重或者用overwrite模式覆盖整个分区。我一般建议结果里带上partition_id和task_attempt_id方便排查重复问题。4.4 结果写出的格式选择结果写出看起来简单其实也有讲究。如果结果是向量embedding写 Parquet 比写 JSON 高效得多因为 Parquet 是列式存储向量这种定长数组压缩率高。如果结果是文本标签JSON 或 CSV 都行但要注意转义问题。写出的并行度要和读入匹配否则会出现读的时候 24 个分区写的时候变成 1 个分区的情况写出阶段变成单点瓶颈。检查方法是看 Spark UI 里写出 Stage 的 Task 数如果远小于读入 Stage就要在写出前repartition。还有一个细节是小文件问题。如果每个分区写一个文件24 个分区就是 24 个文件看起来不多。但如果任务每天跑一次一个月就是 720 个文件HDFS 的 NameNode 压力会累积。解决办法是写出后定期做 compaction或者用spark.sql.files.maxRecordsPerFile控制单文件大小。5. 性能调优与踩坑实录配置跑通只是第一步真正让这套管线跑得快、跑得稳需要不少调优。这一节分享几个我实际踩过的坑和对应的解法。5.1 启动开销为什么这么大第一次跑这个任务的人都会惊讶为什么一个简单的推理任务启动要花好几分钟原因在于 YARN 容器启动、Python 环境解压、模型加载这三块叠加。YARN 容器启动本身就要几秒到几十秒取决于集群负载。Python 环境如果是 venv 打包解压一个几百 MB 的 tar.gz 也要时间。模型加载更是大头一个几 GB 的模型从 HDFS 或本地读进来再传到 GPU几十秒很正常。优化方向有三个。第一是预热用 Spark 的spark.executor.pyspark.memory和spark.python.worker.reuse让 Executor 复用避免每个 Task 重新起 Python 进程。第二是模型缓存把模型文件提前放到所有节点的本地磁盘代码里从本地读而不是从 HDFS 拉。第三是减少 Executor 数量、增大单 Executor 资源因为 Executor 启动开销是固定的数量少则总开销小。我实测下来把 Executor 数从 50 个减到 10 个、每个 Executor 从 1 卡 4G 内存增到 4 卡 16G 内存总耗时反而降了 30%。原因是启动开销从 50 份变成 10 份省下来的时间远超单 Executor 处理时间增加的部分。5.2 显存 OOM 的几种典型场景显存 OOM 是 GPU 推理最常见的报错但原因各不相同需要区别对待。第一种是模型本身太大。比如一个 7B 参数的模型FP32 下光权重就要 28GB单张 24GB 的卡根本放不下。这种情况要么用量化INT8 或 INT4要么用模型并行拆到多卡。量化是最简单的bitsandbytes库几行代码就能把模型压到 4-bit显存占用降到原来的四分之一。第二种是batch size 太大。这个好判断报错信息里会显示尝试分配的显存大小。解法就是调小 batch size或者用梯度累积的思路——但推理没有梯度所以就是单纯调小。第三种是显存碎片。长时间运行后显存被切成很多小块虽然总空闲量够但没有连续的大块可用。这种情况用torch.cuda.empty_cache()能缓解但根治要靠PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True这个环境变量让分配器用可扩展段管理显存。第四种最隐蔽多个 Task 共享一张卡时的竞争。如果task.resource.gpu.amount设成 0.5两个 Task 共享一张卡各自加载模型显存占用翻倍。这种情况要么把 task.amount 调回 1要么让模型在 Task 间共享但前面说过共享有设备匹配问题。我的建议是宁可降低并发也不要共享卡跑大模型。5.3 数据倾斜导致的尾部长尾数据倾斜在 GPU 推理里比在普通 Spark 任务里更致命因为 GPU 任务本来就慢一个倾斜的 Task 可能拖慢整个 Stage 几十分钟。排查方法是看 Spark UI 的 Stage 页面把 Task 按 Duration 排序如果最慢的几个 Task 耗时是平均值的 5 倍以上基本可以确定是倾斜。再看这几个 Task 的 Input Size如果也远大于平均那就是数据分布问题。解法分两种。如果倾斜的 key 数量少比如就几个可以单独把这些 key 的数据抽出来用单独的 RDD 处理最后 union 回去。如果倾斜的 key 很多用加盐打散给每个 key 加一个随机后缀让数据分散到多个分区处理完再去掉后缀聚合。加盐的代码大致是import random def add_salt(row): return (row.key _ str(random.randint(0, 9)), row) rdd.map(add_salt).partitionBy(100).mapPartitions(infer).map(remove_salt)盐的粒度这里是 0-9要根据倾斜程度调倾斜越严重盐的范围越大。5.4 任务失败后的排查链路GPU 任务失败时报错信息往往不直观。我总结了一套排查顺序基本能覆盖 90% 的情况。第一步看 YARN 的容器日志。Spark 的报错经常被截断真正的根因在 YARN 的stderr里。用yarn logs -applicationId app_id拉全量日志搜Error和Exception。第二步区分是 OOM 还是其他错误。如果是CUDA out of memory按 5.2 节的方法处理。如果是CUDA error: an illegal memory access通常是代码 bug比如张量形状不匹配或者设备不一致。第三步检查是不是驱动或 CUDA 版本问题。报错里出现driver version或CUDA runtime字样基本就是版本不兼容。对照宿主机驱动版本和镜像里的 CUDA 版本确认兼容性。第四步看是不是资源申请超限。如果 YARN 报Container killed by YARN for exceeding memory limits说明内存申请少了调大spark.executor.memory或spark.executor.memoryOverhead。GPU 任务的内存开销比普通任务大因为 CUDA context 本身就要占几百 MB。第五步复现。如果以上都排查不出来把失败 Task 的输入数据单独抽出来在单机上跑一遍通常能复现问题。这一步虽然笨但最有效。6. 这套方案适合什么、不适合什么聊了这么多实现细节最后回到选型本身。基于 Spark on YARN 的 GPU 推理不是万能方案它有明确的适用边界。适合的场景离线批量推理、模型评测、数据管线中的模型环节。这些场景的共同特点是数据量大、对延迟不敏感、需要与现有大数据链路复用。特别是当你的数据已经在 HDFS 或 Hive 里用 Spark 直接读、直接推理、直接写回省去了数据搬运的成本这个优势是专门的推理服务框架给不了的。不适合的场景在线低延迟服务、小批量高频调用、需要动态扩缩容的交互式场景。这些场景用专门的推理服务框架比如 Triton、TorchServe 这类更合适它们针对单次请求的延迟做了大量优化而 Spark 的启动开销和调度延迟在这种场景下是硬伤。一个折中的思路是分层在线服务用推理服务框架离线批量用 Spark on YARN两者共享同一份模型文件保证结果一致。这样既拿到了在线服务的低延迟又拿到了离线批处理的吞吐。我个人在实际操作中的体会是这套方案最大的价值不在于快而在于省心。一旦管线搭好后续的模型迭代、数据更新、任务调度都能复用现有的大数据基础设施不需要为推理单独维护一套系统。对于已经有 Spark 和 YARN 的团队来说接入 GPU 推理的边际成本很低这是它相比自建推理平台最大的优势。最后分享一个小技巧在正式跑全量数据前先用 1% 的采样数据跑一遍完整流程。这一步能暴露 90% 的配置问题和代码 bug而且成本极低。我见过太多人直接上全量跑了两个小时才发现模型加载路径写错了白白浪费集群资源。采样验证这个习惯能帮你省下大量调试时间。