Repository navigation
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #484 +/- ##
=======================================
Coverage ? 81.42%
=======================================
Files ? 137
Lines ? 20359
Branches ? 0
=======================================
Hits ? 16578
Misses ? 3781
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
9fb020b to
a657dfd
Compare
…steps Every srun that srtctl launches with a container image ran without --container-name, so Pyxis created a fresh enroot container for each step and extracted the whole image into ENROOT_RUNTIME_PATH (tmpfs) once per step per node. A disaggregated recipe puts a prefill worker, a decode worker, the frontend, the benchmark client and the telemetry exporters on one node from the same image, so a ~40 GiB image cost 150-200 GiB of host RAM per node. That RAM is what the engines' host KV cache (host_cache_size / KV_OFFLOADING) needs; the symptom was the kernel OOM-killing a context worker (SLURM OUT_OF_MEMORY, ~300 GB RSS per rank) 40-60 minutes into a run once the cache filled. Pass --container-name=srtctl_<sha1(image)[:12]>_<job id> on every containerised srun: the first step on a node creates the container, later steps of the same job attach to the same rootfs. Steps launched with ENROOT_REMAP_ROOT=yes get a separate name because a rootfs created without the remap cannot serve them; the job id keeps concurrent jobs on a node apart. Callers can still pass an explicit container_name. Same fix as benchmarking_toolkit's shared-container change; measured on a 1.35 TB-RAM node: four 41 GiB copies -> one, which let a 224 GiB ctx host cache run a full hour where 220-240 GiB were OOM-killed before. Signed-off-by: Iman Tabrizian <itabrizian@nvidia.com>
a657dfd to
318c460
Compare
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
| container_name = container_name or shared_container_name( | ||
| container_image, remap_root=remap_root, job_id=slurm_job_id | ||
| ) | ||
| srun_cmd.append(f"--container-name={container_name}") |
There was a problem hiding this comment.
Blocking: steps that start at the same time race to create the named container, and the loser can start on a half-unpacked rootfs.
srtctl launches the prefill and decode workers, telemetry and services on a node within seconds of each other. Pyxis checks for the name with enroot list -f and then runs enroot create, and nothing locks across steps. The shm mutex only covers the tasks inside one step (pyxis_slurmstepd.c L1304-1337, L1489-1533).
enroot create checks [ -e rootfs ] and then runs unsquashfs -d rootfs straight into the final directory, with no temp directory and rename (runtime.sh L426-438). While a 40 GiB image unpacks, which takes minutes, a second step sees the directory and gets pid == 0. It logs "reusing existing container filesystem" and starts on a partial rootfs. If its check lands slightly earlier, it fails with "File already exists" instead.
The single-step unit tests can't catch this. One fix is to create the container once per node in a dedicated step, and only then launch the steps that attach to it.
| idempotent), and benchmark scripts that `pip install` at run time install | ||
| into the live worker rootfs. The dynamo install lock and sentinel are | ||
| designed for this (see `_serialize_node_install` in `core/schema.py`). | ||
| - `--container-mounts` is still passed on every step; Pyxis applies it to the |
There was a problem hiding this comment.
Blocking: Pyxis does not apply --container-mounts to a step that attaches to a running container. It logs "ignoring --container-mounts when attaching to a running container" and calls remove_all_mounts() (pyxis_slurmstepd.c L1392-1395; v0.20.1 has the same code).
With this PR, whichever step starts first on a node decides the mounts for every later step on that node. The benchmark's runner.get_container_mounts() and anything else that differs from the worker step's mounts would be silently dropped. This paragraph is wrong: the mounts are dropped, so it isn't harmless.
| can be OOM-killed once its host cache fills. | ||
|
|
||
| - The name includes the job id, so two jobs sharing a node never attach to each | ||
| other's container, and Pyxis removes job-scoped containers at job end. |
There was a problem hiding this comment.
Blocking: on a default Pyxis install these named containers are never removed.
- Pyxis defaults to
container_scope = SCOPE_GLOBAL(config.c L37). - The job-end cleanup returns early unless the scope is
job(pyxis_slurmd.c L188). - Only unnamed (
temporary_rootfs) containers are removed when a step ends (pyxis_slurmstepd.c L1753).
On a global-scope cluster, every job therefore leaves an unpacked copy of each image on each node. If ENROOT_DATA_PATH is tmpfs, that leaks host RAM across jobs, which is worse than the problem this PR fixes.
Whether a cluster uses job or global scope differs by cluster. Per the Design Rule "Cluster differences live in srtslurm.yaml", this should be an opt-in ClusterConfig field, documented as requiring container_scope=job, rather than always on.
| remap_root = bool(srun_export_env) and all( | ||
| srun_export_env.get(k) == v for k, v in CONTAINER_REMAP_ROOT_EXPORT.items() | ||
| ) | ||
| container_name = container_name or shared_container_name( |
There was a problem hiding this comment.
Blocking: this changes every containerised srun, but the launch snapshots can't show it. tests/launch_snapshots.py (on main) records the keyword arguments passed to start_srun_process. Because the default name is computed in here, the recorded value is container_name=None and --container-name never shows up in tests/snapshots/launch/.
Option: compute the name once in RuntimeContext, which is the follow-up the description already suggests and also gives every step one name. Pass it in from the callers, then rebase and run make snapshots so the launch diff shows the change.
| # Container options | ||
| if container_image: | ||
| srun_cmd.extend(["--container-image", str(container_image)]) | ||
| remap_root = bool(srun_export_env) and all( |
There was a problem hiding this comment.
Nit: working out whether a step remaps root by reading the values in srun_export_env is fragile. A caller that adds another export, or spells the value y (which enroot also accepts), would change the result without anyone noticing. An explicit remap_root argument from the call sites that already pass CONTAINER_REMAP_ROOT_EXPORT would be clearer.
|
|
||
| - The name includes the job id, so two jobs sharing a node never attach to each | ||
| other's container, and Pyxis removes job-scoped containers at job end. | ||
| - Steps launched with `ENROOT_REMAP_ROOT=yes` (the dynamo cold install: workers |
There was a problem hiding this comment.
Nit: enroot applies ENROOT_REMAP_ROOT when a container starts (it sets up the user namespace), not when the rootfs is created. A rootfs created without the remap could serve a remapped step. The real reason to keep the names separate is that an attaching step joins the running container's namespaces (reuse_ns), including its user namespace. The explanation here and in the shared_container_name docstring should say that.
| attaching step and may warn that the container already exists, which is | ||
| harmless. | ||
| - The `Connection Commands` printed at start-up include the name, so an | ||
| interactive `srun ... --pty bash` lands in the workers' rootfs. |
There was a problem hiding this comment.
Nit: on recipes that install dynamo, workers use the separate remap container. _print_connection_info uses the name without the remap, so there the interactive shell does not land in the workers' rootfs.
| node creates the container; later steps of the same job that run the same image | ||
| on that node attach to it instead of extracting the image again. | ||
|
|
||
| Pyxis/enroot extracts each unnamed container into `ENROOT_RUNTIME_PATH`, which |
There was a problem hiding this comment.
Nit: enroot create unpacks into ENROOT_DATA_PATH, for named and unnamed containers alike (runtime.sh L425), not ENROOT_RUNTIME_PATH. The PR description says the same thing and needs the same fix.
| two containers (remapped workers/frontend, plain client/services/stage steps), | ||
| not one; recipes whose container already has dynamo installed hold one. | ||
| - Steps that write into the rootfs now share it with the engines on that node: | ||
| `setup_script` runs once per step against the same files (make it |
There was a problem hiding this comment.
Nit: being idempotent isn't enough for setup_script and benchmark-side pip install once they share a rootfs. Several steps can run them at the same moment against the same files, and only the dynamo install is serialised (_serialize_node_install). A note that they also need to be safe to run concurrently, or a lock, would help.
| subshell. Distinct FDs keep the two node-local and cross-node locks | ||
| independent and refactor-proof even if that inner subshell is removed. | ||
|
|
||
|
|
There was a problem hiding this comment.
Nit: double blank line in the docstring.
|
Review summary (by REVIEW.md): changes requested Extracting the image once per node is worth doing, but Pyxis and enroot don't behave the way the change assumes. Details are in the inline comments. Blocking
Nits: inline. Also, the description refers to "benchmarking_toolkit", which looks like an internal reference (REVIEW.md asks for none). A safer version: opt-in through Checked locally: |
This branch has not been deployed
3 participants
Problem
Every
srunsrtctl launches with a container image runs without--container-name, so Pyxis creates a fresh enroot container per step and extracts the image intoENROOT_RUNTIME_PATH(tmpfs = host RAM on most clusters) once per step per node. A disaggregated recipe puts a prefill worker, a decode worker, the frontend, the benchmark client and the telemetry exporters on one node from the same image, so a ~40 GiB image costs 150-200 GiB of host RAM per node.That is the RAM the engines' host KV cache (
host_cache_size,KV_OFFLOADING=dram) needs. Symptom: the kernel OOM-kills a context worker (SLURM OUT_OF_MEMORY, ~300 GB RSS per rank) 40-60 minutes into a run, once the host cache fills.Fix
Pass
--container-name=srtctl_<sha1(image)[:12]>_<job id>on every containerised srun (core/slurm.py::start_srun_process, new helpershared_container_name). The first step on a node creates the container; every later step of the job that runs the same image attaches to it.ENROOT_REMAP_ROOT=yes(workers and the dynamo frontend on dynamo-install recipes) get their own name, because a rootfs created without the remap cannot serve them. Known limitation: on such recipes a node holds two containers (remapped workers/frontend, plain client/services/stage steps) rather than one; recipes whose image already has dynamo hold one. Folding the remap decision intoRuntimeContextso every step of a job uses one name is a possible follow-up.setup_script, benchmark-sidepip install) now share it with the engines on the node; documented, and_serialize_node_install's docstring now states that its lock/sentinel rely on the sharing.Connection Commandsinclude the name, so an interactivesrun --pty bashlands in the workers' rootfs instead of extracting a further copy.--container-mountsstill goes on every step; Pyxis applies it to the attaching step (it may warn that the container exists, harmless).container_name.Evidence
Same change as benchmarking_toolkit's shared-container fix. On a 1.35 TB-RAM node running a disaggregated recipe, the per-node container footprint went from four 41 GiB rootfs copies (165 GiB) to one; a 224 GiB context-worker host KV cache then ran full 3600 s windows where 220 and 240 GiB had been OOM-killed.
enroot liston a node shows one container instead of four.Tests / docs
tests/test_slurm.py: same image -> same job-scoped name across steps; remap-root steps get a distinct name; explicitcontainer_namehonoured; no flag without an image.docs/slurm-faq.md: new "One Container Per Image Per Node" section.ruff checkclean.uv run pytest tests/: the 35 failures on my machine are identical with and without this change (rich-console/ANSI and script-path assertions intest_apply_*,test_benchmarksetc.), i.e. pre-existing locally, none from this PR.