[DO NOT MERGE] Create scripts to audit scontrol calls

helm/slurm-cluster/slurm_scripts/SCONTROL_AUDIT.md

@@ -0,0 +1,141 @@
+# Passive check scontrol audit
+
+Temporary dev-only instrumentation for SCHED-1897 can log every `scontrol`
+invocation made by passive checks. It is disabled by default.
+
+Enable on a dev cluster:
+
+```bash
+helm upgrade <release> helm/slurm-cluster \
+  --namespace <namespace> \
+  --reuse-values \
+  --set slurmScripts.scontrolAudit.enabled=true
+```
+
+Optional values:
+
+```yaml
+slurmScripts:
+  scontrolAudit:
+    enabled: true
+    logPath: ""
+    realScontrolPath: /usr/bin/scontrol
+```
+
+When `logPath` is empty, the wrapper writes to
+`/opt/soperator-outputs/slurm_scripts/scontrol_audit.jsonl` inside the jail
+and `/mnt/jail/opt/soperator-outputs/slurm_scripts/scontrol_audit.jsonl`
+outside the jail. Each line is one JSON object.
+
+Disable by setting `slurmScripts.scontrolAudit.enabled=false` and rolling the
+updated scripts out again.
+
+## Useful fields
+
+- `context`: `prolog`, `epilog`, or `hc_program`
+- `check_name`: current check name, or empty for runner-level filtering
+- `check_phase`: usually `run_check` or `filter_by_node_state`
+- `runner_invocation_id`: one `check_runner.py` invocation; useful for per-job
+  or per-HealthCheckProgram-interval grouping
+- `node`, `job_id`, `job_gpus`
+- `command_class`: `show node`, `show job`, `listjobs`, `update drain`,
+  `update resume`, `update comment`, or another first-level command
+- `argv`: full `scontrol` argument array
+- `caller_exe`, `caller_cmdline`: parent-process clue
+
+## Collection examples
+
+From a worker pod:
+
+```bash
+kubectl -n <namespace> exec <worker-pod> -c slurmd -- \
+  cat /mnt/jail/opt/soperator-outputs/slurm_scripts/scontrol_audit.jsonl \
+  > scontrol_audit.<worker-pod>.jsonl
+```
+
+From inside a node container:
+
+```bash
+cat /mnt/jail/opt/soperator-outputs/slurm_scripts/scontrol_audit.jsonl
+cat /opt/soperator-outputs/slurm_scripts/scontrol_audit.jsonl
+```
+
+Combine several workers:
+
+```bash
+cat scontrol_audit.*.jsonl > scontrol_audit.all.jsonl
+```
+
+## Summary commands
+
+Counts by context, check, and command class:
+
+```bash
+jq -r '[.context, (.check_name // ""), .command_class] | @tsv' \
+  scontrol_audit.all.jsonl | sort | uniq -c | sort -nr
+```
+
+Counts by job and node:
+
+```bash
+jq -r 'select(.job_id != "") |
+  [.job_id, .node, .context, (.check_name // ""), .command_class] | @tsv' \
+  scontrol_audit.all.jsonl | sort | uniq -c | sort -nr
+```
+
+Counts per `HealthCheckProgram` runner invocation:
+
+```bash
+jq -r 'select(.context == "hc_program") |
+  [.runner_invocation_id, .node, (.check_name // ""), .command_class] | @tsv' \
+  scontrol_audit.all.jsonl | sort | uniq -c | sort -nr
+```
+
+Background versus job-coupled load:
+
+```bash
+jq -r '[
+  (if .job_id == "" then "background" else "job" end),
+  .context,
+  (.check_name // ""),
+  .command_class
+] | @tsv' scontrol_audit.all.jsonl | sort | uniq -c | sort -nr
+```
+
+Sample table shape:
+
+```text
+count  context     check_name                 command_class
+1240   hc_program  job_tmpfs_delete_leftover  listjobs
+312    prolog      alloc_mem_used             show job
+156    hc_program  alloc_mem_used             show node
+18     prolog      gpu_health_check           update drain
+9      hc_program  alloc_gpus_busy            update resume
+```
+
+## Static expectations
+
+Based on the built-in check configs and scripts:
+
+- `check_runner.py` may call `scontrol show node ... --json` for node-state
+  filtering, `CHECKS_NODE_*` env export, and drain/undrain/comment decisions.
+  The result is cached within one runner invocation until an update clears it.
+- `alloc_mem_used` in `prolog` requests `CHECKS_JOB_ALLOC_MEM_BYTES`, so it is
+  expected to cause one cached `scontrol show job ... --json` per runner
+  invocation where the check runs.
+- `alloc_mem_used` in `hc_program` requests `CHECKS_NODE_REAL_MEM_BYTES`, so it
+  is expected to use `scontrol show node ... --json`.
+- `alloc_gpus_busy` and `alloc_mem_used` undrain checks run only on drained
+  nodes, so `hc_program` node-state filtering is expected to use
+  `scontrol show node ... --json` when those checks are enabled.
+- Failed checks with `on_fail: drain` can call `scontrol update ... State=drain`:
+  `alloc_gpus_busy`, `alloc_mem_used`, `boot_disk_full`, `gpu_health_check`,
+  and `nvme_raid_health` when that optional check is enabled.
+- Successful checks with `on_ok: undrain` can call
+  `scontrol update ... State=resume`: `alloc_gpus_busy`, `alloc_mem_used`, and
+  `boot_disk_full` when the node reason matches the check reason.
+- `job_tmpfs_delete_leftover` in `hc_program` directly calls
+  `scontrol listjobs --json`, but only after it finds existing job tmpfs
+  directories to evaluate.
+- Other built-in passive checks do not contain executable `scontrol` calls;
+  references in their output text are remediation hints only.

