[BUG] - GPU images are incorrectly built from plain Ubuntu base instead of CUDA base

[marcelovilla]

Describe the bug

Our "Build Docker Images" GHA workflow includes logic to select different base images depending on the target architecture. CPU builds use a standard Ubuntu base image, while GPU builds use a CUDA-enabled base image.

See these relevant lines in the workflow:

nebari-docker-images/.github/workflows/build-push-docker.yaml

Lines 26 to 27 in 0e120d5

	GPU_BASE_IMAGE: nvidia/cuda:12.8.1-base-ubuntu24.04
	GPU_IMAGE_SUFFIX: gpu

nebari-docker-images/.github/workflows/build-push-docker.yaml

Lines 84 to 89 in 0e120d5

	- name: "Set BASE_IMAGE and Image Suffix 📷"
	if: ${{ matrix.platform == 'gpu' }}
	run: |
	echo "GPU Platform Matrix"
	echo "BASE_IMAGE=$GPU_BASE_IMAGE" >> $GITHUB_ENV
	echo "IMAGE_SUFFIX=-$GPU_IMAGE_SUFFIX" >> $GITHUB_ENV

nebari-docker-images/.github/workflows/build-push-docker.yaml

Line 122 in 0e120d5

build-args: BASE_IMAGE=${{ env.BASE_IMAGE }}

---

It used to be the case that our Dockerfiles accepted an ARG for the base image. For example:

nebari-docker-images/Dockerfile.jupyterlab

Line 9 in e1418d0

ARG BASE_IMAGE=ubuntu:20.04

However, #211 appears to have changed that behavior unintentionally. I don’t think we noticed the issue until now, while reviewing #229. Looking at some of our recent image builds, I can confirm the CPU and GPU images are indeed using the same base image.

Taken from the jupyterlab-cpu build for the 2025.10.1rc1 tag GHA logs:

 #7 [linux/amd64 builder 1/4] FROM docker.io/library/ubuntu:24.04@sha256:66460d557b25769b102175144d538d88219c077c678a49af4afca6fbfc1b5252

Taken from the jupyterlab-gpu build for the 2025.10.1rc1 tag GHA logs:

#9 [linux/amd64 builder 1/4] FROM docker.io/library/ubuntu:24.04@sha256:66460d557b25769b102175144d538d88219c077c678a49af4afca6fbfc1b5252

Looking at the image digests, they're exactly the same, which is not the expected behavior. I think we should fix this so we make sure we built GPU images on top of the CUDA base image. I know @dcmcand has some thoughts on how to implement this instead of relying on having an ARG on our Dockerfile. We've been building images like these for some months now and haven't noticed any issues when running GPU workloads on Nebari. That makes me wonder why things still work as before if we're not relying on the CUDA images anymore. I guess it might have to do with the fact that we're still relying on the NVIDIA device plugin Daemonset.

Expected behavior

GPU images should build on top of a CUDA base image.

How to Reproduce the problem?

Take a look at the "Build Docker Images" GHA workflow

Command output

Versions and dependencies used.

No response

Anything else?

No response

[Adam-D-Lewis]

This article might be helpful to understand what's going on. Read it a while ago - https://next.redhat.com/2025/08/26/a-developers-guide-to-pytorch-containers-and-nvidia-solving-the-puzzle/

[viniciusdc]

Here's the original info about the daemon set we were installing (https://github.com/nebari-dev/nebari/tree/main/src/_nebari/stages/kubernetes_initialize/template/modules/nvidia-installer) at least for GCP is in here https://github.com/GoogleCloudPlatform/cos-gpu-installer

[viniciusdc]

15:03 $ which nvidia-smi
/usr/local/nvidia/bin/nvidia-smi

[Adam-D-Lewis]

TLDR - Do we need to use the NVIDIA GPU images since things seem to be working? (Yes!)

Yes, at a minimum, using our own ubuntu image is much more fragile than using the NVIDIA GPU base image. Also, some users may even require the NVIDIA devel images (e.g. nvidia/cuda:12.8.1-devel-ubuntu24.04) if they need to compile CUDA code.

Fragility Issues

