Python GPU 任务调度：CUDA 显存管理和多进程隔离

Python GPU 任务调度CUDA 显存管理和多进程隔离一、模型加载失败——CUDA out of memory的锅多个 AI 服务部署在同一台 GPU 机器上各自用各自的 Python 进程。看似井水不犯河水结果第一个服务加载了 7B 模型占 14GB 显存第二个服务加载 1.5B 模型占 3GB第三个服务一加载就报CUDA out of memory。明明 GPU 有 24GB 显存只是两个服务用了 17GB为什么第三个加载 3GB 就报 OOM答案是 CUDA 的内存碎片化。PyTorch 的显存分配不是先申请一大块按需分发而是用多少申请多少。两个进程虽然总占用只有 17GB但内存碎片加上 PyTorch 缓存机制导致剩余空间不足以连续分配 3GB。二、CUDA 显存管理架构flowchart TB subgraph GPU[GPU 显存24GB HBM] direction TB Frag1[进程 A: 模型参数\n10GB 连续] Cache1[进程 A: PyTorch 缓存\n4GB 预留] Frag2[进程 B: 模型参数\n3GB 连续] Cache2[进程 B: PyTorch 缓存\n1GB 预留] Free1[碎片空间\n2GB已碎片化] Free2[可用空间\n4GB] end subgraph Manager[显存管理器] MCP[多进程隔离 (CUDA MPS)] MCP -- Limit1[进程 A 显存限制: 14GB] MCP -- Limit2[进程 B 显存限制: 5GB] CacheClean[内存清理策略] CacheClean -- Empty[torch.cuda.empty_cache()] CacheClean -- IPC[CUDA IPC 共享] end subgraph Schedule[任务调度] Q[任务队列] -- Alloc[显存检查] Alloc --|够用| Run[执行任务] Alloc --|不够| Wait[等待 清理] Wait -- GC[触发 Python GC\n 清理 CUDA 缓存] GC -- Alloc end三、生产级显存管理代码显存管理器import torch import gc import time import threading from typing import Optional, Callable, Dict from dataclasses import dataclass, field from collections import deque import logging logger logging.getLogger(__name__) dataclass class GPUMemoryBudget: GPU 显存预算——每个进程/任务的显存限额 device: int 0 total_memory_mb: int 0 # 为该任务预留的显存MB peak_memory_mb: int 0 # 历史峰值 current_allocated_mb: int 0 # 当前分配 class GPUMemoryManager: GPU 显存管理器——负责任务级别的显存隔离和调度 def __init__(self, device_id: int 0, reserved_memory_mb: int 1024): device_id: GPU 设备编号 reserved_memory_mb: 为系统和 CUDA 上下文预留的显存 self.device_id device_id self.reserved_mb reserved_memory_mb # 检查 GPU 是否可用 if not torch.cuda.is_available(): raise RuntimeError(CUDA 不可用请检查驱动和 PyTorch 版本) self.total_memory_mb torch.cuda.get_device_properties( device_id ).total_memory // (1024 * 1024) self.available_mb self.total_memory_mb - reserved_memory_mb # 任务显存预算表 self._budgets: Dict[str, GPUMemoryBudget] {} self._lock threading.Lock() logger.info( fGPU {device_id}: 总显存 {self.total_memory_mb}MB, f可用 {self.available_mb}MB (预留 {reserved_memory_mb}MB) ) def register_task( self, task_id: str, memory_budget_mb: int, ) - GPUMemoryBudget: 注册一个任务并分配显存预算 如果剩余显存不足抛出异常 with self._lock: # 检查总预算是否超标 allocated_sum sum( b.total_memory_mb for b in self._budgets.values() ) if allocated_sum memory_budget_mb self.available_mb: raise MemoryError( f任务 {task_id} 需要 {memory_budget_mb}MB f但仅剩 {self.available_mb - allocated_sum}MB ) budget GPUMemoryBudget( deviceself.device_id, total_memory_mbmemory_budget_mb, ) self._budgets[task_id] budget logger.info(f任务 {task_id}: 已分配 {memory_budget_mb}MB 预算) return budget def release_task(self, task_id: str): 释放任务的显存预算并清理 CUDA 缓存 with self._lock: if task_id in self._budgets: del self._budgets[task_id] # 触发 Python GC 和 CUDA 缓存清理 gc.collect() torch.cuda.empty_cache() logger.info(f任务 {task_id}: 已释放显存预算) def get_current_usage(self) - Dict[str, int]: 获取当前显存使用情况 allocated torch.cuda.memory_allocated(self.device_id) // (1024 * 1024) cached torch.cuda.memory_reserved(self.device_id) // (1024 * 1024) free self.total_memory_mb - allocated - cached return { total_mb: self.total_memory_mb, allocated_mb: allocated, cached_mb: cached, free_mb: max(0, free), } def safe_execute( self, task_id: str, func: Callable, *args, max_retries: int 3, **kwargs, ): 在显存安全的环境中执行函数 如果 OOM自动清理缓存并重试 for attempt in range(max_retries): try: return func(*args, **kwargs) except torch.cuda.OutOfMemoryError as e: logger.warning( f任务 {task_id} OOM第 {attempt1}/{max_retries} 次尝试 ) # 清理策略先清 Python GC再清 CUDA 缓存 gc.collect() torch.cuda.empty_cache() # 检查显存状态 usage self.get_current_usage() logger.warning(f当前显存: {usage}) if attempt max_retries - 1: raise RuntimeError( f任务 {task_id} 重试 {max_retries} 次后仍然 OOM ) from e # 递增等待给其他释放显存的机会 time.sleep(2 ** attempt)多进程 GPU 隔离方案import multiprocessing as mp from contextlib import contextmanager class GPUProcessPool: GPU 多进程池——每个进程独立管理自己的显存 def __init__(self, gpu_id: int, num_workers: int 2): self.gpu_id gpu_id self.num_workers num_workers self._pool: Optional[mp.Pool] None # 每个 worker 的显存限额 total_mem torch.cuda.get_device_properties(gpu_id).total_memory self.worker_memory_limit ( total_mem // num_workers * 80 // 100 # 每个 worker 用 80% 份额 ) staticmethod def _worker_init(gpu_id: int, memory_limit: int, worker_id: int): 每个 worker 进程的初始化函数 ——在 fork 之后、执行任务之前调用 import os os.environ[CUDA_VISIBLE_DEVICES] str(gpu_id) # 设置 PyTorch 显存限制 torch.cuda.set_per_process_memory_fraction( memory_limit / torch.cuda.get_device_properties(gpu_id).total_memory ) # 设置 PyTorch 显存分配策略 # expandable_segmentsTrue 允许内存段动态扩展减少碎片 os.environ[PYTORCH_CUDA_ALLOC_CONF] expandable_segments:True logger.info(fWorker {worker_id} 初始化完成, 显存限制: {memory_limit} 字节) def start(self): 启动多进程池 self._pool mp.Pool( processesself.num_workers, initializerself._worker_init, initargs(self.gpu_id, self.worker_memory_limit, 0), ) def submit(self, func, *args, **kwargs): 提交任务到进程池 if self._pool is None: raise RuntimeError(请先调用 start()) return self._pool.apply_async(func, args, kwargs)模型加载时的显存优化def load_model_with_memory_budget( model_path: str, budget_mb: int, device: str cuda:0, ) - torch.nn.Module: 在显存预算限制下加载模型 ——如果超出预算自动尝试量化 # 初步估算FP16 模型大约每 1B 参数占用 2GB # 如果预算不够尝试 INT8/INT4 量化 file_size_mb __import__(os).path.getsize(model_path) // (1024 * 1024) if file_size_mb budget_mb * 1.5: # 模型文件超过预算的 1.5 倍必须量化 logger.info(f模型大小 {file_size_mb}MB 超过预算 {budget_mb}MB使用 INT4 量化) from transformers import BitsAndBytesConfig quant_config BitsAndBytesConfig( load_in_4bitTrue, bnb_4bit_compute_dtypetorch.float16, bnb_4bit_use_double_quantTrue, ) model load_model(model_path, device_mapauto, quantization_configquant_config) else: # 预算充足直接加载 torch.cuda.set_per_process_memory_fraction( budget_mb / torch.cuda.get_device_properties(0).total_memory * (1024 * 1024) ) model load_model(model_path) model model.to(device) return model四、边界分析与 Trade-offs多进程 vs MPSMulti-Process Service多进程简单易用隔离性强但有进程间通信开销CUDA MPS共享 GPU 上下文减少碎片但单进程崩溃可能影响其他进程建议推理服务用 MPS训练任务用多进程隔离显存碎片化PyTorch 的expandable_segments可以缓解碎片问题但有轻微性能开销。对于需要频繁分配释放显存的场景如动态 batch size建议开启。任务优先级简单 FIFO 调度可能不够。建议参考 Kubernetes 的 QoS 模型Guaranteed / Burstable / BestEffort设计优先级调度。GPU 共享的经济性一张 A10080GB按 5 个进程平分每个 16GB。如果某个进程实际只用 8GB剩余的 8GB 浪费了。可以使用 MIGMulti-Instance GPU或 MPS 提高利用率。五、总结Python GPU 任务调度的核心挑战是显存碎片化和多进程隔离显存预算制每个任务注册时声明显存需求管理器全局调度多进程隔离每个进程独立管理显存避免相互影响OOM 自愈捕获 OOM → 清理缓存 → 递减重试量化降级模型超过显存预算时自动尝试量化加载一个实用的显存监控命令nvidia-smi --query-gpumemory.used,memory.free --formatcsv -l 1持续观察显存波动。