Pular para o conteúdo principal
Última atualização em

Use custom Docker images with AI Runtime

info

Este recurso está em Beta. Para usá-lo, um administrador do Workspace deve ativar as pré-visualizações de AI Runtime Beta recursos e Databricks Artifact Registry na página Previews do Workspace.

O AI Runtime pode executar uma imagem de container do Docker personalizada armazenada no Artifact Registry. Use uma imagem personalizada quando precisar de:

  • Bibliotecas de sistema ou dependências complexas que você não pode instalar com environment.dependencies.
  • Um ambiente reproduzível entre os ambientes de desenvolvimento, pesquisa e produção.
  • Imagens privadas aprovadas pela organização, criadas pela sua equipe de plataforma ou segurança.

Pré-requisitos​

Push an image to Artifact Registry​

Antes de usar uma imagem personalizada com o AI Runtime, armazene-a no Artifact Registry no mesmo workspace que você usa para enviar a carga de trabalho.

  1. Crie ou selecione um catálogo e um esquema do Unity Catalog para a imagem.

  2. Follow Get started with Artifact Registry to configure Docker authentication, grant the required privileges, and push the image.

  3. Anote o nome do Unity Catalog da imagem no seguinte formato:

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

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

dica

Alternativamente, use o comando auxiliardatabricks air images push na CLI do Databricks.

Use a Docker Image em uma carga de trabalho​

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 Execução uploaded application code.

Ao trazer sua própria Docker Image, environment.dependencies e environment.version não são compatíveis. A especificação de environment.unity_catalog_image em qualquer um dos campos aciona um erro. Se você tiver dependências adicionais, instale os pacotes no Dockerfile.

Envie a carga de trabalho:

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

O perfil deve ser autenticado no mesmo workspace onde a imagem está armazenada.

Variáveis de ambiente injetadas em seu contêiner​

O AI Runtime injeta as seguintes variáveis de ambiente em cada contêiner no runtime:

  • CODE_SOURCE_PATH: caminho para o código do aplicativo de upload, quando code_source estiver configurado.
  • NUM_NODES: número total de nós.
  • LOCAL_WORLD_SIZE: GPUs por nó.
  • WORLD_SIZE: número total de processos.
  • POD_RANK: classificação do nó atual (indexada a partir de 0). Também injetado como NODE_RANK.
  • LOCAL_ADDR: IP do nó local (somente multinó).
  • MASTER_ADDR: endereço de coordenação rank-0 (somente multinó).
  • MASTER_PORT: rank-0 coordination port (multi-node only).

Exemplos​

Os exemplos a seguir mostram como executar código de aplicativo **upload** e treinamento distribuído com uma imagem personalizada.

Executar código de aplicativo upload​

Use code_source para fazer upload do código do aplicativo separadamente da imagem. Você pode editar e reenviar seu código sem reconstruir a imagem. Instale as dependências Python e de sistema do código na imagem.

Coloque train.py em um diretório src local ao lado de workload.yaml. A configuração a seguir faz o upload de src e executa o train.py dela dentro da imagem personalizada:

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 resolve-se em relação a workload.yaml. O AI Runtime define $CODE_SOURCE_PATH para o caminho do diretório upload no contêiner. Consulte code_source para ver as opções de código-fonte.

H100 multinó com RDMA​

Para Jobs H100 multi-nós que precisam de largura de banda de rede completa em instâncias p5 da AWS, baseie sua imagem em uma das imagens base do Databricks com NCCL e EFA pré-configurados:

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}" \
    /app/train.py

Crie sua própria imagem​

Ao criar sua própria imagem, o Databricks recomenda usar a capacidade databricks-ai-runtime com um agente de codificação ou iniciar a partir de uma imagem base Databricks.

Use um agente de codificação​

Instale a capacidade **databricks-ai-runtime** Claude Code para obter orientação passo a passo sobre Dockerfile, incluindo criação do zero, compatibilidade com CUDA/NCCL/EFA, problemas comuns e uma lista de verificação pré-compilação. Esta capacidade exige a CLI do Databricks versão 1.0.0 ou mais recente.

Shell
databricks aitools install --skills databricks-ai-runtime

Imagens base do Databricks​

