SLURM多任务配置无加速：增加ntasks为何耗时不变？

• 使用SLURM调度Python multiprocessing代码时出现异常：
  • 当设置--ntasks=3且--cpus-per-task=40，或--ntasks=1且--cpus-per-task=40时，代码运行耗时均为99秒，无任何差异
  • 但将--cpus-per-task从20调整为40时，耗时从194秒降至99秒，符合CPU核心数翻倍的加速预期
• 当前集群单节点可用CPU为40核，附最小可复现代码及SLURM脚本：

最小可复现代码 #

import multiprocessing as mp
import openpyxl
import os
import time
from multiprocessing import Lock


def write_to_excel(workbook, sheet_name, row, col, data, mylock):
    # 模拟CPU密集型计算
    for k in range(15_000):
        for j in range(15_000):
            a = k + j
            if a % 2 == 0:
                a = a + 1
            else:
                a = a - 1
            if a is None:
                print(a)
    with mylock:
        # 写入Excel
        wb = openpyxl.load_workbook(workbook)
        sheet = wb[sheet_name]
        sheet.cell(row=row, column=col, value=data)
        wb.save(workbook)


if __name__ == "__main__":
    start_time = time.time()
    # 初始化Excel文件
    wb = openpyxl.Workbook()
    wb.save("shared_workbook.xlsx")

    mylock = Lock()

    # 从SLURM环境变量获取资源配置
    num_tasks = int(os.getenv("SLURM_NTASKS", 1))
    cpus_per_task = int(os.getenv("SLURM_CPUS_PER_TASK", 1))

    print(f"num_tasks: {num_tasks}")
    print(f"cpus_per_task: {cpus_per_task}")
    print(f"num_processes: {num_tasks * cpus_per_task}")

    # 固定创建102个进程
    num_processes_to_have = 102

    # 启动进程
    processes = []
    for i in range(num_processes_to_have):
        process = mp.Process(
            target=write_to_excel,
            args=(
                "shared_workbook.xlsx",
                "Sheet",
                i + 1,
                1,
                f"Data from process {i + 1}",
                mylock,
            ),
        )
        processes.append(process)
        process.start()

    # 等待所有进程完成
    for process in processes:
        process.join()

    print("Writing to shared workbook complete.", time.time() - start_time)

SLURM脚本示例 #

#SBATCH --job-name=#####
#SBATCH --output=#####
#SBATCH --time=1:00:00
#SBATCH --mem=8G
#SBATCH --partition=#####
#SBATCH --mail-user=#####
#SBATCH --mail-type=#####
#SBATCH --export=NONE
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=20

1. SLURM资源分配限制：
   • 当指定--ntasks=3且--cpus-per-task=40时，SLURM会尝试为每个任务分配40核，但单节点仅40核，且未指定--nodes参数要求跨节点调度，因此SLURM实际只会在当前节点启动1个任务，占用全部40核，与--ntasks=1且--cpus-per-task=40的资源分配完全一致，所以耗时相同
2. 进程数超过CPU核心数的瓶颈：
   • 代码固定创建102个进程，但单节点最多只能同时运行40个CPU密集型进程（受限于40核），剩余进程会排队等待。无论SLURM指定多少任务数，只要总可用CPU核心数为40，实际并行的进程数上限就是40，因此耗时不会变化
3. CPU密集型任务的特性：
   • 任务核心是嵌套循环的纯CPU计算，无IO等待，并行效率完全受限于可用CPU核心数，超过核心数的进程只会增加调度开销，无法提升运行速度

1. 正确配置SLURM参数：
   • 单节点任务只需设置--ntasks=1，并将--cpus-per-task设为节点可用核心数（40），无需指定多任务数或超过节点核心数的CPU配额
   • 若需跨节点调度，必须添加--nodes参数，但注意当前代码使用本地文件锁，跨节点写入Excel会存在一致性问题，不建议采用
2. 优化进程数设置：
   • 不要固定创建102个进程，应根据SLURM分配的总CPU数动态调整，避免进程数超过核心数：

     # 取分配的CPU数和任务数的最小值
     num_processes_to_have = min(num_tasks * cpus_per_task, 102)
     

   • 或直接使用mp.cpu_count()获取当前可用CPU核心数，确保进程数不超过硬件上限
3. 优化Excel写入逻辑：
   • 目前每个进程单独加载、保存Excel的方式效率极低，可改为：
     • 所有进程先完成计算，将结果存入内存（如共享队列或列表）
     • 最后由主进程统一写入Excel，无需使用锁，大幅减少IO开销

内容的提问来源于stack exchange，提问作者FluidMechanics Potential Flows