如何从中心Python脚本通过Slurm并行运行多个计算任务

我在采用Slurm作为工作负载管理器的HPC集群上运行Python脚本。脚本里有个包含约10000个带ID的对象实例的字典：

d = {id1 : instance1, ...}

中心脚本会筛选出约100个ID执行计算，调用d[idx].calculate()即可启动计算。该方法后台会准备输入文件并执行命令行指令：

mpirun -np <num_cores> calc_executable.exe

这个可执行文件基于MPI，只需设置核心数<num_cores>，无需额外配置。计算耗时较长（最长可能24小时），所有任务相互独立，可使用不同核心数。计算完成后需要收集输出并更新对应对象实例的属性。

我可以修改代码直接调用命令行并自行准备输入文件（非首选），但需要调度器能灵活读取处理结果。我希望通过中心Python脚本在Slurm中调度所有计算任务，且能实时追踪任务状态，以便任务完成后立即处理结果并写入数据库。请问适用的Python库是什么？需要支持任务提交与结果追踪。

已尝试/调研的方案 #

multiprocessing #

最初用multiprocessing库，把主脚本通过Slurm分配N个核心，每个计算任务在独立进程运行：

from multiprocessing import Pool
with Pool(processes=8) as pool:
    results = pool.starmap(run_vasp, inputs)

该方案可行，但没借助Slurm调度计算任务，得手动设置并行进程数，当前代码无法在任务完成后立即处理结果（此问题可忽略）。

未尝试的Python库 #

• pyslurm：可以提交Slurm任务，但需要自行封装类来监控任务状态和获取结果。
• Dask：可能可行，但有资料显示它不适用于长耗时任务。
• snakemake：当前需求下过于复杂，需要用其专属语言配置工作流，我希望全程用Python处理，但未来可能考虑使用工作流管理器。

1. pyslurm + 自定义状态监控 #

pyslurm是对接Slurm API的官方Python绑定，完全匹配你的核心需求：

• 直接通过Python代码提交Slurm作业（指定核心数、输出路径等参数）
• 实时查询作业状态（排队、运行、完成、失败等）
• 无额外工作流框架依赖，全程用Python控制

实现思路 #

1. 封装作业提交函数：为每个目标ID生成Slurm作业参数，提交后记录作业ID与对象ID的映射
2. 循环监控作业状态：定期查询Slurm作业状态，遍历所有待处理任务
3. 即时处理结果：一旦发现作业完成，立即读取输出、更新对象属性并写入数据库，同时移除已处理任务

核心代码示例 #

import pyslurm
import time

# 存储作业ID与对象ID的映射关系
job_obj_map = {}

def submit_slurm_job(obj_id, num_cores):
    # 构建Slurm作业配置
    job_config = {
        "partition": "your_partition",
        "nodes": 1,
        "cpus_per_task": num_cores,
        "output": f"job_{obj_id}.out",
        "error": f"job_{obj_id}.err",
        "wrap": f"mpirun -np {num_cores} calc_executable.exe"
    }
    # 提交作业并记录映射
    job_id = pyslurm.job().create(job_config)
    job_obj_map[job_id] = obj_id
    return job_id

# 提交筛选出的100个任务
selected_obj_ids = get_selected_ids()  # 替换为你的筛选逻辑
for obj_id in selected_obj_ids:
    core_count = get_core_requirement(obj_id)  # 替换为你的核心数获取逻辑
    submit_slurm_job(obj_id, core_count)

# 持续监控作业状态
while job_obj_map:
    # 获取所有作业的当前状态
    active_jobs = pyslurm.job().load()
    for job_id in list(job_obj_map.keys()):
        if job_id not in active_jobs:
            continue
        state = active_jobs[job_id]["job_state"]
        if state == "COMPLETED":
            # 处理完成任务
            obj_id = job_obj_map.pop(job_id)
            instance = d[obj_id]
            # 从输出文件读取结果并更新对象属性
            update_instance_from_output(instance, f"job_{obj_id}.out")
            # 写入数据库
            write_to_db(instance)
        elif state in ["FAILED", "CANCELLED"]:
            # 处理失败任务（重试/日志记录）
            obj_id = job_obj_map.pop(job_id)
            log_job_failure(obj_id)
    # 每30秒检查一次状态
    time.sleep(30)

2. Parsl（备选方案） #

Parsl是专为HPC设计的Python并行任务调度库，原生支持Slurm等集群调度器，无需手动处理Slurm作业细节，同时支持实时状态追踪与结果回收：

• 用Python装饰器定义任务，自动提交到Slurm
• 内置状态监控，任务完成后可立即获取结果
• 支持动态分配资源（不同任务使用不同核心数）

核心代码示例 #

import parsl
from parsl.config import Config
from parsl.executors import HighThroughputExecutor
from parsl.providers import SlurmProvider

# 配置Slurm集群
slurm_config = Config(
    executors=[
        HighThroughputExecutor(
            label='slurm_executor',
            provider=SlurmProvider(
                partition='your_partition',
                nodes_per_block=1,
                init_blocks=0,
                max_blocks=100,  # 最大并发任务数
                walltime='24:00:00',  # 单任务最长运行时间
            ),
        )
    ]
)
parsl.load(slurm_config)

# 定义计算任务
@parsl.python_app
def run_calc(obj_id, num_cores):
    # 调用对象的计算方法（需确保对象可序列化，或改为文件交互）
    instance = d[obj_id]
    instance.calculate(num_cores=num_cores)
    # 返回需写入数据库的属性
    return obj_id, instance.get_result_attrs()

# 提交所有任务
task_futures = []
for obj_id in selected_obj_ids:
    core_count = get_core_requirement(obj_id)
    task_futures.append(run_calc(obj_id, core_count))

# 实时处理完成的任务
for future in parsl.as_completed(task_futures):
    obj_id, result_attrs = future.result()
    instance = d[obj_id]
    # 更新对象属性
    instance.update_attrs(result_attrs)
    # 写入数据库
    write_to_db(instance)

方案对比 #

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