Expose node RealMemory without controller RPC

helm/slurm-cluster/slurm_scripts/check_runner.py

@@ -8,6 +8,9 @@
 import time
 import typing
 
+SOPERATOR_NODE_METADATA_FILE = "/run/soperator/node_metadata.env"
+SOPERATOR_NODE_REAL_MEMORY_BYTES = "SOPERATOR_NODE_REAL_MEMORY_BYTES"
+
 # Set up logging
 try:
     log_stdout = "/dev/stdout"
@@ -114,7 +117,7 @@ class Check(typing.NamedTuple):
     # - CHECKS_NODE_COMMENT - comment field of the Slurm node
     # - CHECKS_NODE_REAL_MEM_BYTES - total allocatable memory in bytes for the Slurm node
     # - CHECKS_JOB_ALLOC_MEM_BYTES - memory in bytes, allocated for this Slurm job on this node
-    # These values are extracted from long-running commands, that's why they aren't exported by default
+    # Some values are extracted from long-running commands, so they aren't exported by default.
     need_env: list[str] = []
 
 class NodeInfo(typing.NamedTuple):
@@ -356,10 +359,39 @@ def export_needed_env(check: Check):
         if env == "CHECKS_NODE_COMMENT":
             os.environ["CHECKS_NODE_COMMENT"] = get_node_info().comment
         if env == "CHECKS_NODE_REAL_MEM_BYTES":
-            os.environ["CHECKS_NODE_REAL_MEM_BYTES"] = str(get_node_info().real_memory_bytes)
+            os.environ["CHECKS_NODE_REAL_MEM_BYTES"] = str(get_node_real_memory_bytes())
         if env == "CHECKS_JOB_ALLOC_MEM_BYTES":
             os.environ["CHECKS_JOB_ALLOC_MEM_BYTES"] = str(get_job_info().allocated_memory_bytes)
 
+# Get node RealMemory from node-local metadata, avoiding a controller RPC in the normal path.
+# Fall back to Slurm node info for compatibility with workers that have not yet been restarted
+# with an image and pod specification that publish the metadata file.
+@functools.lru_cache(maxsize=1)
+def get_node_real_memory_bytes() -> int:
+    try:
+        with open(SOPERATOR_NODE_METADATA_FILE, encoding="utf-8") as metadata_file:
+            for raw_line in metadata_file:
+                key, separator, value = raw_line.rstrip("\n").partition("=")
+                if key != SOPERATOR_NODE_REAL_MEMORY_BYTES:
+                    continue
+                if separator == "" or not value.isdecimal() or int(value) <= 0:
+                    raise ValueError(f"Invalid {SOPERATOR_NODE_REAL_MEMORY_BYTES} value")
+
+                real_memory_bytes = int(value)
+                logging.info(
+                    f"Node RealMemory from {SOPERATOR_NODE_METADATA_FILE}: "
+                    f"{real_memory_bytes} bytes"
+                )
+                return real_memory_bytes
+
+        raise ValueError(f"Missing {SOPERATOR_NODE_REAL_MEMORY_BYTES} value")
+    except Exception as e:
+        logging.warning(
+            f"Failed to get node RealMemory from {SOPERATOR_NODE_METADATA_FILE}: {e}; "
+            "falling back to Slurm node info"
+        )
+        return get_node_info().real_memory_bytes
+
 # Get GPU platform tags, e.g. ["8xH200", "8xGPU] from "nvidia-smi"
 # Please note, this command can be executed from both jail or host rootfs
 # The list starts with more specific tags, and ends with less specific ones

helm/slurm-cluster/slurm_scripts/check_runner_test.py

@@ -116,5 +116,111 @@ def test_uncomment_on_ok_still_runs_in_hc_program(self):
             self.assertEqual(["uncomment"], calls)
 
 
