Aangepaste Docker-installatiekopieën gebruiken

Belangrijk

Aangepaste Docker-images voor AI Runtime CLI-workloads zijn in bèta.

Met Docker Container Services (DCS) kunt u uw eigen Docker-containerinstallatiekopieën overbrengen naar air workloads. Gebruik een aangepaste image wanneer u deze nodig hebt:

  • Specifieke systeembibliotheekversies.
  • Complexe afhankelijkheden die niet schoon in environment.dependenciespassen.
  • Een exacte omgeving om onderzoeksresultaten te reproduceren.
  • Standaardimages die zijn gemaakt door het platform- of beveiligingsteam van uw organisatie.

Prerequisites

Een installatiekopie registreren

Voordat u een workload met een aangepaste installatiekopie uitvoert, moet u deze registreren met air register image. Bij registratie haalt het Databricks-platform de image op en cachet deze. Elke gebruiker moet één keer per imagetag een image registreren. Registreer u alleen opnieuw wanneer u een nieuwe tag pusht of referenties roteert. Registratie duurt 2-6 minuten en blokkeert totdat de afbeelding gereed is.

Openbare afbeeldingen

Registreer openbare installatiekopieën door de URL van de Docker-installatiekopieën en uw Databricks-profiel op te geven:

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

De korte afbeeldingsreferentie werkt ook. Bijvoorbeeld: library/ubuntu:latest.

Privé-Docker Hub-afbeeldingen

Als u eerst een persoonlijk toegangstoken wilt genereren om een privé Docker Hub-image te registreren. Klik in uw Docker Hub accountinstellingen op Persoonlijke toegangstokensNieuw token genereren. Alleen-lezentoegang is voldoende.

Kies een van de volgende verificatiemethoden:

Meld u aan bij Docker Hub op de terminal. U wordt gevraagd om uw Docker Hub gebruikersnaam en persoonlijk toegangstoken:

docker login

Dit slaat uw aanmeldgegevens op in ~/.docker/config.json. Registreer vervolgens de image — air leest de inloggegevens automatisch:

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

Interactieve verificatie gebruiken

Authenticeer en sla aanmeldgegevens in één stap op in een Databricks secret scope:

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

U wordt gevraagd om uw Docker Hub gebruikersnaam en persoonlijk toegangstoken. Inloggegevens worden opgeslagen in de secret scope van uw werkruimte voor toekomstige registraties.

Sla referenties op in een Databricks-geheim en verwijs er rechtstreeks naar:

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

Een Docker-image gebruiken in een workload

Geef de Docker-image op in uw workload-YAML onder 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

Wanneer u uw eigen Docker-image gebruikt, worden environment.dependencies en environment.version niet ondersteund. Het opgeven van environment.docker_image.url bij een van beide velden veroorzaakt een fout. Als u extra afhankelijkheden hebt, installeert u in plaats daarvan de pakketten in het Dockerfile.

De werklast verzenden:

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

Omgevingsvariabelen die zijn geïnjecteerd in uw container

AI Runtime injecteert de volgende omgevingsvariabelen in elke container tijdens runtime:

  • NUM_NODES — totaal aantal knooppunten.
  • LOCAL_WORLD_SIZE — GPU's per knooppunt.
  • WORLD_SIZE — totaal aantal processen.
  • POD_RANK — huidige positie van knooppunt (0 geïndexeerd). Ook ingevoegd als NODE_RANK.
  • LOCAL_ADDR — IP van lokaal knooppunt (alleen voor meerdere knooppunten).
  • MASTER_ADDR — rang-0 coördinatieadres (alleen voor meerdere knooppunten).
  • MASTER_PORT — rank-0-coördinatiepoort (alleen bij configuraties met meerdere knooppunten).

Examples

A10 met één knooppunt

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 met meerdere knooppunten met RDMA

Voor H100-taken met meerdere knooppunten waarvoor volledige netwerkbandbreedte op AWS p5-exemplaren nodig is, baseert u uw installatiekopie op een van de Databricks-basisinstallatiekopieën met NCCL en EFA vooraf geconfigureerd:

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

Maak uw eigen image

Databricks-basisinstallatiekopieën

Databricks publiceert basisimages op Docker Hub op databricksruntime/air met CUDA, NCCL en cloudspecifieke netwerkconfiguratie (AWS EFA of Azure InfiniBand) vooraf geconfigureerd.

Tag Wolk Variant Wanneer gebruiken
dcs-base-aws-runtime AWS Runtime Alleen vooraf gebouwde wielen installeren
dcs-base-aws-devel AWS Devel CUDA-extensies compileren (vereist nvcc)
dcs-base-azure-runtime Azure Runtime Alleen vooraf gebouwde wielen installeren
dcs-base-azure-devel Azure Devel CUDA-extensies compileren (vereist nvcc)

Gebruik de runtimevariant , tenzij uw Dockerfile CUDA-extensies compileert, zoals flash-attn, apex of aangepaste kernels.

Een voorbeeld van een Dockerfile waarin PyTorch wordt toegevoegd aan een Databricks-basisimage. De basisimages bevatten Python op /opt/venv, beheerd door uv. uv pip install is standaard gericht op die omgeving; als u een andere omgeving wilt gebruiken, maakt en activeert u een venv voordat u deze uitvoert 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

Bouwen, uploaden en registreren:

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

Requirements

  • Images moeten worden gehost op Docker Hub. Amazon ECR, Google GCR en GitHub GHCR worden niet ondersteund.
  • De grootte van de afbeelding moet kleiner zijn dan 20 GB.
  • WORKDIR wordt tijdens runtime niet gehonoreerd. Gebruik absolute paden voor bestanden die in de afbeelding zijn gebakken. Gebruik bijvoorbeeld python /app/train.py, niet python train.py.
  • U kunt environment.dependencies of environment.version niet met environment.docker_image.url gebruiken. Als u extra pakketten nodig hebt boven op wat er in de image zit, moet u die aan de Dockerfile toevoegen.

Zie ook