如何为使用torch.multiprocessing的Python多进程程序配置profiler

Python的sys.setprofile配置是进程级独立的，torch.multiprocessing创建子进程时，无论使用fork还是spawn启动模式，子进程都不会自动继承父进程注册的profile回调函数，必须在子进程内部单独初始化才能生效。

方案1：子进程入口函数显式初始化（最直接） #

在你的子进程执行函数run_parameter_server的开头添加profiler启用代码即可：

def run_parameter_server(rank, world_size):
    # 开头先初始化profiler
    import sys
    import tracer
    sys.setprofile(tracer.tracefunc)
    # 原有业务代码
    ...

方案2：使用initializer统一初始化（适合多子进程场景） #

如果子进程数量多，不想逐个修改目标函数，可以通过Process的initializer参数统一执行初始化逻辑：

# 先定义统一的初始化函数
def init_tracer():
    import sys
    import tracer
    sys.setprofile(tracer.tracefunc)

# 创建子进程时传入初始化参数
import torch.multiprocessing as mp
p = mp.Process(
    target=run_parameter_server, 
    args=(0, world_size),
    initializer=init_tracer # 子进程启动后会先执行这个函数
)
p.start()
processes.append(p)

方案3：装饰器包裹目标函数（无侵入修改原函数） #

如果不想改动原有业务函数代码，可以写装饰器包裹目标函数完成初始化：

def with_tracer(func):
    def wrapper(*args, **kwargs):
        import sys
        import tracer
        sys.setprofile(tracer.tracefunc)
        return func(*args, **kwargs)
    return wrapper

# 给目标函数加装饰器即可
@with_tracer
def run_parameter_server(rank, world_size):
    # 原有业务代码
    ...

• 你现有tracefunc中用可变对象indent=[0]作为默认参数，每个子进程拥有独立的内存空间，不会出现跨进程的计数冲突，无需修改原有实现。
• 多进程同时向stdout打印内容可能出现输出乱序，可考虑将每个进程的日志写入独立文件避免混淆。
• 如果使用mp.spawn接口批量启动子进程，同样支持传入initializer参数完成统一初始化。

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