聊聊 AI 调度算法怎么在 GPU 不够用时把活儿分好一、GPU 不够分调度器得先想清楚怎么排做 AI 集群的都知道最缺的就是 GPU。但不同任务对 GPU 的需求完全不一样训练任务一占就是几天推理任务要的是低延迟开发调试还得随时能插队。这三类任务挤在一起调度器要是没安排好要么集群空转要么大家排队等到崩溃。AI 调度和传统作业调度有两个地方特别不一样能抢回来训练任务支持 checkpoint高优先级的任务来了能把低优先级的先挂起来等有空位再恢复接着跑。资源需求不固定同一个训练任务用 1 张卡跑可能要 10 小时用 8 张卡跑 1.5 小时就够了。调度器得在“用多少卡”和“等多久”之间找平衡。二、调度算法到底怎么分调度算法主要看两个维度资源怎么分公平怎么保。flowchart TB subgraph 资源分配策略 A[优先级调度: 按任务优先级分配] B[公平调度: 按用户/团队均分] C[利用率优先: 填充碎片资源] end subgraph 公平性目标 D[最大最小公平: 优先满足最小需求] E[加权公平: 按权重分配份额] F[主导资源公平: 多资源维度均衡] end subgraph 抢占机制 G[优雅抢占: 保存checkpoint后释放] H[强制抢占: 直接终止低优先级任务] I[部分抢占: 只抢占部分GPU] end A -- D E B -- E F C -- G I优先级调度最直观高优先级的先上。但有个坑低优先级的可能永远轮不上饥饿。解决办法是加个“老化”机制——任务等得越久优先级自动往上提一点。公平调度是保证每个用户或团队都能分到资源。比如 Hadoop YARN 的 Capacity Scheduler给每个团队建个队列保证最低资源量。Kubernetes 的 ResourceQuota 也是类似给每个命名空间设上限防止一个人把资源占光了。主导资源公平DRF是多资源调度里的核心算法。集群里既有 GPU、CPU 还有内存不同任务需求不一样——训练主要吃 GPU数据处理主要吃 CPU 和内存。DRF 会算出每个用户的“主导资源”就是占用比例最高的那个然后在主导资源上实现公平。抢占机制在 AI 调度里特别重要。优雅抢占是靠 checkpoint 保存状态恢复时接着跑不浪费之前的计算。部分抢占更灵活只抢一部分 GPU任务可以缩容继续跑不用完全终止。三、代码怎么实现下面这段代码把优先级调度、DRF 和抢占的核心逻辑串起来了。from dataclasses import dataclass, field from typing import Optional from enum import Enum import time class TaskPriority(Enum): 任务优先级枚举 CRITICAL 0 # 在线推理延迟敏感 HIGH 1 # 交互式开发需要即时响应 NORMAL 2 # 常规训练可等待 LOW 3 # 离线批处理可抢占 class TaskStatus(Enum): 任务状态枚举 PENDING pending RUNNING running PREEMPTED preempted COMPLETED completed dataclass class ResourceRequest: 资源请求多维度资源需求 gpu: int 0 cpu: int 0 memory_gb: float 0.0 dataclass class AITask: AI 调度任务 task_id: str user: str priority: TaskPriority resources: ResourceRequest status: TaskStatus TaskStatus.PENDING submit_time: float field(default_factorytime.time) start_time: Optional[float] None checkpoint_supported: bool True # 是否支持优雅抢占 dataclass class ClusterResources: 集群资源总量 total_gpu: int 32 total_cpu: int 256 total_memory_gb: float 1024.0 used_gpu: int 0 used_cpu: int 0 used_memory_gb: float 0.0 def available(self) - ResourceRequest: return ResourceRequest( gpuself.total_gpu - self.used_gpu, cpuself.total_cpu - self.used_cpu, memory_gbself.total_memory_gb - self.used_memory_gb, ) def can_allocate(self, req: ResourceRequest) - bool: avail self.available() return (avail.gpu req.gpu and avail.cpu req.cpu and avail.memory_gb req.memory_gb) def allocate(self, req: ResourceRequest) - None: self.used_gpu req.gpu self.used_cpu req.cpu self.used_memory_gb req.memory_gb def release(self, req: ResourceRequest) - None: self.used_gpu - req.gpu self.used_cpu - req.cpu self.used_memory_gb - req.memory_gb class DRFScheduler: 主导资源公平调度器 在多资源维度上实现公平分配同时支持优先级和抢占 def __init__(self, cluster: ClusterResources): self.cluster cluster self.pending_tasks: list[AITask] [] self.running_tasks: list[AITask] [] # 每个用户的资源消耗记录 self.user_allocation: dict[str, ResourceRequest] {} def submit(self, task: AITask) - None: 提交任务到等待队列 self.pending_tasks.append(task) def _dominant_share(self, user: str) - float: 计算用户的主导资源份额 主导资源 用户在所有资源维度中占比最高的那个 DRF 在主导资源维度上实现公平分配 if user not in self.user_allocation: return 0.0 alloc self.user_allocation[user] gpu_share alloc.gpu / self.cluster.total_gpu if self.cluster.total_gpu 0 else 0 cpu_share alloc.cpu / self.cluster.total_cpu if self.cluster.total_cpu 0 else 0 mem_share alloc.memory_gb / self.cluster.total_memory_gb if self.cluster.total_memory_gb 0 else 0 return max(gpu_share, cpu_share, mem_share) def _effective_priority(self, task: AITask) - float: 计算任务的有效优先级基础优先级 等待时间老化 等待时间越长有效优先级越高防止低优先级任务饥饿 base task.priority.value # 数值越小优先级越高 wait_seconds time.time() - task.submit_time # 每等待 60 秒优先级提升 0.1等效于优先级降级 aging wait_seconds / 600.0 return base - aging def schedule(self) - list[AITask]: 执行一轮调度从等待队列中选择任务分配资源 综合考虑优先级、DRF 公平性和资源可用性 scheduled [] # 按有效优先级排序数值越小越优先 self.pending_tasks.sort(keyself._effective_priority) still_pending [] for task in self.pending_tasks: if self.cluster.can_allocate(task.resources): # 分配资源 self.cluster.allocate(task.resources) task.status TaskStatus.RUNNING task.start_time time.time() self.running_tasks.append(task) # 更新用户资源消耗记录 if task.user not in self.user_allocation: self.user_allocation[task.user] ResourceRequest() self.user_allocation[task.user].gpu task.resources.gpu self.user_allocation[task.user].cpu task.resources.cpu self.user_allocation[task.user].memory_gb task.resources.memory_gb scheduled.append(task) else: still_pending.append(task) self.pending_tasks still_pending return scheduled def preempt(self, victim_task: AITask) - bool: 抢占运行中的任务释放其占用的资源 优先选择支持 checkpoint 的任务进行优雅抢占 if victim_task not in self.running_tasks: return False # 释放资源 self.cluster.release(victim_task.resources) victim_task.status TaskStatus.PREEMPTED self.running_tasks.remove(victim_task) # 更新用户资源消耗 if victim_task.user in self.user_allocation: self.user_allocation[victim_task.user].gpu - victim_task.resources.gpu self.user_allocation[victim_task.user].cpu - victim_task.resources.cpu self.user_allocation[victim_task.user].memory_gb - victim_task.resources.memory_gb # 被抢占的任务重新加入等待队列 self.pending_tasks.append(victim_task) return True def find_preemption_candidate(self, required: ResourceRequest) - Optional[AITask]: 查找可抢占的任务优先抢占低优先级且支持 checkpoint 的任务 # 按优先级从低到高排序同优先级优先抢占支持 checkpoint 的 candidates sorted( self.running_tasks, keylambda t: (t.priority.value, not t.checkpoint_supported) ) for candidate in candidates: if candidate.resources.gpu required.gpu: # 只抢占 GPU 足够的任务 return candidate return None def get_user_fairness_report(self) - dict[str, dict[str, float]]: 生成用户公平性报告展示各用户的主导资源份额 report {} for user in self.user_allocation: report[user] { dominant_share: round(self._dominant_share(user), 4), gpu_allocated: self.user_allocation[user].gpu, cpu_allocated: self.user_allocation[user].cpu, memory_allocated_gb: round(self.user_allocation[user].memory_gb, 1), } return reportDRFScheduler这个类把主导资源公平的核心逻辑实现了。_dominant_share算每个用户在多资源上的主导份额_effective_priority用老化机制防止饥饿preempt和find_preemption_candidate负责抢占。四、实际用的时候得注意这几个坑DRF 的公平不等于高效。DRF 在多资源维度上算公平但一个 GPU 密集的训练任务和一个 CPU 密集的数据处理任务DRF 会给它们相同的主导资源份额。结果可能是训练任务等了几小时集群其实有空闲资源。如果追求利用率Tetris 这种利用率优先的算法可能更合适。抢占有恢复成本。优雅抢占要保存 checkpoint大模型的 checkpoint 能到几十 GB保存和加载都要时间。频繁抢占的话大量时间花在 I/O 上集群有效利用率反而下降。建议加个抢占冷却期同一个任务短时间内最多被抢占一次。优先级老化得调好参数。老化太快优先级就没区分度了老化太慢低优先级任务还是等很久。建议根据业务 SLA 来定低优先级任务的最大等待时间别超过 SLA 的 50%。适用场景有限制。DRF 适合多用户共享 GPU 集群、需要保证公平的场景。单用户或单团队的集群简单的 FIFO 或利用率优先调度更高效。在线推理服务最好用专用集群别和训练任务混在一起。五、最后说两句AI 调度算法的核心就是在资源不够用时平衡效率、公平和响应时间。DRF 保证多资源维度上的公平优先级老化防止低优先级任务饿死抢占机制让高优先级任务能及时拿到资源。落地的时候多用户共享集群可以用 DRF但得把优先级老化参数调好确保低优先级任务在可接受的时间内能跑起来。抢占优先选支持 checkpoint 的任务减少计算浪费。调度策略得根据实际负载持续调没有一劳永逸的最优解。