Run Containers in Jobs
Pyxis adds container options to srun and sbatch. This page covers the flags, how to get your data
and environment into the container, and complete job scripts you can copy.
The Container Flags
Pass these to srun, or as #SBATCH directives / srun arguments inside an sbatch script.
Core flags
These are the flags used by TIR's own reference workloads and the ones to build on.
| Flag | What it does |
|---|---|
--container-image=<ref> | The image to run. A registry reference (nvcr.io/nvidia/pytorch:25.09-py3) or a path to a squash file (/pfs/images/pytorch.sqsh) |
--container-name=<name> | Names the imported container so it is reused by later steps and jobs instead of being re-imported. See Cache and manage images |
--container-mounts=<src>:<dst>[,…] | Bind-mounts host paths into the container. Comma-separate multiple mounts |
--container-workdir=<path> | Working directory inside the container |
--mpi=pmix | Use PMIx for process launch. Required for MPI ranks and recommended for any multi-node container job |
Other Pyxis flags
Pyxis also ships the flags below. They are standard Pyxis rather than TIR-specific, and the exact set
depends on the Pyxis version inside your cluster image — confirm with srun --help | grep -i container on your own cluster before relying on one in a production script.
| Flag | Typical use |
|---|---|
--container-writable / --container-readonly | Allow or forbid writes to the container's root filesystem |
--container-remap-root / --no-container-remap-root | Whether your user appears as root inside the container |
--container-save=<path.sqsh> | Write the container's filesystem back out to a squash file after the step |
--container-env=<VAR[,VAR…]> | Pass specific variables through instead of the default environment handling |
--container-mount-home / --no-container-mount-home | Mount or skip your home directory |
--container-entrypoint / --no-container-entrypoint | Run or skip the image's declared entrypoint |
srun, not #SBATCHInside a batch script, srun --container-image=… is the reliable form. It makes it obvious which step
runs in the container and lets you run set-up commands on the host before it — writing directories,
staging data, printing diagnostics.
Mount Your Data
A container starts from the image's own filesystem. Your cluster storage exists on the node, but you have to pass in the paths the job needs.
What is on the node
| Storage | Typical mount | Writable | Shared across nodes |
|---|---|---|---|
| Parallel File System (PFS) | the path you set when mounting it, e.g. /pfs | Yes | Yes — same filesystem on the login node and every worker |
| Shared File System (SFS) | e.g. /shared or /mnt | Yes | Yes |
| Dataset | the path you set when mounting it | No — read-only | Yes |
/dev/shm | /dev/shm | Yes | No — per pod, backed by memory |
| Node-local scratch | e.g. /tmp | Yes | No — per node, and lost on restart |
Manage these from the Volumes tab.
Mounting them
srun --container-image=nvcr.io/nvidia/pytorch:25.09-py3 \
--container-mounts=/pfs:/pfs,/shared:/shared \
--container-workdir=/pfs/project \
python train.py
You can remap paths — --container-mounts=/pfs/datasets/imagenet:/data makes the dataset appear at
/data inside the container, which is handy when your image or code expects a fixed path.
The container's filesystem is discarded when the step ends, and node-local paths (/tmp, /dev/shm)
do not survive a node restart and are invisible to other nodes. Checkpoints, logs and model outputs
must go to a PFS, SFS or Weka mount. This is the single most common way to lose a training run.
InfiniBand and RDMA
Device nodes are not automatically inside the container. For jobs that use RDMA — most multi-node NCCL training — mount them explicitly:
--container-mounts=/pfs:/pfs,/dev/infiniband:/dev/infiniband
See Multi-node container training for the full picture.
Environment Variables
Variables you export before srun are visible to the process inside the container, which is how the
NCCL and framework tuning in the examples below works:
export NCCL_DEBUG=INFO
srun --container-image=… python train.py
Slurm's own variables (SLURM_JOB_ID, SLURM_NODEID, SLURM_PROCID, SLURM_NTASKS, …) are
available inside the container too — use them to derive ranks and per-job output paths.
Nothing is injected into your job environment. Any NCCL_*, UCX_* or framework tuning is entirely
yours to set, which means defaults come from the image. For multi-node RDMA runs you will normally
need to set several — see
Multi-node container training.
If your image needs a secret at runtime, read it inside the job from a file on your mounted storage. Do not bake credentials into a job script or into the image.
GPUs Inside the Container
Ask Slurm for GPUs the usual way; the container inherits them.
#SBATCH --gres=gpu:8 # 8 GPUs on the node
#SBATCH --gpus-per-node=8 # equivalent, newer syntax
Inside the container, nvidia-smi shows exactly the GPUs the job was allocated — not the whole node.
CUDA_VISIBLE_DEVICES is set by Slurm accordingly.
srun --gres=gpu:2 --container-image=nvcr.io/nvidia/pytorch:25.09-py3 nvidia-smi -L
The cluster advertises GPUs to Slurm as a GRES with the card type, so you can check what is available before you submit:
sinfo -o "%20N %10c %10m %25G %10T"
scontrol show node slinky-0
The container brings its own CUDA runtime, so pick a tag built for your GPU generation. An image built for older CUDA may run poorly, or not at all, on newer hardware such as H200 or B200. NGC tags state the CUDA version they carry.
Worked Examples
Single-node training on 8 GPUs
#!/bin/bash
#SBATCH --job-name=train-8gpu
#SBATCH --partition=all
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --gpus-per-node=8
#SBATCH --cpus-per-task=64
#SBATCH --time=08:00:00
#SBATCH --output=/pfs/logs/%x-%j.out
#SBATCH --error=/pfs/logs/%x-%j.err
IMAGE=nvcr.io/nvidia/pytorch:25.09-py3
export NCCL_DEBUG=WARN
srun --container-image="$IMAGE" \
--container-name=pytorch-train \
--container-mounts=/pfs:/pfs \
--container-workdir=/pfs/project \
torchrun --standalone --nproc_per_node=8 train.py \
--data /pfs/datasets/imagenet \
--checkpoint-dir /pfs/checkpoints/${SLURM_JOB_ID}
A one-off command in a container
srun --gres=gpu:1 --time=00:15:00 \
--container-image=nvcr.io/nvidia/pytorch:25.09-py3 \
--container-mounts=/pfs:/pfs \
python /pfs/scripts/validate_dataset.py
Hyperparameter sweep as a job array
#!/bin/bash
#SBATCH --job-name=sweep
#SBATCH --partition=all
#SBATCH --array=0-15
#SBATCH --nodes=1
#SBATCH --gpus-per-node=1
#SBATCH --time=02:00:00
#SBATCH --output=/pfs/logs/%x-%A_%a.out
LRS=(0.1 0.03 0.01 0.003)
BSS=(32 64 128 256)
LR=${LRS[$((SLURM_ARRAY_TASK_ID / 4))]}
BS=${BSS[$((SLURM_ARRAY_TASK_ID % 4))]}
srun --container-image=/pfs/images/pytorch-25.09.sqsh \
--container-mounts=/pfs:/pfs \
--container-workdir=/pfs/project \
python train.py --lr "$LR" --batch-size "$BS" \
--out /pfs/runs/lr${LR}_bs${BS}
Using a squash file here rather than a registry reference means 16 array tasks do not each pull the image — see Cache and manage images.
Non-GPU pre-processing step
srun --nodes=1 --ntasks=8 --time=01:00:00 \
--container-image=python:3.12-slim \
--container-mounts=/pfs:/pfs \
python /pfs/scripts/shard_dataset.py --workers 8
python:3.12-slim, alpine and similar images may have no bash. Use sh -c instead of
bash -lc, or pick a fuller base image.
Interactive Sessions
For debugging, allocate a node and get a shell inside the container:
srun --partition=all --nodes=1 --gres=gpu:1 --time=01:00:00 --pty \
--container-image=nvcr.io/nvidia/pytorch:25.09-py3 \
--container-mounts=/pfs:/pfs \
--container-workdir=/pfs \
bash
You are now inside the container on a compute node with a GPU. nvidia-smi, python, and your /pfs
data are all there. Exit the shell to release the allocation.
For a longer session that survives disconnects, allocate first and attach steps to it:
salloc --partition=all --nodes=1 --gres=gpu:1 --time=04:00:00
srun --container-image=nvcr.io/nvidia/pytorch:25.09-py3 --container-mounts=/pfs:/pfs --pty bash
An idle salloc keeps its GPUs out of the queue for other jobs. Always set --time and exit when
you are done. Check what you are holding with squeue -u $USER.