Skip to content

Commit 8ba59e7

Browse files
author
Fabrizio Waldner
committed
Expose node RealMemory without controller RPC
1 parent 7b7244b commit 8ba59e7

9 files changed

Lines changed: 340 additions & 2 deletions

File tree

‎helm/slurm-cluster/slurm_scripts/check_runner.py‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@
88
import time
99
import typing
1010

11+
SOPERATOR_NODE_METADATA_FILE = "/run/soperator/node_metadata.env"
12+
SOPERATOR_NODE_REAL_MEMORY_BYTES = "SOPERATOR_NODE_REAL_MEMORY_BYTES"
13+
1114
# Set up logging
1215
try:
1316
log_stdout = "/dev/stdout"
@@ -114,7 +117,7 @@ class Check(typing.NamedTuple):
114117
# - CHECKS_NODE_COMMENT - comment field of the Slurm node
115118
# - CHECKS_NODE_REAL_MEM_BYTES - total allocatable memory in bytes for the Slurm node
116119
# - CHECKS_JOB_ALLOC_MEM_BYTES - memory in bytes, allocated for this Slurm job on this node
117-
# These values are extracted from long-running commands, that's why they aren't exported by default
120+
# Some values are extracted from long-running commands, so they aren't exported by default.
118121
need_env: list[str] = []
119122

120123
class NodeInfo(typing.NamedTuple):
@@ -356,10 +359,39 @@ def export_needed_env(check: Check):
356359
if env == "CHECKS_NODE_COMMENT":
357360
os.environ["CHECKS_NODE_COMMENT"] = get_node_info().comment
358361
if env == "CHECKS_NODE_REAL_MEM_BYTES":
359-
os.environ["CHECKS_NODE_REAL_MEM_BYTES"] = str(get_node_info().real_memory_bytes)
362+
os.environ["CHECKS_NODE_REAL_MEM_BYTES"] = str(get_node_real_memory_bytes())
360363
if env == "CHECKS_JOB_ALLOC_MEM_BYTES":
361364
os.environ["CHECKS_JOB_ALLOC_MEM_BYTES"] = str(get_job_info().allocated_memory_bytes)
362365

366+
# Get node RealMemory from node-local metadata, avoiding a controller RPC in the normal path.
367+
# Fall back to Slurm node info for compatibility with workers that have not yet been restarted
368+
# with an image and pod specification that publish the metadata file.
369+
@functools.lru_cache(maxsize=1)
370+
def get_node_real_memory_bytes() -> int:
371+
try:
372+
with open(SOPERATOR_NODE_METADATA_FILE, encoding="utf-8") as metadata_file:
373+
for raw_line in metadata_file:
374+
key, separator, value = raw_line.rstrip("\n").partition("=")
375+
if key != SOPERATOR_NODE_REAL_MEMORY_BYTES:
376+
continue
377+
if separator == "" or not value.isdecimal() or int(value) <= 0:
378+
raise ValueError(f"Invalid {SOPERATOR_NODE_REAL_MEMORY_BYTES} value")
379+
380+
real_memory_bytes = int(value)
381+
logging.info(
382+
f"Node RealMemory from {SOPERATOR_NODE_METADATA_FILE}: "
383+
f"{real_memory_bytes} bytes"
384+
)
385+
return real_memory_bytes
386+
387+
raise ValueError(f"Missing {SOPERATOR_NODE_REAL_MEMORY_BYTES} value")
388+
except Exception as e:
389+
logging.warning(
390+
f"Failed to get node RealMemory from {SOPERATOR_NODE_METADATA_FILE}: {e}; "
391+
"falling back to Slurm node info"
392+
)
393+
return get_node_info().real_memory_bytes
394+
363395
# Get GPU platform tags, e.g. ["8xH200", "8xGPU] from "nvidia-smi"
364396
# Please note, this command can be executed from both jail or host rootfs
365397
# The list starts with more specific tags, and ends with less specific ones

‎helm/slurm-cluster/slurm_scripts/check_runner_test.py‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,5 +116,111 @@ def test_uncomment_on_ok_still_runs_in_hc_program(self):
116116
self.assertEqual(["uncomment"], calls)
117117

118118