+def load_check_runner():
+    required_env = {
+        "SLURMD_NODENAME": "worker-1",
+        "CHECKS_OUTPUTS_BASE_DIR": "/opt/soperator-outputs",
+        "CHECKS_CONTEXT": "hc_program",
+        "CHECKS_CONFIG": "/opt/slurm_scripts/checks.json",
+    }
+    with mock.patch.dict(os.environ, required_env):
+        spec = importlib.util.spec_from_file_location(
+            "check_runner_under_test", CHECK_RUNNER_PATH
+        )
+        module = importlib.util.module_from_spec(spec)
+        assert spec.loader is not None
+        spec.loader.exec_module(module)
+        return module
+
+
+check_runner = load_check_runner()
+
+
+class NodeRealMemoryMetadataTest(unittest.TestCase):
+    def setUp(self):
+        check_runner.get_node_real_memory_bytes.cache_clear()
+
+    def test_reads_real_memory_from_local_metadata_without_node_rpc(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            metadata_file = Path(tmpdir) / "node_metadata.env"
+            metadata_file.write_text(
+                "IGNORED=value\n"
+                "SOPERATOR_NODE_REAL_MEMORY_BYTES=999292928\n",
+                encoding="utf-8",
+            )
+
+            with (
+                mock.patch.object(
+                    check_runner, "SOPERATOR_NODE_METADATA_FILE", str(metadata_file)
+                ),
+                mock.patch.object(check_runner, "get_node_info") as get_node_info,
+            ):
+                result = check_runner.get_node_real_memory_bytes()
+
+            self.assertEqual(999292928, result)
+            get_node_info.assert_not_called()
+
+    def test_exports_local_real_memory_for_checks(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            metadata_file = Path(tmpdir) / "node_metadata.env"
+            metadata_file.write_text(
+                "SOPERATOR_NODE_REAL_MEMORY_BYTES=999292928\n",
+                encoding="utf-8",
+            )
+            check = check_runner.Check(need_env=["CHECKS_NODE_REAL_MEM_BYTES"])
+
+            with (
+                mock.patch.object(
+                    check_runner, "SOPERATOR_NODE_METADATA_FILE", str(metadata_file)
+                ),
+                mock.patch.object(check_runner, "get_node_info") as get_node_info,
+                mock.patch.dict(os.environ, {}, clear=False),
+            ):
+                check_runner.export_needed_env(check)
+                exported_value = os.environ["CHECKS_NODE_REAL_MEM_BYTES"]
+
+            self.assertEqual("999292928", exported_value)
+            get_node_info.assert_not_called()
+
+    def test_falls_back_to_slurm_when_metadata_is_missing(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            missing_file = Path(tmpdir) / "missing.env"
+            node_info = check_runner.NodeInfo(real_memory_bytes=2147483648)
+
+            with (
+                mock.patch.object(
+                    check_runner, "SOPERATOR_NODE_METADATA_FILE", str(missing_file)
+                ),
+                mock.patch.object(
+                    check_runner, "get_node_info", return_value=node_info
+                ) as get_node_info,
+            ):
+                result = check_runner.get_node_real_memory_bytes()
+
+            self.assertEqual(2147483648, result)
+            get_node_info.assert_called_once_with()
+
+    def test_falls_back_to_slurm_when_metadata_is_invalid(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            metadata_file = Path(tmpdir) / "node_metadata.env"
+            metadata_file.write_text(
+                "SOPERATOR_NODE_REAL_MEMORY_BYTES=invalid\n", encoding="utf-8"
+            )
+            node_info = check_runner.NodeInfo(real_memory_bytes=1073741824)
+
+            with (
+                mock.patch.object(
+                    check_runner, "SOPERATOR_NODE_METADATA_FILE", str(metadata_file)
+                ),
+                mock.patch.object(
+                    check_runner, "get_node_info", return_value=node_info
+                ) as get_node_info,
+            ):
+                result = check_runner.get_node_real_memory_bytes()
+
+            self.assertEqual(1073741824, result)
+            get_node_info.assert_called_once_with()
+
+
 if __name__ == "__main__":
     unittest.main(verbosity=2)

images/worker/slurmd.dockerfile

@@ -157,6 +157,7 @@ RUN mkdir -p /var/log/slurm/multilog && \
 
 # Copy slurmd entrypoint script
 COPY images/worker/slurmd_entrypoint.sh /opt/bin/slurm/
+COPY images/worker/write_soperator_metadata.sh /opt/bin/slurm/
 
 # Copy worker init script (controller readiness + topology for ephemeral nodes)
 COPY images/worker/worker_init.py /opt/bin/slurm/
@@ -167,6 +168,7 @@ COPY images/worker/docker_proxy_nginx_entrypoint.sh /opt/bin/slurm/
 COPY images/worker/dockerd_entrypoint.sh /opt/bin/slurm/
 
 RUN chmod +x /opt/bin/slurm/slurmd_entrypoint.sh && \
+    chmod +x /opt/bin/slurm/write_soperator_metadata.sh && \
     chmod +x /opt/bin/slurm/supervisord_entrypoint.sh && \
     chmod +x /opt/bin/slurm/worker_init.py && \
     chmod +x /opt/bin/slurm/docker_proxy_nginx_entrypoint.sh && \

images/worker/slurmd_entrypoint.sh

@@ -42,6 +42,9 @@ else
     export TOPO_SWITCH_TIER2="unknown"
 fi
 
+echo "Export Soperator node metadata"
+/opt/bin/slurm/write_soperator_metadata.sh
+
 echo "Evaluate variables in the Slurm node 'Extra' field"
 evaluated_extra=$(eval echo "$SLURM_NODE_EXTRA")
 

images/worker/write_soperator_metadata.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+set -euo pipefail
+
+metadata_file="${1:-/run/soperator/node_metadata.env}"
+node_real_memory_bytes="${SOPERATOR_NODE_REAL_MEMORY_BYTES:-}"
+
+if ! [[ "${node_real_memory_bytes}" =~ ^[0-9]+$ ]] || [[ "${node_real_memory_bytes}" == "0" ]]; then
+    rm -f -- "${metadata_file}"
+    echo "SOPERATOR_NODE_REAL_MEMORY_BYTES is unavailable or invalid; skipping node metadata export" >&2
+    exit 0
+fi
+
+metadata_dir="$(dirname -- "${metadata_file}")"
+install -d -m 0755 "${metadata_dir}"
+
+temporary_file="$(mktemp "${metadata_file}.tmp.XXXXXX")"
+trap 'rm -f -- "${temporary_file}"' EXIT
+
+printf 'SOPERATOR_NODE_REAL_MEMORY_BYTES=%s\n' "${node_real_memory_bytes}" > "${temporary_file}"
+chmod 0644 "${temporary_file}"
+mv -f -- "${temporary_file}" "${metadata_file}"
+trap - EXIT
+
+echo "Exported node metadata to ${metadata_file}"

images/worker/write_soperator_metadata_test.py

@@ -0,0 +1,76 @@
+import os
+import stat
+import subprocess
+import tempfile
+import unittest
+from pathlib import Path
+
+
+SCRIPT_PATH = Path(__file__).with_name("write_soperator_metadata.sh")
+
+
+class WriteSoperatorMetadataTest(unittest.TestCase):
+    def run_writer(
+        self, metadata_file: Path, real_memory_bytes: str | None
+    ) -> subprocess.CompletedProcess[str]:
+        env = os.environ.copy()
+        if real_memory_bytes is None:
+            env.pop("SOPERATOR_NODE_REAL_MEMORY_BYTES", None)
+        else:
+            env["SOPERATOR_NODE_REAL_MEMORY_BYTES"] = real_memory_bytes
+
+        return subprocess.run(
+            ["bash", str(SCRIPT_PATH), str(metadata_file)],
+            check=False,
+            env=env,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+
+    def test_writes_real_memory_metadata_atomically(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            metadata_file = Path(tmpdir) / "nested" / "node_metadata.env"
+
+            result = self.run_writer(metadata_file, "999292928")
+
+            self.assertEqual(0, result.returncode)
+            self.assertEqual(
+                "SOPERATOR_NODE_REAL_MEMORY_BYTES=999292928\n",
+                metadata_file.read_text(encoding="utf-8"),
+            )
+            self.assertEqual(0o644, stat.S_IMODE(metadata_file.stat().st_mode))
+            self.assertEqual([], list(metadata_file.parent.glob("*.tmp.*")))
+
+    def test_replaces_existing_metadata(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            metadata_file = Path(tmpdir) / "node_metadata.env"
+            metadata_file.write_text(
+                "SOPERATOR_NODE_REAL_MEMORY_BYTES=1\n", encoding="utf-8"
+            )
+
+            result = self.run_writer(metadata_file, "2147483648")
+
+            self.assertEqual(0, result.returncode)
+            self.assertEqual(
+                "SOPERATOR_NODE_REAL_MEMORY_BYTES=2147483648\n",
+                metadata_file.read_text(encoding="utf-8"),
+            )
+
+    def test_missing_or_invalid_metadata_is_skipped(self):
+        for value in (None, "", "0", "-1", "1.5", "invalid"):
+            with self.subTest(value=value), tempfile.TemporaryDirectory() as tmpdir:
+                metadata_file = Path(tmpdir) / "node_metadata.env"
+                metadata_file.write_text(
+                    "SOPERATOR_NODE_REAL_MEMORY_BYTES=1\n", encoding="utf-8"
+                )
+
+                result = self.run_writer(metadata_file, value)
+
+                self.assertEqual(0, result.returncode)
+                self.assertFalse(metadata_file.exists())
+                self.assertIn("skipping node metadata export", result.stderr)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

internal/consts/node_metadata.go

@@ -0,0 +1,5 @@
+package consts
+
+// EnvNodeRealMemoryBytes carries the byte representation of the RealMemory value
+// rendered into slurm.conf for a worker node.
+const EnvNodeRealMemoryBytes = "SOPERATOR_NODE_REAL_MEMORY_BYTES"

internal/render/worker/container.go

@@ -243,6 +243,14 @@ func renderContainerNodeSetSlurmd(
 		return corev1.Container{}, fmt.Errorf("checking resource requests: %w", err)
 	}
 
+	for _, env := range nodeSet.ContainerSlurmd.CustomEnv {
+		if env.Name == consts.EnvNodeRealMemoryBytes {
+			return corev1.Container{}, fmt.Errorf("environment variable %q is managed by Soperator", consts.EnvNodeRealMemoryBytes)
+		}
+	}
+
+	realMemoryBytes := common.RenderRealMemorySlurmd(resources) * 1024 * 1024
+
 	appArmorProfile := nodeSet.ContainerSlurmd.AppArmorProfile
 	if nodeSet.AppArmorProfileUseDefault {
 		appArmorProfile = fmt.Sprintf("%s/%s", "localhost", naming.BuildAppArmorProfileName(nodeSet.ParentalCluster.Name, nodeSet.ParentalCluster.Namespace))
@@ -265,6 +273,7 @@ func renderContainerNodeSetSlurmd(
 				nodeSet.GPU.Nvidia.GDRCopyEnabled,
 				nodeSet.DockerEnabled,
 				nodeSet.NodeExtra,
+				realMemoryBytes,
 			),
 			nodeSet.ContainerSlurmd.CustomEnv...,
 		),
@@ -337,6 +346,7 @@ func renderNodeSetSlurmdEnv(
 	enableGDRCopy bool,
 	dockerEnabled bool,
 	slurmNodeExtra string,
+	realMemoryBytes int64,
 ) []corev1.EnvVar {
 	envVar := []corev1.EnvVar{
 		{
@@ -360,6 +370,10 @@ func renderNodeSetSlurmdEnv(
 			Name:  consts.EnvDockerEnabled,
 			Value: strconv.FormatBool(dockerEnabled),
 		},
+		{
+			Name:  consts.EnvNodeRealMemoryBytes,
+			Value: strconv.FormatInt(realMemoryBytes, 10),
+		},
 	}
 
 	if len(slurmNodeExtra) > 0 {