Use imagens personalizadas do Docker com o AI Runtime

O AI Runtime pode rodar uma imagem de contêiner Docker personalizada armazenada no Registro de Artefatos. Use uma imagem personalizada quando precisar:

• 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 pela sua plataforma ou equipe de segurança.

Pré-requisitos

• Instale ou atualize para a CLI do Databricks versão 1.19.0 ou superior. Consulte Instalar ou atualizar a CLI do Databricks.
• Peça a um administrador do espaço de trabalho para ativar a versão prévia de AI Runtime Beta Features. Para obter instruções, consulte Gerenciar prévias no nível do espaço de trabalho.
• Configure o Registro de Artefatos no mesmo espaço de trabalho e obtenha permissão para enviar e ler imagens. Veja Comece com o Registro de Artefatos.
• Instale e inicie o Docker na sua máquina local.

Enviar uma imagem para o Registro de Artefatos

Antes de usar uma imagem personalizada com o AI Runtime, armazene-a no Registro de Artefatos no mesmo espaço de trabalho que você usa para enviar a carga de trabalho.

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

2. Siga as instruções em Primeiros passos com o Artifact Registry para configurar a autenticação do Docker, conceder os privilégios necessários e fazer push da imagem.

3. Note o nome do Catálogo Unity da imagem no seguinte formato:

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

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

Use uma imagem Docker em uma carga de trabalho

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

Este exemplo executa train.py direto de /app na imagem. Para enviar o código do aplicativo separadamente sem reconstruir a imagem, veja Executar código de aplicação enviado.

Ao trazer sua própria imagem do Docker, environment.dependencies e environment.version não são compatíveis. Especificar environment.unity_catalog_image com qualquer campo dispara um erro. Se você tiver dependências adicionais, instale os pacotes no Dockerfile.

Envie a carga de trabalho:

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

O perfil deve autenticar no mesmo espaço de trabalho 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 em runtime:

• CODE_SOURCE_PATH: caminho para o código da aplicação enviado, 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 atual do nó (0-indexada). Também é injetado como NODE_RANK.
• LOCAL_ADDR: IP local do nó (somente em multinó).
• MASTER_ADDR: endereço de coordenação de rank-0 (somente para vários nós).
• MASTER_PORT: porta de coordenação de rank-0 (somente em configurações com vários nós).

Exemplos

Os exemplos a seguir mostram como executar código de aplicação carregado e treinamento distribuído com uma imagem personalizada.

Executar código de aplicação carregado

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

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

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 é resolvido em relação a workload.yaml. O runtime de IA define $CODE_SOURCE_PATH como o caminho do diretório enviado no contêiner. Veja code_source para opções de código-fonte.

H100 de vários nós com RDMA

Para execuções H100 de vários nós que precisam de largura de banda total de rede em instâncias p5 da AWS, baseie sua imagem em uma das imagens base do Databricks com NCCL e EFA pré-configurados:

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

Criar sua própria imagem

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

Usar um agente de codificação

Instale a habilidade databricks-ai-runtime do 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. Essa habilidade requer a CLI do Databricks versão 1.0.0 ou mais recente.

databricks aitools install --skills databricks-ai-runtime --experimental

Imagens base do Databricks

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

O seguinte Dockerfile adiciona o PyTorch a uma imagem base do Databricks. As imagens base fornecem Python em /opt/venv, gerenciado por uv. uv pip install é direcionado a esse ambiente por padrão. 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 copia isso para /app/train.py. Se você enviar o código do aplicativo com code_source, ometa a COPY instrução e mantenha train.py no seu diretório fonte.

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

Construa a imagem localmente:

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

Depois, siga o Comece com o Registro de Artefatos para marcar e enviar a imagem para o Registro de Artefatos. Use o nome resultante <catalog>.<schema>.<image>:<tag> como environment.unity_catalog_image no YAML da carga de trabalho.

Limitations

• As imagens devem ser armazenadas no Registro de Artefatos, no espaço de trabalho onde você envia a carga de trabalho.
• O tamanho da imagem deve ter menos de 20 GB.
• WORKDIR não é respeitado em runtime. Use caminhos absolutos para arquivos inseridos na imagem. Por exemplo, use python /app/train.py, e 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, deverá adicioná-los ao Dockerfile.

Resoluçã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 tempo de execução com um erro do OpenSSL quando uma biblioteca tenta criar um contexto SSL, por exemplo:

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

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

Isso ocorre porque cargas de trabalho em tempo de execução de IA rodam em hosts com os Padrões Federais de Processamento de Informações (FIPS) ativados. Quando as bibliotecas criptográficas da imagem não são compatíveis com FIPS, o OpenSSL falha ao inicializar no modo FIPS, portanto, a criação de um contexto SSL falha.

Solução recomendada:

As cargas de trabalho corporativas, governamentais, de saúde e finanças geralmente dependem da conformidade com a FIPS 140-2 ou 140-3 para auditorias FedRAMP, CMMC ou HIPAA. Se sua carga de trabalho precisar permanecer em conformidade com FIPS, crie sua imagem com bibliotecas criptográficas compatíveis com FIPS.

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

Para desativar o modo FIPS, defina isso no YAML da sua carga de trabalho, em env_variables:

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

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

ENV OPENSSL_FORCE_FIPS_MODE=0

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

Recursos adicionais

• Início rápido da CLI do Databricks para Runtime de IA
• Referência de YAML para carga de trabalho
• air grupo de comandos