如何在Slurm集群不同节点调用Julia函数并独占节点全核

需求说明 #

• 从outer函数调用inner函数，每个inner调用运行在不同节点：ij=1时用节点1全部16核，ij=2时用节点2全部16核，依此类推。

用户现有代码 #

Julia 代码 #

using Distributed
addprocs(32)
    
println("Number of processes: ", nprocs())
println("Number of workers: ", nworkers())


@everywhere function inner(a,ij)
   sleep(5);
   println("Inside inner")

   return a*ij;
end

function outer(a,N)
   tt0 = time()
   g(x) = ij -> inner(x, ij);
   arrsum = sum(pmap(g(a), (1:N)));
   tt1 = time()
   println("outer time = $(tt1-tt0)")

   return arrsum
end

println("outer = ",outer(1,5))

Slurm 提交脚本 #

#!/bin/bash

#SBATCH -J m_node
#SBATCH -t 0-04:00:00
#SBATCH --nodes 2
#SBATCH --ntasks-per-node 1
#SBATCH --cpus-per-task=16

srun /home/userdir/julia-1.10.4/bin/julia  /home/userdir/Work/julia_mnode.jl

当前问题 #

未达预期：每个inner调用仅使用1核，分散在2个节点上。

解决方案 #

要实现每个inner调用独占节点全部16核，需调整Slurm配置和Julia进程分配逻辑，核心是让每个节点的Julia worker绑定该节点所有CPU核心，且每个inner任务分配到对应节点的worker执行。

步骤1：修改Slurm脚本 #

调整参数确保每个节点的Julia worker绑定全部核心：

#!/bin/bash

#SBATCH -J m_node
#SBATCH -t 0-04:00:00
#SBATCH --nodes 2
#SBATCH --ntasks-per-node 1
#SBATCH --cpus-per-task=16
# 绑定CPU核心到当前任务，确保worker能使用节点全部核心
#SBATCH --cpu-bind=cores

# 通过--machine-file让Julia自动获取Slurm节点列表，每个节点启动一个worker
srun /home/userdir/julia-1.10.4/bin/julia --project --machine-file $SLURM_JOB_NODELIST /home/userdir/Work/julia_mnode.jl

步骤2：调整Julia代码 #

去掉手动创建进程的逻辑，改为按节点分组分配任务：

using Distributed
# 不再手动addprocs，由srun通过--machine-file自动添加节点worker
    
println("Number of processes: ", nprocs())
println("Number of workers: ", nworkers())

# 按节点分组整理worker列表
node_workers = Dict{String, Vector{Int}}()
for w in workers()
    host = fetch(@spawnat w gethostname())
    node_workers[host] = push!(get(node_workers, host, Int[]), w)
end
# 按节点排序得到固定顺序的worker组
sorted_node_groups = sort(collect(values(node_workers)))

@everywhere function inner(a,ij)
    # 打印当前worker所在节点及可用核心数，验证配置
    println("Worker $(myid()) on host $(gethostname()), using $(Sys.CPU_THREADS) cores")
    sleep(5);
    return a*ij;
end

function outer(a,N)
    tt0 = time()
    total = 0
    for ij in 1:N
        # 循环分配节点，超过节点数则重复使用
        group_idx = ((ij-1) % length(sorted_node_groups)) + 1
        target_worker = sorted_node_groups[group_idx][1]
        # 在指定节点的worker上执行inner，独占该节点全部核心
        result = fetch(@spawnat target_worker inner(a, ij))
        total += result
    end
    tt1 = time()
    println("outer time = $(tt1-tt0)")
    return total
end

println("outer = ",outer(1,5))

关键说明 #

• Slurm侧：--cpu-bind=cores确保每个节点的Julia worker能使用该节点的全部16核；--machine-file让Julia自动识别并添加节点worker。
• Julia侧：
  • 按节点分组worker，确保每个inner任务分配到指定节点的worker执行。
  • 如果inner函数需要多线程并行，可在函数内部添加Threads.@threads逻辑，此时会自动利用节点的16核资源。

内容的提问来源于stack exchange，提问作者evening silver fox