在SLURM集群用submitit运行含ProcessPoolExecutor的函数遇错求助

问题 #

在SLURM集群上使用submitit运行调用ProcessPoolExecutor的函数时，遭遇pickle序列化错误。相关代码及报错信息如下：

业务代码 #

def func():
    print("Hello world!")

def parallel_executor():
    with ProcessPoolExecutor(max_workers=24) as executor_inner:
        j = executor_inner.submit(func)
        j.result()

提交代码 #

executor = submitit.AutoExecutor()
j = executor.submit(parallel_executor)

报错信息 #

Traceback (most recent call last):
  File "/home/X/venvs/smac_env/lib/python3.10/site-packages/submitit/core/submission.py", line 55, in process_job
    result = delayed.result()
  File "/home/X/venvs/smac_env/lib/python3.10/site-packages/submitit/core/utils.py", line 133, in result
    self._result = self.function(*self.args, **self.kwargs)
  File "/home/X/meta_hpo/test_submitit.py", line 33, in parallel_executor
    j.result()
  File "X/Python/3.10.8-GCCcore-12.2.0/lib/python3.10/concurrent/futures/_base.py", line 458, in result
    return self.__get_result()
  File "X/Python/3.10.8-GCCcore-12.2.0/lib/python3.10/concurrent/futures/_base.py", line 403, in __get_result
    raise self._exception
  File "X/Python/3.10.8-GCCcore-12.2.0/lib/python3.10/multiprocessing/queues.py", line 244, in _feed
    obj = _ForkingPickler.dumps(obj)
  File "X/Python/3.10.8-GCCcore-12.2.0/lib/python3.10/multiprocessing/reduction.py", line 51, in dumps
    cls(buf, protocol).dump(obj)
_pickle.PicklingError: Can't pickle <function func at 0x151b9dbed6c0>: attribute lookup func on __main__ failed

解决方法 #

1. 修复pickle序列化问题 #

错误根源是：submitit提交的任务在SLURM节点上重新加载代码时，子进程无法在__main__模块中找到func函数，导致pickle反序列化失败。可通过两种方式修复：

方式1：将函数放入可导入模块

创建独立模块my_tasks.py存放业务函数：

# my_tasks.py
def func():
    print("Hello world!")

def parallel_executor():
    from concurrent.futures import ProcessPoolExecutor
    with ProcessPoolExecutor(max_workers=24) as executor_inner:
        j = executor_inner.submit(func)
        j.result()

主脚本中导入并提交：

import submitit
from my_tasks import parallel_executor

executor = submitit.AutoExecutor()
j = executor.submit(parallel_executor)

方式2：添加主模块保护

在主脚本中用if __name__ == "__main__":包裹提交逻辑，避免子进程重复执行初始化代码，同时确保函数能被正确识别：

from concurrent.futures import ProcessPoolExecutor
import submitit

def func():
    print("Hello world!")

def parallel_executor():
    with ProcessPoolExecutor(max_workers=24) as executor_inner:
        j = executor_inner.submit(func)
        j.result()

if __name__ == "__main__":
    executor = submitit.AutoExecutor()
    j = executor.submit(parallel_executor)

2. 使用submitit原生并行机制替代ProcessPoolExecutor #

既然已经基于submitit调度任务，可直接利用其原生能力管理并行子任务，更适配SLURM集群的资源调度逻辑：

单参数/无参数任务示例

import submitit

def func():
    print("Hello world!")

def parallel_executor():
    # 用LocalExecutor调度节点内子任务，也可根据需求改用SlurmExecutor跨节点调度
    inner_executor = submitit.LocalExecutor()
    # 单任务直接submit，批量任务用map_array
    job = inner_executor.submit(func)
    job.result()

if __name__ == "__main__":
    executor = submitit.AutoExecutor()
    j = executor.submit(parallel_executor)

多参数批量任务示例

import submitit

def func_with_arg(name):
    print(f"Hello {name}!")

def parallel_executor():
    inner_executor = submitit.LocalExecutor()
    # 批量传递参数列表
    args_list = ["Alice", "Bob", "Charlie"]
    jobs = inner_executor.map_array(func_with_arg, args_list)
    for job in jobs:
        job.result()

if __name__ == "__main__":
    executor = submitit.AutoExecutor()
    j = executor.submit(parallel_executor)

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