如何在Slurm作业脚本中根据分配节点自动请求CPU核心数？

方法1：通过节点名直接判断（无需管理员权限） #

不需要修改集群配置，直接在作业脚本里根据分配到的节点名前缀判断，自动设置CPU核心数：

#!/bin/bash
#SBATCH --nodes=1
#SBATCH --ntasks=1
# 批量提交时可添加作业数组：#SBATCH --array=1-100

# 根据节点名前缀确定核心数
if [[ $SLURMD_NODENAME == node12-* ]]; then
    CORE_NUM=2
else
    CORE_NUM=4
fi

# 读取批量参数（如果是作业数组）
# INPUT_PARAM=$(sed -n "${SLURM_ARRAY_TASK_ID}p" your_params_list.txt)

# 用对应核心数执行任务
srun --cpus-per-task=$CORE_NUM your_executable --input $INPUT_PARAM

方法2：给节点添加自定义特性（需管理员权限） #

让管理员给两类节点添加自定义特性，后续脚本通过特性值获取对应核心数，扩展性更强：

1. 管理员执行节点特性配置：

# 给node12系列节点添加特性core_opt=2
scontrol update NodeName=node12-* Features=core_opt:2
# 给node20系列节点添加特性core_opt=4
scontrol update NodeName=node20-* Features=core_opt:4

2. 作业脚本写法：

#!/bin/bash
#SBATCH --nodes=1
#SBATCH --ntasks=1
# 批量提交添加：#SBATCH --array=1-100

# 从节点特性中提取预设的最优核心数
CORE_NUM=$(scontrol show node $SLURMD_NODENAME | grep -oP 'core_opt:\K\d+')

# 读取批量参数
# INPUT_PARAM=$(sed -n "${SLURM_ARRAY_TASK_ID}p" your_params_list.txt)

# 执行任务
srun --cpus-per-task=$CORE_NUM your_executable --input $INPUT_PARAM

方法3：通过Slurm Prolog脚本全局配置（需管理员权限） #

让管理员配置Prolog脚本，在作业启动前自动设置环境变量，脚本直接读取变量即可：

1. 管理员创建Prolog脚本（比如/etc/slurm/prolog.sh）：

#!/bin/bash
# 根据节点类型设置环境变量
if [[ $SLURMD_NODENAME == node12-* ]]; then
    export DESIRED_CORES=2
elif [[ $SLURMD_NODENAME == node20-* ]]; then
    export DESIRED_CORES=4
fi

2. 在slurm.conf中配置Prolog路径：

Prolog=/etc/slurm/prolog.sh

3. 作业脚本写法：

#!/bin/bash
#SBATCH --nodes=1
#SBATCH --ntasks=1
# 批量提交添加：#SBATCH --array=1-100

# 读取批量参数
# INPUT_PARAM=$(sed -n "${SLURM_ARRAY_TASK_ID}p" your_params_list.txt)

# 直接使用环境变量指定核心数
srun --cpus-per-task=$DESIRED_CORES your_executable --input $INPUT_PARAM

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