PyTorch多进程spawn调用类方法遇pickle锁对象错误的技术问询

torch.multiprocessing.spawn采用spawn启动模式，该模式下子进程会重新加载整个Python代码，而非直接复制主进程内存（不同于fork模式）。启动时主进程会把DistributedMission的实例序列化（pickle）后传递给子进程，而threading.Lock是依赖操作系统底层线程原语的对象，本身不支持pickle序列化，因此触发TypeError: cannot pickle '_thread.lock' object错误。

根据锁的使用场景，有两种可行方案：

方案1：延迟初始化线程锁（适用于控制单进程内线程串行） #

不在类的__init__中创建锁，而是在子进程启动后的worker方法里，第一次需要使用锁时再初始化。因为spawn模式下每个子进程会独立执行代码，各自创建自己的锁，完全避开序列化问题：

import torch
import threading

class DistributedMission:
    def __init__(self):
        # 先不初始化锁
        self.lock = None

    def _worker(self, rank):
        # 子进程内初始化锁，每个子进程拥有独立的锁实例
        if self.lock is None:
            self.lock = threading.Lock()
        
        # 用锁控制串行逻辑
        with self.lock:
            print(f"Rank {rank} 进入临界区")
            # 你的业务代码
            print(f"Rank {rank} 退出临界区")

    def start(self):
        torch.multiprocessing.spawn(self._worker, nprocs=2)

if __name__ == "__main__":
    mission = DistributedMission()
    mission.start()

方案2：使用可序列化的跨进程锁（适用于控制多进程间串行） #

如果需要的是跨进程的同步（而非单进程内线程同步），可以用multiprocessing.Manager创建支持pickle的锁对象，这类锁能在进程间安全传递和使用：

import torch
from multiprocessing import Manager

class DistributedMission:
    def __init__(self):
        # 通过Manager创建可序列化的跨进程锁
        manager = Manager()
        self.lock = manager.Lock()

    def _worker(self, rank):
        with self.lock:
            print(f"Rank {rank} 进入跨进程临界区")
            # 你的业务代码
            print(f"Rank {rank} 退出跨进程临界区")

    def start(self):
        torch.multiprocessing.spawn(self._worker, nprocs=2)

if __name__ == "__main__":
    mission = DistributedMission()
    mission.start()

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