Mastering Container Usage on DGX On-Prem

Table of Contents

1. Introduction

2. Understanding the Technology Stack

   • SLURM

   • Enroot

   • Pyxis

   • How It All Works Together

3. Container Deployment Strategies

   • Ephemeral Containers

   • Persistent Containers

4. Finding the Right Container

   • Accessing NGC Containers

   • Setting Up Access to Private Containers

5. Accessing Your Data in Containers

6. Interactive Development and Container Customization

   • Quick Interactive Sessions

   • Customizing Persistent Containers

   • Development Workflows

7. Essential SLURM and Pyxis Options

1. Introduction

Before diving into the technical details, let's understand why this setup is valuable:

• Reproducibility: Your code runs the same way every time, regardless of the compute node

• Environment Management: No more dependency conflicts or "it works on my machine" issues

• Resource Efficiency: Share optimized, pre-built environments across your team

• Flexibility: Switch between different frameworks (PyTorch, TensorFlow, etc.) without system-level changes

2. Understanding the Technology Stack

Running containerized workloads on the DGX On-Prem cluster combines three powerful technologies that work seamlessly together to give you the flexibility of containers with the performance of HPC:

• SLURM: Your job scheduler and resource manager

• Enroot: A container runtime built for HPC environments

• Pyxis: The bridge that connects SLURM and Enroot automatically

This guide will help you understand these tools and use them effectively, whether you're running quick experiments or large-scale training jobs.

SLURM

Think of SLURM (Simple Linux Utility Resource Management) as the traffic controller for your cluster. It:

• Decides which jobs run when and where

• Allocates GPUs, CPUs, and memory fairly among users

• Handles queues when the cluster is busy

• Monitors job progress and resource usage

Enroot

Enroot is like Docker's HPC-savvy cousin. Unlike Docker, it:

• Runs without requiring administrator privileges (no root access needed)

• Focuses on performance rather than security isolation

• Integrates naturally with shared filesystems

• Handles GPU access efficiently

Pyxis

Pyxis is the magic that makes everything work together automatically. When you submit a job, it:

• Downloads your container image if needed

• Sets up the container environment

• Ensures your job has access to GPUs and data

• Cleans up when your job finishes

How It All Works Together

Let's compare three approaches to running the same PyTorch container.

Traditional Docker (Local Machine)

This is how you would normally do it when running containers from your own computer:

Enroot Equivalent (HPC Environments)

These are the Enroot equivalent commands (when compared to Docker) that Pyxis automatically runs behind the scenes:

SLURM + Pyxis (DGX On-Prem)

The SLURM + Pyxis approach is clearly the simplest - you just specify which container you want, and everything else happens automatically!

The following diagram illustrates how these different parts interact to each other from a very high-level perspective.

Open

3. Container Deployment Strategies

Ephemeral Containers

This is a quick and clean approach for on-demand tasks. It's best for experiments, one-off jobs, or when you want the latest version every time. The container is downloaded fresh, your job runs, then everything is cleaned up automatically. Perfect for keeping things tidy!

Persistent Containers

This is a save and reuse approach for regular tasks. It's best for repeated jobs with the same environment, custom setups, or when you want faster job startup. The first time this runs, the container is downloaded and saved with the name pytorch_container. Future jobs with the same --container-name will reuse the existing container, starting much faster.

This approach is a time-saver when you're running multiple related jobs or iterations of the same experiment!

4. Finding the Right Container

You can virtually run any publicly available container image from a container registry, but we strongly advise you to look for container images on the NGC (NVIDIA GPU Cloud) Catalog, which is available here: https://catalog.ngc.nvidia.com . When you find a container you'd like to use, look for the "Get Container" button on the top right corner of the page, and copy the image path for the tag you wish to use.

1. Browse to the container you want on NGC

2. Click "Get Container" (top right)

3. Copy the image path for your desired tag

4. Use it directly in your --container-image parameter

The screenshots below illustrate the process.

Open

Browse to the container you want on NGC

 

Open

