Update submit-bench.sh

Author: sbryngelson

.github/workflows/phoenix/submit-bench.sh

@@ -1,64 +1,101 @@
-#!/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_cpu_opts="\
+sbatch_script="$1"
+device="$2"
+
+job_slug="`basename "$1" | sed 's/\.sh$//' | sed 's/[^a-zA-Z0-9]/-/g'`-$2"
+
+# read the body of the user script
+sbatch_body=$(<"$sbatch_script")
+
+# 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 02:00:00                             # walltime
+#SBATCH -q embers                               # QOS
+#SBATCH -o $job_slug.out                        # stdout+stderr
+#SBATCH --mem-per-cpu=2G                        # default mem (overridden below)
+"
+
+# CPU vs GPU overrides
+if [[ "$device" == "cpu" ]]; then
+  sbatch_device_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\
 "
-
-sbatch_gpu_opts="\
+elif [[ "$device" == "gpu" ]]; then
+  sbatch_device_opts="\
 #SBATCH -CL40S
 #SBATCH --ntasks-per-node=4       # Number of cores per node required
 #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 02: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 SLURM job finishes ──────────
+while :; do
+  # Try sacct first
+  STATE=$(sacct -j "$JOBID" --format=State --noheader --parsable2 | head -n1)
+
+  # Fallback to squeue if sacct is empty
+  if [[ -z "$STATE" ]]; then
+    STATE=$(squeue -j "$JOBID" -h -o "%T" || echo "")
+  fi
+
+  # If it’s one of SLURM’s terminal states, break immediately
+  case "$STATE" in
+    COMPLETED|FAILED|CANCELLED|TIMEOUT)
+      echo "✅ SLURM job $JOBID reached terminal state: $STATE"
+      break
+      ;;
+    "")
+      echo "✅ SLURM job $JOBID no longer in queue; assuming finished"
+      break
+      ;;
+    *)
+      echo "⏳ SLURM job $JOBID state: $STATE"
+      sleep 10
+      ;;
+  esac
+done
 
+# Now retrieve the exit code and exit with it
+EXIT_CODE=$(sacct -j "$JOBID" --noheader --format=ExitCode | head -1 | cut -d: -f1)
+echo "🔚 SLURM job $JOBID exit code: $EXIT_CODE"
+exit "$EXIT_CODE"