helm/slurm-cluster/slurm_scripts/check_runner.py

@@ -1,3 +1,4 @@
+import contextlib
 import functools
 import json
 import logging
@@ -139,6 +140,13 @@ class JobInfo(typing.NamedTuple):
     logging.error(f"Failed to get environment variable '{ke.args[0]}', exiting: {ke}")
     sys.exit(0)
 
+CHECKS_RUNNER_INVOCATION_ID = os.environ.get("CHECKS_RUNNER_INVOCATION_ID")
+if not CHECKS_RUNNER_INVOCATION_ID:
+    job_or_interval = SLURM_JOB_ID if SLURM_JOB_ID else "nojob"
+    CHECKS_RUNNER_INVOCATION_ID = f"{SLURMD_NODENAME}.{CHECKS_CONTEXT}.{job_or_interval}.{os.getpid()}.{int(time.time() * 1000)}"
+os.environ["CHECKS_RUNNER_INVOCATION_ID"] = CHECKS_RUNNER_INVOCATION_ID
+os.environ["CHECKS_RUNNER_PID"] = str(os.getpid())
+
 def main():
     start_time = time.perf_counter()
     logging.info("Started")
@@ -261,7 +269,8 @@ def filter_by_node_state(checks: list[Check]) -> list[Check]:
     # Skip if all checks don't care
     if all("any" in check.node_states for check in checks):
         return checks
