From faed230dca6ff0969466075324a3f5b896c56619 Mon Sep 17 00:00:00 2001
From: Toby Jennings <tobyj@slac.stanford.edu>
Date: Wed, 8 Jul 2026 19:03:41 +0000
Subject: [PATCH] fix: update default provisioner script

- trap SIGTERM in script and perform one final provisioner loop
- keep provisionerJob alive until N consecutive allocateNodes failures
---
 doc/changes/DM-55434.bugfix.rst               |  1 +
 .../bps/htcondor/etc/htcondor_defaults.yaml   | 49 ++++++++++++++-----
 python/lsst/ctrl/bps/htcondor/lssthtc.py      |  5 +-
 python/lsst/ctrl/bps/htcondor/provisioner.py  |  3 ++
 tests/test_htcondor_service.py                |  2 +-
 tests/test_provisioner.py                     |  3 ++
 6 files changed, 48 insertions(+), 15 deletions(-)
 create mode 100644 doc/changes/DM-55434.bugfix.rst

diff --git a/doc/changes/DM-55434.bugfix.rst b/doc/changes/DM-55434.bugfix.rst
new file mode 100644
index 0000000..8b3e626
--- /dev/null
+++ b/doc/changes/DM-55434.bugfix.rst
@@ -0,0 +1 @@
+Add signal handling and resilience to default provisioningScript
diff --git a/python/lsst/ctrl/bps/htcondor/etc/htcondor_defaults.yaml b/python/lsst/ctrl/bps/htcondor/etc/htcondor_defaults.yaml
index 702135b..3060310 100644
--- a/python/lsst/ctrl/bps/htcondor/etc/htcondor_defaults.yaml
+++ b/python/lsst/ctrl/bps/htcondor/etc/htcondor_defaults.yaml
@@ -16,21 +16,44 @@ provisioning:
   provisioningPlatform: "s3df"
   provisioningScript: |
     #!/bin/bash
-    set -e
     set -x
-    while true; do
-        ${CTRL_EXECUTE_DIR}/bin/allocateNodes.py \
-            --account {provisioningAccountingUser} \
-            --auto \
-            --node-count {provisioningNodeCount} \
-            --maximum-wall-clock {provisioningMaxWallTime} \
-            --glidein-shutdown {provisioningMaxIdleTime} \
-            --queue {provisioningQueue} \
-            --nodeset {nodeset} \
-            {provisioningExtraOptions} \
-            {provisioningPlatform}
-        sleep {provisioningCheckInterval}
+
+    allocating=1
+    failures=0
+    max_failures=10
+
+    allocate=("${CTRL_EXECUTE_DIR}/bin/allocateNodes.py" \
+          --account '{provisioningAccountingUser}' \
+          --auto \
+          --node-count '{provisioningNodeCount}' \
+          --maximum-wall-clock '{provisioningMaxWallTime}' \
+          --glidein-shutdown '{provisioningMaxIdleTime}' \
+          --queue '{provisioningQueue}' \
+          --nodeset '{nodeset}' \
+          {provisioningExtraOptions} \
+          '{provisioningPlatform}')
+
+    trap "allocating=0" SIGTERM
+
+    echo "$(date -u): starting provisioningJob"
+    while (( allocating )); do
+        if "${allocate[@]}"; then
+          failures=0
+        else
+          echo "$(date -u): allocatedNodes failed" >&2
+          failures=$((failures + 1))
+        fi
+        if [[ $failures -ge $max_failures ]]; then
+          echo "$(date -u): Too many consecutive allocateNodes failures, exiting" >&2
+          exit 1
+        fi
+        sleep {provisioningCheckInterval} &
+        wait $! || :
     done
+    echo "$(date -u): Shutting down provisioningJob..."
+    sleep 20
+    echo "$(date -u): Running FINAL provisioner"
+    "${allocate[@]}"
     exit 0
   provisioningScriptConfig: |
     config.platform["{provisioningPlatform}"].user.name="${USER}"
diff --git a/python/lsst/ctrl/bps/htcondor/lssthtc.py b/python/lsst/ctrl/bps/htcondor/lssthtc.py
index 1461214..a753409 100644
--- a/python/lsst/ctrl/bps/htcondor/lssthtc.py
+++ b/python/lsst/ctrl/bps/htcondor/lssthtc.py
@@ -222,6 +222,9 @@ class WmsNodeType(IntEnum):
     "periodic_remove",
     "accounting_group",
     "accounting_group_user",
+    "kill_sig",
+    "want_graceful_removal",
+    "job_max_vacate_time",
 }
 HTC_VALID_JOB_DAG_KEYS = {
     "dir",
@@ -1144,7 +1147,7 @@ def add_service_job(self, job):
         Parameters
         ----------
         job : `HTCJob`
-            HTCJob to add to the HTCDag as a FINAL job.
+            HTCJob to add to the HTCDag as a SERVICE job.
         """
         # Add dag level attributes to each job
         job.add_job_attrs(self.graph["attr"])
diff --git a/python/lsst/ctrl/bps/htcondor/provisioner.py b/python/lsst/ctrl/bps/htcondor/provisioner.py
index 634ba70..b853625 100644
--- a/python/lsst/ctrl/bps/htcondor/provisioner.py
+++ b/python/lsst/ctrl/bps/htcondor/provisioner.py
@@ -192,6 +192,9 @@ def provision(self, dag: HTCDag) -> None:
             "executable": f"{self.script_name}",
             "should_transfer_files": "NO",
             "getenv": "True",
+            "kill_sig": "SIGTERM",
+            "want_graceful_removal": "True",
+            "job_max_vacate_time": "60",
         }
         cmds |= {
             "output": str(job.subfile.with_suffix(".$(Cluster).out")),
diff --git a/tests/test_htcondor_service.py b/tests/test_htcondor_service.py
index 31df2de..5022ac2 100644
--- a/tests/test_htcondor_service.py
+++ b/tests/test_htcondor_service.py
@@ -227,7 +227,7 @@ def testPrepareProvision(self, mock_write):
             prov_script = Path(tmpdir) / "provisioningJob.bash"
             self.assertTrue(prov_script.is_file())
             script_contents = prov_script.read_text()
-            self.assertIn(f"--nodeset {timestamp}", script_contents)
+            self.assertIn(f"--nodeset '{timestamp}'", script_contents)
 
     def testSubmitWithConfigPath(self):
         """Only testing value for wms_config_path being passed
diff --git a/tests/test_provisioner.py b/tests/test_provisioner.py
index 37820f9..390590e 100644
--- a/tests/test_provisioner.py
+++ b/tests/test_provisioner.py
@@ -149,6 +149,9 @@ def testProvision(self):
             "output": f"jobs/{script.stem}/{script.stem}.$(Cluster).out",
             "error": f"jobs/{script.stem}/{script.stem}.$(Cluster).out",
             "log": f"jobs/{script.stem}/{script.stem}.$(Cluster).log",
+            "kill_sig": "SIGTERM",
+            "want_graceful_removal": "True",
+            "job_max_vacate_time": "60",
         }
         dag = HTCDag(name="default")