fix: pin privileged Python host-check image by digest (#22)

Author: samharshe

cmax/scripts/1-audit/checks/platform_config.py

@@ -36,7 +36,8 @@
 Environment variables:
 
 * ``CLUSTERMAX_AUDIT_K8S_NAMESPACE`` (default ``default``)
-* ``CLUSTERMAX_AUDIT_K8S_HOST_CHECK_IMAGE`` (default ``python:3.12-alpine``)
+* ``CLUSTERMAX_AUDIT_K8S_HOST_CHECK_IMAGE`` (default: digest-pinned
+  ``python:3.12-alpine``; use a digest-pinned mirror for air-gapped clusters)
 * ``CLUSTERMAX_AUDIT_K8S_HOST_CHECK_PULL_POLICY`` (default ``IfNotPresent``)
 * ``CLUSTERMAX_AUDIT_K8S_MAX_HOST_CHECKS`` (default: all GPU nodes)
 * ``CLUSTERMAX_PYXIS_CHECK_IMAGE`` - container the mount check launches.
@@ -68,6 +69,11 @@
 from typing import Any
 
 SCHEMA_VERSION = 1
+# Pin the multi-architecture index so both x86 and Arm hosts use fixed images.
+DEFAULT_K8S_HOST_CHECK_IMAGE = (
+    "python:3.12-alpine@"
+    "sha256:c4634f578a412db396771b61b064c6e546c9d6414c7fb5b1b05d5871f1885f7b"
+)
 CHECK_KEYS = ("vm_iommu", "arm_smmu_virtualization")
 ALL_CHECK_KEYS = (*CHECK_KEYS, "nccl_topo_file", "nccl_ib_qps")
 CACHE_PATH_ENV = "CLUSTERMAX_PLATFORM_CHECK_CACHE"
@@ -2020,7 +2026,7 @@ def run_k8s_check(
     include_rdma_iommu: bool = True,
 ) -> tuple[list[dict[str, Any]], list[str]]:
     namespace = os.environ.get("CLUSTERMAX_AUDIT_K8S_NAMESPACE", "default")
-    image = os.environ.get("CLUSTERMAX_AUDIT_K8S_HOST_CHECK_IMAGE", "python:3.12-alpine")
+    image = os.environ.get("CLUSTERMAX_AUDIT_K8S_HOST_CHECK_IMAGE", DEFAULT_K8S_HOST_CHECK_IMAGE)
     max_nodes = as_int(os.environ.get("CLUSTERMAX_AUDIT_K8S_MAX_HOST_CHECKS"), default=len(nodes))
     return load_fanout().fan_out_k8s(
         lambda node: run_k8s_host_check(

cmax/scripts/1-audit/checks/system/hbm_memory_exposure.py

@@ -10,9 +10,9 @@
 Kubernetes harness environment variables:
 
 * ``CLUSTERMAX_AUDIT_K8S_NAMESPACE`` (default ``default``)
-* ``CLUSTERMAX_AUDIT_K8S_HOST_CHECK_IMAGE`` (default
-  ``python:3.12-alpine``) - must be reachable from GPU nodes; override
-  for air-gapped clusters.
+* ``CLUSTERMAX_AUDIT_K8S_HOST_CHECK_IMAGE`` (default: digest-pinned
+  ``python:3.12-alpine``) - must be reachable from GPU nodes; use a
+  digest-pinned mirror for air-gapped clusters.
 * ``CLUSTERMAX_AUDIT_K8S_HOST_CHECK_PULL_POLICY`` (default
   ``IfNotPresent``)
 * ``CLUSTERMAX_AUDIT_K8S_MAX_HOST_CHECKS`` (default: all GPU nodes)
@@ -33,6 +33,11 @@
 
 
 GIB_KB = 1024 * 1024
+# Pin the multi-architecture index so both x86 and Arm hosts use fixed images.
+DEFAULT_K8S_HOST_CHECK_IMAGE = (
+    "python:3.12-alpine@"
+    "sha256:c4634f578a412db396771b61b064c6e546c9d6414c7fb5b1b05d5871f1885f7b"
+)
 # Lower bound for what counts as an HBM-sized memory-only NUMA node. 64 GiB
 # catches every coherent NVIDIA platform shipping today (GH200 96 GiB, GB200
 # 192 GiB per GPU, GB300 288 GiB per GPU, H100 SXM 80 GiB if ever exposed
@@ -735,7 +740,7 @@ def run_k8s_check() -> tuple[list[dict[str, Any]], list[str]]:
         return [], [NO_K8S_GPU_NODES]
 
     namespace = os.environ.get("CLUSTERMAX_AUDIT_K8S_NAMESPACE", "default")
-    image = os.environ.get("CLUSTERMAX_AUDIT_K8S_HOST_CHECK_IMAGE", "python:3.12-alpine")
+    image = os.environ.get("CLUSTERMAX_AUDIT_K8S_HOST_CHECK_IMAGE", DEFAULT_K8S_HOST_CHECK_IMAGE)
     max_nodes = as_int(os.environ.get("CLUSTERMAX_AUDIT_K8S_MAX_HOST_CHECKS"), default=len(nodes))
     return load_fanout().fan_out_k8s(_hbm_per_node(namespace, image), nodes=nodes, max_nodes=max_nodes)
 

tests/audit/test_k8s_utility_image_pins.py

@@ -3,10 +3,14 @@
 
 from __future__ import annotations
 
+import importlib.util
+import subprocess
 import sys
 import tempfile
 from pathlib import Path
+from unittest import mock
 
+import pytest
 import yaml
 
 sys.path.insert(0, str(Path(__file__).resolve().parent))
@@ -28,6 +32,36 @@
     "ubuntu:24.04@"
     "sha256:561618e2c15bf2397621dd04f96926663a3b5616c189cf7e38db7e82f5c538ea"
 )
+PYTHON_IMAGE = (
+    "python:3.12-alpine@"
+    "sha256:c4634f578a412db396771b61b064c6e546c9d6414c7fb5b1b05d5871f1885f7b"
+)
+
+
+@pytest.mark.parametrize("check_path", ["platform_config.py", "system/hbm_memory_exposure.py"])
+@pytest.mark.parametrize("override", [None, "registry.example/python@" + PYTHON_IMAGE.split("@")[1]])
+def test_python_host_checks_use_pinned_default_and_preserve_override(check_path, override):
+    spec = importlib.util.spec_from_file_location("host_check", SCRIPT.parent / "checks" / check_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    node = {"name": "gpu-0"}
+    fanout = mock.Mock()
+    fanout.k8s_gpu_nodes.return_value = [node]
+    fanout.fan_out_k8s.side_effect = lambda worker, **kwargs: worker(node)
+    env = {} if override is None else {"CLUSTERMAX_AUDIT_K8S_HOST_CHECK_IMAGE": override}
+    with (
+        mock.patch.dict(module.os.environ, env, clear=True),
+        mock.patch.object(module, "load_fanout", return_value=fanout),
+        mock.patch.object(module, "kubectl", return_value=subprocess.CompletedProcess([], 0, "{}")),
+        mock.patch.object(module, "run_k8s_host_check") as run_host,
+    ):
+        if check_path == "platform_config.py":
+            module.run_k8s_check([node])
+        else:
+            module.run_k8s_check()
+    namespace, selected_node, image = run_host.call_args.args
+    manifest = module.pod_manifest(namespace, selected_node, image, "test-pod")
+    assert manifest["spec"]["containers"][0]["image"] == (override or PYTHON_IMAGE)
 
 
 def _applied_manifest(function_name: str, setup: str) -> dict: