1. 从GPU排队乱象到自研调度内核ax的诞生动机1.1 最早遇到的那个让人抓狂的问题事情得从半年前说起。我们当时部署了一套面向内部多团队共享的推理服务平台底层GPU就那么几块卡上游算法团队、数据分析团队每天都在往平台上丢任务。最初任务量不大大家排队等等也无所谓等模型越训越多、推理请求越变越密冲突就来了A团队一个显存占用非常大的batch推理任务直接卡住了后面十几个轻量任务而轻量任务里又混着几个时效性极高的实时打标请求。结果就是——重任务没跑完轻任务全堵着等到重任务释放显存轻任务里的实时请求早超时了。最难受的是这个问题不是单纯的排队能解决的。因为推理任务不像离线批处理那样只需先来后到它至少有三个维度要同时考虑显存占用大小、请求的紧急程度、预估的执行时长。有的任务只占2GB显存但需要在10秒内出结果有的任务要占40GB显存但可以等到凌晨再跑。拿一个固定的FIFO队列或者简单的优先级队列去扛永远只能满足其中一部分诉求。于是我在想能不能做一个轻量的调度内核单独负责任务编排这一层它不关心模型具体怎么跑、训练脚本怎么写只关心一件事——在有限的GPU资源上按合理的顺序把合适的任务放到合适的卡上。这个调度内核就是后来我们内部代号叫ax的东西。这个名字没啥玄学含义就是一个随手敲出来的缩写后面大家叫着顺口就留下来了如果你非要解读a就当是asynchronous异步x就当是交叉切换——异步交叉的调度倒也算贴切。1.2 为什么我没直接拿现成的调度器来用说实话接到这个需求的第一反应肯定是去翻开源社区看看有没有现成的框架能直接抄作业。包括Kubernetes的调度器、Nomad、以及专门做AI任务调度的Volcano、KubeFlow我都翻过一轮。最终没有直接拿来用原因也简单我们的场景不是在线集群调度而是单机/小规模多卡环境下的推理任务编排。K8s那套东西是为大规模分布式场景设计的引入它等于把整个部署架构全部推倒重来成本不可接受。我们的任务模型很轻不是什么Pod、容器、镜像那一套就是一段Python函数 显存需求 预估时长 优先级。拿容器调度的抽象去套就像拿货车去送外卖能到但每单都亏。需要抢占与退避策略非常场景化比如实时打标请求到来时可以临时挂起正在跑的长任务这种推理场景特有的调度语义通用调度器默认不会给你。于是决定自研一个极小内核。设计目标就三条零外部依赖内部纯Python实现、可嵌入到现有服务进程里、调度策略可以像写业务代码一样自由定制。ax这个项目就是干这个的。1.3 ax的定位和边界不是什么都能干的大调度写代码之前我先给ax划了个清晰的边界。它不负责GPU显存的精确分配和CUDA context管理那是推理框架和模型自己该管的事它不负责跨机器、跨机房的分布式编排那是K8s的活。它只负责一个点在任务进入执行前决定谁先跑、跑到哪张卡上、跑多长时间该让位。换句话说它做的是一个决策者而不是执行者。这个定位意味着ax可以塞进很多现有的服务里。比如你有一个FastAPI推理服务每个请求进来先丢给ax由ax决定现在是不是执行这个推理的好时机再决定是否把请求转给穿CUDA后端或者你有一个批量处理脚本每天要跑几十个离线任务用ax在内部按优先级别排队等等。我的实践感受是把调度从业务代码里到处散落的队列逻辑抽出来集中成内核之后整个系统的复杂度一下子下降了很多。2. ax的核心设计调度器与执行器分离外加一个租约机制2.1 为什么调度器必须和执行器分家最早版本的ax其实没做调度器/执行器分离。任务直接推到一个全局队列一个后台线程循环从队列取任务、直接执行。看起来简单但写到一半就发现两个绕不开的问题第一如果执行逻辑和调度逻辑在同一个循环里那么执行一个长任务会阻塞调度循环导致调度器无法响应新来的紧急任务。第二GPU卡上正在跑的任务状态没有谁能实时掌握你说它正在跑万一进程崩了、卡死了调度器完全感知不到后续任务可能永远等不到资源释放。所以ax从第二个版本起直接参考了分布式系统里常见的思路Scheduler调度器独立线程/进程维护任务队列、资源台账和所有worker的状态。它只做判断和决策不执行任何具体任务。Executor执行器/Worker每一个GPU卡对应一个workerworker向调度器注册自己的能力比如我是第2号卡可用显存24GB然后循环等待调度器下发任务。这样一个任务从提交到跑完路径是客户端提交给scheduler - scheduler放进队列 - scheduler根据资源台账挑一个合适的worker - 下发任务 - worker拿到任务真正执行 - 回报结果和最新显存占用 - scheduler更新台账。这样做的好处是如果某个worker所在的进程因为模型加载失败而崩了scheduler不会跟着崩它只是在心跳超时后把对应的资源标记为离线再把已经下发但没确认执行的任务重新入队。隔离性一下子就出来了。2.2 租约机制谁能保证任务一定在执行调度器给worker下发任务之后怎么知道worker是真的在执行还是已经死悄悄了最简单粗暴的方式是让worker做心跳上报但心跳只能证明进程还活着并不能证明任务最终成功跑完了。于是我在ax里引入了租约lease的思想所谓租约就是调度器给下发的每个任务盖一个有效时间戳。任务下发时带着一个lease_duration比如说120秒。worker必须在每个lease区间内回报心跳而且必须在任务完成时上报当前租约已完结。如果调度器在租约超时后仍然没有收到任务的完成确认它不会傻等而是直接把任务状态重置为待调度重新入队。这个思想本质上借鉴的是一套lease heartbeat的模型在分布式共识算法、对象存储元数据管理里都很常见。它解决的核心矛盾是在异步环境里接收方调度器无法依赖发送方worker的主动汇报来保证进度必须通过超时机制兜底。具体到代码里ax为每个任务维护了一个截止时间字段调度循环每轮都会扫一遍所有running状态的任务凡是当前时间 lease_deadline的一律判定为疑似失败并重新调度。这样做避免了绝大部分任务卡死但没人管的僵尸问题。2.3 资源和任务怎么描述宁可简单不要复杂在资源描述方面ax只认两个维度卡编号gpu_id和可用显存gpu_mem_mb。没有去做细粒度的SM占用估算、NVLink带宽探测因为这些对推理场景来说通常过度设计。我需要的是一个快速判断这张卡能不能放下我这个任务的布尔结果而不是一个精确到小数点后两位的资源画像。任务描述也很朴素只保留四个关键信息task_id全局唯一标识。gpu_mem_required_mb预估显存占用单位MB由提交方填写。priority优先级0-100数字越大越优先。estimated_duration_s预估执行时长单位秒用于计算租约和调度策略。payload真正的执行负载——通常是一个包含函数引用和参数序列化的字典。如果你想让某个任务必须等到凌晨再执行也没问题提交时在payload里加一个scheduled_at字段调度器在取队列时直接跳过那些未到执行时间的任务即可。ax不内置时间表只提供排序接口具体规则完全由调用方自己定义。这是我刻意做的设计选择——调度器的职责是把机制做好而不是替你决定什么任务该优先。3. 从零手写ax调度循环状态机、排序器和任务分发3.1 任务五态比分布式系统教科书还朴素ax里的任务并不是简单地排队-执行-完成三段式而是维护了一个五态模型PENDING排队中 - READY可调度 - RUNNING执行中 - SUCCEEDED成功 / FAILED失败 / CANCELLED取消。每个状态之间允许的转换非常明确PENDING - READY当任务的所有前置条件满足例如执行时间未到、所依赖的上游任务已完成。READY - RUNNING调度器在某一轮循环中选中该任务并成功下发到worker。RUNNING - SUCCEEDEDworker回报成功。RUNNING - FAILEDworker回报失败或者租约超时被强制重试重试次数耗尽后转为FAILED。任何状态 - CANCELLED调用方主动取消。这个状态机说实话一点都不复杂但它最大的价值是让任务到底卡在哪一步这个问题变得可观测。以前排查问题时要靠猜现在只要打开ax自带的metrics接口输入task_id直接看出状态流转耗时分布立刻知道是排队等了太久还是执行时租约超时。3.2 调度循环每一轮扫描都做了什么ax的调度器主体是一个无限循环线程每隔schedule_interval_ms默认500毫秒跑一轮。每一轮做四件事收集worker心跳更新资源台账标记离线worker。扫描running任务检查租约是否超时超时的重新入队。从pending队列里取READY任务进入候选池。对候选任务按优先级排序逐个匹配可用worker的显存能放下去的就下发放不下去的留在队列里等下一轮。核心其实就这么短短几十行逻辑但是有两个陷阱值得提醒第一个陷阱是资源台账的记性不能太好。当一个任务下发到worker但还没开始加载模型时worker报告的可用显存并不会立刻减少因为模型加载是一个渐进的过程。如果调度器在下发任务的那一刻就把这张卡的显存扣掉预估用量并且再也不管它那这个数字会一直虚高——因为实际加载时显存占用往往比预估要高预扣值不够后续就会有任务被误判为放得下而发出去结果在worker本地OOM。我在第二个版本里改成下发任务时按预估占用预扣但每轮循环根据worker实际报告的空闲显存量反推修正预扣值。两者取一个更保守的值作为可真可用显存。第二个陷阱是不要在一个调度循环里反复下发同一个任务。如果下发任务后worker的确认回包因为网络抖动延迟了等到下一轮循环时这个任务还处于RUNNING状态吗不一定——它可能还在READY状态确认没回来于是调度器可能又把它重新下发一遍。这就导致同一个任务在两台worker上同时跑造成双倍显存消耗。解决办法是引入下发记录表凡是下发过但未确认的任务在确认超时窗口比如1.5个心跳周期内不允许被再次调度。3.3 Worker注册与心跳上报一个小型净室协议每个worker在启动时会拿着自己的gpu_id向调度器注册。注册信息除了gpu_id和总显存外还包括心跳周期heartbeat_interval_ms默认3000毫秒和能力标签capabilities可选项用于标记这张卡是A100还是消费级卡。调度器收到注册后把worker加入到可调度列表并在资源台账中初始化它的状态为ONLINE。心跳上报的路径很简单worker每隔固定时间向调度器发送一个dict里面包含gpu_mem_used_mb、gpu_mem_available_mb、current_task_ids三个字段。调度器收到后更新对应worker的台账。这里有一个我和团队反复调过的参数心跳超时阈值必须设成心跳周期的3倍以上。为什么因为推理任务的显存占用经常会出现瞬时尖峰。举个例子一个目标检测模型在batch16时显存占用可能是5GB但某个批次出现了特别密集的锚框显存瞬间跳到8GB持续几百毫秒后又降回去。如果worker的心跳线程刚好在这个瞬时时点采样到了8GB并上报给调度器而调度器马上把可用显存调低那后续本来可以放下的任务就会被拒之门外。更麻烦的是如果这类高占用采样恰好连续两次出现在心跳中调度器甚至会怀疑worker异常。所以我在worker侧加了一个显存平滑器——记录最近N次采样的滑动窗口上报时取P95值而不是瞬时值同时心跳超时阈值放宽到3个周期。这样既不会漏掉真正的OOM风险又不会因为瞬时抖动而误杀。3.4 失败重试的幂等设计重试得多但不能重试出乱子推理任务不同于离线训练很多失败其实是瞬时故障比如CUDA OOM、模型加载时的一次性报错、GPU驱动临时性的ECC错误。ax默认给每个任务配置了max_retries3并且在重试时采取指数退避第一次重试延迟2秒第二次4秒第三次8秒。这个设计是从网络路由协议里借鉴来的核心目的就一个避免在同一时刻所有失败任务同时重新冲击调度器形成重试风暴。但重试还有一个更隐蔽的坑任务可能根本不是失败而是执行了一部分但没回报结果。比如一个任务在worker上已经跑了90%还没有完成上报调度器就判它超时并重新派发到另一张卡上——那这个任务等于被执行了两遍如果它内部本身不是幂等的比如会往数据库重复插入记录就会造成脏数据。应对措施很简单提交方可以在任务定义里声明idempotentTrue或idempotentFalse。对于非幂等任务ax在租约超时后不会自动重试而是先把任务挂到UNACKNOWLEDGED队列并再次向原worker索要执行结果只有确认worker也联系不上了才把任务标记为FAILED_UNKNOWN_STATE并交给人工处理。这个设计牺牲了一点自动化但换来了数据可靠性的保障——在调度系统里宁可慢一点不能错一点。4. 接入实录把ax塞进现有的推理服务里4.1 接入方式一个装饰器就能解决的绝不动核心代码为了让团队内部改造工作量最小化我做了一个很取巧的设计ax提供的是一个装饰器ax.submit(...)装饰到已有的推理函数上函数本身什么都不用改。调用端只需要改动一行——从直接调函数变成通过ax提交并等待结果。当时给这套接口定的原则是调度逻辑不进模型代码只在服务入口那一层包一层即可。为什么因为模型推理函数往往被各处交叉复用离线脚本、实时接口、数据处理流水线你如果要在函数内部加入向调度器请求资源的逻辑那所有调用方都要跟着变。反过来如果你在调用侧统一接入用一个统一入口封装那改造范围就能控制在最小。下面是我们内部一个通用做法的核心伪代码from ax import SchedulerClient, task client SchedulerClient(127.0.0.1:8765) task(gpu_mem_required_mb8192, priority80, estimated_duration_s30) def run_yolo_inference(img_batch): # 原本的推理逻辑原封不动 return model.predict(img_batch) # 调用时从直接调用改为提交并等待 future client.submit(run_yolo_inference, img_batchbatch) result future.result(timeout60)每次提交时ax会把函数引用序列化后连同资源需求、优先级一起打包进任务定义由调度器决定何时把任务放给某个worker去执行。worker端拿到任务后反序列化、执行、再序列化结果返回。这一套流程我自己用下来最大的感受是对业务代码的侵入几乎为零而调度能力是完整的。4.2 关键参数调优哪些值不能照抄必须结合业务实测接入不是终点调参数才是见功力的时候。下面这几个参数我在不同项目里反复调过坑也比较多直接列成表格供大家参考参数默认值调优建议踩过的坑schedule_interval_ms500任务总量小于100时可以设到200总量大于1000时改到800-1000间隔太小调度线程空转吃CPU太大紧急任务等待时间明显增加heartbeat_interval_ms3000对短任务小于10秒场景建议改为1000对长任务3000够了心跳太密网络包开销大太疏租约判定迟钝lease_durationestimated_duration_s * 1.5 15s一定要给模型加载时间和显存瞬时尖峰留出缓冲建议再乘1.2设太紧长任务被频繁误判为超时并重跑白白浪费资源prefetch_task_num1同一worker上可以预取的任务数建议设2-3可以掩盖调度间隙设太高单个worker上挤入过多任务显存瞬间超出预算priority_aging_interval_s60低优先级任务每等待60秒优先级自动1防止饿死不设这个低优先级任务可能永远轮不到这里重点说一下lease_duration。很多人会把estimated_duration_s直接当租约时长结果是模型加载占用了10秒但任务预估时长只有5秒于是租约在第7.5秒就超时了调度器误判任务失败并重新派发——旧的instance还在占着显存跑新的instance又进来了两张卡一起卡死。我们后来统一规范租约时长 预估执行时长 × 1.5 模型加载缓冲15秒 网络往返预留5秒。这个公式虽然粗糙但比拿预估直接顶稳得多。4.3 实测数据从排队2小时到秒级调度改造完成后的那个星期我做了一组对比测试。场景是同一批200个推理任务包含不同显存需求和优先级分别用旧版固定FIFO队列和新版ax调度各跑一遍结果如下指标FIFO队列ax调度全部任务完成时长43分钟29分钟高优先级任务的平均等待时间18.5分钟2.3秒低优先级任务完成率在相同时间内100%但拖累高优任务94%部分让位给高优任务因显存不足导致的失败次数17次3次GPU整体利用率61%85%我自己最关心的其实是高优先级任务的等待时间这一项——从18.5分钟降到2.3秒这意味着实时打标请求不需要再跟着离线大任务排队了。因为ax通过优先级抢占资源匹配能让轻量高优任务插到一个小显存空闲的GPU卡上先跑而不必等待重型任务的释放。GPU利用率从61%涨到85%并不算夸张但我知道这里面的水分低优先级离线任务本来可以更晚跑没必要全部挤压在同一时段。所以后来我又调整了调度策略给离线任务设置窗口时间——只在晚上批量执行白天把它们挂起。这样白天的卡就留给高优在线任务晚间的卡完全归离线任务用整体利用率反而更健康。5. 避坑记录ax调度落地之后踩过的三个大坑5.1 坑一心跳抖动导致的任务抖动我们第一次把心跳周期调到1秒之后很快发现worker上报的显存数据非常神经质调度器不停地改变对某张卡可用显存的判断导致一些任务的执行卡顿、调度频繁重置。分析之后发现问题出在CUDA的内存分配策略上——cuDNN/cuBLAS会在推理中途动态申请workspace这个操作没法预测。我的解决方案是双管齐下一是前面提到的滑动窗口取P95值二是把worker上报显存分为计算可用显存和原始可用显存两个字段。计算可用显存 原始可用显存 - 预留缓冲20%。调度决策只用计算可用显存这样即使cuDNN临时申请了workspace也大概率落在预留缓冲份额内不会突破调度器的判断红线。这个设计的本质是显存管理不能拿理论值做预算必须给真实运行留出安全垫。5.2 坑二优先级队列的饿死问题刚开始实现时我天真地以为只要priority数字大就先跑就行了。结果跑了半天发现一批低优先级的离线任务从早到晚一个都没执行全部被持续涌入的高优实时任务压到队尾。优先级是绝对数值时高优任务不断插队低优任务永远轮不到——这就是经典的饥饿问题。解决办法是引入**优先级老化priority aging**机制任务每在队列中等待60秒priority就自动加1直到达到调度器上限。这样低优先级任务不会永远垫底只是晚一点、再晚一点但一定轮得到。实际观察到的效果是离线任务最迟等15-20分钟就能开始跑而高优任务的等待时间因为老化抬升并没有出现不可控的增长——因为它们的数量本身有限不太可能永远占用调度器。这个模式在任务类型混杂的场景里属于必备手段。5.3 坑三重试风暴差点把GPU搞崩有一次我们排查一个批量推理任务的大规模失败发现原因出在数据文件本身损坏。但ax不知道这一点它看到一个任务失败就按2秒、4秒、8秒的退避策略重试——200个任务齐刷刷重试等于每隔几秒就有几十个prompt重新加载模型、申请显存、跑推理、失败、再重试。GPU资源瞬间被重试任务占满正常任务全部被挤到队尾最终引发连锁超时。后来我们给ax加了一个重试熔断器如果一个worker上同一个任务的连续失败次数达到max_retries调度器会在30秒内不再接收该worker上的任何任务同时将该任务标记为FAILED_NEEDS_ATTENTION交给人工检查数据文件或模型权重。重试只是策略不是信仰——在异常环境下要允许系统停下来等一下而不是条件反射地疯狂重试。6. ax后续的演进思路几个值得探索的方向ax做到现在这个状态已经稳定运行在我们内部的多套推理服务里。但说实话它离一个完整的调度系统还有不少距离。我现在正在琢磨的几个方向一是感知真实显存尖峰。当前是根据历史P95值来预判但不同模型在不同输入下的显存曲线其实有规律可循。如果能做一个显存画像模块——针对每个模型离线跑一批样例输入统计出它的显存使用分布调度时直接使用该分布预估显存而不是用一个静态数字——调度精度会提升一个台阶。二是多卡协同调度。现在的worker绑定单卡但有些大模型的张量并行需要同时占用2-4张卡。这需要ax把卡这一层抽象成资源组调度时以组为单位匹配。目前正在实现中核心难点是组内卡之间的通信带宽预测——不是所有卡都在同一台机器上时跨机通信成本要比单机内高很多。三是调度策略的可观测性。调度这事实在太容易黑盒化了如果某个任务被不同规则反复抢占你很难手工复盘到底是哪条规则发挥了作用。我打算给ax加一个简单的调度决策日志每次下发任务时记录一句为什么选它可能是优先级最高可能是显存匹配也可能是老化机制抬上来的。这样后续做策略优化就有据可依而不是全靠拍脑袋。这些方向暂时还在验证未必每条都值得做成通用能力。但ax这个项目从最开始一个解决GPU排队乱象的小工具能一步步走到内部稳定运行、被多个团队共同使用我觉得最有价值的不是调度器本身而是我总结出的这条思路在一个看似笨拙的排队问题上只要你愿意抽取出调度决策的独立层很多原本纠缠不清的运维难题都会自然解开。如果你们团队也在被任务互相踩脚困扰不妨也试试这种轻量自研的路子从最朴素的优先级资源匹配开始一步步迭代出真正适合自己业务的调度内核。