Python Multiprocessing代码在Jupyter/IPython中无法运行的问题

问题描述 #

学习伯克利《Python数值方法》时，运行书中的多进程对比示例代码，串行版本正常执行，但并行版本在Jupyter Lab和IPython环境中均报错：

AttributeError: Can't get attribute 'random_square' on <module '__main__' (built-in)>

原因分析 #

Jupyter/IPython的交互式环境与普通Python脚本的模块加载逻辑存在差异：

• multiprocessing.Pool默认使用fork方式启动子进程，子进程会继承主进程内存空间，但无法识别在交互式环境中动态定义的random_square函数。
• 子进程尝试从__main__模块加载目标函数时，该函数并不存在于可被导入的__main__模块命名空间中，因此触发属性不存在错误。

解决方案 #

方案1：将函数封装到独立模块 #

把需要并行执行的函数放到单独的Python文件中，让子进程可以正常导入：

1. 创建square_utils.py文件，写入函数定义：

import numpy as np

def random_square(seed):
    np.random.seed(seed)
    random_num = np.random.randint(0, 10)
    return random_num**2

2. 在Jupyter中导入函数并运行修正后的代码：

import multiprocessing as mp
import numpy as np
import time
from square_utils import random_square

# 串行版本
t0 = time.time()
results = []
for i in range(10000000): 
    results.append(random_square(i))
t1 = time.time()
print(f'串行执行时间 {t1 - t0} s')

# 并行版本
t0 = time.time()
n_cpu = mp.cpu_count()

# 使用with语句自动管理进程池资源
with mp.Pool(processes=n_cpu) as pool:
    results = pool.map(random_square, range(10000000))
t1 = time.time()
print(f'并行执行时间 {t1 - t0} s')

方案2：使用spawn启动上下文并包裹代码块 #

采用spawn启动方式（重新启动Python解释器，避免fork带来的命名空间问题），并将多进程执行逻辑放在if __name__ == '__main__':块中：

import multiprocessing as mp
import numpy as np
import time

def random_square(seed):
    np.random.seed(seed)
    random_num = np.random.randint(0, 10)
    return random_num**2

# 串行版本
t0 = time.time()
results = []
for i in range(10000000): 
    results.append(random_square(i))
t1 = time.time()
print(f'串行执行时间 {t1 - t0} s')

# 并行版本（修正后）
if __name__ == '__main__':
    t0 = time.time()
    # 创建spawn类型的进程上下文
    ctx = mp.get_context('spawn')
    n_cpu = ctx.cpu_count()
    
    with ctx.Pool(processes=n_cpu) as pool:
        results = pool.map(random_square, range(10000000))
    t1 = time.time()
    print(f'并行执行时间 {t1 - t0} s')

额外注意点 #

• 原代码中results = [pool.map(...)])会将结果额外包裹一层列表，pool.map本身已返回列表，需去掉外层括号。
• 推荐使用with语句管理Pool资源，避免手动关闭进程池的繁琐操作。

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