SCHED-1897: Add idle node memory health check

helm/slurm-cluster/slurm_scripts/idle_mem_used.sh

@@ -0,0 +1,91 @@
+#!/bin/bash
+
+set -euo pipefail
+
+meminfo_path="${IDLE_MEM_USED_MEMINFO_PATH:-/proc/meminfo}"
+node_name="${SLURMD_NODENAME:-unknown}"
+node_real_memory_bytes="${CHECKS_NODE_REAL_MEM_BYTES:-}"
+
+echo "[$(date)] Check memory usage when node ${node_name} is idle"
+echo "Slurm RealMemory input: ${node_real_memory_bytes:-<unavailable>} bytes"
+echo "Idle state source: local 'scontrol listjobs' (no controller RPC)"
+echo "Idle state rule: exit code 1 with \"No slurmstepd's found on this node\" means the node has no local jobs"
+
+if listjobs_output="$(scontrol listjobs 2>&1)"; then
+    listjobs_rc=0
+else
+    listjobs_rc=$?
+fi
+
+echo "scontrol listjobs exit code: ${listjobs_rc}"
+echo "scontrol listjobs output: ${listjobs_output:-<empty>}"
+
+if (( listjobs_rc == 0 )); then
+    node_is_idle=false
+    echo "Local Slurm jobs are present; treating the node as non-idle"
+elif (( listjobs_rc == 1 )) && [[ "${listjobs_output}" == *"No slurmstepd's found on this node"* ]]; then
+    node_is_idle=true
+    echo "No local slurmstepd was found; treating the node as idle"
+else
+    echo "Could not determine whether the node is idle from 'scontrol listjobs'; skipping memory validation" >&2
+    exit 0
+fi
+
+echo "Node is idle: ${node_is_idle}"
+if [[ "${node_is_idle}" != "true" ]]; then
+    echo "Node has local jobs; skipping memory validation"
+    exit 0
+fi
+
+if ! [[ "${node_real_memory_bytes}" =~ ^[0-9]+$ ]] || [[ "${node_real_memory_bytes}" == "0" ]]; then
+    echo "Invalid or unavailable Slurm RealMemory '${node_real_memory_bytes:-<unavailable>}'; expected a positive byte count, skipping memory validation" >&2
+    exit 0
+fi
+
+meminfo_values="$(
+    awk '
+        /^MemTotal:/ { total = $2 }
+        /^MemAvailable:/ { available = $2 }
+        END { if (total != "" && available != "") print total, available }
+    ' "${meminfo_path}" 2>/dev/null || true
+)"
+read -r mem_total_kib mem_available_kib <<< "${meminfo_values}"
+
+if ! [[ "${mem_total_kib:-}" =~ ^[0-9]+$ ]] ||
+   ! [[ "${mem_available_kib:-}" =~ ^[0-9]+$ ]] ||
+   (( mem_available_kib > mem_total_kib )); then
+    echo "Could not determine valid MemTotal and MemAvailable values from ${meminfo_path}; skipping memory validation" >&2
+    exit 0
+fi
+
+mem_total_bytes=$((mem_total_kib * 1024))
+mem_available_bytes=$((mem_available_kib * 1024))
+mem_used_bytes=$((mem_total_bytes - mem_available_bytes))
+
+if (( node_real_memory_bytes > mem_total_bytes )); then
+    echo "Slurm RealMemory ${node_real_memory_bytes} bytes exceeds MemTotal ${mem_total_bytes} bytes; skipping memory validation" >&2
+    exit 0
+fi
+
+max_idle_used_bytes=$((mem_total_bytes - node_real_memory_bytes))
+
+mem_total_gb="$(awk -v bytes="${mem_total_bytes}" 'BEGIN { printf "%.2f", bytes / 1000000000 }')"
+mem_available_gb="$(awk -v bytes="${mem_available_bytes}" 'BEGIN { printf "%.2f", bytes / 1000000000 }')"
+mem_used_gb="$(awk -v bytes="${mem_used_bytes}" 'BEGIN { printf "%.2f", bytes / 1000000000 }')"
+node_real_memory_gb="$(awk -v bytes="${node_real_memory_bytes}" 'BEGIN { printf "%.2f", bytes / 1000000000 }')"
+max_idle_used_gb="$(awk -v bytes="${max_idle_used_bytes}" 'BEGIN { printf "%.2f", bytes / 1000000000 }')"
+
+echo "Memory source: ${meminfo_path}"
+echo "Memory measurements: total=${mem_total_gb} GB (${mem_total_bytes} bytes), available=${mem_available_gb} GB (${mem_available_bytes} bytes), used=total-available=${mem_used_gb} GB (${mem_used_bytes} bytes)"
+echo "Slurm RealMemory: ${node_real_memory_gb} GB (${node_real_memory_bytes} bytes)"
+echo "Derived maximum idle used memory: MemTotal-RealMemory=${max_idle_used_gb} GB (${max_idle_used_bytes} bytes)"
+
+if (( mem_available_bytes < node_real_memory_bytes )); then
+    memory_deficit_bytes=$((node_real_memory_bytes - mem_available_bytes))
+    memory_deficit_gb="$(awk -v bytes="${memory_deficit_bytes}" 'BEGIN { printf "%.2f", bytes / 1000000000 }')"
+    echo "Node ${node_name} is IDLE but has only ${mem_available_gb} GB of memory available, below Slurm RealMemory ${node_real_memory_gb} GB by ${memory_deficit_gb} GB (used: ${mem_used_gb} GB; derived maximum idle usage: ${max_idle_used_gb} GB; MemTotal: ${mem_total_gb} GB). This may indicate leftover or spurious processes consuming memory." >&3
+    exit 1
+fi
+
+echo "Idle node leaves enough memory available for Slurm RealMemory"
+exit 0

