Use custom Docker images with AI Runtime

Important

This feature is in Public Preview. To use it, a workspace admin must enable the Databricks Artifact Registry preview from the workspace Previews page.

AI Runtime can run your workload in a custom Docker container image instead of a managed environment. You build the image for linux/amd64, push it to Artifact Registry in your workspace, and reference its Unity Catalog name in the workload YAML that you submit with the Databricks CLI. Use a custom image when you need:

  • System libraries or complex dependencies that you cannot install with environment.dependencies.
  • CUDA extensions that you compile from source, such as flash-attn or custom kernels.
  • A reproducible environment across development, research, and production.
  • Private, organization-approved images built by your platform or security team.

Prerequisites

Push an image to Artifact Registry

Before using a custom image with AI Runtime, store it in Artifact Registry in the same workspace that you use to submit the workload. To build an image first, see Build your own image.

  1. Create or select a Unity Catalog catalog and schema for the image, and obtain the required Artifact Registry privileges.

  2. Create or refresh a workspace OAuth profile for the workspace:

    databricks auth login --profile my-databricks-profile

    The Docker credential helper requires a profile created by databricks auth login. It doesn't support profiles that use personal access tokens or OAuth machine-to-machine (M2M) credentials.

  3. To authenticate Docker, tag the image, and push it in one step, run the databricks air images push helper:

    databricks air images push \
  --profile my-databricks-profile \
  --source my-training-image:v1 \
  --catalog main \
  --schema ml \
  --artifact training:v1

    For more options, including copying an image from another registry, see Push images with the Databricks CLI.

    To push with docker push instead, run databricks auth docker configure --profile my-databricks-profile to configure the Docker credential helper, then follow Push an image in the Artifact Registry guide.

  4. Note the image's Unity Catalog name in the following format:

    <catalog>.<schema>.<image>:<tag>

    For example, main.ml.training:v1. Do not include the registry hostname in the workload configuration.

Use a Docker image in a workload

Specify the image's Unity Catalog name in your workload YAML under environment.unity_catalog_image:

experiment_name: my-dcs-training
environment:
  unity_catalog_image: main.ml.training:v1
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: python /app/train.py

This example runs train.py from /app in the image. To upload application code separately without rebuilding the image, see Run uploaded application code.

When bringing your own Docker image, environment.dependencies and environment.version are not supported. Specifying environment.unity_catalog_image with either field triggers an error. If you have additional dependencies, install the packages in the Dockerfile instead.

Submit the workload:

databricks air run -f workload.yaml -p my-databricks-profile

The profile must authenticate to the same workspace where the image is stored.

Environment variables injected into your container

AI Runtime injects the following environment variables into every container at runtime:

  • CODE_SOURCE_PATH: path to uploaded application code, when code_source is configured.
  • NUM_NODES: total number of nodes.
  • LOCAL_WORLD_SIZE: GPUs per node.
  • WORLD_SIZE: total number of processes.
  • POD_RANK: current node rank (0-indexed). Also injected as NODE_RANK.

Examples

The following example shows how to run uploaded application code with a custom image.

Run uploaded application code

Use code_source to upload application code separately from the image. You can edit and resubmit your code without rebuilding the image. Install the code's Python and system dependencies in the image.

Place train.py in a local src directory next to workload.yaml. The following configuration uploads src and runs its train.py inside the custom image:

experiment_name: my-dcs-uploaded-code
environment:
  unity_catalog_image: main.ml.training:v1
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
code_source:
  root_path: ./src
command: |-
  cd "$CODE_SOURCE_PATH"
  python3 train.py

root_path resolves relative to workload.yaml. AI Runtime sets $CODE_SOURCE_PATH to the uploaded directory's path in the container. See code_source for code source options.

Build your own image

When building your own image, Databricks recommends using databricks-ai-runtime skill with a coding agent or starting from a Databricks base image.

You can also start from another base image, such as nvidia/cuda. In that case, the image must include a CUDA runtime that the host GPU driver supports and NCCL for multi-GPU communication. The databricks-ai-runtime skill covers building an image from scratch.

Use a coding agent

Install the databricks-ai-runtime Claude Code skill for step-by-step Dockerfile guidance, including building from scratch, CUDA/NCCL/EFA compatibility, common problems, and a pre-build checklist. This skill requires Databricks CLI version 1.0.0 or newer.

