Run containers in your jobs with Pyxis

Pyxis is a SLURM plugin. It adds a --container-image option to srun and sbatch. Your job then runs inside a container, not directly on the host.

A container is a standardized package. It holds a program together with all the software that the program needs: libraries, tools, and other dependencies. The container runs the same way on any host, because it does not use the software that is installed on the host.

An image is the saved template for a container. You start a container from an image, and the same image always gives you the same software.

Docker containers are the most common kind. The images on public registries, like the ones in the examples below, 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 the same for all users, and it does not change. To add the software that your job needs, run the job inside a container.

There are two ways to provide the container:

• Use a pre-built image. Pre-built images are commonly distributed by software providers like PyTorch, TensorFlow or NVIDIA. Also machine-learning researchers commonly share pre-built images to reproduce their research environments.
• Build your own image. When no pre-built image has what you need, make a customized image. 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:

#!/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())

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

#!/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

1. 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¶

It is possible to modify an image, for example to install extra software that your job needs. You then point Pyxis to the modified image to start the customized container. We cover how to create a custom image in: Build your own container image with Enroot.