#!/bin/bash

获取每个节点分配的GPU数量

gres_per_gpu=$(/usr/bin/scontrol show job $SLURM_JOB_ID | grep -oP 'gres/gpu=\K[^,]+')

获取分配的节点数量

num_nodes=$(/usr/bin/scontrol show job $SLURM_JOB_ID | grep -oP 'NumNodes=\K[^ ]+')

计算总共分配的GPU数量

gpu_count=$((gres_per_gpu / num_nodes))

获取分配的节点列表

allocated_nodes=$(/usr/bin/scontrol show hostnames $SLURM_JOB_NODELIST)

创建包含分配节点和每个节点GPU数量的hostfile文件

host_file='hostfile' echo '$allocated_nodes:$gpu_count' > $host_file

将脚本参数中的'$'替换为计算得到的GPU数量

sed -i 's/$/:$gpu_count/' '$1'