tao-run-on-slurm
Remote SLURM GPU cluster execution over SSH with sbatch/srun, Pyxis/Enroot containers, and Lustre-backed results. Use when running TAO training/eval/inference jobs on an on-prem or DGX SLURM cluster. Trigger phrases include "run on SLURM", "submit sbatch", "DGX SLURM cluster", "Pyxis/Enroot containe
By nvidia · 1,539 installs
npx skills add nvidia/skills --skill tao-run-on-slurm
Source repository · Upstream listing
SLURM
Standalone install? If this session was not initialized by the TAO skill bank plugin, run the tao setup skill first (host preflight, credentials, cross skill discovery).
Remote GPU compute platform for clusters managed by SLURM. Jobs are submitted
from the launch host to a login node over SSH, staged on a shared
filesystem, submitted with sbatch , and executed with srun container support.
When to use
Use SLURM when the user has access to a managed GPU cluster, shared Lustre
storage, and scheduler owned GPU allocation. Do not use SLURM for local files
that exist only on the agent machine; data and outputs must be reachable from
the cluster.
Preflight + SSH
Confirm SLURM USER and SLURM HOSTNAME are exported and passwordless SSH to a
login host works ( ssh o BatchMode=yes ).
The launch host needs ssh , not local sbatch , srun , Enroot, or a Lustre
mount. Preflight those scheduler, Pyxis, Enroot, and shared storage dependencies
on the selected remote login/compute frame. Model specific inspectors may be
streamed from the installed skill over SSH stdin; do not stage an ad hoc source
patch or treat the launch host as the SLURM frame.
For private nvcr.io images, install ~/.config/enroot/.credentials on the
cluster once per (cluster, user): Pyxis/Enroot does not read NGC KEY from the
job env, and without persistent credentials, auth gated pulls fail with "Could
not process JSON input" at job startup. Install it via the printf ssh
heredoc so the NGC KEY value never lands in shell history, intermediate files,
or chat output; never cat / echo the value.
If a preflight check fails, the agent prompts the user to authorize the
install/fix via Bash. Pip installable Python requirements are the exception:
install them automatically, then rerun preflight.
See references/slurm ssh credentials.md for the full preflight script, the
enroot credentials heredoc, prerequisite key setup (keypair, ssh copy id ,
known hosts , container key mounts, 2FA handling), and the SSH failure
remediation prompt.
Execution — the four verbs
tao run on slurm is a platform consumer : it runs a spec bundle over
ssh + sbatch/squeue/sacct/scancel , mutating only the job record. Storage is
tier A (Lustre) — the dataset is staged to a shared path before submit and
read through Pyxis; never fetch S3 inside the allocation (the scheduler idle
timeout kills GPU idle jobs and bills the wasted time). $BANK =
${TAO SKILL BANK PATH} ; $LOGIN = a resolved SLURM HOSTNAME .
submit
1. Reuse what's already staged — never redo (tier A):
Image: @@IMAGE@@ is a Lustre .sqsh — reuse an existing one if present
( ssh $LOGIN ls <sqsh ); only if missing, convert once with enroot import
(cached by name — see references/slurm container execution.md ).
Dataset: confirm it is already on Lustre ( ssh $LOGIN test e … ) and
reference those paths; tao data io stages only a small auxiliary input
that is not there yet — never re stage existing data, and never the training
set inside the allocation.
Then author the spec at <job dir /specs/spec.yaml on Lustre with those paths.
2. Credentials → sidecar (never inline): if the run needs session creds
(e.g. HF TOKEN ), write them to a mode 600 sidecar on Lustre and let the
template shred it on exit; NGC image pulls use the one time
~/.config/enroot/.credentials (see references/slurm ssh credentials.md ),
not the job env:
3. Open the record — mints the id, binds results dir on Lustre, before launch:
4. Consume the optional model lifecycle. If the validated spec bundle has
execution , preserve its order and semantics while mapping distributed
intent to native SLURM/Pyxis. Stage only its checksum closed
supporting files . The full generic lifecycle and staging contract is in
references/slurm container execution.md .
5. Render templates/slurm/singlenode.sbatch.tmpl — substitute every
@@<NAME @@ ( JOB NAME=$JOB ID , NUM GPUS , CPUS PER TASK , TIME , LOG DIR ,
IMAGE , CONTAINER MOUNTS=<RUNTIME SUPPLIED MOUNTS , COMMAND=<bundle command reading the shared storage
spec , SBATCH EXTRA= account/partition lines, ENV FILE= the sidecar path or
empty, EXTRA ENV= any cluster NCCL knobs) → <job dir /sbatch/job $JOB ID.sbatch .
Lint + syntax check before submit: redact secrets.py lint <sbatch must
pass and bash n <sbatch must succeed.
6. Submit + record RUNNING:
A submit that skipped the gate or the open has no id — so it cannot launch.
status
SLURM state vocab
PENDING PENDING
RUNNING / COMPLETING RUNNING
COMPLETED COMPLETE (confirm status.json in results dir )
FAILED / TIMEOUT / OUT OF MEMORY ERROR (infra vs program classify → retry, M6)
NODE FAIL / BOOT FAIL ERROR , err class=ERR INFRA ( requeue re queues these)
CANCELLED / PREEMPTED / REVOKED CANCELED
(not found) UNKNOWN
Native sub state rides in the transition message . Poll at the chosen interval;
long queue waits are normal — do not stop on elapsed time.
logs
cancel
Treat an already terminated SLURM job as a successful cancel.
Multi node (nodes 1)
Same four verbs, with three additions at submit:
1. Render templates/slurm/multinode.sbatch.tmpl instead of the single node
one — it's a strict superset (adds nodes / wait all nodes + the
rendezvous block). WORLD SIZE is the node count (TAO's misnomer); never
change it to a global rank count.
2. NCCL probe first — before the real job, run a cheap 2 node all reduce
( scripts/nccl allreduce probe.py under the container's torchrun) with a
~120s timeout. Before invoking torchrun, preserve the TAO rendezvous values
as TAO NODE COUNT=$WORLD SIZE ,
TAO GPUS PER NODE=$NUM GPU PER NODE , and
TAO NODE RANK=$SLURM PROCID ; torchrun overwrites its standard
WORLD SIZE with the global process count. NCCL PROBE OK → proceed.
Timed out (the collective hung)
→ set the cluster's NCCL knob in EXTRA ENV and re probe — on CS OCI ORD that
is export NCCL P2P DISABLE=1 (the intra node P2P hang), often with
NCCL SOCKET IFNAME=eth0 / NCCL IB DISABLE=1 . Cache the working env per
cluster so later jobs skip the probe. Gate on gpus per node 1 too —
the P2P hang triggers on a single node with 2+ GPUs.
3. Tier A Lustre, sidecar creds, record, and lint are unchanged.
Cosmos backend guardrails
Read [ references/cosmos slurm guardrails.md ](references/cosmos slurm guardrails.md)
before rendering a Cosmos command. It defines image staging, planner
materialization, Framework and Cosmos RL launch contracts, worker/runtime
requirements, and exit/status handling.
Storage
Use shared filesystem URIs, not local or file:// paths; tao core rejects
local/file paths for remote backends.
lustre:///absolute/path for user provided datasets on Lustre.
slurm:// paths may appear in microservices metadata and are converted to
Lustre paths before the container starts.
Accept either dataset roots (model skills map them to required files) or direct
spec key paths. After SSH succeeds and before generating scripts, test e each
required dataset path from the login host; if it fails, stop and ask for
corrected paths or staged data rather than producing scripts that fail in the
first training job. See references/slurm ssh credentials.md for root vs.
direct spec modes, backend details, and the results dir default.
Container execution
tao core runs TAO containers through Pyxis/Enroot:
1. Stage compact JSON files for specs, environment, and cloud metadata under
<job dir /specs , <job dir /env , and <job dir /meta .
2. Convert the Docker image to a cached SQSH image before the GPU job, with
srun n1 p <conversion partition enroot import . This is a one time cost
per image, not an optional optimization — see Acquire the image off the GPU
allocation below.
3. Write an sbatch script under <job dir /sbatch/job <job id .sbatch .
4. Submit sbatch export=ALL <script .
5. Run the container with srun container image=<image container mounts=<RUNTIME SUPPLIED MOUNTS .
Accepted image formats: /path/to/image.sqsh , registry image:tag ,
docker://registry image:tag , and ordinary registry/image:tag (converted to
Pyxis form when needed). SQSH conversion is cached by image name; for :latest
images the cached SQSH is reused unless force reconvert latest is enabled.
Acquire the image off the GPU allocation
The GPU is yours from the moment the allocation starts, not from when compute
begins. Anything the job does before training — pulling a registry image,
converting it, fetching a dataset — runs on GPUs that are idle, billed, and
visible to the cluster's GPU idle reaper. A first time TAO pull plus enroot
conversion is minutes of that, which is long enough to be killed and long enough
to be expensive.
So the image must already be a local .sqsh when the GPU job starts. Passing a
docker:// or registry image:tag URI straight to srun container image=
makes Pyxis pull and convert inside the allocation — the exact trap. Convert
once on a CPU partition , then point every later job at the resulting file:
The same rule governs data: stage it to Lustre before submit (tier A) rather
than fetching inside the allocation.
CS OCI ORD conversion uses cpu long , 4 CPUs, 7200M memory, no exclusive node,
node local Enroot temp paths, and at least 120 minutes. The execution reference
records the evidence and QOSGrpMemLimit recovery contract.
Partial conversions are self detecting: the SQSH is validated by hsqs magic,
so a truncated file is rejected rather than silently used. Conversion runs once
and is then cached by image name.
A failed conversion must not fall back to the registry image. The tempting
recovery — pass docker://… to srun and let Pyxis handle it — puts the pull
back inside the GPU allocation, which is the cost the conversion existed to
avoid, and it does so precisely when something is already wrong. Treat a failed
or truncated conversion as fatal: fix it on the CPU partition and resubmit.
Diagnostic: if a job is unexpectedly slow to produce output, check what
container image= actually received. A registry URI there — rather than a
.sqsh path — means the pull happened on the GPUs.
Monitoring and cancellation
Scheduler status comes from the stored SLURM job id via squeue / sacct ;
TAO terminal status comes from status.json in the shared results folder.
While chat monitoring is enabled, keep polling at the requested interval for
any non terminal job ( PENDING , RUNNING , or otherwise). Do not stop after a
fixed elapsed time such as 30 minutes; long queue waits are normal on shared
GPU partitions.
Do not send a final response for a non terminal SLURM job when chat
monitoring is enabled. A final response is a detach action; use it only if the
user asked to detach/stop or the job reached terminal state.
Logs are read over SSH from
<job dir /slurm logs/<slurm job name <slurm job id /main.out and .err .
Cancel by looking up backend details.slurm metadata.slurm job id and running
scancel <slurm job id over SSH. Treat missing or already terminated jobs as
successful cancellation.
Status mapping:
PENDING Pending
RUNNING or COMPLETING Running
COMPLETED check status.json
FAILED , BOOT FAIL , DEADLINE , OUT OF MEMORY , NODE FAIL retry if
logs match retriable infrastructure patterns, otherwise Error
CANCELLED , PREEMPTED , REVOKED Canceled
TIMEOUT Error
SUSPENDED , STOPPED Running (still scheduler owned and may resume;
the native sub state rides in the transition message — same convention as
docker paused )
Required inputs
Ask for these in the SLURM intak