Gebruik aangepaste Docker-images met AI Runtime

Belangrijk

Deze functie bevindt zich in de bètaversie. Om het te gebruiken, moet een workspace-beheerder de AI Runtime Beta Features en Databricks Artifact Registry previews inschakelen vanaf de pagina Voorlezen van de werkruimte.

AI Runtime kan een aangepaste Docker-containerimage draaien die is opgeslagen in het Artifact Registry. Gebruik een aangepaste image wanneer u deze nodig hebt:

  • Systeembibliotheken of complexe afhankelijkheden die je niet met environment.dependencies kunt installeren.
  • Een reproduceerbare omgeving over ontwikkeling, onderzoek en productie.
  • Privé, door de organisatie goedgekeurde afbeeldingen gebouwd door jouw platform of beveiligingsteam.

Prerequisites

  • Installeer of update naar Databricks CLI versie 1.19.0 of hoger. Zie De Databricks CLI installeren of bijwerken.
  • Vraag een workspace-beheerder om de AI Runtime Beta Features preview in te schakelen. Raadpleeg Voorbeelden op werkruimteniveau beheren voor instructies.
  • Stel Artifact Registry in dezelfde werkruimte in en verkrijg toestemming om afbeeldingen te pushen en te lezen. Zie Begin met het Artefactregister.
  • Installeer en start Docker op je lokale machine.

Stuur een afbeelding naar het Artefactregister

Voordat je een aangepaste image met AI Runtime gebruikt, sla deze op in het Artifact Registry in dezelfde werkruimte die je gebruikt om de workload in te dienen.

  1. Maak of selecteer een Unity Catalog-catalogus en schema voor de afbeelding.

  2. Volg Get Started with Artifact Registry om Docker-authenticatie te configureren, de vereiste rechten te verlenen en de image te pushen.

  3. Noteer de Unity Catalog-naam van de afbeelding in het volgende formaat:

    <catalog>.<schema>.<image>:<tag>
    

    Bijvoorbeeld: main.ml.training:v1. Neem de hostnaam van het register niet op in de workloadconfiguratie.

Tip

Alternatief kun je het databricks air images push helper-commando gebruiken in Databricks CLI.

Gebruik een Docker-image in een workload

Specificeer de naam van de Unity Catalog van de image in je workload YAML onder environment.unity_catalog_image:

experiment_name: my-dcs-training
environment:
  unity_catalog_image: main.ml.training:v1
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: python /app/train.py

Dit voorbeeld voert train.py uit vanuit /app in de afbeelding. Om applicatiecode apart te uploaden zonder de afbeelding opnieuw te bouwen, zie Run uploaded application code.

Wanneer u uw eigen Docker-image gebruikt, worden environment.dependencies en environment.version niet ondersteund. Het opgeven van environment.unity_catalog_image 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:

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

Het profiel moet authenticeren bij dezelfde werkruimte waar de afbeelding wordt opgeslagen.

Omgevingsvariabelen die zijn geïnjecteerd in uw container

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

  • CODE_SOURCE_PATH: pad naar geüploade applicatiecode, wanneer code_source is geconfigureerd.
  • NUM_NODES: totaal aantal knooppunten.
  • LOCAL_WORLD_SIZE: GPU's per node.
  • WORLD_SIZE: totaal aantal processen.
  • POD_RANK: huidige knooppuntrang (0-geïndexeerd). Ook ingevoegd als NODE_RANK.
  • LOCAL_ADDR: IP-adres van de lokale node (alleen multi-node).
  • MASTER_ADDR: rank-0 coördinatieadres (alleen voor multi-node).
  • MASTER_PORT: rank-0-coördinatiepoort (alleen voor meerdere nodes).

Examples

De volgende voorbeelden laten zien hoe je geüploade applicatiecode en gedistribueerde training kunt uitvoeren met een aangepaste afbeelding.

Voer geüploade applicatiecode uit

Gebruik code_source het om applicatiecode apart van de afbeelding te uploaden. Je kunt je code bewerken en opnieuw indienen zonder de afbeelding opnieuw op te bouwen. Installeer de Python- en systeemafhankelijkheden van de code in de afbeelding.

Plaats train.py in een lokale src directory naast workload.yaml. De volgende configuratie uploadt src en voert train.py uit binnen de aangepaste image:

experiment_name: my-dcs-uploaded-code
environment:
  unity_catalog_image: main.ml.training:v1
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
code_source:
  root_path: ./src
command: |-
  cd "$CODE_SOURCE_PATH"
  python3 train.py