Click Get Container and copy the Tag

 

Accessing NGC Containers

You should be able to pick any container you'd like to use. Some of them are public, some of them are private. For private containers, you need to have an API Key setup on DGX Head. The next section will further elaborate on that. You're probably asking yourself now: how do I know a container image is private? If the container page shows the "NVIDIA AI Enterprise Essentials" tag, then that container is private.

Here's the thing: Enroot tries to authenticate you with the registry regardless of whether a container is public or private. So for public containers, there's a clever way to bypass this authentication step. To skip authentication for public containers, you can replace docker://nvcr.io/nvidia/ with docker://nvcr.io#nvidia/. Notice the # instead of / after nvcr.io - this tells Enroot "don't bother authenticating, just grab the public container!" If you need to access private registries or NGC private collections, then look at the next section.

In the meantime, these are some popular NGC containers for AI/ML:

• nvcr.io/nvidia/pytorch:24.11-py3 - PyTorch with CUDA support

• nvcr.io/nvidia/tensorflow:24.11-tf2-py3 - TensorFlow 2.x

• nvcr.io/nvidia/cuda:12.3-devel-ubuntu22.04 - CUDA development environment

• nvcr.io/nvidia/rapids:24.10-cuda12.0-runtime-ubuntu22.04 - RAPIDS data science

• nvcr.io/nvidia/tritonserver:24.11-py3 - Triton Inference Server

Setting Up Access to Private Containers

This is a straightforward step that should only be done once. First, you generate an API Key on NGC. Then, you copy that Key to your home folder on DGX On-Prem Head node.

Getting an NGC API Key

1. Login at NGC

2. Go to Setup → Generate API Key

3. Scroll down to “Legacy Keys” and click on “Generate Legacy Key”

4. Copy the API Key and keep it securely as it will be used in the next step

The screenshots below illustrate the process.

Open

Scroll down to Legacy Keys

Now that you have generated the API Key, we just need to tell Enroot to use them.

Open

Click on Generate Legacy Key

 

Setting Up Enroot

First, connect to the DGX On-Prem cluster.

You'll need to enter your NetID password. If you're off-campus, make sure you're connected to the VPN first. Now, copy the commands right next, replace YOUR_NGC_API_KEY with the Key you created in the previous step, paste it to the terminal and run them.

If you want to double check the .credentials file have your API Key, you can run cat ~/.config/enroot/.credentials and expect to see something like the following.

5. Accessing Your Data in Containers

Containers have their own isolated filesystem, so you need to explicitly grant access to your files. The DGX On-Prem automatically mounts your home folder, but lab directories have to be manually mounted. So if you need access to other directories, just use the --container-mounts option.

6. Interactive Development and Container Customization

Now that you know how to deploy containers and access your data, let's explore how to use them for interactive development. This is where containers really shine - you can experiment, install packages, and develop code without worrying about breaking anything!

Quick Interactive Sessions

Sometimes you just want to jump into a container and try things out. No need for a job script - srun gives you instant interactive access:

Need to work with your lab's data? Just add the mount:

Pro tip: Use --container-name to avoid re-downloading the container every time:

Customizing Persistent Containers

Here's where things get really interesting. By default, containers are read-only - you can't install new packages or modify files. But with --container-writable, you can customize your environment and the changes persist!

Next time you use --container-name='my_custom_pytorch', all your installed packages are still there:

Important notes about writable containers:

• Changes are saved to the named container, not the base image

• Each user's named containers are separate (your 'my_custom_pytorch' is different from someone else's)

• Writable containers use more disk space - clean up old ones you're not using

• System package installations (apt/yum) might need --container-writable even in interactive mode

Development Workflows

Here's a typical development workflow that many users find effective:

1. Start with exploration:

2. Customize for your project:

3. Develop and test interactively:

4. Submit batch jobs:

This workflow keeps development fast and interactive while ensuring your production runs use the exact same environment!

7. Essential SLURM and Pyxis Options

These are the most common tags you will likely be using, but feel free to browse the Pyxis Documentation.