-    node_info = get_node_info()
+    with scontrol_audit_scope(phase="filter_by_node_state"):
+        node_info = get_node_info()
     return [
         check for check in checks
         if (
@@ -272,69 +281,108 @@ def filter_by_node_state(checks: list[Check]) -> list[Check]:
 
 # Run a specific check
 def run_check(check: Check, in_jail=False):
-    # Export environment variables requested by this check
-    export_needed_env(check)
-
-    # Print info about the running check
-    log_rel_path = string.Template(check.log).safe_substitute(
-        worker=SLURMD_NODENAME, context=CHECKS_CONTEXT, name=check.name
-    )
-    log_abs_path = os.path.join(CHECKS_OUTPUTS_BASE_DIR, log_rel_path)
-    if not in_jail:
-        log_abs_path = "/mnt/jail" + log_abs_path
-    start_time = time.perf_counter()
-    logging.info(f"Running check {check.name} ({check.command}), logging to {log_abs_path}")
-    logging.info(f"Check spec: {json.dumps(check._asdict(), indent=2)}")
-
-    # Create parent dirs with full permissions
-    os.makedirs(os.path.dirname(log_abs_path), mode=0o777, exist_ok=True)
+    with scontrol_audit_scope(check=check, phase="run_check", in_jail=in_jail):
+        # Export environment variables requested by this check
+        export_needed_env(check)
 
-    # Execute the check command
-    cmd = ["bash", "-l", "-c", f"{check.command} 3>&1 1>\"{log_abs_path}\" 2>&1"]
-    result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-
-    # Build the reason message
-    reason_base = string.Template(check.reason_base.rstrip()).safe_substitute(
-        context=CHECKS_CONTEXT, name=check.name
+        # Print info about the running check
+        log_rel_path = string.Template(check.log).safe_substitute(
+            worker=SLURMD_NODENAME, context=CHECKS_CONTEXT, name=check.name
+        )
+        log_abs_path = os.path.join(CHECKS_OUTPUTS_BASE_DIR, log_rel_path)
+        if not in_jail:
+            log_abs_path = "/mnt/jail" + log_abs_path
+        start_time = time.perf_counter()
+        logging.info(f"Running check {check.name} ({check.command}), logging to {log_abs_path}")
+        logging.info(f"Check spec: {json.dumps(check._asdict(), indent=2)}")
+
+        # Create parent dirs with full permissions
+        os.makedirs(os.path.dirname(log_abs_path), mode=0o777, exist_ok=True)
+
+        # Execute the check command
+        check_command = f"{check.command} 3>&1 1>\"{log_abs_path}\" 2>&1"
+        if scontrol_audit_enabled():
+            # bash -l may reset PATH via profile scripts; restore the audit wrapper first.
+            check_command = f"export PATH=\"/opt/slurm_scripts:$PATH\"; {check_command}"
+        cmd = ["bash", "-l", "-c", check_command]
+        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+
+        # Build the reason message
+        reason_base = string.Template(check.reason_base.rstrip()).safe_substitute(
+            context=CHECKS_CONTEXT, name=check.name
+        )
+        reason = reason_base
+        details = result.stdout.strip().encode('unicode_escape').decode()
+        if check.reason_append_details and details:
+            reason += f": {details}"
+        reason += f" [{CHECKS_CONTEXT}]"
+
+        # Log check running time
+        end_time = time.perf_counter()
+        logging.info(f"Completed check {check.name} in {end_time - start_time:.3f} seconds")
+
+        # React to the check result
+        if result.returncode != 0:
+            logging.info(f"Check {check.name}: FAIL ({details})")
+
+            # Drain / comment the Slurm node
+            if check.on_fail == "drain" and "DRAIN" not in get_node_info().state_flags:
+                drain_node(reason)
+            elif check.on_fail == "comment":
+                comment_node(reason)
+
+            # Exit, if the check failed
+            # In prolog context, restart the job and print a message to the job output
+            # (Slurm passes lines from Prolog stdout that start with "print " to job output)
+            if CHECKS_CONTEXT == "prolog":
+                print(f"print Slurm healthcheck failed on node {SLURMD_NODENAME}, trying to automatically requeue")
+                sys.exit(1)
+
+            sys.exit(0)
+
+        logging.info(f"Check {check.name}: OK")
+
+        # Undrain / uncomment the Slurm node, if it was marked with the same reason
+        if check.on_ok == "undrain" and "DRAIN" in get_node_info().state_flags:
+            if get_node_info().reason and get_node_info().reason.startswith(reason_base):
+                undrain_node()
+        elif check.on_ok == "uncomment":
+            if get_node_info().comment and get_node_info().comment.startswith(reason_base):
+                uncomment_node()
+
+def scontrol_audit_enabled() -> bool:
+    return os.environ.get("SCONTROL_AUDIT_ENABLED", "").lower() in ("1", "true", "yes", "on")
+
+@contextlib.contextmanager
+def scontrol_audit_scope(
+    check: typing.Optional[Check] = None,
+    phase: str = "",
+    in_jail: typing.Optional[bool] = None,
+):
+    keys = (
+        "CHECKS_CURRENT_CHECK_NAME",
+        "CHECKS_CURRENT_CHECK_COMMAND",
+        "CHECKS_CURRENT_CHECK_PHASE",
+        "CHECKS_CURRENT_CHECK_IN_JAIL",
     )
-    reason = reason_base
-    details = result.stdout.strip().encode('unicode_escape').decode()
-    if check.reason_append_details and details:
-        reason += f": {details}"
-    reason += f" [{CHECKS_CONTEXT}]"
-
-    # Log check running time
-    end_time = time.perf_counter()
-    logging.info(f"Completed check {check.name} in {end_time - start_time:.3f} seconds")
+    old_values = {key: os.environ.get(key) for key in keys}
 
-    # React to the check result
-    if result.returncode != 0:
-        logging.info(f"Check {check.name}: FAIL ({details})")
+    os.environ["CHECKS_CURRENT_CHECK_NAME"] = check.name if check else ""
+    os.environ["CHECKS_CURRENT_CHECK_COMMAND"] = check.command if check else ""
+    os.environ["CHECKS_CURRENT_CHECK_PHASE"] = phase
+    if in_jail is not None:
+        os.environ["CHECKS_CURRENT_CHECK_IN_JAIL"] = "1" if in_jail else "0"
+    else:
+        os.environ["CHECKS_CURRENT_CHECK_IN_JAIL"] = ""
 
-        # Drain / comment the Slurm node
-        if check.on_fail == "drain" and "DRAIN" not in get_node_info().state_flags:
-            drain_node(reason)
-        elif check.on_fail == "comment":
-            comment_node(reason)
-
-        # Exit, if the check failed
-        # In prolog context, restart the job and print a message to the job output
-        # (Slurm passes lines from Prolog stdout that start with "print " to job output)
-        if CHECKS_CONTEXT == "prolog":
-            print(f"print Slurm healthcheck failed on node {SLURMD_NODENAME}, trying to automatically requeue")
-            sys.exit(1)
-
-        sys.exit(0)
-
-    logging.info(f"Check {check.name}: OK")
-
-    # Undrain / uncomment the Slurm node, if it was marked with the same reason
-    if check.on_ok == "undrain" and "DRAIN" in get_node_info().state_flags:
-        if get_node_info().reason and get_node_info().reason.startswith(reason_base):
-            undrain_node()
-    elif check.on_ok == "uncomment":
-        if get_node_info().comment and get_node_info().comment.startswith(reason_base):
-            uncomment_node()
+    try:
+        yield
+    finally:
+        for key, old_value in old_values.items():
+            if old_value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = old_value
 
 # Export additional environment variables requested by the check
 # Their values are obtained from long-running commands, that's why they aren't exported by default
@@ -401,8 +449,6 @@ def get_platform_tags() -> list[str]:
 
 # Get info about the Slurm node from "scontrol show node"
 # Please note, this command can be executed from both jail or host rootfs
-# This function returns the cached value for subsequent calls
-@functools.lru_cache(maxsize=1)
 def get_node_info() -> NodeInfo:
     try:
         result = subprocess.run(
@@ -433,8 +479,6 @@ def get_node_info() -> NodeInfo:
 
 # Get info about the Slurm job from "scontrol show job"
 # Please note, this command can be executed from both jail or host rootfs
-# This function returns the cached value for subsequent calls
-@functools.lru_cache(maxsize=1)
 def get_job_info() -> JobInfo:
     try:
         if CHECKS_CONTEXT not in ("prolog", "epilog"):
@@ -503,8 +547,6 @@ def drain_node(reason):
             ["scontrol", "update", f"NodeName={SLURMD_NODENAME}", "State=drain", f"Reason=\"{reason}\""],
             check=False, stderr=subprocess.DEVNULL
         )
-        # Invalidate cache for the Slurm node info
-        get_node_info.cache_clear()
     except Exception as e:
         logging.warning(f"Failed to drain Slurm node {SLURMD_NODENAME}: {e}")
 
@@ -516,8 +558,6 @@ def undrain_node():
             ["scontrol", "update", f"NodeName={SLURMD_NODENAME}", "State=resume"],
             check=False, stderr=subprocess.DEVNULL
         )
-        # Invalidate cache for the Slurm node info
-        get_node_info.cache_clear()
     except Exception as e:
         logging.warning(f"Failed to undrain Slurm node {SLURMD_NODENAME}: {e}")
 
@@ -529,8 +569,6 @@ def comment_node(comment):
             ["scontrol", "update", f"NodeName={SLURMD_NODENAME}", f"Comment=\"{comment}\""],
             check=False, stderr=subprocess.DEVNULL
         )
-        # Invalidate cache for the Slurm node info
-        get_node_info.cache_clear()
     except Exception as e:
         logging.warning(f"Failed to comment Slurm node {SLURMD_NODENAME}: {e}")
 
@@ -542,8 +580,6 @@ def uncomment_node():
             ["scontrol", "update", f"NodeName={SLURMD_NODENAME}", "Comment=\"\""],
             check=False, stderr=subprocess.DEVNULL
         )
-        # Invalidate cache for the Slurm node info
-        get_node_info.cache_clear()
     except Exception as e:
         logging.warning(f"Failed to uncomment Slurm node {SLURMD_NODENAME}: {e}")