119+
def load_check_runner():
120+
required_env = {
121+
"SLURMD_NODENAME": "worker-1",
122+
"CHECKS_OUTPUTS_BASE_DIR": "/opt/soperator-outputs",
123+
"CHECKS_CONTEXT": "hc_program",
124+
"CHECKS_CONFIG": "/opt/slurm_scripts/checks.json",
125+
}
126+
with mock.patch.dict(os.environ, required_env):
127+
spec = importlib.util.spec_from_file_location(
128+
"check_runner_under_test", CHECK_RUNNER_PATH
129+
)
130+
module = importlib.util.module_from_spec(spec)
131+
assert spec.loader is not None
132+
spec.loader.exec_module(module)
133+
return module
134+
135+
136+
check_runner = load_check_runner()
137+
138+
139+
class NodeRealMemoryMetadataTest(unittest.TestCase):
140+
def setUp(self):
141+
check_runner.get_node_real_memory_bytes.cache_clear()
142+
143+
def test_reads_real_memory_from_local_metadata_without_node_rpc(self):
144+
with tempfile.TemporaryDirectory() as tmpdir:
145+
metadata_file = Path(tmpdir) / "node_metadata.env"
146+
metadata_file.write_text(
147+
"IGNORED=value\n"
148+
"SOPERATOR_NODE_REAL_MEMORY_BYTES=999292928\n",
149+
encoding="utf-8",
150+
)
151+
152+
with (
153+
mock.patch.object(
154+
check_runner, "SOPERATOR_NODE_METADATA_FILE", str(metadata_file)
155+
),
156+
mock.patch.object(check_runner, "get_node_info") as get_node_info,
157+
):
158+
result = check_runner.get_node_real_memory_bytes()
159+
160+
self.assertEqual(999292928, result)
161+
get_node_info.assert_not_called()
162+
163+
def test_exports_local_real_memory_for_checks(self):
164+
with tempfile.TemporaryDirectory() as tmpdir:
165+
metadata_file = Path(tmpdir) / "node_metadata.env"
166+
metadata_file.write_text(
167+
"SOPERATOR_NODE_REAL_MEMORY_BYTES=999292928\n",
168+
encoding="utf-8",
169+
)
170+
check = check_runner.Check(need_env=["CHECKS_NODE_REAL_MEM_BYTES"])
171+
172+
with (
173+
mock.patch.object(
174+
check_runner, "SOPERATOR_NODE_METADATA_FILE", str(metadata_file)
175+
),
176+
mock.patch.object(check_runner, "get_node_info") as get_node_info,
177+
mock.patch.dict(os.environ, {}, clear=False),
178+
):
179+
check_runner.export_needed_env(check)
180+
exported_value = os.environ["CHECKS_NODE_REAL_MEM_BYTES"]
181+
182+
self.assertEqual("999292928", exported_value)
183+
get_node_info.assert_not_called()
184+
185+
def test_falls_back_to_slurm_when_metadata_is_missing(self):
186+
with tempfile.TemporaryDirectory() as tmpdir:
187+
missing_file = Path(tmpdir) / "missing.env"
188+
node_info = check_runner.NodeInfo(real_memory_bytes=2147483648)
189+
190+
with (
191+
mock.patch.object(
192+
check_runner, "SOPERATOR_NODE_METADATA_FILE", str(missing_file)
193+
),
194+
mock.patch.object(
195+
check_runner, "get_node_info", return_value=node_info
196+
) as get_node_info,
197+
):
198+
result = check_runner.get_node_real_memory_bytes()
199+
200+
self.assertEqual(2147483648, result)
201+
get_node_info.assert_called_once_with()
202+
203+
def test_falls_back_to_slurm_when_metadata_is_invalid(self):
204+
with tempfile.TemporaryDirectory() as tmpdir:
205+
metadata_file = Path(tmpdir) / "node_metadata.env"
206+
metadata_file.write_text(
207+
"SOPERATOR_NODE_REAL_MEMORY_BYTES=invalid\n", encoding="utf-8"
208+
)
209+
node_info = check_runner.NodeInfo(real_memory_bytes=1073741824)
210+
211+
with (
212+
mock.patch.object(
213+
check_runner, "SOPERATOR_NODE_METADATA_FILE", str(metadata_file)
214+
),
215+
mock.patch.object(
216+
check_runner, "get_node_info", return_value=node_info
217+
) as get_node_info,
218+
):
219+
result = check_runner.get_node_real_memory_bytes()
220+
221+
self.assertEqual(1073741824, result)
222+
get_node_info.assert_called_once_with()
223+
224+
119225
if __name__ == "__main__":
120226
unittest.main(verbosity=2)

