Använd anpassade Docker-bilder med AI Runtime

AI Runtime kan köra en anpassad Docker-containerbild lagrad i Artifact Registry. Använd en anpassad avbildning när du behöver:

• Systembibliotek eller komplexa beroenden som du inte kan installera med environment.dependencies.
• En reproducerbar miljö inom utveckling, forskning och produktion.
• Privata, av organisationen godkända bilder byggda av din plattform eller säkerhetsteam.

Prerequisites

• Installera eller uppdatera till Databricks CLI version 1.19.0 eller högre. Se Installera eller uppdatera Databricks CLI.
• Be en workspace-administratör att aktivera AI Runtime Beta Features-förhandsvisningen. För instruktioner, se Hantera förhandsvisningar på arbetsytsnivå.
• Ställ in Artifact Registry i samma arbetsyta och skaffa behörighet att pusha och läsa bilder. Se Kom igång med Artefaktregistret.
• Installera och starta Docker på din lokala dator.

Skicka en bild till Artefaktregistret

Innan du använder en anpassad avbild med AI Runtime ska du lagra den i Artifact Registry i samma arbetsyta som du använder för att skicka in arbetsbelastningen.

1. Skapa eller välj en Unity Catalog-katalog och ett schema för bilden.

2. Följ Get Started with Artifact Registry för att konfigurera Docker-autentisering, ge de nödvändiga rättigheterna och skicka ut bilden.

3. Notera bildens Unity-katalognamn i följande format:

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

   Till exempel main.ml.training:v1. Inkludera inte registervärdnamnet i arbetsbelastningskonfigurationen.

Använd en Docker-image i en arbetsbelastning

Ange bildens Unity-katalognamn i din arbetsbelastnings-YAML under 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

Detta exempel kör train.py från /app i bilden. För att ladda upp applikationskod separat utan att bygga om bilden, se Kör uppladdad applikationskod.

När du använder en egen Docker-avbild stöds inte environment.dependencies och environment.version. Om du anger environment.unity_catalog_image med något av fälten utlöses ett fel. Om du har ytterligare beroenden installerar du paketen i Dockerfile i stället.

Skicka arbetsbelastningen:

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

Profilen måste autentiseras till samma arbetsyta där bilden lagras.

Miljövariabler som matas in i containern

AI Runtime matar in följande miljövariabler i varje container vid körning:

• CODE_SOURCE_PATH: sökväg till uppladdad applikationskod, när code_source konfigureras.
• NUM_NODES: totalt antal noder.
• LOCAL_WORLD_SIZE: GPU:er per nod.
• WORLD_SIZE: totalt antal processer.
• POD_RANK: nuvarande nodrank (0-indexerad). Injiceras också som NODE_RANK.
• LOCAL_ADDR: lokal nod-IP (endast flernoder).
• MASTER_ADDR: rank-0-koordineringsadress (endast för flera noder).
• MASTER_PORT: rank-0-koordinationsport (endast för flera noder).

Exempel

Följande exempel visar hur man kör uppladdad applikationskod och distribuerad träning med en anpassad bild.

Kör uppladdad applikationskod

Använd code_source den för att ladda upp applikationskod separat från bilden. Du kan redigera och skicka in din kod igen utan att bygga om bilden. Installera kodens Python- och systemberoenden i bilden.

Placera train.py i en lokal src katalog bredvid workload.yaml. Följande konfiguration laddar upp src och kör dess train.py i den anpassade avbilden:

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 tolkas relativt till workload.yaml. AI Runtime sätter $CODE_SOURCE_PATH den uppladdade katalogens sökväg i containern. Se code_source för alternativ för kodkälla.

H100 med flera noder med RDMA

För H100-jobb med flera noder som behöver fullständig nätverksbandbredd på AWS p5-instanser baserar du avbildningen på en av Databricks-basavbildningarna med NCCL och EFA förkonfigurerade:

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

Skapa en egen avbildning

När du skapar en egen avbildning rekommenderar Databricks att du använder databricks-ai-runtime kunskaper med en kodningsagent eller startar från en Databricks-basavbildning.