helm/slurm-cluster/slurm_scripts/idle_mem_used.sh.json

@@ -0,0 +1,16 @@
+{
+  "name": "idle_mem_used",
+  "command": "./idle_mem_used.sh",
+  "platforms": ["any"],
+  "skip_for_cpu_jobs": false,
+  "skip_for_partial_gpu_jobs": false,
+  "contexts": ["hc_program"],
+  "node_states": ["any"],
+  "on_fail": "drain",
+  "on_ok": "none",
+  "reason_base": "[node_problem] $name",
+  "reason_append_details": true,
+  "run_in_jail": false,
+  "log": "slurm_scripts/$worker.$name.$context.out",
+  "need_env": ["CHECKS_NODE_REAL_MEM_BYTES"]
+}

helm/slurm-cluster/slurm_scripts/idle_mem_used_test.py

@@ -0,0 +1,171 @@
+import os
+import subprocess
+import tempfile
+import unittest
+from pathlib import Path
+
+
+SCRIPT_PATH = Path(__file__).with_name("idle_mem_used.sh")
+GIB = 1024 * 1024 * 1024
+
+
+class IdleMemUsedTest(unittest.TestCase):
+    def run_check(
+        self,
+        *,
+        listjobs_rc: int,
+        listjobs_output: str,
+        total_bytes: int | None,
+        available_bytes: int | None,
+        node_real_memory_bytes: int | None = 56 * GIB,
+    ) -> subprocess.CompletedProcess[str]:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            scontrol_path = Path(tmpdir) / "scontrol"
+            scontrol_path.write_text(
+                "#!/bin/bash\n"
+                "printf '%s\\n' \"${MOCK_SCONTROL_LISTJOBS_OUTPUT}\"\n"
+                "exit \"${MOCK_SCONTROL_LISTJOBS_RC}\"\n",
+                encoding="utf-8",
+            )
+            scontrol_path.chmod(0o755)
+
+            meminfo_path = Path(tmpdir) / "meminfo"
+            if total_bytes is not None and available_bytes is not None:
+                meminfo_path.write_text(
+                    f"MemTotal:       {total_bytes // 1024} kB\n"
+                    f"MemAvailable:   {available_bytes // 1024} kB\n",
+                    encoding="utf-8",
+                )
+
+            env = os.environ.copy()
+            if node_real_memory_bytes is None:
+                env.pop("CHECKS_NODE_REAL_MEM_BYTES", None)
+            else:
+                env["CHECKS_NODE_REAL_MEM_BYTES"] = str(node_real_memory_bytes)
+            env.update(
+                {
+                    "PATH": f"{tmpdir}:{env['PATH']}",
+                    "SLURMD_NODENAME": "worker-1",
+                    "IDLE_MEM_USED_MEMINFO_PATH": str(meminfo_path),
+                    "MOCK_SCONTROL_LISTJOBS_RC": str(listjobs_rc),
+                    "MOCK_SCONTROL_LISTJOBS_OUTPUT": listjobs_output,
+                }
+            )
+
+            return subprocess.run(
+                [
+                    "bash",
+                    "-c",
+                    'exec 3>&1; exec bash "$1"',
+                    "idle-mem-used-test",
+                    str(SCRIPT_PATH),
+                ],
+                check=False,
+                env=env,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+            )
+
+    def test_non_idle_node_skips_memory_validation(self):
+        result = self.run_check(
+            listjobs_rc=0,
+            listjobs_output="JOBID\n1011",
+            total_bytes=None,
+            available_bytes=None,
+        )
+
+        self.assertEqual(0, result.returncode)
+        self.assertIn("Node is idle: false", result.stdout)
+        self.assertIn("Local Slurm jobs are present", result.stdout)
+        self.assertIn("Node has local jobs; skipping memory validation", result.stdout)
+        self.assertNotIn("Memory measurements:", result.stdout)
+
+    def test_idle_node_with_enough_available_memory_passes(self):
+        result = self.run_check(
+            listjobs_rc=1,
+            listjobs_output="No slurmstepd's found on this node",
+            total_bytes=64 * GIB,
+            available_bytes=60 * GIB,
+        )
+
+        self.assertEqual(0, result.returncode)
+        self.assertIn("Node is idle: true", result.stdout)
+        self.assertIn("No local slurmstepd was found", result.stdout)
+        self.assertIn(f"used=total-available=4.29 GB ({4 * GIB} bytes)", result.stdout)
+        self.assertIn(
+            f"Slurm RealMemory: 60.13 GB ({56 * GIB} bytes)", result.stdout
+        )
+        self.assertIn(
+            f"Derived maximum idle used memory: MemTotal-RealMemory=8.59 GB ({8 * GIB} bytes)",
+            result.stdout,
+        )
+
+    def test_idle_node_with_insufficient_available_memory_fails(self):
+        result = self.run_check(
+            listjobs_rc=1,
+            listjobs_output="No slurmstepd's found on this node",
+            total_bytes=64 * GIB,
+            available_bytes=22 * GIB,
+        )
+
+        self.assertEqual(1, result.returncode)
+        self.assertIn("Node is idle: true", result.stdout)
+        self.assertIn("Node worker-1 is IDLE but has only 23.62 GB", result.stdout)
+        self.assertIn("below Slurm RealMemory 60.13 GB by 36.51 GB", result.stdout)
+        self.assertIn("derived maximum idle usage: 8.59 GB", result.stdout)
+        self.assertIn("leftover or spurious processes", result.stdout)
+
+    def test_unavailable_memory_data_does_not_drain_node(self):
+        result = self.run_check(
+            listjobs_rc=1,
+            listjobs_output="No slurmstepd's found on this node",
+            total_bytes=None,
+            available_bytes=None,
+        )
+
+        self.assertEqual(0, result.returncode)
+        self.assertIn("Could not determine valid MemTotal and MemAvailable", result.stderr)
+
+    def test_unavailable_real_memory_does_not_drain_node(self):
+        result = self.run_check(
+            listjobs_rc=1,
+            listjobs_output="No slurmstepd's found on this node",
+            total_bytes=None,
+            available_bytes=None,
+            node_real_memory_bytes=None,
+        )
+
+        self.assertEqual(0, result.returncode)
+        self.assertIn("Invalid or unavailable Slurm RealMemory", result.stderr)
+        self.assertNotIn("Memory measurements:", result.stdout)
+
+    def test_real_memory_larger_than_memtotal_does_not_drain_node(self):
+        result = self.run_check(
+            listjobs_rc=1,
+            listjobs_output="No slurmstepd's found on this node",
+            total_bytes=64 * GIB,
+            available_bytes=60 * GIB,
+            node_real_memory_bytes=72 * GIB,
+        )
+
+        self.assertEqual(0, result.returncode)
+        self.assertIn("exceeds MemTotal", result.stderr)
+        self.assertNotIn("Derived maximum idle used memory", result.stdout)
+
+    def test_unexpected_listjobs_failure_does_not_validate_memory(self):
+        result = self.run_check(
+            listjobs_rc=2,
+            listjobs_output="Unable to inspect local jobs",
+            total_bytes=None,
+            available_bytes=None,
+        )
+
+        self.assertEqual(0, result.returncode)
+        self.assertIn("scontrol listjobs exit code: 2", result.stdout)
+        self.assertIn("Could not determine whether the node is idle", result.stderr)
+        self.assertNotIn("Memory measurements:", result.stdout)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

