Custom images for AI Workbench

AI & Data Science NVIDIA AI Workbench
1

Hi,

I’m confused about the features available in AI Workbench, specifically regarding base images. The available base images do not seem to offer the most recent CUDA major version (13) or Pytorch (2.9 or later). These documentation pages:

provide instructions on generating AI Workbench compatible images, which appear to mainly be adding the required metadata labels. I did this for the most recent Nvidia released pytorch (25.12-py3) and uploaded this container to a private container registry on a self hosted Gitlab. AI workbench has the Gitlab integration configured with a personal access token, and is able to push projects to repositories on this Gitlab. However, I’m getting an error “registry ‘%s’ is not supported yet”, (formatted to the registry hostname of the container I provided), when I try to create a project with the custom image. The container is available to read publicly without authentication. The above documentation pages reference using custom registry URLs, and the following posts imply that this should be supported since it’s a Gitlab based registry and is a public container.

The post says that supported registries are Nvidia NGC, Github, and Gitlab. I wasn’t sure if by Gitlab it was actually gitlab.com, so I looked into the AI Workbench service binary, wb-svc. From what I can tell, it’s comparing the container registry hostname with nvcr.io, ghcr.io, and docker.io when it tries to create the project, but I’m not entirely sure what the logic is. I also gained no further information from the AI workbench logs at /home/<user>/.nvwb/logs/workbench.log, which displayed the same error message and no additional details.

  1. Is there any way to get AI Workbench to pull images from local container registries? The container image is also pulled (and built) locally on the device so it shouldn’t technically need to pull anything at all.
  2. Otherwise, is Nvidia planning on pushing more updated versions of CUDA and Pytorch containers that have the required AI Workbench metadata labels?

Thanks

2

For the short term in case this helps anyone else:

If you select a base container and create the project, you can then publish to a git server and update the .project/spec.yaml file to a custom registry and image. When you pull the changes in Workbench, it will rebuild the container. I was able to do this to get CUDA 13 and Pytorch 2.10.

I’m still curious in the long term support for newer images release by Nvidia.

3

It helps, but I get the strong feeling that it’s a feature that will be dropped very quickly, and we will all likely be better off adapting to using docker natively rather than this wrapped version.

I hope I’m wrong, it’s super useful having the UI

4

howdy - sorry for taking so long to reply.

Current issue

Publishing workbench containers to ngc is constrained by a fairly “robust” release process that winds through various kinds of “scrutiny”.

As such it takes us a while to publish them. We will have new images out for the impending release in the next few days.

Regardless, we are changing up how workbench handles base images over the next few releases to eliminate this bottleneck and make it easier for users to use their own containers.

How it “works”

Workbench needs some basic info to know how to handle the container. That info can be in image labels, or it can be in the spec.yaml file. In fact, once the info is in the spec.yaml file Workbench doesn’t pay attention to the image labels.

In addition, we currently only support ngc, github and gitlab (.com and self-hosted) as private registries which means pulling from docker hub requires the container to be public….

This has to do with how we handle auth.

In the meantime, there are two klunky approaches:

  1. Pull the container you want to use; add a relatively small set of “image labels” via a container file; push to a public registry; then use the “BYOC” feature in the create project flow.
    1. See here for prepping image: Use a Custom Container Image — NVIDIA AI Workbench User Guide
    2. See here for BYOC: Use a Custom Container Image — NVIDIA AI Workbench User Guide
  2. OR you can edit the spec.yaml file as mentioned above.

Where we are heading in the near term

Removing need for existing labels on base image

There are two approaches we are thinking through:

  • Let users or an agent specify the relevant info up front during the clone/creation process and add it to the spec.yaml file
  • AND/OR have workbench check the container can be run, and then pull/run it and have workbench deterministically fish out the info and add it to the spec.yaml file.

Moving to a “post pull” modification approach

The bureaucratic bottleneck in our image publishing is due to us adding some extra software bits. If we move towards a “config pack” approach that does that post pull, i.e. as part of the local build process, then we can just provide the most recent nvidia containers as they come out… kind of.

The issue would be that we will need to potentially update the “config pack” for that image or image class as they arise, which means a workbench release.

I’m not convinced that’s the case, but we are looking into it.

Supporting more private registries, starting with Docker Hub

Basically we just need to add an “integration” that supports auth with docker hub, etc. See here: How to Connect Integrations — NVIDIA AI Workbench User Guide

This is perpetually backlogged but I will escalate it.

Mentioned Feature Requests

Local container registry

I’ve been advocating for ability to add a local container registry as an image source for a while, but there are a few “reproducibility” issues with that. It’s possible we are too hung up on reproducibility…

Pointing at an existing local image

Also been advocating for this. Again, reproducibility issues.

It looks like we’ve glued ourselves too closely to reproducibility and we may want to take a more of a batteries not included approach to how users handle images themselves.