Python multiprocess.spawn传递对象及GPU设备动态复用问题求助

我们可以用multiprocessing.Queue维护可用GPU设备列表，让任务进程从队列中获取空闲设备，执行完成后立即归还设备，实现GPU资源的即时复用，无需等待整组进程结束。

以下是完整实现代码：

import multiprocessing
import subprocess

def worker(device_queue):
    # 从队列获取空闲GPU设备
    device = device_queue.get()
    try:
        # 执行目标进程，沿用原有subprocess命令
        subprocess.run(f'python special_process.py {device}'.split(), check=True)
    finally:
        # 无论任务成功/失败，都归还设备到队列
        device_queue.put(device)

if __name__ == '__main__':
    # 设置spawn启动方式，适配GPU环境（避免fork模式的CUDA资源冲突）
    multiprocessing.set_start_method('spawn')
    
    # 初始化可用GPU队列（设备编号0-7）
    gpu_count = 8
    device_queue = multiprocessing.Queue()
    for dev in range(gpu_count):
        device_queue.put(dev)
    
    # 创建并启动50个任务进程
    total_tasks = 50
    processes = []
    for _ in range(total_tasks):
        p = multiprocessing.Process(target=worker, args=(device_queue,))
        processes.append(p)
        p.start()
    
    # 等待所有进程执行完成
    for p in processes:
        p.join()

核心细节说明 #

• 设备队列机制：队列作为共享资源，初始存入所有GPU编号。每个进程启动时先抢占空闲设备，确保同一时间单GPU仅运行一个进程。
• spawn启动模式：显式指定该模式是因为CUDA等GPU库通常不兼容fork模式，可避免进程间GPU资源初始化冲突。
• 设备安全归还：finally块保证设备无论任务结果如何都会被归还，不会出现资源泄漏导致后续任务无设备可用。
• 动态调度效率：任意任务完成后立即释放GPU，等待中的进程会马上获取该设备启动，彻底解决原分组模式下的等待瓶颈。

如果special_process.py支持直接作为模块导入，还可以优化为直接调用函数，减少子进程开销：

import multiprocessing
from special_process import main  # 假设脚本入口函数为main

def worker(device_queue):
    device = device_queue.get()
    try:
        main(device)  # 直接调用业务逻辑函数
    finally:
        device_queue.put(device)

# 后续初始化队列、创建进程的代码与上述一致

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