diff --git a/helm/slurm-cluster/slurm_scripts/idle_mem_used.drain.py b/helm/slurm-cluster/slurm_scripts/idle_mem_used.drain.py new file mode 100644 index 000000000..124da9502 --- /dev/null +++ b/helm/slurm-cluster/slurm_scripts/idle_mem_used.drain.py @@ -0,0 +1,243 @@ +import json +import logging +import os +import subprocess +import sys +import time + + +BYTES_PER_GB = 1_000_000_000 +LOGGER = logging.getLogger("idle_mem_used") + + +class InfoFilter(logging.Filter): + def filter(self, record: logging.LogRecord) -> bool: + return record.levelno < logging.WARNING + + +def configure_logging() -> None: + stdout_handler = logging.StreamHandler(sys.stdout) + stdout_handler.setLevel(logging.DEBUG) + stdout_handler.addFilter(InfoFilter()) + stdout_handler.setFormatter(logging.Formatter("%(message)s")) + + stderr_handler = logging.StreamHandler(sys.stderr) + stderr_handler.setLevel(logging.WARNING) + stderr_handler.setFormatter(logging.Formatter("%(message)s")) + + LOGGER.setLevel(logging.INFO) + LOGGER.handlers.clear() + LOGGER.addHandler(stdout_handler) + LOGGER.addHandler(stderr_handler) + LOGGER.propagate = False + + +def run_capture(command: list[str]) -> subprocess.CompletedProcess[str]: + try: + return subprocess.run( + command, + check=False, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + except OSError as e: + return subprocess.CompletedProcess( + args=command, + returncode=127, + stdout=f"{command[0]}: {e}\n", + ) + + +def run_capture_c_locale(command: list[str]) -> subprocess.CompletedProcess[str]: + env = os.environ.copy() + env["LC_ALL"] = "C" + try: + return subprocess.run( + command, + check=False, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + env=env, + ) + except OSError as e: + return subprocess.CompletedProcess( + args=command, + returncode=127, + stdout=f"{command[0]}: {e}\n", + ) + + +def get_local_job_count(listjobs_output: str) -> int | None: + try: + data = json.loads(listjobs_output) + except json.JSONDecodeError: + return None + + jobs = data.get("jobs") + if not isinstance(jobs, list): + return None + + return len(jobs) + + +def parse_memory_values(free_bytes_output: str) -> tuple[int, int] | None: + for line in free_bytes_output.splitlines(): + fields = line.split() + if not fields or fields[0] != "Mem:": + continue + + if len(fields) < 7: + return None + + if not fields[1].isdecimal() or not fields[6].isdecimal(): + return None + + mem_total_bytes = int(fields[1]) + mem_available_bytes = int(fields[6]) + if mem_available_bytes > mem_total_bytes: + return None + + return mem_total_bytes, mem_available_bytes + + return None + + +def bytes_to_gb(value: int) -> str: + return f"{value / BYTES_PER_GB:.2f}" + + +def log_human_memory_snapshot() -> None: + LOGGER.info("Memory snapshot (free -hw):") + result = run_capture_c_locale(["free", "-hw"]) + if result.stdout: + LOGGER.info(result.stdout.rstrip("\n")) + + if result.returncode != 0: + LOGGER.warning( + "Could not print the human-readable memory snapshot with 'free -hw'" + ) + + +def write_drain_reason(message: str) -> None: + os.write(3, f"{message}\n".encode("utf-8", errors="backslashreplace")) + + +def main() -> int: + configure_logging() + + node_name = os.environ.get("SLURMD_NODENAME", "unknown") + node_real_memory_raw = os.environ.get("CHECKS_NODE_REAL_MEM_BYTES", "") + + LOGGER.info("[%s] Check memory usage when node %s is idle", time.ctime(), node_name) + LOGGER.info( + "Slurm RealMemory input: " + "%s bytes", + node_real_memory_raw if node_real_memory_raw else "", + ) + + listjobs_result = run_capture(["scontrol", "listjobs", "--json"]) + listjobs_output = listjobs_result.stdout.rstrip("\n") + + LOGGER.info("scontrol listjobs --json exit code: %s", listjobs_result.returncode) + LOGGER.info( + "scontrol listjobs --json output: " + "%s", + listjobs_output if listjobs_output else "", + ) + + if listjobs_result.returncode != 0: + LOGGER.warning( + "Could not determine whether the node is idle because " + "'scontrol listjobs --json' failed; skipping memory validation" + ) + return 0 + + local_job_count = get_local_job_count(listjobs_output) + if local_job_count is None: + LOGGER.warning( + "Could not determine whether the node is idle because " + "'scontrol listjobs --json' returned invalid job data; " + "skipping memory validation" + ) + return 0 + + LOGGER.info("Local Slurm job count from JSON .jobs array: %s", local_job_count) + node_is_idle = local_job_count == 0 + if node_is_idle: + LOGGER.info("The JSON .jobs array is empty; treating the node as idle") + else: + LOGGER.info( + "The JSON .jobs array contains local jobs; treating the node as non-idle" + ) + + LOGGER.info("Node is idle: %s", str(node_is_idle).lower()) + if not node_is_idle: + LOGGER.info("Node has local jobs; skipping memory validation") + return 0 + + if not node_real_memory_raw.isdecimal() or int(node_real_memory_raw) <= 0: + real_memory_display = node_real_memory_raw or "" + LOGGER.warning( + "Invalid or unavailable Slurm RealMemory '%s'; expected a positive " + "byte count, skipping memory validation", + real_memory_display, + ) + return 0 + + node_real_memory_bytes = int(node_real_memory_raw) + free_bytes_result = run_capture_c_locale(["free", "-b"]) + if free_bytes_result.returncode != 0: + LOGGER.warning( + "Could not read local memory information with 'free -b'; " + "skipping memory validation" + ) + return 0 + + memory_values = parse_memory_values(free_bytes_result.stdout) + if memory_values is None: + LOGGER.warning( + "Could not determine valid total and available memory from 'free -b'; " + "skipping memory validation" + ) + return 0 + + mem_total_bytes, mem_available_bytes = memory_values + + if node_real_memory_bytes > mem_total_bytes: + LOGGER.warning( + "Slurm RealMemory %s bytes exceeds MemTotal %s bytes; " + "skipping memory validation", + node_real_memory_bytes, + mem_total_bytes, + ) + return 0 + + mem_available_gb = bytes_to_gb(mem_available_bytes) + node_real_memory_gb = bytes_to_gb(node_real_memory_bytes) + + LOGGER.info( + "Memory comparison: " + "available=%s GB (%s bytes), " + "Slurm RealMemory=%s GB (%s bytes)", + mem_available_gb, + mem_available_bytes, + node_real_memory_gb, + node_real_memory_bytes, + ) + log_human_memory_snapshot() + + if mem_available_bytes < node_real_memory_bytes: + write_drain_reason( + f"available memory {mem_available_gb} GB < configured " + f"{node_real_memory_gb} GB; stop leftover processes or reboot" + ) + return 1 + + LOGGER.info("Idle node leaves enough memory available for Slurm RealMemory") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/helm/slurm-cluster/slurm_scripts/idle_mem_used.drain.py.json b/helm/slurm-cluster/slurm_scripts/idle_mem_used.drain.py.json new file mode 100644 index 000000000..06bdcb8af --- /dev/null +++ b/helm/slurm-cluster/slurm_scripts/idle_mem_used.drain.py.json @@ -0,0 +1,16 @@ +{ + "name": "idle_mem_used", + "command": "/usr/bin/python3 ./idle_mem_used.drain.py", + "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.py b/helm/slurm-cluster/slurm_scripts/idle_mem_used.undrain.py new file mode 100644 index 000000000..fe8b393b7 --- /dev/null +++ b/helm/slurm-cluster/slurm_scripts/idle_mem_used.undrain.py @@ -0,0 +1,177 @@ +import logging +import os +import subprocess +import sys +import time + + +BYTES_PER_GB = 1_000_000_000 +LOGGER = logging.getLogger("idle_mem_used") + + +class InfoFilter(logging.Filter): + def filter(self, record: logging.LogRecord) -> bool: + return record.levelno < logging.WARNING + + +def configure_logging() -> None: + stdout_handler = logging.StreamHandler(sys.stdout) + stdout_handler.setLevel(logging.DEBUG) + stdout_handler.addFilter(InfoFilter()) + stdout_handler.setFormatter(logging.Formatter("%(message)s")) + + stderr_handler = logging.StreamHandler(sys.stderr) + stderr_handler.setLevel(logging.WARNING) + stderr_handler.setFormatter(logging.Formatter("%(message)s")) + + LOGGER.setLevel(logging.INFO) + LOGGER.handlers.clear() + LOGGER.addHandler(stdout_handler) + LOGGER.addHandler(stderr_handler) + LOGGER.propagate = False + + +def run_capture_c_locale(command: list[str]) -> subprocess.CompletedProcess[str]: + env = os.environ.copy() + env["LC_ALL"] = "C" + try: + return subprocess.run( + command, + check=False, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + env=env, + ) + except OSError as e: + return subprocess.CompletedProcess( + args=command, + returncode=127, + stdout=f"{command[0]}: {e}\n", + ) + + +def parse_memory_values(free_bytes_output: str) -> tuple[int, int] | None: + for line in free_bytes_output.splitlines(): + fields = line.split() + if not fields or fields[0] != "Mem:": + continue + + if len(fields) < 7: + return None + + if not fields[1].isdecimal() or not fields[6].isdecimal(): + return None + + mem_total_bytes = int(fields[1]) + mem_available_bytes = int(fields[6]) + if mem_available_bytes > mem_total_bytes: + return None + + return mem_total_bytes, mem_available_bytes + + return None + + +def bytes_to_gb(value: int) -> str: + return f"{value / BYTES_PER_GB:.2f}" + + +def log_human_memory_snapshot() -> None: + LOGGER.info("Memory snapshot (free -hw):") + result = run_capture_c_locale(["free", "-hw"]) + if result.stdout: + LOGGER.info(result.stdout.rstrip("\n")) + + if result.returncode != 0: + LOGGER.warning( + "Could not print the human-readable memory snapshot with 'free -hw'" + ) + + +def main() -> int: + configure_logging() + + node_name = os.environ.get("SLURMD_NODENAME", "unknown") + node_real_memory_raw = os.environ.get("CHECKS_NODE_REAL_MEM_BYTES", "") + + LOGGER.info( + "[%s] Check whether memory usage has recovered on drained node %s", + time.ctime(), + node_name, + ) + LOGGER.info( + "Slurm RealMemory input: %s bytes", + node_real_memory_raw if node_real_memory_raw else "", + ) + LOGGER.info( + "Node eligibility source: this check is scheduled only for drained nodes" + ) + + if not node_real_memory_raw.isdecimal() or int(node_real_memory_raw) <= 0: + real_memory_display = node_real_memory_raw or "" + LOGGER.warning( + "Invalid or unavailable Slurm RealMemory '%s'; keeping node drained", + real_memory_display, + ) + return 1 + + node_real_memory_bytes = int(node_real_memory_raw) + free_bytes_result = run_capture_c_locale(["free", "-b"]) + if free_bytes_result.returncode != 0: + LOGGER.warning( + "Could not read local memory information with 'free -b'; " + "keeping node drained" + ) + return 1 + + memory_values = parse_memory_values(free_bytes_result.stdout) + if memory_values is None: + LOGGER.warning( + "Could not determine valid total and available memory from 'free -b'; " + "keeping node drained" + ) + return 1 + + mem_total_bytes, mem_available_bytes = memory_values + if node_real_memory_bytes > mem_total_bytes: + LOGGER.warning( + "Slurm RealMemory %s bytes exceeds MemTotal %s bytes; " + "keeping node drained", + node_real_memory_bytes, + mem_total_bytes, + ) + return 1 + + mem_available_gb = bytes_to_gb(mem_available_bytes) + node_real_memory_gb = bytes_to_gb(node_real_memory_bytes) + + LOGGER.info( + "Memory comparison: " + "available=%s GB (%s bytes), " + "Slurm RealMemory=%s GB (%s bytes)", + mem_available_gb, + mem_available_bytes, + node_real_memory_gb, + node_real_memory_bytes, + ) + log_human_memory_snapshot() + + if mem_available_bytes < node_real_memory_bytes: + LOGGER.warning( + "Available memory %s GB is still below configured memory %s GB; " + "keeping node drained", + mem_available_gb, + node_real_memory_gb, + ) + return 1 + + LOGGER.info( + "Drained node leaves enough memory available for Slurm RealMemory; " + "memory recovery confirmed" + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/helm/slurm-cluster/slurm_scripts/idle_mem_used.undrain.py.json b/helm/slurm-cluster/slurm_scripts/idle_mem_used.undrain.py.json new file mode 100644 index 000000000..b392f427c --- /dev/null +++ b/helm/slurm-cluster/slurm_scripts/idle_mem_used.undrain.py.json @@ -0,0 +1,16 @@ +{ + "name": "idle_mem_used", + "command": "/usr/bin/python3 ./idle_mem_used.undrain.py", + "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..6b38c89e2 --- /dev/null +++ b/helm/slurm-cluster/slurm_scripts/idle_mem_used_test.py @@ -0,0 +1,307 @@ +import os +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + + +PYTHON_DRAIN_SCRIPT_PATH = Path(__file__).with_name("idle_mem_used.drain.py") +PYTHON_UNDRAIN_SCRIPT_PATH = Path(__file__).with_name("idle_mem_used.undrain.py") +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 = PYTHON_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 "$1" "$2"', + "idle-mem-used-test", + sys.executable, + 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=PYTHON_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=PYTHON_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=PYTHON_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) + + def test_undrain_keeps_node_drained_when_real_memory_is_unavailable(self): + result = self.run_check( + listjobs_rc=64, + listjobs_output="undrain must not query local jobs", + total_bytes=None, + available_bytes=None, + node_real_memory_bytes=None, + script_path=PYTHON_UNDRAIN_SCRIPT_PATH, + ) + + self.assertEqual(1, result.returncode) + self.assertNotIn("listjobs", result.stdout) + self.assertNotIn("listjobs", result.stderr) + self.assertIn("Invalid or unavailable Slurm RealMemory", result.stderr) + self.assertIn("keeping node drained", result.stderr) + + +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..ff30a481b --- /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.py"] + - isNotNull: + path: data["idle_mem_used.undrain.py"] + - matchRegex: + path: data["checks.json"] + pattern: '"name": "idle_mem_used"' + - matchRegex: + path: data["checks.json"] + pattern: '"command": "/usr/bin/python3 ./idle_mem_used.drain.py"' + - matchRegex: + path: data["checks.json"] + pattern: '"command": "/usr/bin/python3 ./idle_mem_used.undrain.py"' + - 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.py: + enabled: false + idle_mem_used.undrain.py: + enabled: false + asserts: + - isNull: + path: data["idle_mem_used.drain.py"] + - isNull: + path: data["idle_mem_used.undrain.py"] + - notMatchRegex: + path: data["checks.json"] + pattern: '"command": "/usr/bin/python3 ./idle_mem_used.(drain|undrain).py"' diff --git a/helm/slurm-cluster/values.yaml b/helm/slurm-cluster/values.yaml index edc0b73aa..d1685c485 100644 --- a/helm/slurm-cluster/values.yaml +++ b/helm/slurm-cluster/values.yaml @@ -812,6 +812,14 @@ slurmScripts: enabled: true customContent: null customConfig: null + idle_mem_used.drain.py: + enabled: true + customContent: null + customConfig: null + idle_mem_used.undrain.py: + enabled: true + customContent: null + customConfig: null job_tmpfs_delete.sh: enabled: true customContent: null