fix(slurm): make the generated job script portable across clusters - #163
Open
mkuznet1 wants to merge 6 commits into
Open
fix(slurm): make the generated job script portable across clusters#163mkuznet1 wants to merge 6 commits into
mkuznet1 wants to merge 6 commits into
Conversation
Clusters that do not advertise GPU GRES reject any job script carrying --gpus-per-node, so the generated sbatch fails before launch. Add slurm.skip_gpus_directive (default false) to omit the directive and rely on exclusive/nproc_per_node instead. Co-authored-by: Cursor <cursoragent@cursor.com>
The pre-submission check ran `madengine --version` with a 5s timeout, which a cold interpreter start off shared/NFS storage exceeds, aborting submission on a perfectly healthy environment. Raise it so the probe only catches a hang. Co-authored-by: Cursor <cursoragent@cursor.com>
The node-info step shelled out to the host package manager with no time limit. On a node where yum wants to import a repo GPG key the command waits on a prompt that never arrives, so the whole multi-node run hangs before the workload starts. Co-authored-by: Cursor <cursoragent@cursor.com>
A batch job is not guaranteed to inherit the submitter's PATH: a site can default sbatch to --export=NONE, and the module loads in the job body can rewrite it. The pre-submission check then passes on the login node while the compute node aborts with "madengine not found in PATH". Render the per-user bin directory and the directory the madengine console script was resolved from at submission time into the generated script, so the job puts the same interpreter back on PATH instead of relying on inheritance.
The single-node workspace probe matched \bnfs\b only, but df -T reports nfs4 on most modern NFS mounts. A shared submission directory was therefore classified as node-local and the job copied the whole project into /tmp instead of using the shared path. Match \bnfs[0-9]*\b so nfs, nfs3 and nfs4 are all recognized. The rendered job script now also has coverage for the --gpus-per-node opt-out it grew earlier in this batch: skip_gpus_directive shipped without tests, so nothing failed if the directive crept back into the template. Both states of the flag are asserted against the rendered script.
The shared-filesystem probe grepped the entire `df -T` output line, which carries the mount point as well as the type. A local disk mounted at a path such as /mnt/nfs-scratch therefore matched, the submission directory was classified as shared, and the single-node job worked out of storage the other side of the run could not see. Read the type column alone via `df --output=fstype` and anchor the pattern to it. The option is GNU coreutils 8.21 and up, so an awk fallback over `df -T` covers older systems. beegfs and panfs join the list of shared types while the pattern is being rewritten; both are common enough on HPC sites to be worth recognizing.
There was a problem hiding this comment.
Pull request overview
This PR improves SLURM portability in madengine by removing cluster-specific assumptions from the generated job script and hardening environment/probing logic so jobs submit and run consistently across more SLURM sites.
Changes:
- Adds
slurm.skip_gpus_directiveto optionally omit#SBATCH --gpus-per-nodefor clusters that reject it. - Makes the generated SLURM job script more robust by re-establishing
PATHinside the job and improving the shared-filesystem probe (fstype-only parsing and wider shared FS matching). - Caps the “informational” rocm-libs package queries to avoid hangs.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| tests/unit/test_slurm_job_template.py | Adds unit coverage for PATH exports, shared-fstype probe behavior, and skip_gpus_directive. |
| src/madengine/orchestration/run_orchestrator.py | Caps rocm-libs package-manager queries to prevent interactive hangs. |
| src/madengine/deployment/templates/slurm/job.sh.j2 | Updates SLURM job template to support skip_gpus_directive, re-add PATH inside the job, and harden shared-filesystem detection. |
| src/madengine/deployment/slurm.py | Plumbs skip_gpus_directive, increases CLI probe timeout, and passes submission-time madengine bin dir into the template context. |
Suppressed comments (3)
src/madengine/orchestration/run_orchestrator.py:773
Console.sh()already supports atimeoutparameter; prefer it over the externaltimeoutcommand to avoid depending on coreutils being present in the runtime environment.
print(self.console.sh("timeout 10 yum info rocm-libs", canFail=True))
src/madengine/orchestration/run_orchestrator.py:775
- Same as above: use
Console.sh(..., timeout=10)rather than shelling out totimeoutso the informational query is capped without adding a dependency on thetimeoutbinary.
print(self.console.sh("timeout 10 zypper info rocm-libs", canFail=True))
src/madengine/orchestration/run_orchestrator.py:777
- Same as above: use
Console.sh(..., timeout=10)rather thantimeout 10 ...to keep this capped on distros/environments wheretimeoutmay be unavailable.
print(self.console.sh("timeout 10 tdnf info rocm-libs", canFail=True))
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
mkuznet1
marked this pull request as ready for review
August 7, 2026 14:32
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Three clusters in a row have had to patch madengine's source before they could submit a
job. The RCCL CI integration
(ROCm/rocm-systems#9055) carries a
patch_madengine_for_cluster()helper that rewritesjob.sh.j2,slurm.pyandrun_orchestrator.pyat runtime, and the cluster we brought up most recently needed thesame edits by hand. Every one of those patches is a cluster-specific assumption baked
into madengine, not something the cluster did wrong. This PR removes five of the six, and
hardens the filesystem probe that the last of those fixes exposed.
Ticket: AICOMNET-366.
What is in here
feat(slurm): allow opting out of the --gpus-per-node sbatch directiveGresTypes=(null)rejects any job script carrying--gpus-per-node, so submission fails before launch.slurm.skip_gpus_directive(defaultfalse) omits the directive; the run relies on--exclusive+nproc_per_node.fix(slurm): raise the madengine availability probe timeoutmadengine --versionprobe had a 5 s timeout, which a cold interpreter start off shared/NFS storage exceeds — a healthy environment was reported as broken.fix(run): cap the informational rocm-libs package queryyum info rocm-libscan block forever on an interactive GPG-key prompt, hanging the run on a purely informational query.fix(slurm): inherit the submitter's PATH in the sbatch job--export=NONE, andmodule loadcan rewrite it), so the compute node aborts with "madengine not found in PATH" after the login-node check passed.$HOME/.local/binand the directory the madengine console script was resolved from at submission time.fix(slurm): match nfs4 in the shared-filesystem probe\bnfs\bonly, butdf -Treportsnfs4on most modern mounts, so a shared submission directory was treated as node-local and the whole project was copied into/tmp.\bnfs[0-9]*\b.fix(slurm): read the filesystem type, not the whole df linedf -Tline, mount point included, so a local disk at a path such as/mnt/nfs-scratchwas classified as shared storage — the opposite mistake, and the more dangerous one.df --output=fstype(with an awk fallback for pre-8.21 coreutils) and the pattern is anchored to it.beegfsandpanfsjoin the list of shared types.The
nfs4commit also backfills the coverageskip_gpus_directiveshipped without, sonothing fails silently if the directive creeps back into the template.
What is deliberately not in here
patch_madengine_for_cluster()has a sixth patch: a per-node venv bootstrap for clusterswhose compute nodes cannot use the submitter's interpreter. That is a design change to
how madengine reaches compute nodes, not a portability one-liner, and it is tracked
separately. After this PR the CI helper is down to that single patch.
Testing
tests/unit: 565 passed (544 before, plus 21 new intests/unit/test_slurm_job_template.pycovering the PATH exports, the fstype probe —
nfs/nfs3/nfs4/lustre/gpfs/ceph/beegfs/panfsaccepted,ext4/xfs/overlay/tmpfsand anext4disk under/mnt/nfs-scratchrejected — and both states ofskip_gpus_directive).tests/integrationon a GPU node: 149 passed, 1 failed, 2 skipped. The failure(
test_renderD_count_matches_gpu_count,assert 8 == 17) is a property of that node's KFDtopology and reproduces without these changes.
GresTypes=(null)cluster (Broadcom Thor2 / RoCE), aPrimus/Megatron training workload on an RCCL overlay image:
job.shcarries both PATH exports and no--gpus-per-node;bnxt_re*NICs;pre-change run on the same node pair.