Utilize imagens Docker personalizadas com o AI Runtime

O AI Runtime pode executar a sua carga de trabalho numa imagem de contentor Docker personalizada em vez de um ambiente gerido. Constróis a imagem para linux/amd64, envia-a para o Registo de Artefactos no teu espaço de trabalho e referencias o nome do Catálogo Unity no YAML da carga de trabalho que submetes com a linha de comando Databricks. Use uma imagem personalizada quando precisar:

• Bibliotecas do sistema ou dependências complexas que não pode instalar com environment.dependencies.
• Extensões CUDA compiladas a partir da fonte, como flash-attn ou kernels personalizados.
• 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.

Prerequisites

• 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 que ative a pré-visualização do Databricks Artifact Registry. 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 Artifact Registry

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. Para criar primeiro uma imagem, veja Crie a sua própria imagem.

1. Crie ou selecione um catálogo e esquema do Catálogo Unity para a imagem, e obtenha os privilégios necessários do Registo de Artefactos.

2. Crie ou atualize um perfil OAuth de espaço de trabalho para o espaço de trabalho:

   databricks auth login --profile my-databricks-profile
   

   O assistente de credenciais Docker requer um perfil criado por databricks auth login. Não suporta perfis que utilizem tokens de acesso pessoais ou credenciais OAuth machine-to-machine (M2M).

3. Para autenticar o Docker, marca a imagem e faz push numa só etapa, executa o databricks air images push helper:

   databricks air images push \
     --profile my-databricks-profile \
     --source my-training-image:v1 \
     --catalog main \
     --schema ml \
     --artifact training:v1
   

   Para mais opções, incluindo copiar uma imagem de outro registo, veja Push imagens com a CLI Databricks.

   Para enviar com docker push em vez disso, execute databricks auth docker configure --profile my-databricks-profile para configurar o assistente de credenciais Docker e siga Enviar uma imagem no guia do Registo de Artefactos.

4. 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 numa carga de trabalho

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: rank atual do nó (indexado a partir de 0). Também é injetado como NODE_RANK.

Examples

O exemplo seguinte mostra como executar código de aplicação carregado 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.

Crie a sua 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.

Também pode começar a partir de outra imagem base, como nvidia/cuda. Nesse caso, a imagem deve incluir um runtime CUDA que o driver da GPU do host suporta e NCCL para comunicação multi-GPU. A databricks-ai-runtime competência abrange construir uma imagem do zero.

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

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

As imagens devem usar a linux/amd64 plataforma. Num Mac com Apple Silicon, o Docker cria imagens arm64 por defeito, por isso passa sempre --platform linux/amd64. Para construir a imagem localmente, execute o seguinte:

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

Depois , envia a imagem para o Registo de Artefactos e usa 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 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 é honrado 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.

O Docker credential helper requer um perfil criado com databricks auth login

Falha ao enviar uma imagem com um erro como o seguinte:

Error: profile "my-profile" uses auth_type "". Docker credential helper requires a profile created by databricks auth login

Este erro ocorre quando o perfil utiliza um token de acesso pessoal ou outro método de autenticação que o assistente de credenciais Docker não suporta. Crie um perfil OAuth de workspace e depois volte a executar o push com esse perfil:

databricks auth login --profile my-databricks-profile

Múltiplos perfis correspondem ao ID do espaço de trabalho

A configuração da autenticação do Docker ou o envio de uma imagem falham com um erro como o seguinte:

Error: multiple Databricks profiles match workspace ID 1111111111111111: profile-a and profile-b. Remove duplicate workspace_id entries before using Docker credential helper

Este erro ocorre quando mais do que um perfil no seu ficheiro de configuração (~/.databrickscfg por defeito) aponta para o mesmo espaço de trabalho. Remove o perfil que não usas para o Docker, depois volta a executar o comando com o perfil restante. Para sair de um perfil e removê-lo do ficheiro de configuração, execute o seguinte:

databricks auth logout --profile profile-a --delete

Para todas as opções de logout, veja o grupo de comandos auth.

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

• Início rápido do AI Runtime
• Referência YAML da workload
• air grupo de comandos