Update submit.sh

Author: sbryngelson

.github/workflows/phoenix/submit.sh

@@ -1,64 +1,83 @@
-#!/bin/bash
-
-set -e
+#!/usr/bin/env bash
+set -euo pipefail
 
 usage() {
     echo "Usage: $0 [script.sh] [cpu|gpu]"
+    exit 1
 }
 
-if [ ! -z "$1" ]; then
-    sbatch_script_contents=`cat $1`
-else
-    usage
-    exit 1
-fi
+[[ $# -eq 2 ]] || usage
+
+sbatch_script="$1"
+device="$2"
+
+# read the body of the user script
+sbatch_body=$(<"$sbatch_script")
 
-sbatch_cpu_opts="\
-#SBATCH -p cpu-small               # partition
-#SBATCH --ntasks-per-node=24       # Number of cores per node required
-#SBATCH --mem-per-cpu=2G           # Memory per core\
+# common SBATCH directives
+sbatch_common_opts="\
+#SBATCH -J shb-${sbatch_script%%.sh}-$device   # job name
+#SBATCH --account=gts-sbryngelson3             # account
+#SBATCH -N1                                     # nodes
+#SBATCH -t 03:00:00                             # walltime
+#SBATCH -q embers                               # QOS
+#SBATCH -o ${sbatch_script%%.sh}-$device.%j.out  # stdout+stderr
+#SBATCH --mem-per-cpu=2G                        # default mem (overridden below)
 "
 
-sbatch_gpu_opts="\
+# CPU vs GPU overrides
+if [[ "$device" == "cpu" ]]; then
+  sbatch_device_opts="\
+#SBATCH -p cpu-small
+#SBATCH --ntasks-per-node=24
+"
+elif [[ "$device" == "gpu" ]]; then
+  sbatch_device_opts="\
 #SBATCH -p gpu-v100,gpu-a100,gpu-h100,gpu-l40s
-#SBATCH --ntasks-per-node=4       # Number of cores per node required
-#SBATCH -G2\
+#SBATCH --ntasks-per-node=4
+#SBATCH -G2
 "
-
-if [ "$2" = "cpu" ]; then
-    sbatch_device_opts="$sbatch_cpu_opts"
-elif [ "$2" = "gpu" ]; then
-    sbatch_device_opts="$sbatch_gpu_opts"
 else
-    usage
-    exit 1
+  usage
 fi
 
-job_slug="`basename "$1" | sed 's/\.sh$//' | sed 's/[^a-zA-Z0-9]/-/g'`-$2"
+# submit and capture the JobID
+JOBID=$(sbatch <<-EOT | awk '{print $4}'
+	#!/usr/bin/env bash
+	${sbatch_common_opts}
+	${sbatch_device_opts}
 
-sbatch <<EOT
-#!/bin/bash
-#SBATCH -Jshb-$job_slug            # Job name
-#SBATCH --account=gts-sbryngelson3 # charge account
-#SBATCH -N1                        # Number of nodes required
-$sbatch_device_opts
-#SBATCH -t 03:00:00                # Duration of the job (Ex: 15 mins)
-#SBATCH -q embers                  # QOS Name
-#SBATCH -o$job_slug.out            # Combined output and error messages file
-#SBATCH -W                         # Do not exit until the submitted job terminates.
+	set -e -x
 
-set -e
-set -x
+	cd "\$SLURM_SUBMIT_DIR"
+	echo "Running in \$(pwd):"
 
-cd "\$SLURM_SUBMIT_DIR"
-echo "Running in $(pwd):"
+	# load your modules & env
+	. ./mfc.sh load -c p -m $device
 
-job_slug="$job_slug"
-job_device="$2"
+	# user script contents
+	${sbatch_body}
+EOT
+)
 
-. ./mfc.sh load -c p -m $2
+echo "🚀 Submitted SLURM job $JOBID"
 
-$sbatch_script_contents
+# if this wrapper is killed/canceled, make sure SLURM job is cleaned up
+trap '[[ -n "${JOBID:-}" ]] && scancel "$JOBID" >/dev/null 2>&1 || :' EXIT
 
-EOT
+# poll until job finishes
+while true; do
+  STATE=$(sacct -j "$JOBID" --noheader --format=State | head -1 | cut -d' ' -f1)
+  echo "⏳ SLURM job $JOBID state: $STATE"
+  case "$STATE" in
+    COMPLETED|FAILED|CANCELLED|TIMEOUT) break ;;
+    *) sleep 10 ;;
+  esac
+done
+
+# show final report
+sacct -j "$JOBID" --format=JobID,State,ExitCode,Elapsed,MaxRSS
 
+# exit with the job's real code (left of the colon)
+EXIT_CODE=$(sacct -j "$JOBID" --noheader --format=ExitCode | head -1 | cut -d: -f1)
+exit "$EXIT_CODE"