helm/slurm-cluster/tests/idle_mem_used_test.yaml

@@ -0,0 +1,36 @@
+suite: test idle memory health check
+templates:
+  - slurm-scripts-cm.yaml
+tests:
+  - it: should enable the idle memory check with node RealMemory
+    asserts:
+      - isNotNull:
+          path: data["idle_mem_used.sh"]
+      - matchRegex:
+          path: data["checks.json"]
+          pattern: '"name": "idle_mem_used"'
+      - matchRegex:
+          path: data["checks.json"]
+          pattern: '"command": "./idle_mem_used.sh"'
+      - matchRegex:
+          path: data["checks.json"]
+          pattern: '"contexts": \[\s*"hc_program"\s*\]'
+      - matchRegex:
+          path: data["checks.json"]
+          pattern: '"on_fail": "drain"'
+      - matchRegex:
+          path: data["checks.json"]
+          pattern: '"need_env": \[\s*"CHECKS_NODE_REAL_MEM_BYTES"\s*\]'
+
+  - it: should allow the idle memory check to be disabled
+    set:
+      slurmScripts:
+        builtIn:
+          idle_mem_used.sh:
+            enabled: false
+    asserts:
+      - isNull:
+          path: data["idle_mem_used.sh"]
+      - notMatchRegex:
+          path: data["checks.json"]
+          pattern: '"name": "idle_mem_used"'

helm/slurm-cluster/values.yaml

@@ -774,6 +774,10 @@ slurmScripts:
       enabled: true
       customContent: null
       customConfig: null
+    idle_mem_used.sh:
+      enabled: true
+      customContent: null
+      customConfig: null
     job_tmpfs_delete.sh:
       enabled: true
       customContent: null