‎images/worker/slurmd.dockerfile‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,7 @@ RUN mkdir -p /var/log/slurm/multilog && \
157157

158158
# Copy slurmd entrypoint script
159159
COPY images/worker/slurmd_entrypoint.sh /opt/bin/slurm/
160+
COPY images/worker/write_soperator_metadata.sh /opt/bin/slurm/
160161

161162
# Copy worker init script (controller readiness + topology for ephemeral nodes)
162163
COPY images/worker/worker_init.py /opt/bin/slurm/
@@ -167,6 +168,7 @@ COPY images/worker/docker_proxy_nginx_entrypoint.sh /opt/bin/slurm/
167168
COPY images/worker/dockerd_entrypoint.sh /opt/bin/slurm/
168169

169170
RUN chmod +x /opt/bin/slurm/slurmd_entrypoint.sh && \
171+
chmod +x /opt/bin/slurm/write_soperator_metadata.sh && \
170172
chmod +x /opt/bin/slurm/supervisord_entrypoint.sh && \
171173
chmod +x /opt/bin/slurm/worker_init.py && \
172174
chmod +x /opt/bin/slurm/docker_proxy_nginx_entrypoint.sh && \

‎images/worker/slurmd_entrypoint.sh‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,9 @@ else
4242
export TOPO_SWITCH_TIER2="unknown"
4343
fi
4444

45+
echo "Export Soperator node metadata"
46+
/opt/bin/slurm/write_soperator_metadata.sh
47+
4548
echo "Evaluate variables in the Slurm node 'Extra' field"
4649
evaluated_extra=$(eval echo "$SLURM_NODE_EXTRA")
4750

‎images/worker/write_soperator_metadata.sh‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
#!/bin/bash
2+
3+
set -euo pipefail
4+
5+
metadata_file="${1:-/run/soperator/node_metadata.env}"
6+
node_real_memory_bytes="${SOPERATOR_NODE_REAL_MEMORY_BYTES:-}"
7+
8+
if ! [[ "${node_real_memory_bytes}" =~ ^[0-9]+$ ]] || [[ "${node_real_memory_bytes}" == "0" ]]; then
9+
rm -f -- "${metadata_file}"
10+
echo "SOPERATOR_NODE_REAL_MEMORY_BYTES is unavailable or invalid; skipping node metadata export" >&2
11+
exit 0
12+
fi
13+
14+
metadata_dir="$(dirname -- "${metadata_file}")"
15+
install -d -m 0755 "${metadata_dir}"
16+
17+
temporary_file="$(mktemp "${metadata_file}.tmp.XXXXXX")"
18+
trap 'rm -f -- "${temporary_file}"' EXIT
19+
20+
printf 'SOPERATOR_NODE_REAL_MEMORY_BYTES=%s\n' "${node_real_memory_bytes}" > "${temporary_file}"
21+
chmod 0644 "${temporary_file}"
22+
mv -f -- "${temporary_file}" "${metadata_file}"
23+
trap - EXIT
24+
25+
echo "Exported node metadata to ${metadata_file}"

