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

Usar imagens personalizadas do Docker com AI Runtime

info

Esse recurso está em Beta. Para usá-lo, um administrador do workspace deve ativar as prévias de Recursos Beta do AI Runtime e Registro de Artefatos do Databricks na página de Pré-visualizações do workspace.

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

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

Pré-requisitos​

Push de uma imagem para o 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. Siga Começar com o Registro de Artefatos para configurar a autenticação do Docker, conceder os privilégios necessários e enviar a imagem.

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

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

    Por exemplo, main.ml.training:v1. Não inclua o hostname do registro na configuração da carga de trabalho.

dica

Como alternativa, use o databricks air images push comando auxiliar na CLI do Databricks.

Use uma Docker Image em uma carga de trabalho​

Especifique o nome do Unity Catalog da imagem no seu YAML de carga de trabalho em 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

Este exemplo execução train.py de /app na imagem. Para fazer upload do código do aplicativo separadamente sem reconstruir a imagem, consulte execução código do aplicativo upload.

Ao trazer sua própria Docker Image, environment.dependencies e environment.version não são suportados. Especificar environment.unity_catalog_image com qualquer campo Trigger 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 se autenticar no mesmo workspace onde a imagem está armazenada.

Variáveis de ambiente injetadas no seu container​

O AI Runtime injeta as seguintes variáveis de ambiente em cada container no tempo de execução:

  • CODE_SOURCE_PATH: caminho para o código da aplicação upload, quando code_source está 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 (baseada em 0). Também injetado como NODE_RANK.
  • LOCAL_ADDR: IP do nó local (apenas multinó).
  • MASTER_ADDR: endereço de coordenação rank-0 (somente multinó).
  • MASTER_PORT: porta de coordenação de rank-0 (apenas multinó).

Exemplos​

Os exemplos a seguir mostram como fazer a execução do código do aplicativo upload e o treinamento distribuído com uma imagem personalizada.

execução código do 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 do 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 upload src e execução seu train.py 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 em relação a workload.yaml. O AI Runtime define $CODE_SOURCE_PATH como o caminho do diretório upload no contêiner. Consulte code_source para opções de código-fonte.

H100 multinó com RDMA​

Para jobs H100 de vários nós que precisam de largura de banda de rede completa em instâncias AWS p5, 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 o uso da habilidade databricks-ai-runtime com um agente de codificação ou começar a partir de uma imagem base do Databricks.

Usar um agente de programação​

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

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

Imagens base do Databricks​

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

Etiqueta

Variante

CUDA

Use ao

dcs-base-aws-runtime

Runtime

12

Instalando apenas Python wheels pré-configuradas

dcs-base-aws-devel

Dev

12

Compilando extensões CUDA (exige nvcc)

dcs-base-aws-runtime-cu13

Runtime

13

Instalando apenas rodas pré-construídas no CUDA 13

dcs-base-aws-devel-cu13

Dev

13

Compilando extensões CUDA no CUDA 13 (requer nvcc)

Etiqueta

Variante

CUDA

Use ao

dcs-base-aws-runtime

Runtime

12

Instalando apenas Python wheels pré-configuradas

dcs-base-aws-devel

Dev

12

Compilando extensões CUDA (exige nvcc)

dcs-base-aws-runtime-cu13

Runtime

13

Instalando apenas rodas pré-construídas no CUDA 13

dcs-base-aws-devel-cu13

Dev

13

Compilando extensões CUDA no CUDA 13 (requer nvcc)

O seguinte Dockerfile adiciona PyTorch a uma imagem de base do Databricks. As imagens de base fornecem Python em /opt/venv, gerenciado por uv. uv pip install tem esse ambiente como destino por default. Para usar um ambiente diferente, crie e ative um venv antes de executar uv pip install.

Para incluir seu script de treinamento na imagem, coloque train.py ao lado do 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

Criar a imagem localmente:

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

Em seguida, siga Começar com o Artifact Registry para adicionar tags e enviar a imagem para o Artifact Registry. Use o nome do <catalog>.<schema>.<image>:<tag> resultante como environment.unity_catalog_image no YAML da carga de trabalho.

dica

Alternativamente, use o comando auxiliar databricks air images push na CLI do Databricks e siga as instruções interativas.

Limitações​

  • As imagens devem ser armazenadas no Artifact Registry no workspace onde você envia a carga de trabalho.
  • O tamanho da imagem deve ser menor que 20 GB.
  • WORKDIR não é respeitado durante o Runtime. Use caminhos absolutos para arquivos incluídos na imagem. Por exemplo, use python /app/train.py, não python train.py.
  • Você não pode usar environment.dependencies ou environment.version com environment.unity_catalog_image. Se você precisar de pacotes extras além do que está na imagem, você deve adicioná-los ao Dockerfile.

Solução de problemas​

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

ssl.SSLError ao carregar dependências​

Uma imagem personalizada pode falhar no Runtime com um erro de 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 elas 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) habilitados. Quando as bibliotecas criptográficas da imagem não estão em conformidade com FIPS, o OpenSSL falha ao inicializar no modo FIPS, então a criação de um contexto SSL falha.

Solução recomendada:

Cargas de trabalho corporativas, governamentais, de saúde e financeiras geralmente dependem da compliance com FIPS 140-2 ou 140-3 para auditorias FedRAMP, CMMC ou HIPAA. Se a sua carga de trabalho deve permanecer compatível com FIPS, crie a sua imagem com bibliotecas criptográficas compatíveis com FIPS.

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

Para desabilitar o modo FIPS, defina-o no YAML da sua 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 ela se aplique a todas as cargas de trabalho que usam a imagem:

Dockerfile
ENV OPENSSL_FORCE_FIPS_MODE=0

Envie novamente a carga de trabalho e confirme se o erro de SSL não aparece mais ao carregar as dependências.

Outros recursos​

Nesta página