子类化mp.Process实现进程级GPU ID时遇递归错误，求解决方法

问题背景 #

需要子类化multiprocessing.Process，为每个进程分配唯一的GPU ID作为进程级常量，用于指定不同GPU运行任务。编写的代码尝试通过元类自动分配GPU ID，但运行时触发RecursionError: maximum recursion depth exceeded，定位到是_Popen方法引发的无限递归。

用户初始代码如下：

import multiprocessing as mp
from typing import Type

# Attempt to override default behaviour of the `Process`
# to add a gpu_id parameter during construction.
class GPUProcessMeta(type):
    def __call__(cls, *args, **kwargs):
        obj = cls.__new__(cls, *args, **kwargs)
        gpu_id = cls.next_id
        if gpu_id in cls.used_ids:
            raise RuntimeError(
                f"Attempt to reserve reserved processor {gpu_id} {cls.used_ids=}"
            )
        cls.next_id += 1
        cls.used_ids.append(gpu_id)
        kwargs["gpu_id"] = gpu_id
        obj.__init__(*args, **kwargs)
        return obj

class GPUProcess(mp.Process, metaclass=GPUProcessMeta):
    used_ids: list[int] = []
    next_id: int = 0

    def __init__(
        self,
        group=None,
        target=None,
        name=None,
        args=(),
        kwargs={},
        *,
        daemon=None,
        gpu_id=None,
    ):
        super(GPUProcess, self).__init__(
            group,
            target,
            name,
            args,
            kwargs,
            daemon=daemon,
        )
        self._gpu_id = gpu_id

    @property
    def gpu_id(self):
        return self._gpu_id

    def __del__(self):
        GPUProcess.used_ids.remove(self.gpu_id)

    def __repr__(self) -> str:
        return f"<{type(self)} gpu_id={self.gpu_id} hash={hash(self)}>"

    @classmethod
    def create_gpu_context(cls) -> Type[mp.context.DefaultContext]:
        context = mp.get_context()
        context.Process = cls
        return context

def test_gpu_pool():    
    ctx = GPUProcess.create_gpu_context()
    with ctx.Pool(2) as pool:
        payload = (tuple(range(3)) for _ in range(10))
        response = pool.starmap(
            _dummy_func,
            payload,
        )
    assert response == ((0, 1, 2),) * 10

def _dummy_func(*args, **kwargs):
    return args, kwargs

if __name__ == "__main__":
    test_gpu_pool()

查看Python标准库源码发现，multiprocessing.Process的_Popen静态方法会委托给当前上下文的Process类实现。当把上下文的Process替换为自定义的GPUProcess后，调用GPUProcess._Popen时会再次触发相同逻辑，形成无限递归。此外，类级别的used_ids和next_id在多进程环境下无法同步，__del__方法执行时机不可靠，也会导致ID管理混乱。

解决方案 #

方案1：直接继承multiprocessing.process.BaseProcess #

绕过标准库的Process包装类，直接继承底层BaseProcess，从根源避免递归问题：

import multiprocessing as mp
from multiprocessing.process import BaseProcess
import itertools

class GPUProcess(BaseProcess):
    # 使用循环迭代器分配固定数量的GPU ID（假设有2个GPU）
    _gpu_ids = itertools.cycle(range(2))

    def __init__(
        self,
        group=None,
        target=None,
        name=None,
        args=(),
        kwargs={},
        *,
        daemon=None,
        gpu_id=None,
    ):
        super().__init__(group, target, name, args, kwargs, daemon=daemon)
        # 自动分配GPU ID，也支持手动指定
        self._gpu_id = gpu_id if gpu_id is not None else next(self._gpu_ids)

    @property
    def gpu_id(self):
        return self._gpu_id

    def run(self):
        # 在子进程启动时设置GPU环境变量
        import os
        os.environ["CUDA_VISIBLE_DEVICES"] = str(self.gpu_id)
        # 执行原run逻辑
        super().run()

    @classmethod
    def create_gpu_context(cls):
        context = mp.get_context()
        context.Process = cls
        return context

def test_gpu_pool():    
    ctx = GPUProcess.create_gpu_context()
    with ctx.Pool(2) as pool:
        payload = (tuple(range(3)) for _ in range(10))
        response = pool.starmap(_dummy_func, payload)
    assert response == ((0, 1, 2),) * 10

def _dummy_func(*args, **kwargs):
    # 在任务函数中获取当前进程的GPU ID
    current_process = mp.current_process()
    print(f"当前进程使用GPU ID: {current_process.gpu_id}")
    return args, kwargs

if __name__ == "__main__":
    test_gpu_pool()

方案2：重写_Popen方法避免递归 #

如果必须继承mp.Process，可以重写_Popen方法，直接调用底层上下文的实现：

import multiprocessing as mp
import itertools

class GPUProcess(mp.Process):
    _gpu_ids = itertools.cycle(range(2))

    def __init__(
        self,
        group=None,
        target=None,
        name=None,
        args=(),
        kwargs={},
        *,
        daemon=None,
        gpu_id=None,
    ):
        super().__init__(group, target, name, args, kwargs, daemon=daemon)
        self._gpu_id = gpu_id if gpu_id is not None else next(self._gpu_ids)

    @property
    def gpu_id(self):
        return self._gpu_id

    @staticmethod
    def _Popen(process_obj):
        # 直接调用默认上下文的底层Process实现，跳过递归逻辑
        return mp.get_context()._default_context.Process._Popen(process_obj)

    def run(self):
        import os
        os.environ["CUDA_VISIBLE_DEVICES"] = str(self.gpu_id)
        super().run()

    @classmethod
    def create_gpu_context(cls):
        context = mp.get_context()
        context.Process = cls
        return context

关键改进点 #

• 避免递归：要么直接继承BaseProcess，要么重写_Popen方法跳过标准库的递归委托逻辑。
• 进程安全的ID分配：使用itertools.cycle（适合固定GPU数量场景）或multiprocessing.Manager（适合动态ID池场景）管理GPU ID，避免类变量在多进程中的同步问题。
• 可靠的GPU设置：在run方法中设置GPU环境变量，确保子进程启动时立即生效。
• 移除不可靠逻辑：删除__del__方法，改用循环分配或管理器实现ID复用，避免多进程中析构函数执行时机不确定的问题。

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