Verwenden Sie benutzerdefinierte Docker-Images mit AI Runtime

AI Runtime kann ein benutzerdefiniertes Docker-Container-Image ausführen, das in der Artifact Registry gespeichert ist. Verwenden Sie bei Bedarf ein benutzerdefiniertes Bild:

• Systembibliotheken oder komplexe Abhängigkeiten, die Sie nicht mit environment.dependencies installieren können.
• Eine reproduzierbare Umgebung über Entwicklung, Forschung und Produktion hinweg.
• Private, von der Organisation genehmigte Bilder, erstellt von Ihrer Plattform oder Ihrem Sicherheitsteam.

Voraussetzungen

• Installieren oder aktualisieren Sie auf Databricks CLI Version 1.19.0 oder höher. Siehe Installieren oder Aktualisieren der Databricks CLI.
• Bitten Sie einen Workspace-Administrator, die AI Runtime Beta Features-Vorschau zu aktivieren. Für Anleitungen siehe Vorschauen auf Arbeitsflächenebene verwalten.
• Richte das Artefaktregister im selben Arbeitsbereich ein und erhalte die Erlaubnis, Bilder zu pushen und zu lesen. Siehe Starten Sie mit dem Artefaktregister.
• Installiere und starte Docker auf deinem lokalen Rechner.

Schiebe ein Bild ins Artefaktregister

Bevor Sie ein benutzerdefiniertes Image mit AI Runtime verwenden, speichern Sie es im Artifact Registry im selben Arbeitsbereich, mit dem Sie die Arbeitslast einreichen.

1. Erstellen oder wählen Sie einen Unity-Katalogkatalog und ein Schema für das Image aus.

2. Folgen Sie Erste Schritte mit Artifact Registry, um die Docker-Authentifizierung zu konfigurieren, die erforderlichen Berechtigungen zu gewähren und das Image zu pushen.

3. Beachten Sie den Unity Catalog-Namen des Image im folgenden Format:

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

   Beispiel: main.ml.training:v1 Fügen Sie den Registry-Hostnamen nicht in die Workload-Konfiguration ein.

Verwenden Sie ein Docker-Image in einer Workload

Geben Sie den Unity Catalog-Namen des Images in Ihrer Workload-YAML-Datei unter environment.unity_catalog_image an:

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

Dieses Beispiel verläuft von train.py bis /app im Image. Um Anwendungscode separat hochzuladen, ohne das Bild neu zu erstellen, siehe Ausführen hochgeladener Anwendungscode.

Wenn Sie ein eigenes Docker-Image verwenden, werden environment.dependencies und environment.version nicht unterstützt. Die Angabe environment.unity_catalog_image mit beiden Feldern löst einen Fehler aus. Wenn Sie über zusätzliche Abhängigkeiten verfügen, installieren Sie stattdessen die Pakete in der Dockerfile-Datei.

Übermitteln Sie die Arbeitsauslastung:

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

Das Profil muss sich im selben Arbeitsbereich authentifizieren, in dem das Bild gespeichert ist.

In Ihren Container eingefügte Umgebungsvariablen

KI-Runtime fügt die folgenden Umgebungsvariablen in jeden Container zur Laufzeit ein:

• CODE_SOURCE_PATH: Pfad zum hochgeladenen Anwendungscode, wenn code_source konfiguriert ist.
• NUM_NODES: Gesamtzahl der Knoten.
• LOCAL_WORLD_SIZE: GPUs je Knoten.
• WORLD_SIZE: Gesamtzahl der Prozesse.
• POD_RANK: aktueller Knotenrang (0-indexiert). Auch eingefügt als NODE_RANK.
• LOCAL_ADDR: lokale IP-Adresse des Knotens (nur bei mehreren Knoten).
• MASTER_ADDR: Rank-0-Koordinierungsadresse (nur bei Multi-Node).
• MASTER_PORT: Rank-0-Koordinationsport (nur bei mehreren Knoten).

Examples

Die folgenden Beispiele zeigen, wie hochgeladener Anwendungscode und verteiltes Training mit einem benutzerdefinierten Image ausgeführt werden können.

Ausführen Sie hochgeladenen Anwendungscode

Verwenden Sie code_source es, um den Anwendungscode separat vom Bild hochzuladen. Du kannst deinen Code bearbeiten und erneut einreichen, ohne das Bild neu zu erstellen. Installiere die Python- und Systemabhängigkeiten des Codes im Image.

Platziere train.py in einem lokalen src-Verzeichnis neben workload.yaml. Die folgende Konfiguration lädt src hoch und führt train.py im benutzerdefinierten Image aus:

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 resolviert relativ zu workload.yaml. AI Runtime setzt $CODE_SOURCE_PATH auf den Pfad des hochgeladenen Verzeichnisses im Container. Siehe code_source für Code-Quelloptionen.

H100 mit mehreren Knoten und RDMA

Für H100-Jobs mit mehreren Knoten, die auf AWS p5-Instances die volle Netzwerkbandbreite benötigen, verwenden Sie eines der Databricks-Basis-Images als Basis, in denen NCCL und EFA vorkonfiguriert sind:

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

