Usare immagini Docker personalizzate

Important

Le immagini Docker personalizzate per i carichi di lavoro dell'interfaccia della riga di comando di Ai Runtime sono in versione beta.

Docker Container Services (DCS) consente di usare un'immagine del contenitore Docker personalizzata nei air carichi di lavoro. Usare un'immagine personalizzata quando è necessario:

  • Versioni specifiche della libreria di sistema.
  • Dipendenze complesse che non rientrano in modo pulito in environment.dependencies.
  • Un ambiente esatto per riprodurre i risultati della ricerca.
  • Immagini standard create dal team di sicurezza o dalla piattaforma dell'organizzazione.

Prerequisites

Registrare un'immagine

Prima di eseguire un carico di lavoro con un'immagine personalizzata, registrarlo con air register image. La registrazione scarica l'immagine e la memorizza nella cache della piattaforma Databricks. Ogni utente deve registrare un'immagine una volta per ogni tag immagine. Ripetere la registrazione solo quando si esegue il push di un nuovo tag o si ruotano le credenziali. La registrazione richiede 2-6 minuti e blocca fino a quando l'immagine non è pronta.

Immagini pubbliche

Registrare immagini pubbliche specificando l'URL dell'immagine Docker e il profilo di Databricks:

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

Funziona anche il riferimento abbreviato dell'immagine. Ad esempio: library/ubuntu:latest.

Immagini private di Docker Hub

Per registrare un'immagine di Docker Hub privata, generare prima un token di accesso personale. Nelle impostazioni dell'account Docker Hub fare clic su Token di accesso personali → Genera nuovo token. L'accesso in sola lettura è sufficiente.

Scegliere uno dei metodi di autenticazione seguenti:

Accedere a Docker Hub al terminale. Verrà richiesto il nome utente Docker Hub e il token di accesso personale:

docker login

In questo modo le credenziali vengono archiviate in ~/.docker/config.json. Registrare quindi l'immagine: air legge automaticamente le credenziali:

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

Uso dell’autenticazione interattiva

Autenticare e archiviare le credenziali in un ambito segreto databricks in un unico passaggio:

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

Ti verrà richiesto di inserire il tuo nome utente di Docker Hub e il tuo token di accesso personale. Le credenziali vengono archiviate nell'ambito del segreto dell'area di lavoro per le registrazioni future.

Archivia le credenziali in un secret di Databricks e facci riferimento direttamente:

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

Usare un'immagine Docker in un carico di lavoro

Specifica l'immagine Docker nel file YAML del carico di lavoro sotto 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

Quando si utilizza una propria immagine Docker, environment.dependencies e environment.version non sono supportati. Specificare environment.docker_image.url in uno dei due campi genera un errore. Se si hanno dipendenze aggiuntive, installare invece i pacchetti nel Dockerfile.

Inviare il carico di lavoro:

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

Variabili di ambiente inserite nel contenitore

Il runtime di intelligenza artificiale inserisce le variabili di ambiente seguenti in ogni contenitore in fase di esecuzione:

  • NUM_NODES — numero totale di nodi.
  • LOCAL_WORLD_SIZE — GPU per nodo.
  • WORLD_SIZE — numero totale di processi.
  • POD_RANK — posizione corrente del nodo (con indice a partire da 0). Viene anche iniettato come NODE_RANK.
  • LOCAL_ADDR — IP del nodo locale (solo multinodo).
  • MASTER_ADDR — indirizzo di coordinamento rank-0 (solo per configurazioni multinodo).
  • MASTER_PORT — porta di coordinamento rank-0 (solo multinodo).

Examples

A10 a nodo singolo

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 multinodo con RDMA

Per i processi H100 multinodo che richiedono una larghezza di banda di rete completa nelle istanze di AWS p5, basare l'immagine su una delle immagini di base di Databricks con NCCL ed EFA preconfigurate:

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

Creare un'immagine personalizzata

Immagini di base di Databricks

Databricks pubblica immagini di base su Docker Hub in databricksruntime/air con CUDA, NCCL e rete specifica del cloud (AWS EFA o Azure InfiniBand) preconfigurata.

Tag Cloud Variant Usa quando
dcs-base-aws-runtime AWS Runtime Installazione solo ruote predefinite
dcs-base-aws-devel AWS Devel Compilazione di estensioni CUDA (richiede nvcc)
dcs-base-azure-runtime Azure Runtime Installazione solo ruote predefinite
dcs-base-azure-devel Azure Devel Compilazione di estensioni CUDA (richiede nvcc)

Usare la variante runtime, a meno che il Dockerfile non compili estensioni CUDA come flash-attn, apex o kernel personalizzati.

Dockerfile di esempio che aggiunge PyTorch a un'immagine di base di Databricks. Le immagini di base forniscono Python in /opt/venv, gestito da uv. uv pip install è destinato a tale ambiente per impostazione predefinita; per usare un ambiente diverso, creare e attivare un venv prima di eseguire 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

Compilare, eseguire il push e registrare:

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

Requirements

  • Le immagini devono essere ospitate in Docker Hub. Amazon ECR, Google GCR e GitHub GHCR non sono supportati.
  • Le dimensioni dell'immagine devono essere minori di 20 GB.
  • WORKDIR non viene rispettato in fase di esecuzione. Usare percorsi assoluti per i file inseriti nell'immagine. Ad esempio, usare python /app/train.py, non python train.py.
  • Non è possibile usare environment.dependencies o environment.version con environment.docker_image.url. Se sono necessari pacchetti aggiuntivi oltre a ciò che si trova nell'immagine, è necessario aggiungerli al Dockerfile.

Vedere anche

  • Last updated on