O Databricks publica imagens base no Docker Hub em databricksruntime/air com CUDA, NCCL e rede específica da cloud (AWS EFA ou Azure InfiniBand) pré-configurados.

Etiqueta

Variante

CUDA

Quando usar

dcs-base-aws-runtime

Runtime

12

Instalando somente Python wheels pré-configuradas

dcs-base-aws-devel

Devel

12

Compilando extensões CUDA (requer nvcc)

dcs-base-aws-runtime-cu13

Runtime

13

Instalando apenas Python wheels pré-compiladas, no CUDA 13

dcs-base-aws-devel-cu13

Devel

13

Compilação de extensões CUDA no CUDA 13 (requer nvcc)

Etiqueta

Variante

CUDA

Quando usar

dcs-base-aws-runtime

Runtime

12

Instalando somente Python wheels pré-configuradas

dcs-base-aws-devel

Devel

12

Compilando extensões CUDA (requer nvcc)

dcs-base-aws-runtime-cu13

Runtime

13

Instalando apenas Python wheels pré-compiladas, no CUDA 13

dcs-base-aws-devel-cu13

Devel

13

Compilação de extensões CUDA no CUDA 13 (requer 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.

Para incluir seu script de treinamento na imagem, coloque train.py próximo ao Dockerfile. O Dockerfile o copia para /app/train.py. Se você fizer upload do código do aplicativo com code_source, omita a instrução COPY e mantenha train.py no seu diretório de origem.

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

Build the image locally:

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

Em seguida, siga Primeiros passos com o Artifact Registry para marcar e enviar a imagem para o Artifact Registry. Use o nome <catalog>.<schema>.<image>:<tag> resultante como environment.unity_catalog_image no YAML da workload.

dica

Alternativamente, use o comando auxiliardatabricks air images push na CLI do Databricks e siga os prompts interativos.

Limitações​

  • As imagens devem ser armazenadas no Artifact Registry no workspace onde você envia a workload.
  • O tamanho da imagem deve ser inferior a 20 GB.
  • WORKDIR não é honrado em tempo de execução. Usar caminhos absolutos para arquivos incorporados na imagem. Por exemplo, use python /app/train.py, não python train.py.
  • Não é possível usar environment.dependencies ou environment.version com environment.unity_catalog_image. Se você precisar de pacotes extras além do que está na imagem, adicione-os ao Dockerfile.

Solução de problemas​

Para erros de autenticação, permissão, envio de imagem ou descoberta de imagem relacionados ao registro, consulte Solucionar problemas do Artifact Registry.

ssl.SSLError ao carregar dependências​

Uma imagem personalizada pode falhar em Runtime com um erro OpenSSL quando uma biblioteca tenta criar um contexto SSL, por exemplo:

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

O erro aparece ao importar bibliotecas que abrem conexões de rede, como huggingface_hub, e impede que sejam carregadas.

Isso ocorre porque as cargas de trabalho do AI Runtime são executadas em hosts com os Padrões de Processamento de Informações Federais (FIPS) ativados. Quando as bibliotecas criptográficas da imagem não estão em conformidade com o FIPS, o OpenSSL falha ao inicializar no modo FIPS, portanto, a criação de um contexto SSL falha.

Solução recomendada:

Cargas de trabalho empresariais, governamentais, de saúde e financeiras frequentemente dependem da compliance com FIPS 140-2 ou 140-3 para auditorias FedRAMP, CMMC ou HIPAA. Se sua carga de trabalho deve permanecer em conformidade com FIPS, crie sua imagem com bibliotecas criptográficas em conformidade com FIPS.

Se sua carga de trabalho não exige compliance com FIPS, você pode desabilitar o modo FIPS definindo a variável de ambiente OPENSSL_FORCE_FIPS_MODE como 0. Fazer isso pode silenciosamente violar os requisitos de compliance.

Para desabilitar o modo FIPS, defina-o em seu YAML de carga de trabalho em env_variables:

YAML
env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

Como alternativa, defina a variável em seu Dockerfile para que se aplique a cada carga de trabalho que use a imagem:

Dockerfile
ENV OPENSSL_FORCE_FIPS_MODE=0

Reenvie a carga de trabalho e confirme que o erro SSL não aparece mais quando as dependências são carregadas.

Outros recursos​

Nesta página