Use imagens Docker personalizadas

Importante

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

O Docker Container Services (DCS) permite-lhe trazer a sua própria imagem de contentor 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 encaixam perfeitamente em environment.dependencies.
  • Um ambiente exato para reproduzir resultados de investigação.
  • Imagens padrão criadas pela plataforma ou equipa de segurança da sua organização.

Pré-requisitos

Registar uma imagem

Antes de executar uma carga de trabalho com uma imagem personalizada, regista-a com air register image. O registo obtém e armazena em cache a imagem na plataforma Databricks. Cada utilizador deve registar uma imagem uma vez por etiqueta de imagem. Faça um novo registo apenas quando enviar uma nova tag ou alterar as credenciais. O registo demora 2 a 6 minutos e fica bloqueado até a imagem estar pronta.

Imagens públicas

Registe imagens públicas fornecendo a URL da imagem Docker e o seu perfil Databricks:

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

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

Imagens privadas do Docker Hub

Para registar uma imagem privada do Docker Hub, gere primeiro um token de acesso pessoal. Nas definições da sua conta Docker Hub, clique em Tokens de Acesso PessoalGerar novo token. O acesso apenas de leitura é suficiente.

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

Inicie sessão no Docker Hub no terminal. Ser-lhe-á solicitado o seu nome de utilizador do Docker Hub e o token de acesso pessoal:

docker login

Isto armazena as suas credenciais em ~/.docker/config.json. Depois regista a imagem — air lê automaticamente as credenciais:

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

Usando autenticação interativa

Autentique e armazene credenciais num escopo secreto Databricks num único passo:

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

Será solicitado o seu nome de utilizador do Docker Hub e o token de acesso pessoal. As credenciais são armazenadas no escopo secreto do seu espaço de trabalho para registos futuros.

Armazene credenciais num segredo Databricks e faça referência direta:

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

Utilizar uma imagem do Docker num workload

Especifique a imagem 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 utilizar uma imagem Docker própria, environment.dependencies e environment.version não são suportados. Especificar environment.docker_image.url 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:

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

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:

  • 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 do nó local (apenas para vários nós).
  • MASTER_ADDR — endereço de coordenação do rank 0 (apenas em configurações multi-nó).
  • MASTER_PORT — porta de coordenação rank-0 (apenas para multinó).

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 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:
  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

Constrói a tua própria imagem

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.

Etiqueta Nuvem Variant Utilizar quando
dcs-base-aws-runtime AWS Runtime Instalação apenas de jantes pré-montadas
dcs-base-aws-devel AWS Devel Compilação de extensões CUDA (requer nvcc)
dcs-base-azure-runtime Azure Runtime Instalação apenas de jantes pré-montadas
dcs-base-azure-devel Azure Devel Compilação de extensões CUDA (requer nvcc)

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

Exemplo de Dockerfile que adiciona o PyTorch a uma imagem base do Databricks. As imagens base fornecem Python em /opt/venv, gerido por uv. uv pip install atinge esse ambiente por defeito; Para usar um ambiente diferente, crie e ative 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 registar:

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 alojadas no Docker Hub. Amazon ECR, Google GCR e GitHub GHCR não são suportados.
  • 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.docker_image.url. Se precisares de pacotes extra para além do que está na imagem, tens de os adicionar ao Dockerfile.

Consulte também