root_path resolueert ten opzichte van workload.yaml. AI Runtime stelt het pad van de geüploade map in de container in $CODE_SOURCE_PATH . Zie code_source voor broncodeopties.

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:
  unity_catalog_image: main.ml.training:v1
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

Wanneer je je eigen image bouwt, raadt Databricks aan om de databricks-ai-runtime skill met een codeeragent te gebruiken of te beginnen met een Databricks-basisimage.

Een coderingsagent gebruiken

Installeer de databricks-ai-runtime Claude Code-skill voor stapsgewijze Dockerfile-richtlijnen, waaronder vanaf nul bouwen, compatibiliteit met CUDA/NCCL/EFA, veelvoorkomende problemen en een checklist vóór het bouwen. Voor deze vaardigheid is Databricks CLI versie 1.0.0 of hoger vereist.

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

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 Variant CUDA Wanneer gebruiken
dcs-base-azure-runtime Runtime 12 Alleen vooraf gebouwde wielen installeren
dcs-base-azure-devel Devel 12 CUDA-extensies compileren (vereist nvcc)

De volgende Dockerfile voegt PyTorch toe aan een Databricks-basisimage. De basisimages bevatten Python op /opt/venv, beheerd door uv. uv pip install Richt zich standaard op die omgeving. Om een andere omgeving te gebruiken, maak en activeer je een venv voordat je start uv pip install.

Om je trainingsscript in de image op te nemen, plaats je train.py naast de Dockerfile. Het Dockerfile kopieert het naar /app/train.py. Als je applicatiecode uploadt met code_source, laat de COPY instructie weg en bewaar train.py in plaats daarvan in je bronmap.

FROM databricksruntime/air:dcs-base-azure-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.py /app/train.py

Bouw het beeld lokaal:

docker build --platform linux/amd64 -t my-training-image:v1 .

Volg daarna Get Started with Artifact Registry om de afbeelding te taggen en naar Artifact Registry te pushen. Gebruik de resulterende <catalog>.<schema>.<image>:<tag> naam als environment.unity_catalog_image in de workload YAML.

Tip

Of gebruik het databricks air images push helper-commando in Databricks CLI en volg de interactieve prompts.

Limitations

  • Afbeeldingen moeten worden opgeslagen in het Artifact Registry in de werkruimte waar je de werklast indient.
  • 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.unity_catalog_image gebruiken. Als u extra pakketten nodig hebt boven op wat er in de image zit, moet u die aan de Dockerfile toevoegen.

Probleemoplossingsproces

Voor fouten bij registergerelateerde authenticatie, machtigingen, het pushen van images of het detecteren van images, zie Troubleshoot Artifact Registry.

ssl.SSLError bij het laden van afhankelijkheden

Een aangepaste installatiekopie kan tijdens runtime mislukken met een OpenSSL-fout wanneer een bibliotheek probeert een SSL-context te maken, bijvoorbeeld:

ssl.SSLError: [CRYPTO] unknown error (_ssl.c:3076)

De fout wordt weergegeven bij het importeren van bibliotheken die netwerkverbindingen openen, zoals huggingface_huben voorkomt dat ze worden geladen.

Dit gebeurt omdat AI-runtime-workloads draaien op hosts met Federal Information Processing Standards (FIPS) ingeschakeld. Wanneer de cryptografische bibliotheken van de image niet FIPS-conform zijn, kan OpenSSL niet initialiseren in de FIPS-modus, waardoor het maken van een SSL-context mislukt.

Aanbevolen oplossing:

Workloads voor ondernemingen, overheid, gezondheidszorg en financiën zijn vaak afhankelijk van FIPS 140-2- of 140-3-naleving voor FedRAMP-, CMMC- of HIPAA-audits. Als uw workload aan FIPS moet blijven voldoen, bouw dan uw image met cryptografische bibliotheken die aan FIPS voldoen.

Als uw workload geen FIPS-naleving vereist, kunt u de FIPS-modus uitschakelen door de OPENSSL_FORCE_FIPS_MODE omgevingsvariabele in te stellen op 0. Daardoor kunnen nalevingseisen ongemerkt worden geschonden.

Als u de FIPS-modus wilt uitschakelen, stelt u deze in uw workload YAML in onder env_variables:

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

U kunt de variabele ook instellen in uw Dockerfile, zodat deze geldt voor elke workload die de image gebruikt:

ENV OPENSSL_FORCE_FIPS_MODE=0

Verzend de workload opnieuw en bevestig dat de SSL-fout niet meer wordt weergegeven wanneer afhankelijkheden worden geladen.

Aanvullende bronnen