diff --git a/helm/slurm-cluster/slurm_scripts/check_runner.py b/helm/slurm-cluster/slurm_scripts/check_runner.py index bde75b08d..4b486fc77 100644 --- a/helm/slurm-cluster/slurm_scripts/check_runner.py +++ b/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 diff --git a/helm/slurm-cluster/slurm_scripts/check_runner_test.py b/helm/slurm-cluster/slurm_scripts/check_runner_test.py index 32ce7621e..8b18425cc 100644 --- a/helm/slurm-cluster/slurm_scripts/check_runner_test.py +++ b/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) diff --git a/images/worker/slurmd.dockerfile b/images/worker/slurmd.dockerfile index de94eaacc..e16297eba 100644 --- a/images/worker/slurmd.dockerfile +++ b/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 && \ diff --git a/images/worker/slurmd_entrypoint.sh b/images/worker/slurmd_entrypoint.sh index 88d5026bf..3236034b5 100755 --- a/images/worker/slurmd_entrypoint.sh +++ b/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") diff --git a/images/worker/write_soperator_metadata.sh b/images/worker/write_soperator_metadata.sh new file mode 100755 index 000000000..66a00a9f3 --- /dev/null +++ b/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}" diff --git a/images/worker/write_soperator_metadata_test.py b/images/worker/write_soperator_metadata_test.py new file mode 100644 index 000000000..1ea93c49f --- /dev/null +++ b/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) diff --git a/internal/consts/node_metadata.go b/internal/consts/node_metadata.go new file mode 100644 index 000000000..3acfe58b6 --- /dev/null +++ b/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" diff --git a/internal/render/worker/container.go b/internal/render/worker/container.go index 00eeb2142..f1b720a47 100644 --- a/internal/render/worker/container.go +++ b/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 { diff --git a/internal/render/worker/statefulset_test.go b/internal/render/worker/statefulset_test.go index 8d17c0f6a..5abbcd661 100644 --- a/internal/render/worker/statefulset_test.go +++ b/internal/render/worker/statefulset_test.go @@ -664,6 +664,81 @@ func TestRenderNodeSetStatefulSet_DockerEnabled(t *testing.T) { } } +func TestRenderNodeSetStatefulSet_NodeRealMemoryMetadata(t *testing.T) { + createNodeSet := func() *values.SlurmNodeSet { + return &values.SlurmNodeSet{ + Name: "test-nodeset", + ParentalCluster: client.ObjectKey{ + Namespace: "test-namespace", + Name: "test-cluster", + }, + ContainerSlurmd: values.Container{ + NodeContainer: slurmv1.NodeContainer{ + Image: "test-image", + Resources: corev1.ResourceList{ + corev1.ResourceMemory: resource.MustParse("1G"), + corev1.ResourceCPU: resource.MustParse("1"), + }, + }, + }, + ContainerMunge: values.Container{ + NodeContainer: slurmv1.NodeContainer{Image: "munge-image"}, + }, + VolumeSpool: corev1.VolumeSource{ + HostPath: &corev1.HostPathVolumeSource{Path: "/tmp/spool"}, + }, + VolumeJail: corev1.VolumeSource{ + HostPath: &corev1.HostPathVolumeSource{Path: "/tmp/jail"}, + }, + StatefulSet: values.StatefulSet{Replicas: 1}, + SupervisorDConfigMapName: "supervisord-config", + SSHDConfigMapName: "sshd-config", + GPU: &slurmv1alpha1.GPUSpec{Enabled: false}, + } + } + + t.Run("exports the rendered Slurm RealMemory in bytes", func(t *testing.T) { + result, err := worker.RenderNodeSetStatefulSet( + "test-cluster", + createNodeSet(), + &slurmv1.Secrets{}, + consts.CGroupV2, + false, + false, + "", + ) + assert.NoError(t, err) + + for _, container := range result.Spec.Template.Spec.Containers { + if container.Name == consts.ContainerNameSlurmd { + // Slurm RealMemory is expressed in whole MiB: 1G becomes 953 MiB. + assertEnvValue(t, container.Env, consts.EnvNodeRealMemoryBytes, "999292928") + return + } + } + t.Fatal("slurmd container not found") + }) + + t.Run("rejects a custom override of Soperator metadata", func(t *testing.T) { + nodeSet := createNodeSet() + nodeSet.ContainerSlurmd.CustomEnv = []corev1.EnvVar{{ + Name: consts.EnvNodeRealMemoryBytes, + Value: "1", + }} + + _, err := worker.RenderNodeSetStatefulSet( + "test-cluster", + nodeSet, + &slurmv1.Secrets{}, + consts.CGroupV2, + false, + false, + "", + ) + assert.ErrorContains(t, err, "is managed by Soperator") + }) +} + func TestRenderNodeSetStatefulSet_PersistentVolumeClaimRetentionPolicy(t *testing.T) { createNodeSet := func(ephemeralNodes *bool, jailSubMounts []slurmv1alpha1.NodeVolumeMount, customVolumeMounts []slurmv1alpha1.NodeVolumeMount) *values.SlurmNodeSet { return &values.SlurmNodeSet{