Python多进程场景下的GPU资源约束调度实现方案咨询

你写的初步逻辑目前跑不通预期效果，主要有三个硬伤：

• 循环里直接启动全部100个进程，没有任何准入控制，会瞬间挤满系统进程列表，甚至触发系统内存溢出
• 共享的GPU资源列表没有加互斥锁，多个进程同时读写列表会产生竞态，出现多个进程抢占同一块GPU、导致显存溢出的问题
• 完全没有GPU释放逻辑，进程跑完后不会把占用的GPU号放回资源池，后续进程根本拿不到空闲资源

方案1：低改造成本，基于你原有代码逻辑升级 #

这个方案不需要推翻你现有的写法，只需要补全信号量控制、互斥锁、资源归还三个核心逻辑就行，适配度最高：

• 用多进程信号量做并发闸门，初始值等于GPU总数量（4），拿不到信号量的进程会自动阻塞等待，直到有其他进程释放资源
• 加全局互斥锁，保证从共享资源池取、放GPU的操作是原子的，不会出现抢资源的问题
• 进程跑完（哪怕任务异常报错），必须把GPU号放回资源池、释放信号量，实现资源复用

完整可运行代码如下：

import multiprocessing
import time

def worker(proc_num, gpu_pool, semaphore, lock):
    # 阻塞等待GPU准入资格
    semaphore.acquire()
    try:
        # 加锁取空闲GPU
        with lock:
            assigned_gpu = gpu_pool.pop()
        # 替换为你的实际业务逻辑，设置当前进程可见的GPU
        print(f"进程{proc_num} 启动，占用GPU{assigned_gpu}")
        time.sleep(2) # 模拟任务运行时长
    finally:
        # 无论任务成功失败，都要把GPU放回资源池
        with lock:
            gpu_pool.append(assigned_gpu)
        # 释放信号量，允许下一个排队进程启动
        semaphore.release()

if __name__ == "__main__":
    GPU_TOTAL = 4
    # 信号量控制最大并行数=GPU数量
    sem = multiprocessing.Semaphore(GPU_TOTAL)
    # 资源池操作锁
    lock = multiprocessing.Lock()
    manager = multiprocessing.Manager()
    gpu_pool = manager.list(range(GPU_TOTAL))
    jobs = []

    for i in range(100):
        p = multiprocessing.Process(
            target=worker,
            args=(i, gpu_pool, sem, lock)
        )
        jobs.append(p)
        p.start()

    for job in jobs:
        job.join()

    print(f"全部任务执行完成，剩余可用GPU：{list(gpu_pool)}")

方案2：工程化更常用的轻量调度方案 #

这个方案不需要自己手动维护资源池，代码更简洁，稳定性更高：
原理是直接把进程池的最大并行数设为GPU总数量（4）——进程池本身就会控制最多同时跑4个进程，天然匹配4块GPU的承载能力。每个进程只需要根据自己在池中的序号，绑定对应GPU就行，跑完的进程会被池自动回收，给下一个排队的任务腾位置。

参考实现代码：

import multiprocessing
import os
import time

def worker(proc_num):
    # 获取当前进程在池中的序号，映射到对应GPU
    process_id = multiprocessing.current_process()._identity[0]
    assigned_gpu = (process_id - 1) % 4
    # 设置CUDA可见的GPU，深度学习框架会自动识别
    os.environ["CUDA_VISIBLE_DEVICES"] = str(assigned_gpu)
    print(f"进程{proc_num} 运行，绑定GPU{assigned_gpu}")
    time.sleep(2) # 替换为实际业务逻辑

if __name__ == "__main__":
    # 进程池最大并行数=GPU总数，不需要额外写调度逻辑
    pool = multiprocessing.Pool(processes=4)
    pool.map(worker, range(100))
    pool.close()
    pool.join()
    print("全部100个任务执行完成")

• 不要为了图方便跳过等待直接一次性启动所有进程，进程创建本身会占用不小的系统内存，100个进程哪怕还没开始跑GPU任务，也可能先把系统内存占满
• 资源池的取、放操作必须加锁，没有互斥锁的情况下，多个进程同时执行pop操作，会拿到同一个GPU号，直接导致单卡显存被多个进程打满崩溃
• 一定要把GPU归还的逻辑放在finally代码块里，哪怕任务运行时抛异常退出，也要确保GPU资源被释放，不然会出现资源泄漏，后续任务永远等不到空闲GPU

内容的提问来源于stack exchange，提问作者william007