Run as a nested job step when launched inside an existing allocation - #71
Open
bvanessen wants to merge 4 commits into
Open
Run as a nested job step when launched inside an existing allocation#71bvanessen wants to merge 4 commits into
bvanessen wants to merge 4 commits into
Conversation
When invoked from inside a scheduler allocation (salloc/sbatch/flux alloc), a blocking launch previously forced allocation-level flags (-N, --partition, ...) that made SLURM treat the srun as a request for a brand-new allocation instead of a job step in the current one. CLI: - With no job-size flag (-N/-g/--gpumem-at-least), inherit the node count from the enclosing allocation via a new scheduler-agnostic probe (hpc_launcher.schedulers.num_nodes_in_current_allocation) that checks the Flux, SLURM, and LSF environment markers. Explicit flags still win; --local is excluded. - Add -c/--cpus-per-task to give a job/step an explicit CPU footprint (SLURM --cpus-per-task, Flux --cores-per-task/--cores-per-slot). SLURM (blocking launch with SLURM_JOB_ID set): - Drop allocation-selection flags (--partition/--account/--reservation) with a warning; they are fixed by the enclosing allocation and would otherwise force a new one. - Reject a node request that exceeds the allocation upfront, instead of letting srun silently queue it as a second allocation. - Add --overlap by default so concurrent steps in one allocation don't hang on "step creation temporarily disabled, retrying (Requested nodes are busy)" (SLURM >= 20.11 gives steps exclusive CPUs). --overlap is omitted when --cpus-per-task (steps pack onto disjoint CPUs via the implied --exact) or --exclusive is given, and can be removed with -x ~--overlap. Non-blocking (sbatch/flux batch) submissions are unchanged: submitting a batch job from inside an allocation deliberately creates a new job. Adds tests/nested_job_step_test.py (17 tests, no scheduler binaries required) and documents the behavior in launch_cli.md and torchrun-hpc_cli.md.
…e tests Validating the nested-job-step feature on CTS-2 (Slurm) and El Capitan family (Flux) surfaced several issues that kept concurrent steps from packing onto one allocation, plus a CLI bug. All were confirmed and the fixes re-validated on hardware. Slurm (slurm.py): - A step with a stated footprint gets --exact plus --mem=0 (share the job's memory instead of implicitly consuming all of it) and, for CPU-only steps, --gres=none (stop implicitly holding every GPU in the job). Without these, concurrent steps serialize on "step creation temporarily disabled, retrying (Requested nodes are busy)". CTS-2 (cts2.py): - Apply --mpibind=off / --gpu-bind=none only when creating an allocation, not to nested steps: together they disabled both mechanisms that export per-step CUDA_VISIBLE_DEVICES (mpibind SPANK plugin and Slurm's GPU binding), leaving the variable unset inside steps. El Capitan family (el_capitan_family.py): - Do not add --exclusive to nested Flux jobs; it forced each nested `flux run` to demand the whole node, serializing concurrent launches. - Use Flux's own affinity plugins (-ompibind=off, -ogpu-affinity= per-task, -ocpu-affinity=per-task) for nested jobs and for any multi-GPU-per-task job: mpibind assigns GPUs by the NUMA locality of a task's cores rather than the granted resource set, which both caps a task at one GPU per domain and can hand two concurrent jobs the same GPU. Whole-node single-GPU-per-task launches keep mpibind. CLI (configure.py): - An explicit --gpus-per-proc 0 (CPU-only launch) is honored instead of being silently promoted to 1 on GPU systems; only an unset value now takes the default. Also guard the system-params info log when no active params exist. Step resource reporting (schedulers/): - Split the report script into slurm_step_resources.sh and flux_step_resources.sh over a shared step_resources_common.sh (CPU affinity, memory/NUMA, cgroup limits, GPU device-open tests, runtime view), with per-scheduler identity and verbose sections. Tests: - tests/slurm_step_resources_test.py + tests/run_slurm_step_packing.sh: from inside a 1-node CTS-2 allocation, pack 7 concurrent steps (2/1/1-GPU x 8 cores + 4 CPU-only x 16 cores), asserting distinct step IDs, overlapping lifetimes, disjoint CPU sets (SMT- and core-specialization-aware), cgroup-verified disjoint GPUs, and CUDA_VISIBLE_DEVICES consistency where exported. - tests/flux_step_resources_test.py + tests/run_flux_step_packing.sh: the Flux equivalent, sized to the node's core count, judging GPU visibility via ROCR/CUDA_VISIBLE_DEVICES. - nested_job_step_test.py: pin the new flag behavior (--exact/--mem=0/ --gres=none, CTS-2 binding flags, El Cap --exclusive and affinity selection); resource_validation_test.py: pin --gpus-per-proc 0. - Document in-allocation behavior updates in launch_cli.md.
bvanessen
marked this pull request as ready for review
August 25, 2026 22:03
The manual packing scripts (tests/run_{slurm,flux}_step_packing.sh) and
their pytest twins had grown into near-duplicates whose parsing and
verification logic had already been fixed in two places more than once.
Extending them to multi-node workloads made the duplication worse, so
the common structure now lives in shared modules and each scheduler
supplies only its policy.
Workloads (specs are nodes:gpus_per_task:cores_per_task, 1 task/node):
- 1-node (7 steps/jobs, as before): 2/1/1-GPU x 8 cores + 4 CPU-only
(16 cores on Slurm/matrix; sized to the node on Flux).
- 2-node (8 steps/jobs): adds a node-spanning GPU step (1 GPU + 8 cores
per task) and a spanning CPU step, mixed with the singles; sized so
the worst-case per-node packing (4 GPUs) still fits.
Multi-node support required per-task report files (a spanning step's
tasks would interleave on shared stdout) and per-host verification:
CPU/GPU disjointness is now checked per node, since distinct nodes
legitimately reuse the same core and GPU numbering, and a spanning
step's tasks must agree on their step/job ID.
New shared pieces:
- tests/step_packing_common.sh: shell harness (temp dir, wrapper
generation, launch/wait loops), parameterized by policy variables.
- tests/verify_step_packing.py: the single verifier, usable as a CLI
(shell harness) and as a module (pytest). Policies: gpu-mode
device-open (Slurm cgroup ground truth + CUDA_VISIBLE_DEVICES
consistency) vs visible-env (Flux ROCR/CUDA_VISIBLE_DEVICES);
cpu-mode exact (Slurm whole-core x SMT siblings) vs normalized-range
(Flux, physical-core normalization with core-specialization
tolerance). Reports all violations, not just the first.
- tests/step_packing_pytest.py: shared pytest harness plus a
shared_tmp_path fixture.
The run_* wrappers and *_test.py files shrink to skip guards, workload
specs, and the policy knobs.
Fixes a real failure found on a 2-node tuolumne run: the work dir was
created under TMPDIR (/var/tmp), which is node-local on LC systems, so
a spanning job's second task failed with "step_wrapper.sh: No such
file or directory". The wrapper script and reports are now written to
a shared filesystem ($HOME by default, override with STEP_PACKING_DIR)
-- same fix in the pytest fixture, which skips cleanly where home is
not writable.
tbennun
requested changes
Aug 28, 2026
| class SlurmScheduler(Scheduler): | ||
|
|
||
| @staticmethod | ||
| def in_slurm_allocation() -> bool: |
Collaborator
There was a problem hiding this comment.
why not make a generic in_allocation static method, so you have FluxScheduler.in_allocation() and slurm...
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
2 participants
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.
When invoked from inside a scheduler allocation (salloc/sbatch/flux alloc), a blocking launch previously forced allocation-level flags (-N, --partition, ...) that made SLURM treat the srun as a request for a brand-new allocation instead of a job step in the current one.
CLI:
SLURM (blocking launch with SLURM_JOB_ID set):
Non-blocking (sbatch/flux batch) submissions are unchanged: submitting a batch job from inside an allocation deliberately creates a new job.
Adds tests/nested_job_step_test.py (17 tests, no scheduler binaries required) and documents the behavior in launch_cli.md and torchrun-hpc_cli.md.