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 TAO service or SDK 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
Run all five steps in order before generating any launcher or submitting any job:
1. SSH connectivity — confirm SLURM USER and SLURM HOSTNAME are set and
passwordless SSH to the login host works ( ssh o BatchMode=yes ).
2. SDK install — optionally install the TAO SDK wrapper for Job handles + S3
wrapping ( nvidia tao sdk[slurm] , on public PyPI).
3. SLURM account — always resolve SLURM ACCOUNT via
sacctmgr show associations before generating any scripts. If unset and
only one account exists for the user, auto select it. If multiple exist, list
them and require the user to export one. Never submit a job without a verified
account — an invalid account causes "Invalid account or account/partition
combination" after SQSH conversion has already run.
4. GPU partition — verify at least one partition in SLURM PARTITION exists
via sinfo . The packaged default polar,polar3,polar4,grizzly is valid on
CS OCI ORD but may not exist on other clusters.
5. Enroot credentials — 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; without persistent
credentials, auth gated pulls fail with "Could not process JSON input". Use
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,
account/partition discovery commands, the enroot credentials heredoc,
prerequisite key setup (keypair, ssh copy id , known hosts , container key
mounts, 2FA handling), and the SSH failure remediation prompt.
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. Optionally convert the Docker image to a cached SQSH image with
srun n1 p <conversion partition enroot import . Do NOT use the cpu
partition for this step — cpu has a ~30 min wall time limit that is
shorter than the conversion time for large TAO images (9+ layers, 30 min).
Use cpu long (or another partition with ≥2 h limit) and set
SLURM CONVERSION PARTITION=cpu long and
SLURM CONVERSION TIMEOUT MINUTES=120 before constructing SlurmSDK .
The SDK validates the SQSH via SquashFS magic bytes before reusing it, so
partial files from failed conversions are automatically rejected and
reconverted — no manual cleanup needed. See the SQSH Conversion And
Caching section of references/slurm container execution.md for the full
env knob table ( SLURM ENROOT TEMP PATH for xattr restricted filesystems,
memory, force reconvert), cache/dedup semantics, live monitoring commands,
and manual pre staging.
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=/lustre .
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.
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 Paused
Required inputs
Ask for these in the SLURM intake; see references/slurm ssh credentials.md
for the full credential list, microservices schema keys, and defaults.
SLURM USER (required): SSH username for the login node.
SLURM HOSTNAME (required): Comma separated login hostnames for failover.
SLURM PARTITION (required): Partition list for GPU submission. Packaged
default polar,polar3,polar4,grizzly , treated as 4 hour queues.
SSH KEY PATH (preferred, expected before launch): private key for
non interactive public key auth. Ask for this first in remediation; prefer it
over the SSH AUTH SOCK agent socket fallback.
SLURM BASE RESULTS DIR (optional): base shared filesystem path; default
/lustre/fsw/portfolios/edgeai/users/<your dir (your per user Lustre dir).
SLURM ACCOUNT (resolve at preflight, not in initial intake): account for
SBATCH account . Auto discovered via sacctmgr during preflight step 3;
only ask the user if multiple accounts are found.
Do not ask for SLURM BASE RESULTS DIR in the initial intake unless the user
wants a custom results root.
Resource defaults
Defaults from tao core :
num nodes : 1
num gpus : 4
max num gpus per node : 8
cpus per task : 16
time hours : 4
timeout hours : 3.8
max time hours : 4
container mounts : /lustre
use requeue : true
use sqsh : true
When generating launchers or wrapper scripts for SLURM, set the wall time
defaults explicitly from the packaged platform resource defaults:
Do not default to 12 hours on SLURM. If the user supplies a longer
SLURM TIME HOURS , verify that the selected partition supports it before
submitting. For the packaged default partition list
polar,polar3,polar4,grizzly , reject requests above 4 hours and ask for a
different partition only if the user actually wants a longer wall time.
When num gpus is greater than or equal to max num gpus per node , the
handler treats the request as exclusive per node and computes additional nodes
from total GPU count when necessary.
Multi node, SDK, and retries
For multi node jobs ( num nodes 1 ), the SDK builds the sbatch directives and
exports the PyTorch distributed rendezvous env vars automatically: WORLD SIZE ,
NUM GPU PER NODE , NODE RANK , MASTER ADDR , and MASTER PORT (29500).
TAO entrypoints read WORLD SIZE + NUM GPU PER NODE and build torchrun
internally. Cosmos RL has special multi node role handling for controller,
policy, and rollout workers.
Use Lustre, not S3, for SLURM job inputs. The GPU allocation starts the
moment the job is dispatched, so a long s3:// download at the top of the
script burns the allocation, can get the job killed for GPU idle, and is billed
either way. Stage training data on the shared filesystem first and reference it
as lustre:///... . S3/HF/NGC pre fetch is fine for small auxiliary inputs
(checkpoints, configs), not training datasets. K8s/Brev do not share this
scheduler idle constraint.
Auto retry of infrastructure failures ( NODE FAIL , BOOT FAIL , NCCL transport
timeouts, CUDA driver init failures, GPU/IB link down, OOM killer node reaping,
Xid errors) is automatic in the SDK, with a stable user facing Job.id across
retries. Plain training failures surface immediately so a broken spec does not
consume the retry budget. SBATCH requeue is enabled by default via
SLURM USE REQUEUE=true .
See references/slurm container execution.md for the full multi node
env var/sbatch directive detail and table, cluster requirements, the optional
TAO SDK path ( SlurmSDK , build entrypoint , ActionWorkflow ) with code, the
Lustre not S3 rule in full, and the failure mode checklist;
references/slurm execution sdk.md covers the MAX JOB RETRIES retry budget.
When the SDK is in scope, read tao skill bank:tao run platform for the
SlurmSDK kwarg reference.
References
references/slurm ssh credentials.md — preflight script, SSH/key setup,
enroot credentials, full credential list, backend details, storage rules,
SSH remediation prompt.
references/slurm container execution.md — container execution steps,
monitoring, status mapping, cancellation, multi node detail, SDK use,
Lustre not S3, auto retry, failure modes.
references/slurm preflight storage.md — extended preflight/storage notes.
references/slurm execution sdk.md — extended execution/SDK notes.
references/detailed guide.md — navigation map for the split references.