Erstellen Eines eigenen Images

Beim Erstellen eines eigenen Bilds empfiehlt Databricks, databricks-ai-runtime mit einem Coding-Agent zu verwenden oder mit einem Databricks-Basis-Bild zu beginnen.

Verwenden eines Codierungs-Agents

Installieren Sie den databricks-ai-runtime-Claude-Code-Skill für eine schrittweise Anleitung zu Dockerfiles, einschließlich der Erstellung von Grund auf, der CUDA/NCCL/EFA-Kompatibilität, häufiger Probleme und einer Checkliste vor dem Build. Diese Fähigkeit erfordert Databricks CLI Version 1.0.0 oder höher.

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

Databricks-Basisbilder

Databricks veröffentlicht Basis-Images auf Docker Hub unter databricksruntime/air mit vorkonfiguriertem CUDA, NCCL und cloud-spezifischer Netzwerkkonfiguration (AWS EFA oder Azure InfiniBand).

Das folgende Dockerfile fügt PyTorch einem Databricks-Basisimage hinzu. Die Basisbilder stellen Python unter /opt/venv bereit, das von uv verwaltet wird. uv pip install Standardmäßig zielt sie auf diese Umgebung ab. Um eine andere Umgebung zu verwenden, erstellen und aktivieren Sie eine venv, bevor Sie uv pip install ausführen.

Um dein Trainingsskript ins Bild einzubinden, platziere train.py es neben die Dockerfile. Das Dockerfile kopiert es nach /app/train.py. Wenn du Anwendungscode mit code_sourcehochlädst, lasse die COPY Anleitung weg und behalte train.py stattdessen im Quellverzeichnis.

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

Erstellen Sie das Bild lokal:

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

Folgen Sie dann Get Start with Artifact Registry , um das Bild zu taggen und ins Artefact Registry zu pushen. Verwenden Sie den resultierenden <catalog>.<schema>.<image>:<tag> Namen als environment.unity_catalog_image im Workload YAML.

Einschränkungen

• Bilder müssen im Artefaktregister im Arbeitsbereich gespeichert werden, in dem Sie die Workload einreichen.
• Die Bildgröße muss unter 20 GB betragen.
• WORKDIR wird zur Laufzeit nicht berücksichtigt. Verwenden Sie absolute Pfade für Dateien, die in das Bild integriert sind. Verwenden Sie beispielsweise python /app/train.py, nicht python train.py.
• Sie können environment.dependencies oder environment.version nicht mit environment.unity_catalog_image verwenden. Wenn Sie zusätzliche Pakete benötigen, die über das Image hinausgehen, müssen Sie sie der Dockerfile-Datei hinzufügen.

Troubleshooting

Bei Authentifizierungs-, Berechtigungs-, Image-Push- oder Image-Erkennungsfehlern im Zusammenhang mit der Registry siehe Fehlerbehebung für Artifact Registry.

ssl.SSLError beim Laden von Abhängigkeiten

Ein benutzerdefiniertes Bild kann zur Laufzeit mit einem OpenSSL-Fehler fehlschlagen, wenn eine Bibliothek versucht, einen SSL-Kontext zu erstellen, z. B.:

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

Der Fehler tritt beim Importieren von Bibliotheken auf, die Netzwerkverbindungen öffnen, wie z. B. huggingface_hub, und verhindert, dass sie geladen werden.

Dies geschieht, weil KI-Laufzeit-Workloads auf Hosts mit Federal Information Processing Standards (FIPS) aktiviert laufen. Wenn die kryptografischen Bibliotheken des Bilds nicht FIPS-kompatibel sind, kann OpenSSL nicht im FIPS-Modus initialisiert werden, sodass das Erstellen eines SSL-Kontexts fehlschlägt.

Empfohlene Lösung:

Unternehmens-, Regierungs-, Gesundheits- und Finanzarbeitslasten hängen häufig von FIPS 140-2- oder 140-3-Compliance für FedRAMP-, CMMC- oder HIPAA-Audits ab. Wenn Ihre Workload FIPS-kompatibel bleiben muss, erstellen Sie Ihr Image mit FIPS-kompatiblen kryptografischen Bibliotheken.

Wenn Ihre Workload keine FIPS-Compliance erfordert, können Sie den FIPS-Modus deaktivieren, indem Sie die OPENSSL_FORCE_FIPS_MODE Umgebungsvariable auf 0festlegen. Dies kann unbemerkt gegen Compliance-Anforderungen verstoßen.

Um den FIPS-Modus zu deaktivieren, legen Sie ihn in Ihrer Workload YAML unter env_variables:

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

Alternativ können Sie die Variable in Ihrer Dockerfile-Datei so festlegen, dass sie für jede Workload gilt, die das Image verwendet:

ENV OPENSSL_FORCE_FIPS_MODE=0

Übermitteln Sie die Workload erneut, und bestätigen Sie, dass der SSL-Fehler nicht mehr angezeigt wird, wenn Abhängigkeiten geladen werden.

Weitere Ressourcen

• Databricks CLI-Schnellstart für AI Runtime
• Workload YAML-Referenz
• air Befehlsgruppe