Använda en kodningsagent

Installera Claude Code-skickligheten databricks-ai-runtime för stegvisa Dockerfile-riktlinjer, inklusive skapande från grunden, CUDA/NCCL/EFA-kompatibilitet, vanliga problem och en checklista för förhandsversion. Den här färdigheten kräver Databricks CLI version 1.0.0 eller senare.

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

Databricks-basavbildningar

Databricks publicerar basavbildningar på Docker Hub med databricksruntime/air CUDA, NCCL och molnspecifika nätverk (AWS EFA eller Azure InfiniBand) förkonfigurerade.

Följande Dockerfile lägger till PyTorch i en Databricks-basbild. Basavbilderna innehåller Python på /opt/venv, som hanteras av uv. uv pip install Siktar på den miljön som standard. För att använda en annan miljö skapa och aktivera en venv innan du kör uv pip install.

För att inkludera ditt träningsskript i bilden, placera train.py bredvid Dockerfile. Dockerfilen kopierar det till /app/train.py. Om du laddar upp applikationskod med code_source, utelämna instruktionen COPY och behåll train.py istället i din källkatalog.

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

Bygg bilden lokalt:

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

Följ sedan Get Started with Artifact Registry för att tagga och skicka bilden till Artifact Registry. Använd det resulterande <catalog>.<schema>.<image>:<tag> namnet som environment.unity_catalog_image i arbetsbelastnings-YAML.

Limitations

• Bilder måste lagras i Artifact Registry i arbetsytan där du skickar in arbetsbelastningen.
• Bildstorleken måste vara under 20 GB.
• WORKDIR respekteras inte vid körning. Använd absoluta sökvägar för filer som har bakats in i bilden. Använd till exempel python /app/train.py, inte python train.py.
• Du kan inte använda environment.dependencies eller environment.version med environment.unity_catalog_image. Om du behöver extra paket utöver vad som finns i avbildningen måste du lägga till dem i Dockerfile.

Felsökning

Information om felsökning av fel som rör autentisering, behörigheter, överföring av avbildningar eller upptäckt av avbildningar i registret finns i Felsök Artifact Registry.

SSL. SSLError vid inläsning av beroenden

En anpassad avbild kan misslyckas under körning med ett OpenSSL-fel när ett bibliotek försöker skapa en SSL-kontext, till exempel:

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

Felet visas vid import av bibliotek som öppnar nätverksanslutningar, till exempel huggingface_hub, och förhindrar att de läses in.

Detta sker eftersom AI-körningsarbetsbelastningar körs på värdar med Federal Information Processing Standards (FIPS) aktiverade. När avbildningens kryptografiska bibliotek inte är FIPS-kompatibla kan OpenSSL inte initieras i FIPS-läge, så det går inte att skapa en SSL-kontext.

Rekommenderad lösning:

Företags-, myndighets-, sjukvårds- och ekonomiarbetsbelastningar är ofta beroende av FIPS 140-2- eller 140-3-efterlevnad för FedRAMP-, CMMC- eller HIPAA-granskningar. Om din arbetsbelastning måste förbli FIPS-kompatibel skapar du avbildningen med FIPS-kompatibla kryptografiska bibliotek.

Om din arbetsbelastning inte kräver FIPS-kompatibilitet kan du inaktivera FIPS-läget genom att ange OPENSSL_FORCE_FIPS_MODE miljövariabeln till 0. Att göra det kan göra att krav på efterlevnad bryts utan att det märks.

Inaktivera FIPS-läge genom att ange detta i arbetsbelastnings-YAML under env_variables:

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

Du kan också ange variabeln i Dockerfile så att den gäller för varje arbetsbelastning som använder avbildningen:

ENV OPENSSL_FORCE_FIPS_MODE=0

Skicka in arbetsbelastningen igen och bekräfta att SSL-felet inte längre visas när beroendena läses in.

Ytterligare resurser

• Databricks CLI-snabbstart för AI Runtime
• YAML-referens för arbetsbelastning
• air kommandogrupp