Skip to main content
Last updated on

Use custom Docker images with AI Runtime

Preview

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:

    Shell
    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:

    Shell
    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:

    Text
    <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:

YAML
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:

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

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

note

If AI Runtime runs your workload outside your workspace's home region, your custom image moves with the workload to the region where it runs. See AI Runtime availability by Geo.

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.
  • MASTER_ADDR: rank-0 coordination address (multi-node only).
  • MASTER_PORT: rank-0 coordination port (multi-node only).

Examples​

The following examples show how to run uploaded application code and distributed training 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:

YAML
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.

Multi-node H100 with RDMA​

For multi-node H100 jobs that need full network bandwidth on AWS p5 instances, base your image on one of the Databricks base images with NCCL and EFA preconfigured:

YAML
experiment_name: my-dcs-distributed
environment:
  unity_catalog_image: main.ml.training:v1
compute:
  num_accelerators: 16 # 2 nodes × 8 H100
  accelerator_type: GPU_8xH100
command: |-
  torchrun \
    --nnodes="${NUM_NODES}" \
    --nproc_per_node="${LOCAL_WORLD_SIZE}" \
    --node_rank="${POD_RANK}" \
    --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \
    --local_addr="$(hostname -I | awk '{print $1}')" \
    /app/train.py

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.

For multi-node GPU_8xH100 workloads, an image built from another base image must also include EFA.

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.

Shell
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-aws-runtime

Runtime

12

Installing pre-built wheels only

dcs-base-aws-devel

Devel

12

Compiling CUDA extensions (requires nvcc)

dcs-base-aws-runtime-cu13

Runtime

13

Installing pre-built wheels only, on CUDA 13

dcs-base-aws-devel-cu13

Devel

13

Compiling CUDA extensions on CUDA 13 (requires nvcc)

Tag

Variant

CUDA

Use when

dcs-base-aws-runtime

Runtime

12

Installing pre-built wheels only

dcs-base-aws-devel

Devel

12

Compiling CUDA extensions (requires nvcc)

dcs-base-aws-runtime-cu13

Runtime

13

Installing pre-built wheels only, on CUDA 13

dcs-base-aws-devel-cu13

Devel

13

Compiling CUDA extensions on CUDA 13 (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.

Dockerfile
FROM databricksruntime/air:dcs-base-aws-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:

Shell
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:

Text
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:

Shell
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:

Text
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:

Shell
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:

Text
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:

YAML
env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

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

Dockerfile
ENV OPENSSL_FORCE_FIPS_MODE=0

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

Additional resources​

On this page