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