Run containers in your jobs with Pyxis

Pyxis is a SLURM plugin that runs your job inside a container instead of directly on the host. Add its --container-image option to srun or sbatch, and you can use the software your job needs — a specific version of PyTorch, for example — without installing anything on the shared system.

A container is a standardized package that holds a program together with everything it needs to run: libraries, tools, and other dependencies. Because a container does not rely on software installed on the host, it runs the same way on any host.

An image is the saved template that a container starts from. The same image always produces the same container, with the same software inside.

Docker images are the most common kind. The images on public registries — including the examples on this page — are Docker images.

Why you must use a container¶

You cannot install or change libraries and packages on the base system: the host software is shared by everyone and stays the same for all users. To add software that your job needs, run the job inside a container instead.

There are two ways to provide the container:

• Use a pre-built image. Software providers such as PyTorch, TensorFlow, and NVIDIA publish ready-to-use images. Researchers also often share pre-built images so that others can reproduce their results.
• Build your own image. When no pre-built image has what you need, build a customized one. See Build your own container image with Enroot.

Run a job with a container image¶

Add --container-image to your srun command or to an #SBATCH directive. Point the option at a pre-built image or at a registry image.

Use a pre-built image from a registry:

srun --gres=gpu:1 --cpus-per-task=8 --mem=8G --container-image=nvcr.io/nvidia/pytorch:24.01-py3 \
    python train.py

In a batch job, add --container-image as an #SBATCH directive, next to your other directives:

#!/bin/bash
#SBATCH --job-name=my-training-run
#SBATCH --partition=batch
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=8
#SBATCH --mem=32G
#SBATCH --time=01:00:00
#SBATCH --container-image=nvcr.io/nvidia/pytorch:24.01-py3
#SBATCH --output=%x-%j.log

python train.py

Submit it with sbatch job.sh.

Image caching¶

Container images are often large. The first download can take several minutes and use a large amount of disk space. To avoid repeat downloads, Pyxis caches each image.

The first time that anyone pulls a registry image, Pyxis stores its layers in /var/cache/enroot. All users share this directory. Later pulls of the same image use the cache. Pyxis does not download the image again.

srun --partition=debug --container-image=nvcr.io/nvidia/pytorch:24.01-py3 true

Container filesystem is read-only¶

You cannot write to a job's container filesystem. Any change inside the container is lost when the job ends. Keep your code, data, and output on the host, and mount the host directory into the container.

Mount your working directory¶

Host directories are not available inside the container by default. Use --container-mounts=SOURCE:TARGET to bind-mount a host directory. To use your current directory inside the container, mount it onto the same path and start the job there:

srun --gres=gpu:1 \
    --container-image=nvcr.io/nvidia/pytorch:24.01-py3 \
    --container-mounts=$PWD:$PWD \
    --container-workdir=$PWD \
    python train.py

• --container-mounts=$PWD:$PWD makes your working directory available at the same path inside the container. Writes to this path go to the host.
• --container-workdir=$PWD starts the job in that directory.

Mount more than one directory with a comma-separated list, for example --container-mounts=$HOME/data:/data,$PWD:$PWD.

Complete example: a GPU smoke test¶

This example runs a small job that trains a tiny model on the GPU inside a container. It checks that the GPU, SLURM, and Pyxis work together. Run it to test your setup, or use it as a template for your own jobs.

1. Create a health_check directory in your home and add the training script:

   /home/<username>/health_check/train_smoke_test.py

   #!/usr/bin/env python3
   """GPU/SLURM/container smoke test: trains a tiny MLP on synthetic data.
   
   No dataset download required — only needs the NGC PyTorch image and a GPU.
   """
   
   import sys
   
   import torch
   from torch import nn
   
   
   def main() -> int:
       if not torch.cuda.is_available():
           print("FAIL: CUDA not available inside container", file=sys.stderr)
           return 1
   
       device = torch.device("cuda")
       print(f"GPU: {torch.cuda.get_device_name(device)}")
   
       torch.manual_seed(0)
       x = torch.randn(512, 20, device=device)
       y = torch.randn(512, 1, device=device)
   
       model = nn.Sequential(
           nn.Linear(20, 64),
           nn.ReLU(),
           nn.Linear(64, 32),
           nn.ReLU(),
           nn.Linear(32, 1),
       ).to(device)
   
       optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
       loss_fn = nn.MSELoss()
   
       losses = []
       for _ in range(50):
           optimizer.zero_grad()
           loss = loss_fn(model(x), y)
           loss.backward()
           optimizer.step()
           losses.append(loss.item())
   
       print(f"Initial loss: {losses[0]:.4f}, final loss: {losses[-1]:.4f}")
       if losses[-1] >= losses[0]:
           print("FAIL: loss did not decrease", file=sys.stderr)
           return 1
   
       print("HEALTH CHECK PASSED")
       return 0
   
   
   if __name__ == "__main__":
       sys.exit(main())
   

2. Add a job script next to it. The --container-mounts path is absolute, so it works in an #SBATCH directive:

   /home/<username>/health_check/health_check.sbatch

   #!/bin/bash
   # Trains a tiny neural network inside an NVIDIA NGC container to smoke-test the
   # GPU + SLURM + Pyxis/Enroot stack end-to-end.
   
   #SBATCH --job-name=health_check
   #SBATCH --partition=debug
   #SBATCH --gres=gpu:1
   #SBATCH --cpus-per-task=8
   #SBATCH --time=00:30:00
   #SBATCH --container-image=nvcr.io/nvidia/pytorch:26.07-py3
   #SBATCH --container-mounts=/home/<username>/health_check/:/home/<username>/health_check/
   #SBATCH --container-workdir=/home/<username>/health_check
   
   set -euo pipefail
   
   start_time=$(date +%s)
   trap 'echo "Job runtime: $(( $(date +%s) - start_time ))s"' EXIT
   
   python3 train_smoke_test.py
   

3. Submit the job:

   sbatch /home/<username>/health_check/health_check.sbatch
   

The job mounts your health_check directory, pulls the container image, and runs the script on the GPU. When it passes, the output file ends with HEALTH CHECK PASSED.

Use a custom image¶

You can modify an image — for example, to add extra software that your job needs — then point Pyxis at the modified image to start the customized container. See Build your own container image with Enroot for how to create one.