#!/bin/bash

PILOT_BIN=$1
PILOT_DIR=$2

function cleanup() {
    echo "Cleaning up temporary directory (${PILOT_DIR})..."
    rm -fr ${PILOT_DIR}
}
trap cleanup EXIT

# Which nodes are in this job?
SCONTROL_LOGGING=`scontrol show hostname ${SLURM_NODELIST} > "${PILOT_DIR}/nodelist"`
if [[ $? != 0 ]]; then
    echo "Failed to get node list, aborting."
    echo ${SCONTROL_LOGGING}
    exit 7
fi

# Start the pilot on each node in this job.
for host in `cat "${PILOT_DIR}/nodelist"`; do
    # On Perlmutter, you can only (reliably) SSH to nodes in your job if you
    # use the 10.* address, but the external IP is sometimes listed first,
    # most commonly when SSH'ing to current host.
    ip=`host $host | grep '10\.' | head -n 1 | awk '{print $4}'`
    ssh ${ip} "${PILOT_BIN}" "${PILOT_DIR}" "${CUDA_VISIBLE_DEVICES}" &
done

# Wait for every node's pilot to finish.
wait

echo "Exiting after all nodes finished."
exit 0