• The NVIDIA images set env vars like NVIDIA_DRIVER_CAPABILITIES that are read and used by the NVIDIA Container Toolkit that we aren't setting.
• I saw # Keep apt from auto upgrading the cublas and nccl packages. in one of the nvidia devel images which prevents apt from upgrading certain packages when a user runs apt upgrade. That isn't allowed on Nebari, but gives an idea of the kinds of fragility we avoid when using NVIDIA images.
• I don't think we have that cuda-compat stuff (more info below) installed (might be installed also by CUDA ContainerToolkit, but I wouldn't bet on it since the image installs it explicitly)

Full Notes

Previously (Before we stopped using CUDA container by accident)

We were using nvidia/cuda:12.8.1-base-ubuntu24.04
We can see the Dockerfile at https://gitlab.com/nvidia/container-images/cuda/-/blob/29a0c4b8b69a4f274e51ed22340b497a83ffa82d/dist/12.8.1/ubuntu2404/base/Dockerfile

From https://next.redhat.com/2025/08/26/a-developers-guide-to-pytorch-containers-and-nvidia-solving-the-puzzle/ we learn the required pieces

┌──────────────────────────────────────────────────┐
│               Your PyTorch Project               │
├──────────────────────────────────────────────────┤
│                     PyTorch                      │
├──────────────────────────────────────────────────┤
│                  CUDA Toolkit (SDK)              │  <-- Inside the Container
├──────────────────────────────────────────────────┤
│              Container's Base OS                 │
└┬────────────────────────────────────────────────┬┘
 │          NVIDIA Container Toolkit Bridge       │
┌┴────────────────────────────────────────────────┴┐
│         Container Runtime (Podman/Docker)        │
├──────────────────────────────────────────────────┤
│           NVIDIA Driver + Host OS                │  <-- On the Host Machine
├──────────────────────────────────────────────────┤
│                  NVIDIA GPU Hardware             │
└──────────────────────────────────────────────────┘

The ones I'll emphasize are

• NVIDIA GPU Drivers
  • installed by the daemonset
• NVIDIA Container Toolkit Bridge (makes GPU visible from processes inside container)
  • installed on the host system as well
  • activated by the gpu daemon set
• CUDA Toolkit (SDK)
  • Is installed via conda env, may also be in the CUDA docker image

Where does nvidia-smi come from on Nebari GPU instances?

The nvidia-smi gets put there by NVIDIA Container Toolkit. It looks at env vars like NVIDIA_PATH and automounts nvidia-smi, etc. if set (I think). So if you run docker run -it quay.io/nebari/nebari-jupyterlab:pytorch-2025-630231e-20251021 bash then nvidia-smi isn't present, but on the k8s cluster with the gpu driver installer daemonset using the same image, nvidia-smi is present.

Which CUDA version does nvidia-smi display?

Another snippet from the article is that when you run nvidia-smi the version of CUDA displayed is "the maximum version of CUDA that the installed driver supports." NOT the CUDA TOOLKIT version.

What's in the nvidia/cuda:12.8.1-base-ubuntu24.04 image?

• install CUDA runtime (NOT complete CUDA Toolkit). Without CUDA toolkit no CUDA code can be compiled. (nvcc will be present if you can compile CUDA code)
• also installs cuda-compat. This allows applications compiled against a newer version of CUDA than what's on the host, to still run successfully.

When might a user need to compile CUDA code?

From https://next.redhat.com/2025/08/26/a-developers-guide-to-pytorch-containers-and-nvidia-solving-the-puzzle/:

1. What kinds of projects need to be compiled? Any project that has a dependency that performs Just-In-Time (JIT) compilation to create optimized code for your specific GPU will require the full CUDA Toolkit.

• Examples: Projects using libraries like vLLM, DeepSpeed, CuPy, or custom CUDA kernels written in C++.
• Things to look for: Check your project’s requirements.txt. If you see packages like triton, xformers, flash-attn, or paged-attention, it’s a huge red flag that you’ll need a full development environment.

What image would be needed if user wants to compile GPU code?

• a devel version (e.g. 12.8.1-devel-ubuntu24.04)

High Level (What we Use)

Google GPU Driver Installer Daemonset

Instructions/docs - https://cloud.google.com/kubernetes-engine/docs/how-to/gpus#terraform_1 (We use ContainerOS on GKE by default)
K8s manifest - https://github.com/GoogleCloudPlatform/container-engine-accelerators/blob/master/nvidia-driver-installer/ubuntu/daemonset-preloaded.yaml