#!/bin/bash

#echo $1

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

Retrieve the number of allocated nodes

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

#echo $num_nodes

Calculate the number of allocated GPUs

gpu_count=$((gres_per_gpu / num_nodes))

#echo $gpu_count

Get the list of allocated hosts

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

Loop through each host and append the GPU count to the host string

for host in $hosts; do sed -i "s/$host/$host:$gpu_count/" "$1" done