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

Usar imagens personalizadas do Docker com AI Runtime

info

Visualização

Esse recurso está em Prévia Pública. Para usá-lo, um administrador do workspace deve habilitar a prévia do Databricks Artifact Registry na página Previews do workspace.

O Runtime de AI pode execução a sua workload em uma imagem de container Docker personalizada em vez de um ambiente gerenciado. Você constrói a imagem para linux/amd64, envia-a para o Artifact Registry no seu workspace e referencia o nome da imagem no Unity Catalog no YAML da workload que você envia com a CLI do Databricks. Use uma imagem personalizada quando precisar de:

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

Enviar 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 o workload. Para criar uma imagem primeiro, consulte Criar sua própria imagem.

  1. Crie ou selecione um catálogo e um esquema do Unity Catalog para a imagem e obtenha as privilégios do Artifact Registry necessários.

  2. Crie ou faça refresh de um perfil OAuth de workspace para o workspace:

    Shell
    databricks auth login --profile my-databricks-profile

    O auxiliar de credenciais do Docker requer um perfil criado por databricks auth login. Ele não oferece suporte a perfis que usam access tokens pessoais ou credenciais OAuth machine-to-machine (M2M).

  3. Para autenticar o Docker, fazer o tag da imagem e enviá-la em um passo, execução o helper databricks air images push:

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

    Para obter mais opções, incluindo copiar uma imagem de outro registro, consulte Enviar imagens com a CLI do Databricks.

    Para enviar com docker push em vez disso, execução databricks auth docker configure --profile my-databricks-profile para configurar o helper de credenciais do Docker, então siga Enviar uma imagem no guia do Artifact Registry.

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

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.

  • 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}" \
    --local_addr="$(hostname -I | awk '{print $1}')" \
    /app/train.py

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

Você também pode começar de outra imagem base, como nvidia/cuda. Nesse caso, a imagem deve incluir um runtime CUDA suportado pelo driver da GPU do host e NCCL para comunicação com múltiplas GPUs. A skill databricks-ai-runtime abrange a criação de uma imagem do zero.

Para workloads de GPU_8xH100 de vários nós, uma imagem criada a partir de outra imagem base também deve incluir EFA.

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

As imagens devem usar a plataforma linux/amd64. Em um Mac com Apple Silicon, as builds do Docker arm64 imagens por default, então sempre passe --platform linux/amd64. Para criar a imagem localmente, execute o seguinte:

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

Em seguida, envie a imagem para o Artifact Registry e use o nome <catalog>.<schema>.<image>:<tag> resultante como environment.unity_catalog_image no YAML da carga de trabalho.

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.

O auxiliar de credenciais do Docker requer um perfil criado por databricks auth login​

O envio de uma imagem falha com um erro como o seguinte:

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

Este erro ocorre quando o perfil usa um access token pessoal ou outro método de autenticação que o auxiliar de credenciais do Docker não oferece suporte. Crie um perfil OAuth de workspace e execute o push novamente com esse perfil:

Shell
databricks auth login --profile my-databricks-profile

Vários perfis correspondem ao ID do workspace​

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

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

Esse erro ocorre quando mais de um perfil no seu arquivo de configuração (~/.databrickscfg por default) aponta para o mesmo workspace. Remova o perfil que você não usa para o Docker e execute o comando novamente com o perfil restante. Para sair de um perfil e removê-lo do arquivo de configuração, execução o seguinte:

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

Para todas as opções de logout, consulte auth grupo de comando.

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