databricks aitools install --skills databricks-ai-runtime

Databricks base images

Databricks publishes base images on Docker Hub at databricksruntime/air with CUDA, NCCL, and cloud-specific networking (AWS EFA or Azure InfiniBand) preconfigured.

Tag Variant CUDA Use when
dcs-base-azure-runtime Runtime 12 Installing pre-built wheels only
dcs-base-azure-devel Devel 12 Compiling CUDA extensions (requires nvcc)

The following Dockerfile adds PyTorch to a Databricks base image. The base images provide Python at /opt/venv, managed by uv. uv pip install targets that environment by default. To use a different environment, create and activate a venv before running uv pip install.

To include your training script in the image, place train.py next to the Dockerfile. The Dockerfile copies it to /app/train.py. If you upload application code with code_source, omit the COPY instruction and keep train.py in your source directory instead.

FROM databricksruntime/air:dcs-base-azure-runtime

RUN uv pip install --no-cache \
    torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0

RUN uv pip install --no-cache \
    transformers==4.45.0 \
    accelerate==0.34.0 \
    'mlflow>=3.6'

COPY ./train.py /app/train.py

Images must use the linux/amd64 platform. On a Mac with Apple Silicon, Docker builds arm64 images by default, so always pass --platform linux/amd64. To build the image locally, run the following:

docker build --platform linux/amd64 -t my-training-image:v1 .

Then push the image to Artifact Registry and use the resulting <catalog>.<schema>.<image>:<tag> name as environment.unity_catalog_image in the workload YAML.

Limitations

  • Images must be stored in Artifact Registry in the workspace where you submit the workload.
  • Image size must be under 20 GB.
  • WORKDIR is not honored at runtime. Use absolute paths for files baked into the image. For example, use python /app/train.py, not python train.py.
  • You cannot use environment.dependencies or environment.version with environment.unity_catalog_image. If you need extra packages beyond what is in the image, you must add them to the Dockerfile.

Troubleshooting

For registry-related authentication, permission, image push, or image discovery errors, see Troubleshoot Artifact Registry.

Docker credential helper requires a profile created by databricks auth login

Pushing an image fails with an error like the following:

Error: profile "my-profile" uses auth_type "". Docker credential helper requires a profile created by databricks auth login

This error occurs when the profile uses a personal access token or another authentication method that the Docker credential helper doesn't support. Create a workspace OAuth profile, then rerun the push with that profile:

databricks auth login --profile my-databricks-profile

Multiple profiles match the workspace ID

Configuring Docker authentication or pushing an image fails with an error like the following:

Error: multiple Databricks profiles match workspace ID 1111111111111111: profile-a and profile-b. Remove duplicate workspace_id entries before using Docker credential helper

This error occurs when more than one profile in your configuration file (~/.databrickscfg by default) points to the same workspace. Remove the profile that you don't use for Docker, then rerun the command with the remaining profile. To log out of a profile and remove it from the configuration file, run the following:

databricks auth logout --profile profile-a --delete

For all logout options, see auth command group.

ssl.SSLError when loading dependencies

A custom image can fail at runtime with an OpenSSL error when a library tries to create an SSL context, for example:

ssl.SSLError: [CRYPTO] unknown error (_ssl.c:3076)

The error appears while importing libraries that open network connections, such as huggingface_hub, and prevents them from loading.

This occurs because AI Runtime workloads run on hosts with Federal Information Processing Standards (FIPS) enabled. When the image's cryptographic libraries are not FIPS-compliant, OpenSSL fails to initialize in FIPS mode, so creating an SSL context fails.

Recommended solution:

Enterprise, government, healthcare, and finance workloads often depend on FIPS 140-2 or 140-3 compliance for FedRAMP, CMMC, or HIPAA audits. If your workload must remain FIPS-compliant, build your image with FIPS-compliant cryptographic libraries.

If your workload does not require FIPS compliance, you can disable FIPS mode by setting the OPENSSL_FORCE_FIPS_MODE environment variable to 0. Doing so can silently break compliance requirements.

To disable FIPS mode, set it in your workload YAML under env_variables:

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

Alternatively, set the variable in your Dockerfile so it applies to every workload that uses the image:

ENV OPENSSL_FORCE_FIPS_MODE=0

Resubmit the workload and confirm the SSL error no longer appears when dependencies load.

Additional resources

  • Last updated on