如果你跟我一样手里有几台带GPU的服务器每周要跑几十个训练任务那你大概率也经历过这种场景张三在群里吼“谁把我卡占了”李四说“我的模型先跑你往后排”最后谁也说不清集群里到底是哪张卡空闲、哪个任务应该被调度起来。我今年把团队集群里的任务管理统一切到了ax调度才算是把这件事彻底理顺了。ax不是那种一上来就感觉很重的老牌调度系统而是一套以自动化执行为核心的轻量级算力调度框架圈子里通常就叫“ax调度”。它解决的事情很直接把任务提交、排队、资源分配、失败重试这些本应该自动化的环节从手动管理里解放出来。如果你是负责GPU集群运维的同学或者组里几个人共享机器、经常因为资源分配吵架这篇文章值得看完我会把部署、配置、调优、踩坑全部摊开讲。1. ax调度到底是什么整体设计与核心思路1.1 为什么我会盯上ax调度一个真实的集群管理痛点先说背景。我们团队一开始只有两台8卡机器谁要训练就自己上去nvidia-smi看看有空卡就CUDA_VISIBLE_DEVICES0,1 python train.py完事。这种模式在人和任务都少的时候没问题但一旦超过五六个人问题就暴露了有人起了任务忘了关连着跑了好几天有人抢卡导致别人OOM还有人两三个任务轮着跑手动切换特别浪费时间。我最早也想过去上Slurm或者Kubernetes但评估一圈之后放弃了。Slurm功能确实强可是它面向的是超算中心那种动辄上百节点的环境配置复杂我们几个人维护起来不划算。Kubernetes更不用说先要搞明白容器、Pod、CRD、调度器那一套对纯做训练算法的人来说学习成本太高。ax调度恰恰卡在这个位置它能提供排队、资源分配、失败重试这些核心能力又不像前两者那样需要专门的运维团队维护。ax调度底层的思想并不复杂就是把“任务想要什么资源”和“集群现在有什么资源”这两件事分开管理。你提交任务时声明需要几张卡、多少显存、跑什么命令调度器会告诉你在哪个节点上用哪几块GPU并且按你配置的策略决定谁先跑。这种模式让团队里每个人都不需要关心别人的任务怎么排只需要把精力放在模型本身的训练上。1.2 ax的核心设计三件事分开管你可以把ax调度想象成一个“医院挂号系统”任务队列是挂号的人资源池是出诊的科室调度策略决定叫号的顺序。这三件事在ax里是深度解耦的这也是它和很多临时脚本方案最大的区别。任务队列解决的是“任务来了放哪里”的问题。ax允许你配置多个队列比如默认队列、高优队列、开发队列每个队列有独立的并发上限和优先级权重。资源池解决的是“哪些机器可以接任务”的问题。ax会定期和每个节点通信收集GPU使用率、显存占用、卡的状态维护一张实时资源表。调度策略则是真正的“大脑”它决定新任务是立即调度、排队等待还是抢占别人的资源。这套设计的核心好处是可控。你想给某个重要任务插队不用去跟同事口头协商调整队列优先级或者提前抢占策略就行你想限制某个人只能使用特定的两张卡改一下资源池配置就行。所有操作都有日志出了矛盾翻日志就好不用靠回忆。1.3 和其他调度工具的定位差异我在选型的时候做过一张对比表贴出来给大家参考维度SlurmKubernetesax调度定位超算中心作业调度容器编排与调度轻量算力任务调度资源粒度节点/分区/整卡Pod/容器/GPU设备任务/GPU卡/显存切片调度对象批处理作业容器应用训练/推理/脚本任务学习成本高很高低维护成本高高低适合场景大型机房微服务AI混合中小型GPU集群对大多数AI团队来说ax调度的定位恰好是“够用且不折腾”。它不是要替代Slurm和K8s而是填补它们在小团队场景下的空档。如果你以后集群规模真的上去了ax调度也设计了对接外部调度器的能力迁移不会太痛苦。2. 部署与配置实操从零到跑通第一个任务2.1 环境准备与安装ax调度对系统要求很低一台普通的Linux服务器就能当控制节点计算节点只需要能通过网络访问即可。我在Ubuntu 20.04和CentOS 7.6上都跑过没遇到什么兼容性问题。需要的环境大概有这几样Linux操作系统内核版本没有特殊要求Python 3.8以上安装时会自动带上对应的CLI工具计算节点上有NVIDIA驱动和nvidia-smi这是GPU发现的唯一依赖节点间网络互通默认走TCP 8080端口可以自行修改安装过程非常简单我直接用的官方预编译包# 控制节点 wget https://example.com/ax-scheduler/ax-controller-latest.tar.gz tar -xzf ax-controller-latest.tar.gz cd ax-controller python install.py # 计算节点 wget https://example.com/ax-scheduler/ax-agent-latest.tar.gz tar -xzf ax-agent-latest.tar.gz cd ax-agent python install.py安装完先别急着启动要有一个控制节点和至少一个计算节点的基本概念。控制节点负责收集信息和分发任务计算节点负责实际执行。我在一台旧的8卡机器上同时跑了控制端和代理端作为试点验证确认稳定后才扩展到其余机器。启动顺序也有讲究先启动控制节点systemctl start ax-controller systemctl enable ax-controller再启动各计算节点systemctl start ax-agent systemctl enable ax-agent启动之后用ax node list检查节点是否都被纳管。这一步别跳过经常有人装完了发现任务提交不进去第一反应是任务排队配置有问题结果其实是节点没接入成功。2.2 核心配置解析一条一条说清楚ax调度的主配置文件是/etc/ax/config.yaml刚上手的人看到一堆字段容易懵我挑几个真正影响调度的关键项讲。cluster: name: lab-cluster heartbeat_timeout: 60 resources: gpu_nodes: - host: gpu-01 gpus: 8 mem_per_gpu: 80 - host: gpu-02 gpus: 4 mem_per_gpu: 40 queues: - name: default max_running: 4 priority: 5 - name: high max_running: 2 priority: 10 policy: scheduler: fair_share preemption: true task_timeout: 86400 retry_count: 2cluster.heartbeat_timeout是节点心跳超时时间单位秒。每个计算节点默认每15秒上报一次心跳如果控制节点连续超过60秒没收到某个节点的心跳就会把它标记为不可用不再往上面派任务。这个值不要太短我试过设成20秒结果节点稍微忙一点就被误判下线任务反复迁移白白浪费了很多时间。resources.gpu_nodes里声明每台机器有几张卡、每张卡多大显存。这里的mem_per_gpu单位是GB默认值建议按实际显存填比如A100 80G就填80V100 32G就填32。这个信息会用于任务调度时的显存校验填错了最直接的后果就是任务被分配到了显存不够的节点上跑起来直接OOM。queues是队列的核心配置。max_running限制该队列同时运行的任务数priority决定队列间的优先级。我习惯建三个队列high给重要且紧急的训练任务default给日常任务dev给调试用的小任务。调试任务的优先级设最低但并发数可以放宽这样大家跑小实验不用排队也不会影响正式训练。policy.scheduler支持fifo和fair_share两种分别对应先来先服务和公平调度。我在下面第三章详细讲。preemption是抢占开关开启后高优先级任务可以挤掉低优先级任务。这个功能要谨慎我一开始开着结果一个挂起的低优任务反复被抢占重启最后只能关掉改成手动处理。配置改完之后执行ax reload就能重新加载不用重启服务。这个设计很实用团队里调整队列权重是常事频繁重启不现实。2.3 命令行工具与API速览ax调度自带命令行工具日常操作基本靠它完成。我把最常用的几个命令列一下# 提交一个普通训练任务 ax submit --queue default --gpus 2 --mem 64 \ --cmd python train.py --config configs/resnet50.yaml # 提交一个高优先级任务并指定节点 ax submit --queue high --gpus 8 --nodes gpu-01 \ --cmd bash run_distributed.sh # 查看所有任务的排队和运行状态 ax queue # 查看某个任务的详细信息 ax status --job-id ax-2024-0001 # 取消任务 ax cancel --job-id ax-2024-0001 # 查看集群GPU总览 ax resource轴调度每次提交都会返回一个类似于ax-2024-0001的任务ID这是你在集群里唯一定位任务的标识日志、状态查询、取消操作都依赖它。我要求团队所有人在提交脚本里必须记录这个ID最朴素的理由你一次性提交五六个任务没有ID根本分不清哪个跑完哪个没跑。--mem参数很多人不理解它的含义它其实是你期望每张卡分配到的显存上限。调度器会拿它乘以gpus来估算总资源占用量。需要注意这只是调度层的容量预占并不会真正限制任务能用的显存真正限制还得靠CUDA环境变量或者容器限制。除了命令行ax也暴露了一套HTTP API适合嵌入你现有的平台系统。最简单的用法是提交任务时直接POSTcurl -X POST http://localhost:8080/api/v1/jobs \ -H Content-Type: application/json \ -d { queue: high, gpus: 4, mem: 80, cmd: python train.py }我们后来做了一个简单的Web发布页面后端就是调这套API团队成员不再需要摸到服务器上敲命令只需要在页面上传参数就行。3. 调度策略与参数调优让每一张卡都动起来3.1 优先级与公平性策略别让所有人都抢高优调度策略直接影响任务池的吞吐时间和公平性这块我最开始没认真研究导致踩了不少坑。ax调度默认的是fifo也就是先来先服务。这个策略最简单但缺点很明显一个低优任务如果排在前面后面的高优任务必须干等哪怕高优任务非常紧急。所以我切到了fair_share公平调度。它的逻辑很像银行窗口排队每个队列都有一个权重调度时按权重比例决定从哪个队列取任务。举个例子high队列权重10default权重5dev权重1那么平均每分配10个任务大概有6个去high、3个去default、1个去dev。这能有效防止低优队列把高优队列完全堵死。有一个比较隐蔽的参数叫scheduling_interval默认是5秒。它决定了调度器每隔多久做一次资源匹配。如果你提交的任务很小、跑得很快可以把这个值调低到2秒如果你跑的是十几个小时的长训练没必要调低反而应该适当调高让调度器少做无用功。这个参数我调过几次最明显的体感是任务提交后从“等几秒才有反应”变成“几乎瞬间就有反应”但每次调低都以牺牲调度器CPU为代价要根据集群规模实测。配置优先级的时候有个常见误区把重要任务全部塞进高优队列。我见过有团队所有任务都走高优结果高优队列排了上百个任务调度策略直接失效。正确做法是只把真正影响线上或者实验进度的任务放进高优其他任务老老实实在默认队列里排队这样才能让高优队列恢复插队能力。3.2 GPU分配粒度从整卡到显存切片ax调度的资源分配支持整卡和显存切片两种粒度。整卡很好理解一张卡只能分配给一个任务显存切片则是把一张大显存卡分成几份让多个小任务共享。这个功能对小团队特别实用尤其当你只有几张卡但大家都要跑小实验的时候。整卡分配用--gpus 1就行调度器会自动选择有空闲卡的节点。显存切片则需要结合显存上限参数使用# 一张80G的A100分成两个40G的切片 ax submit --queue default --gpus 1 --mem 40 \ --cmd python train.py --batch-size 64注意这里--gpus 1仍然是一张卡但调度器会按40G显存来预留容量也就是说这张卡还能再接一个--mem 40的任务。切片本身要求你的训练脚本能够适配显存限制我一般配合torch.cuda.set_per_process_memory_fraction或者CUDA_VISIBLE_DEVICES来约束进程实际占用的显存否则虽然调度员认为它在40G范围内运行实际却可能把整张卡撑爆。显存切片有个隐藏的坑显存实际上不像内存频繁的申请和释放会产生碎片。哪怕理论上你只用了40G运行一段时间后显存碎片可能导致后续分配失败。这个问题我后面会在常见问题章节详细讲这里先说结论切片任务建议定期重启或者尽量不用切片优先采用整卡分配不确定性更少。3.3 失败重试与断点续训高可用不是玄学训练任务跑十几个小时中间网络抖动、GPU驱动崩溃、机器掉电都有可能发生如果没有失败重试机制你只能自己盯着日志手动重启。ax调度的retry_count参数就是干这个的我一般设置为2也就是允许在同一份资源规格下最多重试两次。但光有重试不够任务重启后从哪里继续训练这才是真正影响效率的问题。我要求所有训练代码里必须加入断点保存逻辑并用环境变量从调度器接收checkpoint目录# 提交脚本里加上这个参数 --env CHECKPOINT_DIR/mnt/checkpoints/$JOB_ID训练脚本里这样读取import os ckpt_dir os.environ[CHECKPOINT_DIR] ckpt_path os.path.join(ckpt_dir, last.pt) if os.path.exists(ckpt_path): model.load_state_dict(torch.load(ckpt_path)[model])这样一旦调度器检测到任务异常退出重试时就能自动加载最近的checkpoint继续跑而不是从头再来。我实际观测下来一个大模型任务在retry_count2的加持下有效训练时间提升了差不多30%左右因为省掉了大量重复计算。checkpoint的存储路径一定要放在共享文件系统上不能放在计算节点本地盘。节点挂了本地文件也就没了。我们用的是NFS和JuiceFS后者并发读写表现更好推荐给需要频繁存checkpoint的团队。4. 常见问题与排查实录这些坑我替你踩过了4.1 队列卡死资源不足时的死锁现象最典型的现象是队列里明明有任务在等但调度器一直不调度任务都处于PENDING状态。我第一次遇到时以为配置写错了折腾了半天才发现是资源碎片导致的隐性死锁。举个例子集群里有两张卡任务A占了卡1任务B占了卡2这时候来了任务C它申请两张卡调度器发现凑不出两张空闲卡于是C只能等待。但问题是A和B也没说什么时候结束C就这么一直等着。这还不算最坏更麻烦的是如果C占据了队列头部后面那些只需要一张卡的D、E、F也全被堵死明明集群里有充足的单卡能力。解决这类问题的思路有几个一是给队列设置max_running上限避免太多任务同时占着资源二是配合preemption抢占机制让高优任务能把低优任务赶下来三是合理设置task_timeout训练任务不可能无限期运行超时释放资源。我后来把默认队列的max_running调到了和GPU卡总量相等然后给所有队列统一设置8小时超时死锁现象基本消失了。4.2 显存碎片与OOM误判显存切片功能上线后团队里开始有人反馈任务明明没跑多久就报OOM而且每次报OOM的都是固定节点。我上去nvidia-smi一看显存芯片总量还有不少空闲但每张卡上都有好几段细碎的残留占用加起来的浪费空间非常可观。这是典型的显存碎片问题这跟操作系统内存碎片是同一个道理任务反复申请、释放后显存空间被切碎成小块新任务需要连续分配较大显存时就失败。我用了几种手段缓解第一在计算节点上配置定时检测发现显存碎片率超过30%时触发一次GPU驱动重置强制清空残留进程。第二要求所有切片任务设置PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True这个参数能有效减少pytorch显存碎片化实测对长时间的Transformer训练效果比较明显。第三如果某个任务报OOM但实际峰值显存远低于申请值先上去查一下节点上是不是有死亡进程残留占用。经常是之前某个任务异常退出后没清理干净。4.3 节点掉线后的任务漂移集群里一台机器因为电源或网络问题突然掉线正在运行的任务怎么办ax调度默认的行为是等心跳超时把节点标记为Lost然后对被影响的任务做一次重新调度。这个机制本身没问题但细节上需要注意任务本身并没有真正退出节点只是失联了所以调度器重启任务前要确保旧进程确实已经被清理。我在实际运维中遇到过反例某台机器网络抖动调度器认为节点挂了把任务重新调度到另一台机器并启动了新进程。结果旧节点网络恢复后旧进程还在跑同一份数据被两个进程同时处理写出了重复的数据集。解决的办法是给任务设置lease_timeout相当于租约机制。每个任务运行前会向调度器申请一个租约节点上的代理进程定期续租。一旦节点失联导致续租失败代理进程就会自动杀掉任务等节点恢复后旧任务已经停止调度器再重启也就不用担心重复执行。这个机制特别适合训练这类需要幂等、可恢复的任务。4.4 排查问题的通用思路我把常用的排查思路整理成了固定的三步首先看ax queue确认任务到底在PENDING、RUNNING还是FAILED状态如果PENDING看ax status输出的等待原因一般会写“资源不足”或者“等待队列额度”如果RUNNING登录对应节点看GPU利用率如果FAILED直接去查任务输出日志ax调度会把每次重试的日志都保留下来。90%的调度问题都能通过这三步定位。剩下10%可能是控制节点自身的故障这类问题我一般直接看控制节点日志tail -f /var/log/ax-controller.log5. 一点实战心得与后续想做的事前面写的这些是我从两台机器十个人用到扩展到八台机器几十个任务之后沉淀下来的经验。调度工具归根结底是为人服务的如果它让你多操心、多填表那就是本末倒置了。ax调度最让我满意的一点是它足够轻团队成员不需要学习一套全新的概念体系就能上手提交任务只需要一条ax submit命令。我自己在实操中最大的感受是调度策略的微调永远比新功能更重要。先把资源池配置准确再把队列权重调合理最后才考虑抢占和迁移这些进阶功能。优先级、超时时间、重试次数这三个参数建议每次只调整一个观察两三天再动下一个不要一次性全部改掉。接下来我打算把ax调度和团队的增量训练平台对接起来让模型上线后可以自动根据业务流量调度训练任务再做一套简单的Web面板把资源趋势和任务甘特图展示出来。等跑通之后再回来写一篇新的分享。如果你也正在折腾GPU集群的任务管理希望这篇能帮你少走一些弯路。