从零构建GPU Job Scheduler：轻量级共享显卡调度实战

1. 从一张显卡到一支调度器GPU Job Scheduler 到底在解决什么问题手里有一块 RTX 4060 Laptop GPU或者实验室里摆着几台插了 A100、昇腾 910B 的服务器最朴素的想法通常是写个 PyTorch 脚本python train.py一跑显卡就开始干活了。一个人用没问题但当三个人、五个任务同时想用这块卡的时候事情就变得微妙起来——有人跑大模型微调有人跑 ComfyUI 出图有人只是想验证一下 Paddle GPU 是否安装成功结果显存直接爆掉CUDA out of memory一屏一屏地刷谁也跑不完。这就是GPU Job Scheduler存在的意义。它不是什么高深的东西本质就是一个“排队叫号 资源分配”的中间层把用户提交的任务放进队列按照某种策略决定谁先用卡、用几张卡、用多久任务结束后释放资源给下一个。听起来像操作系统里的进程调度只不过调度的对象从 CPU 时间片换成了 GPU 显存和算力。我最初接触这类需求是在一个混合环境里一台机器上既有 Intel UHD Graphics 核显又有 NVIDIA GeForce RTX 4060 Laptop GPU。核显负责显示输出独显负责计算但总有人不小心把任务跑到了核显上或者两个任务同时抢独显。后来任务多了干脆自己写了一个轻量级的调度器跑了大半年踩了不少坑也总结出一些门道。这篇文章就把这套东西从头到尾拆开讲清楚为什么需要它、核心设计怎么选、代码怎么写、遇到问题怎么排查。适合谁看如果你手头有 GPU 服务器、实验室共享显卡、或者只是想在自己的笔记本上更优雅地管理多个 GPU 任务这篇内容都能直接参考。不需要你是 CUDA 专家但至少要能看懂 Python 和基本的命令行操作。2. 调度器的整体设计与核心思路拆解2.1 为什么不用现成的 Slurm 或 K8s而要自己写一个市面上成熟的调度系统不少Slurm 是超算中心的老牌选手Kubernetes 配合 device plugin 也能做 GPU 调度还有各种云原生方案。但这些东西有一个共同特点——重。Slurm 要配置 slurmctld、slurmd、munge 认证K8s 要搭集群、装 CNI、配 GPU Operator对于一台机器或者一个小型实验室来说属于杀鸡用牛刀。我当时的场景很简单一台工作站两块 GPU五六个人用任务类型主要是 PyTorch 训练、推理服务和偶尔的 ComfyUI 出图。需求就三条第一任务不能互相抢显存第二排队要公平不能一个人霸占一整天第三出问题了能快速看到日志。用 Slurm 当然能实现但维护成本太高而且很多人不熟悉sbatch的语法学习成本也不低。自己写一个调度器的好处在于完全贴合自己的使用习惯想加什么功能就加什么出问题也好排查。坏处也很明显稳定性要靠自己保证边界情况要自己处理。我的建议是如果你的集群超过 10 个节点或者任务类型非常复杂还是老老实实用成熟方案但如果只是小规模共享自己写一个几百行的调度器完全够用而且能学到很多东西。2.2 核心架构队列、资源池、执行器三件套一个最简可用的 GPU Job Scheduler核心就三个模块任务队列存放待执行的任务每个任务包含命令、所需 GPU 数量、显存需求、优先级、提交者等信息。资源池记录当前所有 GPU 的状态包括型号、总显存、已用显存、当前运行的任务列表。执行器从队列里取出任务检查资源是否满足满足就启动进程不满足就等待或者放回队列。这个架构和操作系统的进程调度非常像。队列相当于就绪队列资源池相当于内存管理执行器相当于调度器本身。区别在于GPU 的“时间片”概念很弱任务一旦开始通常要跑完才释放所以调度策略更偏向“先来先服务 优先级抢占”的混合模式。我选择用 Python 来实现原因有三第一GPU 相关的库pynvml、torch.cuda都是 Python 生态最完善第二写起来快改起来也快第三用户提交任务的接口可以用简单的 HTTP 或者命令行不需要额外学一门语言。如果你追求极致性能用 Go 或 Rust 也可以但对于调度这种低频操作秒级Python 完全够用。2.3 调度策略的选择FIFO、优先级还是公平分享调度策略是调度器的灵魂。我试过三种FIFO先来先服务最简单谁先提交谁先跑。问题是如果有人提交了一个要跑 12 小时的大任务后面的人就得干等。对于共享环境来说体验很差。优先级调度给每个任务打一个优先级高优先级的先跑。问题是优先级谁来定如果用户自己填那所有人都会填最高优先级等于没有优先级。如果管理员定那管理员就成了瓶颈。公平分享Fair Share是我最终采用的方案。核心思想是记录每个用户最近一段时间内占用的 GPU 时长占用少的人优先。这样既能防止一个人霸占资源又不需要人工干预。实现上可以用一个简单的滑动窗口比如统计过去 24 小时内每个用户的 GPU 使用时长调度时优先选择使用时长最短的用户的任务。具体实现时我用了一个加权分数score 用户历史使用时长 / 用户权重分数越低优先级越高。权重可以按需分配比如管理员权重高一些普通用户权重相同。这个方案跑了大半年基本没有出现过资源争抢的投诉。3. 核心细节解析与实操要点3.1 GPU 状态监控pynvml 的正确用法与坑点要调度 GPU首先得知道 GPU 当前的状态。NVIDIA 提供了 NVMLNVIDIA Management LibraryPython 绑定是pynvml。基本用法很简单import pynvml pynvml.nvmlInit() device_count pynvml.nvmlDeviceGetCount() for i in range(device_count): handle pynvml.nvmlDeviceGetHandleByIndex(i) mem_info pynvml.nvmlDeviceGetMemoryInfo(handle) util pynvml.nvmlDeviceGetUtilizationRates(handle) print(fGPU {i}: {mem_info.used / 1024**2:.0f}MB / {mem_info.total / 1024**2:.0f}MB, util {util.gpu}%) pynvml.nvmlShutdown()看起来很简单但有几个坑我踩过坑一nvmlDeviceGetMemoryInfo返回的 used 显存包含其他进程。如果你想知道某个任务用了多少显存不能只看总量还要结合进程列表。nvmlDeviceGetComputeRunningProcesses可以拿到每个进程的显存占用但这个接口在某些驱动版本上不稳定偶尔会返回空列表。我的做法是调度器自己记录每个任务启动时的显存基线任务结束后再查一次差值就是该任务的显存占用。坑二核显和独显的区分。在同时有 Intel UHD Graphics 和 NVIDIA RTX 4060 的机器上pynvml只能看到 NVIDIA 的卡Intel 核显不在 NVML 的管理范围内。这其实是好事因为核显本来就不适合做计算。但如果你用的是 AMD 或者昇腾的卡就需要换对应的监控库比如rocm-smi或npu-smi。坑三驱动版本兼容性。我遇到过xid 79: GPU has fallen off the bus这种错误通常是因为驱动崩溃或者硬件接触不良。调度器需要能检测到这种状态并把对应的 GPU 标记为不可用避免继续往上面派任务。检测方法很简单如果nvmlDeviceGetMemoryInfo抛异常就认为该卡异常。3.2 任务提交接口设计命令行还是 HTTP任务提交接口决定了用户的使用体验。我试过两种方案命令行方案用户通过gpu-submit python train.py提交任务调度器返回一个任务 ID用户可以用gpu-status id查看状态。优点是简单直接缺点是跨机器使用不方便而且不容易做权限控制。HTTP 方案调度器跑一个 Web 服务用户通过 POST 请求提交任务GET 请求查看状态。优点是跨平台、易集成缺点是要多写一些代码而且要考虑并发安全。我最终选择了 HTTP 方案因为实验室里有人用 Windows、有人用 macOS、有人用 Linux命令行方案在 Windows 上体验很差。HTTP 接口用 Flask 写核心就两个路由from flask import Flask, request, jsonify app Flask(__name__) app.route(/submit, methods[POST]) def submit(): data request.json cmd data[cmd] gpus data.get(gpus, 1) user data[user] task_id scheduler.submit(cmd, gpus, user) return jsonify({task_id: task_id}) app.route(/status/task_id, methods[GET]) def status(task_id): return jsonify(scheduler.get_status(task_id))这里有一个细节任务提交后调度器不能阻塞等待任务完成而是立即返回任务 ID后台异步执行。我用了一个线程池来跑任务每个任务一个线程线程里用subprocess.Popen启动用户命令。3.3 显存预估与资源匹配怎么判断一个任务能不能跑调度器最核心的逻辑是判断当前 GPU 是否有足够资源运行新任务。这里的关键是“显存预估”。用户提交任务时可以指定需要的显存大小比如--mem 8000表示需要 8GB。如果用户不指定调度器就要自己估算。估算方法有三种方法一按 GPU 数量粗估。比如一张 8GB 的卡默认认为每个任务最多用 6GB留 2GB 给系统和其他进程。这种方法简单但不够精确容易浪费资源。方法二让用户指定。用户最了解自己的任务需要多少显存提交时填一个值。问题是很多用户不知道自己的任务需要多少显存填少了会 OOM填多了会浪费。方法三历史学习。调度器记录每个用户、每种任务类型的历史显存占用下次提交时自动推荐一个值。这个方法最智能但实现复杂而且冷启动阶段没有数据。我采用的是“方法二为主方法一兜底”的策略用户提交时可以指定显存需求如果不指定就按 GPU 总显存的 80% 来分配。同时调度器会记录实际使用量如果实际使用量远小于申请量会在日志里提醒用户下次可以少申请一点。资源匹配的伪代码大概是这样def find_available_gpu(required_mem): for gpu in gpu_pool: if gpu.status ! available: continue if gpu.free_mem required_mem: return gpu return None如果找不到可用的 GPU任务就留在队列里等待下一次调度。调度循环每隔几秒跑一次检查队列里是否有任务可以执行。4. 实操过程与核心环节实现4.1 环境准备从零搭建一个可用的调度器假设你有一台 Ubuntu 22.04 的机器装好了 NVIDIA 驱动和 CUDAPython 版本 3.10。第一步是安装依赖pip install pynvml flask如果你要用 PyTorch 做任务还需要装 PyTorch 的 GPU 版本。这里有一个常见的坑PyTorch 安装教程里通常会让你去官网复制命令但不同 CUDA 版本对应的命令不一样。我的建议是先用nvidia-smi看一下驱动支持的 CUDA 版本然后去 PyTorch 官网找对应的安装命令。比如 CUDA 12.1 对应的命令是pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121装完之后用一个小脚本验证 GPU 是否可用import torch print(torch.cuda.is_available()) print(torch.cuda.device_count()) print(torch.cuda.get_device_name(0))如果输出True和显卡型号说明环境没问题。如果输出False先检查驱动再检查 PyTorch 版本是否匹配。4.2 调度器核心代码从任务提交到执行完成下面是一个简化版的调度器核心代码去掉了 HTTP 部分只保留调度逻辑import threading import subprocess import time import pynvml from queue import PriorityQueue class GPUScheduler: def __init__(self): pynvml.nvmlInit() self.device_count pynvml.nvmlDeviceGetCount() self.gpu_pool [] for i in range(self.device_count): handle pynvml.nvmlDeviceGetHandleByIndex(i) mem_info pynvml.nvmlDeviceGetMemoryInfo(handle) self.gpu_pool.append({ index: i, handle: handle, total_mem: mem_info.total, free_mem: mem_info.total, status: available, task_id: None }) self.task_queue PriorityQueue() self.tasks {} self.lock threading.Lock() self.task_counter 0 def submit(self, cmd, required_mem, user, priority0): with self.lock: self.task_counter 1 task_id self.task_counter self.tasks[task_id] { id: task_id, cmd: cmd, required_mem: required_mem, user: user, status: queued, gpu_index: None, process: None } self.task_queue.put((priority, task_id)) return task_id def find_gpu(self, required_mem): for gpu in self.gpu_pool: if gpu[status] available and gpu[free_mem] required_mem: return gpu return None def schedule_loop(self): while True: if not self.task_queue.empty(): priority, task_id self.task_queue.get() task self.tasks[task_id] gpu self.find_gpu(task[required_mem]) if gpu: gpu[status] busy gpu[free_mem] - task[required_mem] gpu[task_id] task_id task[status] running task[gpu_index] gpu[index] thread threading.Thread(targetself.run_task, args(task, gpu)) thread.start() else: self.task_queue.put((priority, task_id)) time.sleep(2) def run_task(self, task, gpu): env {CUDA_VISIBLE_DEVICES: str(gpu[index])} process subprocess.Popen(task[cmd], shellTrue, envenv) task[process] process process.wait() with self.lock: task[status] finished gpu[status] available gpu[free_mem] task[required_mem] gpu[task_id] None这段代码的核心逻辑是提交任务时放入优先队列调度循环每隔 2 秒检查一次队列找到合适的 GPU 就启动任务任务结束后释放资源。CUDA_VISIBLE_DEVICES环境变量是关键它让任务只能看到分配给它的那块 GPU避免误用其他卡。4.3 参数计算与选择显存、优先级、超时时间怎么定调度器有几个关键参数需要根据实际情况调整显存预留量我通常留 10% 的显存给系统和其他进程。比如 8GB 的卡实际可分配 7.2GB。如果卡上还跑着显示输出比如你用它接显示器预留量要提高到 20%。调度间隔2 秒是一个比较平衡的值。太短会增加 CPU 开销太长会让任务等待时间变长。如果任务提交频率很高可以缩短到 1 秒如果任务很少可以延长到 5 秒。任务超时时间有些任务可能会卡死比如死循环或者等待输入。调度器应该支持超时自动杀死任务。我设置的是 24 小时超过就发邮件提醒用户再给 1 小时宽限期然后强制结束。优先级范围我用了 0 到 10 的整数0 最高10 最低。普通用户默认 5管理员可以指定 0 到 3。实际使用中大部分人用默认值就够了只有紧急任务才会调高优先级。5. 常见问题与排查技巧实录5.1 任务启动失败从 CUDA 版本到权限问题任务启动失败是最常见的问题原因五花八门。我整理了一个排查表现象可能原因排查方法解决方案CUDA out of memory显存不足检查nvidia-smi显存占用减少 batch size 或等待其他任务结束CUDA driver version is insufficient驱动版本太低nvidia-smi查看驱动版本升级驱动或降级 CUDA 版本no CUDA-capable device is detected环境变量错误检查CUDA_VISIBLE_DEVICES确保调度器正确设置环境变量Permission denied文件权限问题ls -l查看文件权限修改权限或换用户运行任务卡住无输出等待输入或死锁strace -p pid查看系统调用设置超时或检查代码逻辑其中CUDA driver version is insufficient这个错误我遇到过好几次。原因是 PyTorch 的 CUDA 版本和驱动版本不匹配。比如驱动只支持 CUDA 11.8但你装了 CUDA 12.1 的 PyTorch就会报这个错。解决方法是用nvidia-smi看驱动支持的 CUDA 版本然后装对应版本的 PyTorch。5.2 显存泄漏与僵尸进程调度器必须处理的边界情况显存泄漏是共享环境的大敌。有些任务结束后进程虽然退出了但显存没有释放导致后续任务无法启动。这种情况通常是因为程序没有正确调用torch.cuda.empty_cache()或者有子进程还在运行。调度器的处理方法是任务结束后检查该 GPU 上的进程列表如果还有残留进程就强制杀死。pynvml.nvmlDeviceGetComputeRunningProcesses可以拿到进程 PID然后用os.kill发送 SIGKILL。僵尸进程是另一个问题。如果用户的任务 fork 了子进程而子进程没有正确退出就会变成僵尸进程。调度器需要定期清理用ps命令找到父进程是调度器的僵尸进程然后wait掉。我踩过的一个坑是调度器本身崩溃后正在运行的任务变成了孤儿进程继续占用 GPU。后来我加了一个启动时的清理逻辑调度器启动时扫描所有 GPU 上的进程如果进程的父进程不是当前调度器就认为它是孤儿进程可以选择杀死或者接管。5.3 多用户环境下的权限与隔离多用户环境下权限和隔离是必须考虑的问题。最简单的做法是调度器以 root 运行任务以提交者的身份运行。这样既能管理 GPU又能保证用户只能访问自己的文件。实现方法是用subprocess.Popen的preexec_fn参数在启动任务前切换用户import pwd import os def demote(user_uid, user_gid): def result(): os.setgid(user_gid) os.setuid(user_uid) return result user_info pwd.getpwnam(task[user]) process subprocess.Popen( task[cmd], shellTrue, preexec_fndemote(user_info.pw_uid, user_info.pw_gid) )这样任务就以普通用户身份运行不能访问其他用户的文件也不能修改系统配置。但有一个问题如果用户的任务需要访问 GPU 设备文件/dev/nvidia*需要确保该用户有权限。通常 NVIDIA 驱动会把这些设备文件的权限设置为 666所以普通用户也能访问。另一个隔离问题是网络。如果多个任务同时运行可能会争抢端口。我的做法是给每个任务分配一个随机端口范围通过环境变量传给任务让任务自己选择端口。5.4 调度器自身的监控与告警调度器本身也需要监控。如果调度器挂了所有任务都会卡住。我用了一个简单的守护脚本每隔 30 秒检查调度器进程是否存活如果不在就重新启动。同时调度器会记录关键事件到日志文件比如任务提交、开始、结束、失败等。告警方面我用了一个简单的邮件通知任务失败时发邮件给提交者GPU 异常时发邮件给管理员。邮件发送用smtplib配置一个 SMTP 服务器即可。如果不想用邮件也可以用 Webhook 推送到聊天工具。日志格式我推荐用 JSON方便后续分析import json import logging logger logging.getLogger(scheduler) handler logging.FileHandler(scheduler.log) handler.setFormatter(logging.Formatter(%(message)s)) logger.addHandler(handler) def log_event(event_type, **kwargs): logger.info(json.dumps({event: event_type, **kwargs}))这样每条日志都是一行 JSON可以用jq或者 Python 脚本快速分析。比如统计每个用户的任务数量cat scheduler.log | jq -r select(.eventtask_finished) | .user | sort | uniq -c6. 从单机到多机调度器的扩展思路单机调度器跑通之后如果 GPU 不够用自然会想到扩展到多机。多机调度的核心问题是如何让多个节点上的调度器协同工作。最简单的方案是选一个主节点其他节点作为工作节点主节点负责任务分配工作节点负责执行。主节点维护一个全局的任务队列和资源池工作节点定期向主节点汇报 GPU 状态。任务提交到主节点后主节点根据各节点的资源情况选择合适的工作节点然后把任务下发过去。工作节点执行任务把结果和日志回传给主节点。这个架构的难点在于网络通信和故障处理。如果工作节点失联主节点需要把该节点上的任务重新调度到其他节点。如果主节点挂了整个系统就瘫痪了。所以生产环境通常需要主节点高可用比如用 etcd 或 ZooKeeper 做选主。不过对于大多数小规模场景单机调度器已经够用了。如果确实需要多机我建议先用 SSH 做简单的远程执行不要一上来就搞复杂的分布式架构。比如主节点通过 SSH 在工作节点上启动任务用nvidia-smi远程查询 GPU 状态。这种方法简单粗暴但很有效。ssh worker1 nvidia-smi --query-gpumemory.used,memory.total --formatcsv解析这个输出就能得到远程节点的 GPU 状态。任务执行也可以用 SSHssh worker1 CUDA_VISIBLE_DEVICES0 python train.py当然这种方式没有考虑并发和故障恢复只适合临时使用。如果要长期运行还是建议用成熟的分布式调度框架或者自己实现一套基于消息队列的架构。7. 一些实操心得与避坑建议调度器跑了大半年最大的体会是简单比智能更重要。我一开始想做一个“智能调度器”能自动预测任务运行时间、自动调整优先级、自动扩缩容。结果发现预测不准、调整不及时、扩缩容出问题反而增加了维护成本。后来回归简单FIFO 加公平分享用户自己指定显存需求调度器只做最基本的资源匹配。这样虽然不够“智能”但稳定可靠用户也容易理解。另一个心得是日志要详细但不要太多。日志太少出问题没法排查日志太多磁盘很快就满了。我的做法是关键事件任务提交、开始、结束、失败必须记录调试信息只在开发阶段开启生产环境关闭。日志文件按天切割保留 30 天。还有一个坑是不要相信用户的输入。用户提交任务时可能会填错显存需求、写错命令、甚至提交恶意命令。调度器需要做基本的校验命令不能为空显存需求不能超过 GPU 总显存用户不能提交超过配额的任务。这些校验虽然简单但能避免很多问题。最后分享一个小技巧给调度器加一个“干跑”模式。用户提交任务时可以选择“干跑”调度器只检查资源是否满足不实际执行任务。这样用户可以在提交前确认自己的任务能不能跑避免提交后才发现资源不够。实现方法很简单在run_task里加一个判断如果是干跑模式就直接返回成功不启动进程。这个调度器后来还扩展了一些功能比如支持任务依赖任务 B 依赖任务 A 完成、支持定时任务每天凌晨跑一次、支持任务模板常用命令保存为模板一键提交。这些功能都不复杂但确实提升了使用体验。如果你也在做类似的东西建议先从最核心的调度逻辑开始跑通了再逐步加功能不要一开始就追求大而全。