Usar imagens personalizadas do Docker

Importante

Esse recurso está em Beta. Não é compatível com espaços de trabalho que utilizam controle de saída serverless (SEG).

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

Ao criar sua própria imagem, a Databricks recomenda usar databricks-ai-runtime habilidade com um agente de programação ou começar com uma imagem base da Databricks.

Usar um agente de codificação

Instale a habilidade databricks-ai-runtime do Claude Code para obter orientação passo a passo sobre Dockerfile, incluindo criação do zero, compatibilidade com CUDA/NCCL/EFA, problemas comuns e uma lista de verificação pré-compilação. Essa habilidade requer a CLI do Databricks versão 1.0.0 ou mais recente.

databricks aitools install --skills databricks-ai-runtime --experimental

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

Resolução de problemas

ssl.SSLError: [CRYPTO] erro desconhecido (_ssl.c) ao carregar dependências

Uma imagem personalizada pode falhar em tempo de execução com um erro do 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 conexões de rede, como huggingface_hub, por exemplo, e impede que elas sejam carregadas.

Isso ocorre porque air as cargas de trabalho são executadas em hosts habilitados para FIPS. Quando as bibliotecas criptográficas da imagem não são compatíveis com FIPS, o OpenSSL falha ao inicializar no modo FIPS, portanto, a criação de um contexto SSL falha.

Solução recomendada:

As cargas de trabalho corporativas, governamentais, de saúde e finanças geralmente dependem da conformidade com a FIPS 140-2 ou 140-3 para auditorias FedRAMP, CMMC ou HIPAA. Se sua carga de trabalho precisar permanecer em conformidade com FIPS, crie sua imagem com bibliotecas criptográficas compatíveis com FIPS.

Se sua carga de trabalho não exigir conformidade com FIPS, você poderá desativar o modo FIPS definindo a variável de ambiente OPENSSL_FORCE_FIPS_MODE como 0. Isso pode interromper silenciosamente os requisitos de conformidade.

Para desativar o modo FIPS, defina isso no YAML da sua carga de trabalho, em env_variables:

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

Como alternativa, defina a variável em seu Dockerfile para que ela se aplique a cada carga de trabalho que usa a imagem:

ENV OPENSSL_FORCE_FIPS_MODE=0

Reenvie a carga de trabalho e confirme se o erro SSL não aparece mais quando as dependências são carregadas.

Consulte também