Usar imagens personalizadas do Docker

Importante

As imagens personalizadas do Docker para cargas de trabalho da CLI do AI Runtime estão em Beta.

Os Serviços de Contêiner do Docker (DCS) permitem que você traga sua própria imagem de contêiner do Docker para air cargas de trabalho. Use uma imagem personalizada quando precisar:

  • Versões específicas da biblioteca do sistema.
  • Dependências complexas que não se encaixam claramente em environment.dependencies.
  • Um ambiente exato para reproduzir os resultados da pesquisa.
  • Imagens padrão criadas pela plataforma ou equipe de segurança da sua organização.

Pré-requisitos

Registrar uma imagem

Antes de executar uma carga de trabalho com uma imagem personalizada, registre-a com air register image. O registro extrai e armazena em cache a imagem na plataforma Databricks. Cada usuário deve registrar uma imagem uma vez por marca de imagem. Registre novamente apenas ao enviar uma nova tag ou alternar credenciais. O registro leva de 2 a 6 minutos e fica bloqueado até que a imagem esteja pronta.

Imagens públicas

Registre imagens públicas fornecendo a URL de imagem do Docker e seu perfil do Databricks:

air register image docker.io/nvidia/cuda:12.9.0-devel-ubuntu24.04 -p my-databricks-profile

A referência de imagem no formato curto também funciona. Por exemplo, library/ubuntu:latest.

Imagens de Docker Hub privadas

Para registrar uma imagem de Docker Hub privada, gere um token de acesso pessoal primeiro. Em suas configurações de conta Docker Hub, clique em Tokens de acesso pessoal → Gerar novo token. O acesso somente de leitura é suficiente.

Escolha um dos seguintes métodos de autenticação:

Faça logon no Docker Hub no terminal. Você será solicitado a fornecer seu nome de usuário Docker Hub e o token de acesso pessoal:

docker login

Isso armazena suas credenciais em ~/.docker/config.json. Em seguida, registre a imagem — air lê as credenciais automaticamente:

air register image myorg/myrepo:mytag -p my-databricks-profile

Usando a autenticação interativa

Autentique e armazene credenciais em um escopo de segredo do Databricks em uma etapa:

air register image myorg/myrepo:mytag --interactive-authenticate -p my-databricks-profile

Você será solicitado a fornecer seu nome de usuário Docker Hub e o token de acesso pessoal. As credenciais são armazenadas no escopo de segredo do workspace para registros futuros.

Armazene credenciais em um segredo do Databricks e faça referência diretamente a ele:

air register image myorg/myrepo:mytag --scope my-secret-scope --key my-docker-key -p my-databricks-profile

Usar uma imagem do Docker em uma carga de trabalho

Especifique a imagem do Docker no YAML da carga de trabalho em environment.docker_image.url:

experiment_name: my-dcs-training
environment:
  docker_image:
    url: myorg/myrepo:mytag
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: python /app/train.py

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

Envie a carga de trabalho:

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

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:

  • 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 (indexada em 0). Também é injetado como NODE_RANK.
  • LOCAL_ADDR — IP do nó local (somente vários nós).
  • MASTER_ADDR — endereço de coordenação rank-0 (somente vários nós).
  • MASTER_PORT — porta de coordenação do rank 0 (somente para múltiplos nós).

Exemplos

A10 de nó único

experiment_name: my-dcs-single-node
environment:
  docker_image:
    url: myorg/myrepo:mytag
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: python3 /app/train.py

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:
  docker_image:
    url: myorg/myrepo:mytag
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

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.

Etiqueta Nuvem Variant Usar quando
dcs-base-aws-runtime AWS Runtime Instalar apenas wheels pré-criados
dcs-base-aws-devel AWS Devel Compilando extensões CUDA (requer nvcc)
dcs-base-azure-runtime Azure Runtime Instalar apenas wheels pré-criados
dcs-base-azure-devel Azure Devel Compilação de extensões CUDA (requer nvcc)

Use a variante runtime, a menos que seu Dockerfile compile extensões CUDA, como flash-attn, apex ou kernels personalizados.

Exemplo de Dockerfile adicionando PyTorch a uma imagem base do Databricks. As imagens base fornecem Python em /opt/venv, gerenciado por uv. uv pip install direciona esse ambiente por padrão; para usar um ambiente diferente, criar e ativar um venv antes de executar uv pip install.

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 /app/train

Compilar, enviar e registrar:

docker build -t myorg/myrepo:mytag .
docker push myorg/myrepo:mytag
air register image myorg/myrepo:mytag --interactive-authenticate -p my-databricks-profile

Requirements

  • As imagens devem ser hospedadas no Docker Hub. Não há suporte para o Amazon ECR, Google GCR e GITHUB GHCR.
  • 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.docker_image.url. Se você precisar de pacotes extras além do que está na imagem, deverá adicioná-los ao Dockerfile.

Consulte também

  • Last updated on