Add a batch_wait script

lib/ramble/ramble/test/workflow_manager_functionality/slurm_workflow_manager.py

@@ -72,6 +72,7 @@ def test_slurm_workflow():
         assert "batch_submit" not in files
         assert "batch_query" not in files
         assert "batch_cancel" not in files
+        assert "batch_wait" not in files
 
         # Assert on slurm workflow manager
         path = os.path.join(ws.experiment_dir, "hostname", "local", "test_slurm")
@@ -80,6 +81,7 @@ def test_slurm_workflow():
         assert "batch_submit" in files
         assert "batch_query" in files
         assert "batch_cancel" in files
+        assert "batch_wait" in files
         with open(os.path.join(path, "batch_submit")) as f:
             content = f.read()
             assert "slurm_execute_experiment" in content

var/ramble/repos/builtin/workflow_managers/slurm/batch_wait.tpl

@@ -0,0 +1,22 @@
+#!/bin/bash
+
+# Ensure job_id is present
+job_id=$(< {experiment_run_dir}/.slurm_job)
+if [ -z "${job_id:-}" ]; then
+    echo "No valid job_id found" 1>&2
+    exit 1
+fi
+
+echo "Waiting for job {job_name} with id ${job_id} to complete..."
+
+while true; do
+    status=$(squeue -h -o "%t" -j "${job_id}" 2>/dev/null)
+    # The absence of the job from squeue indicates it's completed.
+    if [ -z "$status" ]; then
+        break
+    else
+        sleep 10
+    fi
+done
+
+echo "job {job_name} with id ${job_id} is done"

var/ramble/repos/builtin/workflow_managers/slurm/workflow_manager.py

@@ -102,6 +102,12 @@ def __init__(self, file_path):
         dest_path="batch_cancel",
     )
 
+    register_template(
+        name="batch_wait",
+        src_name="batch_wait.tpl",
+        dest_path="batch_wait",
+    )
+
     register_template(
         name="slurm_execute_experiment",
         src_name="slurm_execute_experiment.tpl",