From 7b32d5e178ba7e4f76ebbff6dbaa8978708c564a Mon Sep 17 00:00:00 2001
From: gadorlhiac <gadorlhiac@gmail.com>
Date: Mon, 15 Apr 2024 14:36:24 -0700
Subject: [PATCH] ENH Launch wrapper accounts for partition and account without
 hard coding

---
 launch_scripts/submit_launch_airflow.sh | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/launch_scripts/submit_launch_airflow.sh b/launch_scripts/submit_launch_airflow.sh
index a906ec8e..33fd69f5 100755
--- a/launch_scripts/submit_launch_airflow.sh
+++ b/launch_scripts/submit_launch_airflow.sh
@@ -1,7 +1,29 @@
 #!/bin/bash
+
+# Need to capture partition and account for SLURM
+while [[ $# -gt 0 ]]
+do
+    case "$1" in
+        --partition=*)
+            PARTITION="${1#*=}"
+            shift
+            ;;
+        --account=*)
+            ACCOUNT="${1#*=}"
+            shift
+            ;;
+        *)
+            POS+=("$1")
+            shift
+            ;;
+    esac
+done
+set -- "${POS[@]}"
+
 CMD="${@}"
+CMD="${CMD} --partition=${PARTITION} --account=${ACCOUNT}"
 echo $CMD
 CMD="/sdf/group/lcls/ds/tools/lute/lute_launcher ${CMD}"
-SLURM_ARGS="--partition=milano --account=lcls:data --ntasks=1"
+SLURM_ARGS="--partition=${PARTITION} --account=${ACCOUNT} --ntasks=1"
 echo "Running ${CMD} with ${SLURM_ARGS}"
 sbatch $SLURM_ARGS --wrap "${CMD}"