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

Use Docker Image personalizadas

info

Beta

Imagens Docker personalizadas para workloads da CLI do AI Runtime estão em Beta.

O Docker Container Services (DCS) permite que você utilize sua própria imagem de contêiner Docker para workloads air. Use uma imagem personalizada quando precisar:

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

Pré-requisitos​

Registro uma imagem​

Antes de executar uma carga de trabalho com uma imagem personalizada, registre-a com air register image. O registro baixa e armazena em cache a imagem na plataforma Databricks. Cada usuário deve fazer o registro de uma imagem uma vez por tag de imagem. Faça o registro novamente somente quando você enviar uma nova tag ou rotacionar credenciais. O registro leva 2–6 minutos e bloqueia até que a imagem esteja pronta.

Imagens públicas​

Registre imagens públicas fornecendo o URL da Docker Image e seu perfil do Databricks:

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

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

Imagens privadas do Docker Hub​

Para registrar uma imagem privada do Docker Hub, gere um access token primeiro. Nas configurações da sua account do Docker Hub, clique em tokens de acesso pessoal → Gerar novo token . O acesso somente para leitura é suficiente.

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

Usando docker login (recomendado para uso interativo)​

Log in no Docker Hub no terminal. Será solicitado seu nome de usuário do Docker Hub e access token pessoal:

Shell
docker login

Isso armazena suas credenciais em ~/.docker/config.json. Em seguida, faça o registro da imagem — air lê as credenciais automaticamente:

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

Utilizando autenticação interativa​

Autentique e armazene as credenciais em um Secret Scope do Databricks em uma única etapa:

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

Você será solicitado a informar seu nome de usuário do Docker Hub e personal access token. As credenciais são armazenadas em seu Secret Scope do workspace para registros futuros.

Usando um segredo do Databricks pré-armazenado (recomendado para CI/scripts)​

Armazene as credenciais em um secret do Databricks e referencie-o diretamente:

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

Usar uma Docker Image em uma carga de trabalho​

Especifique a Docker Image no seu YAML de workload em environment.docker_image.url:

YAML
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 Docker Image, environment.dependencies e environment.version não são compatíveis. A especificação de environment.docker_image.url em qualquer um dos campos aciona um erro. Se você tiver dependências adicionais, instale os pacotes no Dockerfile.

Envie a carga de trabalho:

Shell
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 no runtime:

  • 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ó (com índice 0). Também injetado como NODE_RANK.
  • LOCAL_ADDR — IP do nó local (somente multinós).
  • MASTER_ADDR — endereço de coordenação rank-0 (somente multinó).
  • MASTER_PORT — porta de coordenação rank-0 (somente multinó).

Exemplos​

Nó único A10​

YAML
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 multinó com RDMA​

Para Jobs H100 multi-nós que precisam de largura de banda de rede completa em instâncias p5 da AWS, baseie sua imagem em uma das imagens base do Databricks com NCCL e EFA pré-configurados:

YAML
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

Crie sua própria imagem​

Imagens base do Databricks​

O Databricks publica imagens base no Docker Hub em databricksruntime/air com CUDA, NCCL e rede específica da cloud (AWS EFA ou Azure InfiniBand) pré-configurados.

Etiqueta

Nuvem

Variante

Quando usar

dcs-base-aws-runtime

AWS

Runtime

Instalando somente Python wheels pré-configuradas

dcs-base-aws-devel

AWS

Devel

Compilando extensões CUDA (requer nvcc)

dcs-base-azure-runtime

Azure

Runtime

Instalando somente Python wheels pré-configuradas

dcs-base-azure-devel

Azure

Devel

Compilando 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 destina-se a esse ambiente por default; para usar um ambiente diferente, crie e ative um venv antes de executar uv pip install.

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

Crie, envie e registro:

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

Requisitos​

  • Imagens devem ser hospedadas 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 é honrado em tempo de execução. Usar caminhos absolutos para arquivos incorporados na imagem. Por exemplo, use python /app/train.py, não python train.py.
  • Não é possível usar environment.dependencies ou environment.version com environment.docker_image.url. Se você precisar de pacotes extras além do que está na imagem, adicione-os ao Dockerfile.

Veja também​

Nesta página