diff --git a/helm/slurm-cluster/slurm_scripts/idle_mem_used.drain.sh b/helm/slurm-cluster/slurm_scripts/idle_mem_used.drain.sh new file mode 100644 index 000000000..c748c562f --- /dev/null +++ b/helm/slurm-cluster/slurm_scripts/idle_mem_used.drain.sh @@ -0,0 +1,96 @@ +#!/bin/bash + +set -euo pipefail + +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:-} bytes" +echo "Idle state source: local 'scontrol listjobs --json' (no controller RPC)" +echo "Idle state rule: an empty JSON .jobs array means the node has no local jobs" + +if listjobs_output="$(scontrol listjobs --json 2>&1)"; then + listjobs_rc=0 +else + listjobs_rc=$? +fi + +echo "scontrol listjobs --json exit code: ${listjobs_rc}" +echo "scontrol listjobs --json output: ${listjobs_output:-}" + +if (( listjobs_rc != 0 )); then + echo "Could not determine whether the node is idle because 'scontrol listjobs --json' failed; skipping memory validation" >&2 + exit 0 +fi + +if ! local_job_count="$( + jq -er ' + if (.jobs | type) == "array" then + .jobs | length + else + error(".jobs must be an array") + end + ' <<<"${listjobs_output}" 2>/dev/null +)"; then + echo "Could not determine whether the node is idle because 'scontrol listjobs --json' returned invalid job data; skipping memory validation" >&2 + exit 0 +fi + +echo "Local Slurm job count from JSON .jobs array: ${local_job_count}" +if (( local_job_count == 0 )); then + node_is_idle=true + echo "The JSON .jobs array is empty; treating the node as idle" +else + node_is_idle=false + echo "The JSON .jobs array contains local jobs; treating the node as non-idle" +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:-}'; expected a positive byte count, skipping memory validation" >&2 + exit 0 +fi + +if ! free_bytes_output="$(LC_ALL=C free -b 2>&1)"; then + echo "Could not read local memory information with 'free -b'; skipping memory validation" >&2 + exit 0 +fi + +memory_values="$(awk '/^Mem:/ { print $2, $7; exit }' <<<"${free_bytes_output}")" +read -r mem_total_bytes mem_available_bytes <<<"${memory_values}" + +if ! [[ "${mem_total_bytes:-}" =~ ^[0-9]+$ ]] || + ! [[ "${mem_available_bytes:-}" =~ ^[0-9]+$ ]] || + (( mem_available_bytes > mem_total_bytes )); then + echo "Could not determine valid total and available memory from 'free -b'; skipping memory validation" >&2 + exit 0 +fi + +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 + +mem_available_gb="$(awk -v bytes="${mem_available_bytes}" 'BEGIN { printf "%.2f", bytes / 1000000000 }')" +node_real_memory_gb="$(awk -v bytes="${node_real_memory_bytes}" 'BEGIN { printf "%.2f", bytes / 1000000000 }')" + +echo "Memory source: local 'free'" +echo "Memory comparison: available=${mem_available_gb} GB (${mem_available_bytes} bytes), Slurm RealMemory=${node_real_memory_gb} GB (${node_real_memory_bytes} bytes)" +echo "Memory snapshot (free -hw):" +if ! LC_ALL=C free -hw; then + echo "Could not print the human-readable memory snapshot with 'free -hw'" >&2 +fi + +if (( mem_available_bytes < node_real_memory_bytes )); then + echo "available memory ${mem_available_gb} GB < configured ${node_real_memory_gb} GB; stop leftover processes or reboot" >&3 + exit 1 +fi + +echo "Idle node leaves enough memory available for Slurm RealMemory" +exit 0 diff --git a/helm/slurm-cluster/slurm_scripts/idle_mem_used.drain.sh.json b/helm/slurm-cluster/slurm_scripts/idle_mem_used.drain.sh.json new file mode 100644 index 000000000..6887d5b39 --- /dev/null +++ b/helm/slurm-cluster/slurm_scripts/idle_mem_used.drain.sh.json @@ -0,0 +1,16 @@ +{ + "name": "idle_mem_used", + "command": "./idle_mem_used.drain.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": "[user_problem] $name", + "reason_append_details": true, + "run_in_jail": false, + "log": "slurm_scripts/$worker.$name.drain.$context.out", + "need_env": ["CHECKS_NODE_REAL_MEM_BYTES"] +} diff --git a/helm/slurm-cluster/slurm_scripts/idle_mem_used.undrain.sh b/helm/slurm-cluster/slurm_scripts/idle_mem_used.undrain.sh new file mode 100644 index 000000000..9f40517b2 --- /dev/null +++ b/helm/slurm-cluster/slurm_scripts/idle_mem_used.undrain.sh @@ -0,0 +1,55 @@ +#!/bin/bash + +set -euo pipefail + +# This recovery check must fail closed: only a fully validated healthy result +# may allow check_runner to undrain the node. +node_name="${SLURMD_NODENAME:-unknown}" +node_real_memory_bytes="${CHECKS_NODE_REAL_MEM_BYTES:-}" + +echo "[$(date)] Check whether memory usage has recovered on drained node ${node_name}" +echo "Slurm RealMemory input: ${node_real_memory_bytes:-} bytes" +echo "Node eligibility source: this check is scheduled only for drained nodes" + +if ! [[ "${node_real_memory_bytes}" =~ ^[0-9]+$ ]] || [[ "${node_real_memory_bytes}" == "0" ]]; then + echo "Invalid or unavailable Slurm RealMemory '${node_real_memory_bytes:-}'; keeping node drained" >&2 + exit 1 +fi + +if ! free_bytes_output="$(LC_ALL=C free -b 2>&1)"; then + echo "Could not read local memory information with 'free -b'; keeping node drained" >&2 + exit 1 +fi + +memory_values="$(awk '/^Mem:/ { print $2, $7; exit }' <<<"${free_bytes_output}")" +read -r mem_total_bytes mem_available_bytes <<<"${memory_values}" + +if ! [[ "${mem_total_bytes:-}" =~ ^[0-9]+$ ]] || + ! [[ "${mem_available_bytes:-}" =~ ^[0-9]+$ ]] || + (( mem_available_bytes > mem_total_bytes )); then + echo "Could not determine valid total and available memory from 'free -b'; keeping node drained" >&2 + exit 1 +fi + +if (( node_real_memory_bytes > mem_total_bytes )); then + echo "Slurm RealMemory ${node_real_memory_bytes} bytes exceeds MemTotal ${mem_total_bytes} bytes; keeping node drained" >&2 + exit 1 +fi + +mem_available_gb="$(awk -v bytes="${mem_available_bytes}" 'BEGIN { printf "%.2f", bytes / 1000000000 }')" +node_real_memory_gb="$(awk -v bytes="${node_real_memory_bytes}" 'BEGIN { printf "%.2f", bytes / 1000000000 }')" + +echo "Memory source: local 'free'" +echo "Memory comparison: available=${mem_available_gb} GB (${mem_available_bytes} bytes), Slurm RealMemory=${node_real_memory_gb} GB (${node_real_memory_bytes} bytes)" +echo "Memory snapshot (free -hw):" +if ! LC_ALL=C free -hw; then + echo "Could not print the human-readable memory snapshot with 'free -hw'" >&2 +fi + +if (( mem_available_bytes < node_real_memory_bytes )); then + echo "Available memory ${mem_available_gb} GB is still below configured memory ${node_real_memory_gb} GB; keeping node drained" >&2 + exit 1 +fi + +echo "Drained node leaves enough memory available for Slurm RealMemory; memory recovery confirmed" +exit 0 diff --git a/helm/slurm-cluster/slurm_scripts/idle_mem_used.undrain.sh.json b/helm/slurm-cluster/slurm_scripts/idle_mem_used.undrain.sh.json new file mode 100644 index 000000000..d772f693f --- /dev/null +++ b/helm/slurm-cluster/slurm_scripts/idle_mem_used.undrain.sh.json @@ -0,0 +1,16 @@ +{ + "name": "idle_mem_used", + "command": "./idle_mem_used.undrain.sh", + "platforms": ["any"], + "skip_for_cpu_jobs": false, + "skip_for_partial_gpu_jobs": false, + "contexts": ["hc_program"], + "node_states": ["drain"], + "on_fail": "none", + "on_ok": "undrain", + "reason_base": "[user_problem] $name", + "reason_append_details": false, + "run_in_jail": false, + "log": "slurm_scripts/$worker.$name.undrain.$context.out", + "need_env": ["CHECKS_NODE_REAL_MEM_BYTES"] +} diff --git a/helm/slurm-cluster/slurm_scripts/idle_mem_used_test.py b/helm/slurm-cluster/slurm_scripts/idle_mem_used_test.py new file mode 100644 index 000000000..7060dab98 --- /dev/null +++ b/helm/slurm-cluster/slurm_scripts/idle_mem_used_test.py @@ -0,0 +1,289 @@ +import os +import subprocess +import tempfile +import unittest +from pathlib import Path + + +DRAIN_SCRIPT_PATH = Path(__file__).with_name("idle_mem_used.drain.sh") +UNDRAIN_SCRIPT_PATH = Path(__file__).with_name("idle_mem_used.undrain.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, + free_bytes_rc: int = 0, + script_path: Path = DRAIN_SCRIPT_PATH, + ) -> subprocess.CompletedProcess[str]: + with tempfile.TemporaryDirectory() as tmpdir: + scontrol_path = Path(tmpdir) / "scontrol" + scontrol_path.write_text( + "#!/bin/bash\n" + 'if [[ "$*" != "listjobs --json" ]]; then\n' + ' printf \'Unexpected arguments: %s\\n\' "$*" >&2\n' + " exit 64\n" + "fi\n" + "printf '%s\\n' \"${MOCK_SCONTROL_LISTJOBS_OUTPUT}\"\n" + "exit \"${MOCK_SCONTROL_LISTJOBS_RC}\"\n", + encoding="utf-8", + ) + scontrol_path.chmod(0o755) + + free_path = Path(tmpdir) / "free" + free_path.write_text( + "#!/bin/bash\n" + 'case "$*" in\n' + ' "-b")\n' + ' if (( MOCK_FREE_BYTES_RC != 0 )); then\n' + ' printf "Unable to read local memory\\n" >&2\n' + ' exit "${MOCK_FREE_BYTES_RC}"\n' + " fi\n" + ' printf " total used free shared buff/cache available\\n"\n' + ' printf "Mem: %s 0 0 0 0 %s\\n" \\\n' + ' "${MOCK_FREE_TOTAL_BYTES}" "${MOCK_FREE_AVAILABLE_BYTES}"\n' + " ;;\n" + ' "-hw")\n' + ' printf " total used free shared buffers cache available\\n"\n' + ' printf "Mem: mock mock mock mock mock mock mock\\n"\n' + " ;;\n" + " *)\n" + ' printf \'Unexpected arguments: %s\\n\' "$*" >&2\n' + " exit 64\n" + " ;;\n" + "esac\n", + encoding="utf-8", + ) + free_path.chmod(0o755) + + 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", + "MOCK_SCONTROL_LISTJOBS_RC": str(listjobs_rc), + "MOCK_SCONTROL_LISTJOBS_OUTPUT": listjobs_output, + "MOCK_FREE_TOTAL_BYTES": ( + str(total_bytes) if total_bytes is not None else "" + ), + "MOCK_FREE_AVAILABLE_BYTES": ( + str(available_bytes) if available_bytes is not None else "" + ), + "MOCK_FREE_BYTES_RC": str(free_bytes_rc), + } + ) + + 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='{"jobs":[{"job_id":1011}],"errors":[]}', + total_bytes=None, + available_bytes=None, + ) + + self.assertEqual(0, result.returncode) + self.assertIn("Node is idle: false", result.stdout) + self.assertIn("Local Slurm job count from JSON .jobs array: 1", result.stdout) + self.assertIn("JSON .jobs array contains local jobs", result.stdout) + self.assertIn("Node has local jobs; skipping memory validation", result.stdout) + self.assertNotIn("Memory comparison:", result.stdout) + + def test_idle_node_with_enough_available_memory_passes(self): + result = self.run_check( + listjobs_rc=0, + listjobs_output='{"jobs":[],"errors":[]}', + total_bytes=64 * GIB, + available_bytes=60 * GIB, + ) + + self.assertEqual(0, result.returncode) + self.assertIn("Node is idle: true", result.stdout) + self.assertIn("Local Slurm job count from JSON .jobs array: 0", result.stdout) + self.assertIn("JSON .jobs array is empty", result.stdout) + self.assertIn( + f"available=64.42 GB ({60 * GIB} bytes)", result.stdout + ) + self.assertIn( + f"Slurm RealMemory=60.13 GB ({56 * GIB} bytes)", + result.stdout, + ) + self.assertIn("Memory snapshot (free -hw):", result.stdout) + self.assertIn("Mem: mock", result.stdout) + + def test_idle_node_with_insufficient_available_memory_fails(self): + result = self.run_check( + listjobs_rc=0, + listjobs_output='{"jobs":[],"errors":[]}', + total_bytes=64 * GIB, + available_bytes=22 * GIB, + ) + + self.assertEqual(1, result.returncode) + self.assertIn("Node is idle: true", result.stdout) + self.assertIn( + "available memory 23.62 GB < configured 60.13 GB", result.stdout + ) + self.assertIn("stop leftover processes or reboot", result.stdout) + + def test_unavailable_memory_data_does_not_drain_node(self): + result = self.run_check( + listjobs_rc=0, + listjobs_output='{"jobs":[],"errors":[]}', + total_bytes=None, + available_bytes=None, + ) + + self.assertEqual(0, result.returncode) + self.assertIn( + "Could not determine valid total and available memory", result.stderr + ) + + def test_free_failure_does_not_drain_node(self): + result = self.run_check( + listjobs_rc=0, + listjobs_output='{"jobs":[],"errors":[]}', + total_bytes=None, + available_bytes=None, + free_bytes_rc=2, + ) + + self.assertEqual(0, result.returncode) + self.assertIn("Could not read local memory information", result.stderr) + self.assertNotIn("Memory comparison:", result.stdout) + + def test_unavailable_real_memory_does_not_drain_node(self): + result = self.run_check( + listjobs_rc=0, + listjobs_output='{"jobs":[],"errors":[]}', + 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 comparison:", result.stdout) + + def test_real_memory_larger_than_memtotal_does_not_drain_node(self): + result = self.run_check( + listjobs_rc=0, + listjobs_output='{"jobs":[],"errors":[]}', + 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("Memory comparison:", 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 --json exit code: 2", result.stdout) + self.assertIn("scontrol listjobs --json' failed", result.stderr) + self.assertNotIn("Memory comparison:", result.stdout) + + def test_invalid_listjobs_json_does_not_validate_memory(self): + result = self.run_check( + listjobs_rc=0, + listjobs_output="not JSON", + total_bytes=None, + available_bytes=None, + ) + + self.assertEqual(0, result.returncode) + self.assertIn("returned invalid job data", result.stderr) + self.assertNotIn("Node is idle:", result.stdout) + self.assertNotIn("Memory comparison:", result.stdout) + + def test_missing_jobs_array_does_not_validate_memory(self): + result = self.run_check( + listjobs_rc=0, + listjobs_output='{"jobs":null,"errors":[]}', + total_bytes=None, + available_bytes=None, + ) + + self.assertEqual(0, result.returncode) + self.assertIn("returned invalid job data", result.stderr) + self.assertNotIn("Node is idle:", result.stdout) + self.assertNotIn("Memory comparison:", result.stdout) + + def test_undrain_confirms_recovery_without_querying_local_jobs(self): + result = self.run_check( + listjobs_rc=64, + listjobs_output="undrain must not query local jobs", + total_bytes=64 * GIB, + available_bytes=60 * GIB, + script_path=UNDRAIN_SCRIPT_PATH, + ) + + self.assertEqual(0, result.returncode) + self.assertNotIn("listjobs", result.stdout) + self.assertNotIn("listjobs", result.stderr) + self.assertIn("scheduled only for drained nodes", result.stdout) + self.assertIn("memory recovery confirmed", result.stdout) + + def test_undrain_keeps_node_drained_when_available_memory_is_insufficient(self): + result = self.run_check( + listjobs_rc=64, + listjobs_output="undrain must not query local jobs", + total_bytes=64 * GIB, + available_bytes=22 * GIB, + script_path=UNDRAIN_SCRIPT_PATH, + ) + + self.assertEqual(1, result.returncode) + self.assertIn("is still below configured memory", result.stderr) + self.assertIn("keeping node drained", result.stderr) + + def test_undrain_keeps_node_drained_when_memory_data_is_unavailable(self): + result = self.run_check( + listjobs_rc=64, + listjobs_output="undrain must not query local jobs", + total_bytes=None, + available_bytes=None, + free_bytes_rc=2, + script_path=UNDRAIN_SCRIPT_PATH, + ) + + self.assertEqual(1, result.returncode) + self.assertIn("Could not read local memory information", result.stderr) + self.assertIn("keeping node drained", result.stderr) + self.assertNotIn("Memory comparison:", result.stdout) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/helm/slurm-cluster/tests/idle_mem_used_test.yaml b/helm/slurm-cluster/tests/idle_mem_used_test.yaml new file mode 100644 index 000000000..b70c982c6 --- /dev/null +++ b/helm/slurm-cluster/tests/idle_mem_used_test.yaml @@ -0,0 +1,51 @@ +suite: test idle memory health check +templates: + - slurm-scripts-cm.yaml +tests: + - it: should enable the idle memory drain and undrain checks with node RealMemory + asserts: + - isNotNull: + path: data["idle_mem_used.drain.sh"] + - isNotNull: + path: data["idle_mem_used.undrain.sh"] + - matchRegex: + path: data["checks.json"] + pattern: '"name": "idle_mem_used"' + - matchRegex: + path: data["checks.json"] + pattern: '"command": "./idle_mem_used.drain.sh"' + - matchRegex: + path: data["checks.json"] + pattern: '"command": "./idle_mem_used.undrain.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: '"node_states": \[\s*"drain"\s*\]' + - matchRegex: + path: data["checks.json"] + pattern: '"on_ok": "undrain"' + - 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.drain.sh: + enabled: false + idle_mem_used.undrain.sh: + enabled: false + asserts: + - isNull: + path: data["idle_mem_used.drain.sh"] + - isNull: + path: data["idle_mem_used.undrain.sh"] + - notMatchRegex: + path: data["checks.json"] + pattern: '"command": "./idle_mem_used.(drain|undrain).sh"' diff --git a/helm/slurm-cluster/values.yaml b/helm/slurm-cluster/values.yaml index 551844cdf..a0ff18e4c 100644 --- a/helm/slurm-cluster/values.yaml +++ b/helm/slurm-cluster/values.yaml @@ -774,6 +774,14 @@ slurmScripts: enabled: true customContent: null customConfig: null + idle_mem_used.drain.sh: + enabled: true + customContent: null + customConfig: null + idle_mem_used.undrain.sh: + enabled: true + customContent: null + customConfig: null job_tmpfs_delete.sh: enabled: true customContent: null