如何通过Python脚本利用多CPU与节点并行运行Singularity容器

完全可以仅通过Python内置库（subprocessing、concurrent.futures等）实现多节点并行启动Singularity容器，充分利用SLURM分配的硬件资源。以下是具体实现思路和代码：

核心逻辑 #

利用SLURM的srun命令将容器启动任务分发到指定节点，结合Python的并行工具管理多个容器进程，同时保留宿主脚本在主节点的控制能力。

1. 获取SLURM分配的节点列表 #

从SLURM环境变量中解析出所有可用节点：

import subprocess
import os

def get_slurm_nodes():
    nodelist = os.environ.get('SLURM_NODELIST')
    if not nodelist:
        raise RuntimeError("当前未在SLURM作业环境中运行")
    result = subprocess.run(
        ['scontrol', 'show', 'hostnames', nodelist],
        capture_output=True, text=True, check=True
    )
    return result.stdout.strip().split('\n')

2. 计算节点容器分配量 #

根据节点核数、宿主脚本占用资源，计算每个节点可运行的容器数量（以你给出的2个76核节点、40个2核容器为例）：

def calculate_container_allocations(nodes, total_containers, cores_per_container):
    allocations = []
    remaining = total_containers
    
    # 主节点预留1核给宿主脚本
    first_node_available_cores = 76 - 1
    max_first_node = first_node_available_cores // cores_per_container
    alloc = min(max_first_node, remaining)
    allocations.append((nodes[0], alloc))
    remaining -= alloc
    
    # 分配剩余容器到其他节点
    for node in nodes[1:]:
        if remaining <= 0:
            break
        max_node = 76 // cores_per_container
        alloc = min(max_node, remaining)
        allocations.append((node, alloc))
        remaining -= alloc
    return allocations

3. 并行启动容器 #

用concurrent.futures.ThreadPoolExecutor并行在各节点启动容器，通过srun指定节点和CPU资源：

from concurrent.futures import ThreadPoolExecutor

def start_container(node, container_path, task_cmd, cores):
    # 构造srun+singularity启动命令
    cmd = [
        'srun',
        '--nodes=1', '--ntasks=1',
        f'--cpus-per-task={cores}',
        '--nodelist=' + node,
        'singularity', 'exec',
        container_path,
        *task_cmd.split()
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"节点{node}上容器执行失败：{result.stderr}")
    return result.returncode == 0

def run_containers(allocations, container_path, task_cmd, cores_per_container):
    with ThreadPoolExecutor(max_workers=len(allocations)) as executor:
        tasks = []
        for node, count in allocations:
            for _ in range(count):
                tasks.append(executor.submit(
                    start_container, node, container_path, task_cmd, cores_per_container
                ))
        # 等待所有容器执行完成
        for task in tasks:
            task.result()

4. 完整流程封装 #

整合上述模块，实现"启动容器→等待完成→分析→重复"的循环：

def main():
    # 配置参数
    container_path = "/path/to/your/container.sif"
    container_task = "your_task_command_here"
    total_containers = 40
    cores_per_container = 2

    # 获取节点并分配容器
    nodes = get_slurm_nodes()
    allocations = calculate_container_allocations(nodes, total_containers, cores_per_container)
    
    # 启动并等待容器完成
    run_containers(allocations, container_path, container_task, cores_per_container)
    
    # 执行分析逻辑（替换为你的实际代码）
    print("所有容器执行完成，开始分析结果...")
    
    # 如需重复流程，可添加循环
    # while your_condition:
    #     run_containers(...)
    #     analyze(...)

if __name__ == "__main__":
    main()

注意事项 #

• SLURM作业资源申请：需确保宿主脚本的SLURM作业申请了足够的总核数（例中为40*2+1=81核），避免资源不足导致任务排队。
• 日志管理：可修改start_container中的subprocess.run，将容器的标准输出/错误重定向到文件，方便后续排查问题。
• 错误处理：可扩展任务的异常捕获逻辑，实现容器启动失败后的重试机制。

内容的提问来源于stack exchange，提问作者theo-brown