‎images/worker/write_soperator_metadata_test.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
import os
2+
import stat
3+
import subprocess
4+
import tempfile
5+
import unittest
6+
from pathlib import Path
7+
8+
9+
SCRIPT_PATH = Path(__file__).with_name("write_soperator_metadata.sh")
10+
11+
12+
class WriteSoperatorMetadataTest(unittest.TestCase):
13+
def run_writer(
14+
self, metadata_file: Path, real_memory_bytes: str | None
15+
) -> subprocess.CompletedProcess[str]:
16+
env = os.environ.copy()
17+
if real_memory_bytes is None:
18+
env.pop("SOPERATOR_NODE_REAL_MEMORY_BYTES", None)
19+
else:
20+
env["SOPERATOR_NODE_REAL_MEMORY_BYTES"] = real_memory_bytes
21+
22+
return subprocess.run(
23+
["bash", str(SCRIPT_PATH), str(metadata_file)],
24+
check=False,
25+
env=env,
26+
stdout=subprocess.PIPE,
27+
stderr=subprocess.PIPE,
28+
text=True,
29+
)
30+
31+
def test_writes_real_memory_metadata_atomically(self):
32+
with tempfile.TemporaryDirectory() as tmpdir:
33+
metadata_file = Path(tmpdir) / "nested" / "node_metadata.env"
34+
35+
result = self.run_writer(metadata_file, "999292928")
36+
37+
self.assertEqual(0, result.returncode)
38+
self.assertEqual(
39+
"SOPERATOR_NODE_REAL_MEMORY_BYTES=999292928\n",
40+
metadata_file.read_text(encoding="utf-8"),
41+
)
42+
self.assertEqual(0o644, stat.S_IMODE(metadata_file.stat().st_mode))
43+
self.assertEqual([], list(metadata_file.parent.glob("*.tmp.*")))
44+
45+
def test_replaces_existing_metadata(self):
46+
with tempfile.TemporaryDirectory() as tmpdir:
47+
metadata_file = Path(tmpdir) / "node_metadata.env"
48+
metadata_file.write_text(
49+
"SOPERATOR_NODE_REAL_MEMORY_BYTES=1\n", encoding="utf-8"
50+
)
51+
52+
result = self.run_writer(metadata_file, "2147483648")
53+
54+
self.assertEqual(0, result.returncode)
55+
self.assertEqual(
56+
"SOPERATOR_NODE_REAL_MEMORY_BYTES=2147483648\n",
57+
metadata_file.read_text(encoding="utf-8"),
58+
)
59+
60+
def test_missing_or_invalid_metadata_is_skipped(self):
61+
for value in (None, "", "0", "-1", "1.5", "invalid"):
62+
with self.subTest(value=value), tempfile.TemporaryDirectory() as tmpdir:
63+
metadata_file = Path(tmpdir) / "node_metadata.env"
64+
metadata_file.write_text(
65+
"SOPERATOR_NODE_REAL_MEMORY_BYTES=1\n", encoding="utf-8"
66+
)
67+
68+
result = self.run_writer(metadata_file, value)
69+
70+
self.assertEqual(0, result.returncode)
71+
self.assertFalse(metadata_file.exists())
72+
self.assertIn("skipping node metadata export", result.stderr)
73+
74+
75+
if __name__ == "__main__":
76+
unittest.main(verbosity=2)

‎internal/consts/node_metadata.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
package consts
2+
3+
// EnvNodeRealMemoryBytes carries the byte representation of the RealMemory value
4+
// rendered into slurm.conf for a worker node.
5+
const EnvNodeRealMemoryBytes = "SOPERATOR_NODE_REAL_MEMORY_BYTES"

‎internal/render/worker/container.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,14 @@ func renderContainerNodeSetSlurmd(
243243
return corev1.Container{}, fmt.Errorf("checking resource requests: %w", err)
244244
}
245245

246+
for _, env := range nodeSet.ContainerSlurmd.CustomEnv {
247+
if env.Name == consts.EnvNodeRealMemoryBytes {
248+
return corev1.Container{}, fmt.Errorf("environment variable %q is managed by Soperator", consts.EnvNodeRealMemoryBytes)
249+
}
250+
}
251+
252+
realMemoryBytes := common.RenderRealMemorySlurmd(resources) * 1024 * 1024
253+
246254
appArmorProfile := nodeSet.ContainerSlurmd.AppArmorProfile
247255
if nodeSet.AppArmorProfileUseDefault {
248256
appArmorProfile = fmt.Sprintf("%s/%s", "localhost", naming.BuildAppArmorProfileName(nodeSet.ParentalCluster.Name, nodeSet.ParentalCluster.Namespace))
@@ -265,6 +273,7 @@ func renderContainerNodeSetSlurmd(
265273
nodeSet.GPU.Nvidia.GDRCopyEnabled,
266274
nodeSet.DockerEnabled,
267275
nodeSet.NodeExtra,
276+
realMemoryBytes,
268277
),
269278
nodeSet.ContainerSlurmd.CustomEnv...,
270279
),
@@ -337,6 +346,7 @@ func renderNodeSetSlurmdEnv(
337346
enableGDRCopy bool,
338347
dockerEnabled bool,
339348
slurmNodeExtra string,
349+
realMemoryBytes int64,
340350
) []corev1.EnvVar {
341351
envVar := []corev1.EnvVar{
342352
{
@@ -360,6 +370,10 @@ func renderNodeSetSlurmdEnv(
360370
Name: consts.EnvDockerEnabled,
361371
Value: strconv.FormatBool(dockerEnabled),
362372
},
373+
{
374+
Name: consts.EnvNodeRealMemoryBytes,
375+
Value: strconv.FormatInt(realMemoryBytes, 10),
376+
},
363377
}
364378

365379
if len(slurmNodeExtra) > 0 {

0 commit comments

Comments
 (0)