Utilize imagens Docker personalizadas com o AI Runtime

O AI Runtime pode executar uma imagem de contentor Docker personalizada armazenada no Registo de Artefactos. Use uma imagem personalizada quando precisar:

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

Pré-requisitos

• Instale ou atualize para a CLI Databricks versão 1.19.0 ou superior. Consulte Instalar ou atualizar a CLI do Databricks.
• Pede a um administrador de espaço de trabalho para ativar a pré-visualização de funcionalidades beta do AI Runtime . Para instruções, consulte Gerir pré-visualizações ao nível do espaço de trabalho.
• Configura o Registo de Artefactos no mesmo espaço de trabalho e obtém permissão para enviar e ler imagens. Veja Introdução ao Artifact Registry.
• Instala e inicia o Docker na tua máquina local.

Enviar uma imagem para o Registo de Artefactos

Antes de usar uma imagem personalizada com AI Runtime, armazene-a no Registo de Artefactos no mesmo espaço de trabalho que usa para submeter 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 Introdução ao Artifact Registry para configurar a autenticação do Docker, conceder os privilégios necessários e enviar a 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 registo na configuração da carga de trabalho.

Usar uma imagem Docker num workload

Especifique o nome da imagem no Unity Catalog no YAML da sua 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 a partir de /app na imagem. Para carregar o código da aplicação separadamente sem reconstruir a imagem, veja Executar código da aplicação carregado.

Ao utilizar uma imagem Docker própria, environment.dependencies e environment.version não são suportados. Especificar environment.unity_catalog_image com qualquer um dos campos desencadeia um erro. Se tiveres dependências adicionais, instala os pacotes no Dockerfile em vez disso.

Enviar a carga de trabalho:

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

O perfil deve autenticar-se no mesmo espaço de trabalho onde a imagem está armazenada.

Variáveis de ambiente injetadas no seu contentor

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

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

Exemplos

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

Executar código de aplicação carregado

Use code_source para carregar o código da aplicação separadamente da imagem. Pode editar e submeter novamente o seu código sem reconstruir a imagem. Instala o código em Python e as dependências do sistema na imagem.

Coloque train.py num diretório local src ao lado de workload.yaml. A configuração seguinte envia src e executa o 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 relativamente a workload.yaml. O AI Runtime define $CODE_SOURCE_PATH o caminho do diretório carregado no contentor. Consulte code_source para opções de origem do código.

H100 de múltiplos nós com RDMA

Para tarefas H100 multinó que requerem largura de banda de rede completa em instâncias AWS p5, baseie a imagem numa das imagens base da 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

Constrói a tua própria imagem

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

Usar um agente de codificação

Instale a extensão Claude Code databricks-ai-runtime para obter orientação passo a passo sobre Dockerfiles, incluindo criação de raiz, compatibilidade com CUDA/NCCL/EFA, problemas comuns e uma lista de verificação pré-compilação. Esta competência requer a versão 1.0.0 da CLI do Databricks ou mais recente.

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

Imagens de base do Databricks

O Databricks publica imagens base no Docker Hub at databricksruntime/air com CUDA, NCCL e redes específicas da cloud (AWS EFA ou Azure InfiniBand) pré-configuradas.

O seguinte Dockerfile adiciona o PyTorch a uma imagem base de Databricks. As imagens base fornecem Python em /opt/venv, gerido por uv. uv pip install visa esse ambiente por predefinição. Para usar um ambiente diferente, cria e ativa um venv antes de executar uv pip install.

Para incluir o seu script de treino na imagem, coloque train.py ao lado do Dockerfile. O Dockerfile copia-o para /app/train.py. Se carregar o código da aplicação com code_source, omita a instrução COPY e mantenha train.py no teu diretório de origem em vez disso.

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

Constrói a imagem localmente:

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

Depois, siga Introdução ao Artifact Registry para atribuir uma etiqueta à imagem e enviá-la para o Artifact Registry. Use o nome resultante <catalog>.<schema>.<image>:<tag> como environment.unity_catalog_image na carga de trabalho YAML.

Limitations

• As imagens devem ser armazenadas no Registo de Artefactos no espaço de trabalho onde submete a carga de trabalho.
• O tamanho da imagem deve ser inferior a 20 GB.
• WORKDIR não é respeitado durante a execução. Usa caminhos absolutos para ficheiros incorporados na imagem. Por exemplo, use python /app/train.py, não python train.py.
• Não pode usar environment.dependencies nem environment.version com environment.unity_catalog_image. Se precisares de pacotes extra para além do que está na imagem, tens de os adicionar ao Dockerfile.

Troubleshooting

Para erros de autenticação, permissões, envio de imagens ou deteção de imagens relacionados com o registo, consulte Troubleshoot Artifact Registry.

ssl. SSLError ao carregar dependências

Uma imagem personalizada pode falhar em tempo de execução com um erro 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 ligações de rede, como huggingface_hub, e impede o seu carregamento.

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

Solução recomendada:

As cargas de trabalho em empresas, governo, saúde e finanças dependem frequentemente da conformidade com o FIPS 140-2 ou 140-3 para auditorias FedRAMP, CMMC ou HIPAA. Se a sua carga de trabalho tiver de se manter compatível com FIPS, construa a sua imagem com bibliotecas criptográficas compatíveis com FIPS.

Se a sua carga de trabalho não exigir conformidade com FIPS, pode desativar o modo FIPS definindo a OPENSSL_FORCE_FIPS_MODE variável de ambiente para 0. Fazê-lo pode violar silenciosamente os requisitos de conformidade.

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

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

Alternativamente, defina a variável no seu Dockerfile para que se aplique a todas as cargas de trabalho que usem a imagem:

ENV OPENSSL_FORCE_FIPS_MODE=0

Reenvie a carga de trabalho e confirme que o erro SSL já não aparece quando as dependências carregam.

Recursos adicionais

• Guia de iniciação rápida da Databricks CLI para o AI Runtime
• Referência YAML de carga de trabalho
• air grupo de comando