Używaj niestandardowych obrazów Dockera z AI Runtime

Important

Ta funkcja jest dostępna w wersji beta. Aby z niego korzystać, administrator przestrzeni roboczej musi włączyć podglądy AI Runtime Beta Features oraz Databricks Artifact Registry na stronie Podglądów przestrzeni roboczej.

AI Runtime może uruchomić niestandardowy obraz kontenera Dockera przechowywany w rejestrze artefaktów. Jeśli potrzebujesz, użyj obrazu niestandardowego:

  • Biblioteki systemowe lub złożone zależności, których nie da się zainstalować za pomocą environment.dependencies.
  • Powtarzalne środowisko w zakresie rozwoju, badań i produkcji.
  • Prywatne, zatwierdzone przez organizację obrazy przygotowane przez zespół platformy lub zespół ds. bezpieczeństwa.

Prerequisites

Przesuń obraz do Rejestru Artefaktów

Przed użyciem niestandardowego obrazu w AI Runtime zapisz go w Artifact Registry w tym samym obszarze roboczym, którego używasz do przesyłania zadania.

  1. Stwórz lub wybierz katalog i schemat Unity Catalog dla obrazu.

  2. Następuj: Zacznij pracę z Artifact Registry , aby skonfigurować uwierzytelnianie Docker, przydzielić wymagane uprawnienia i wypchnąć obraz.

  3. Zwróć uwagę na nazwę obrazu w katalogu Unity w następującym formacie:

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

    Na przykład main.ml.training:v1. Nie uwzględniaj nazwy hosta rejestru w konfiguracji obciążenia roboczego.

Wskazówka

Alternatywnie, użyj poleceniadatabricks air images push pomocniczego w Databricks CLI.

Użyj obrazu Dockera w obciążeniu roboczym

Określ nazwę obrazu w katalogu Unity Catalog w pliku YAML definicji obciążenia pod 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

Ten przykład jest uruchamiany z train.py na obrazie /app. Aby przesłać kod aplikacji osobno bez przebudowywania obrazu, zobacz Uruchom przesłany kod aplikacji.

W przypadku korzystania z własnego obrazu Docker elementy environment.dependencies i environment.version nie są obsługiwane. Określanie environment.unity_catalog_image za pomocą dowolnego pola wyzwala błąd. Jeśli masz dodatkowe zależności, zainstaluj zamiast tego pakiety w pliku Dockerfile.

Prześlij obciążenie:

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

Profil musi uwierzytelnić się w tym samym obszarze roboczym, w którym przechowywany jest obraz.

Zmienne środowiskowe wprowadzone do kontenera

Środowisko uruchomieniowe sztucznej inteligencji wprowadza następujące zmienne środowiskowe do każdego kontenera w czasie wykonywania:

  • CODE_SOURCE_PATH: ścieżka do przesłanego kodu aplikacji, gdy code_source jest skonfigurowany.
  • NUM_NODES: łączna liczba węzłów.
  • LOCAL_WORLD_SIZE: GPU na węzeł.
  • WORLD_SIZE: całkowita liczba procesów.
  • POD_RANK: aktualna ranga węzła (0-indeksowana). Wstrzyknięto również jako NODE_RANK.
  • LOCAL_ADDR: lokalny adres IP węzła (tylko wielowęzłowy).
  • MASTER_ADDR: adres koordynacji rangi 0 (tylko wielowęzłowe).
  • MASTER_PORT: port koordynacyjny rangi 0 (tylko wielowęzłowy).

Examples

Poniższe przykłady pokazują, jak uruchamiać przesłany kod aplikacji oraz rozproszone szkolenie z niestandardowym obrazem.

Uruchamiaj przesłany kod aplikacji

Używaj do code_source przesyłania kodu aplikacji oddzielnie od obrazu. Możesz edytować i przesyłać swój kod bez ponownego budowania obrazu. Zainstaluj Python i zależności systemowe kodu w obrazie.

Umieść train.py w lokalnym katalogu src obok workload.yaml. Poniższa konfiguracja przesyła src i uruchamia jego train.py wewnątrz niestandardowego obrazu:

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 jest rozpoznawany względem workload.yaml. AI Runtime ustawia $CODE_SOURCE_PATH ścieżkę przesłanego katalogu w kontenerze. Zobacz code_source, aby poznać opcje źródła kodu.

Wielowęzłowy H100 z RDMA

W przypadku wielowęzłowych zadań H100, które wymagają pełnej przepustowości sieci w instancjach AWS p5, oprzyj obraz na jednym z obrazów bazowych Databricks z wstępnie skonfigurowanymi NCCL i EFA:

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

Tworzenie własnego obrazu

Podczas tworzenia własnego obrazu usługa Databricks zaleca używanie databricks-ai-runtime umiejętności z agentem kodowania lub rozpoczynanie pracy z obrazem podstawowym usługi Databricks.

Używanie agenta kodowania

Zainstaluj umiejętność Claude Code databricks-ai-runtime, aby uzyskać wskazówki krok po kroku dotyczące pliku Dockerfile, w tym tworzenia od podstaw, zgodności z CUDA/NCCL/EFA, typowych problemów i listy kontrolnej przed budowaniem. Ta umiejętność wymaga interfejsu wiersza polecenia usługi Databricks w wersji 1.0.0 lub nowszej.

databricks aitools install --skills databricks-ai-runtime

Obrazy bazowe Databricks

Databricks publikuje obrazy bazowe w Docker Hub pod adresem databricksruntime/air, z wstępnie skonfigurowanymi CUDA, NCCL oraz siecią specyficzną dla danej chmury (AWS EFA lub Azure InfiniBand).

Tag Variant CUDA Użyj, gdy
dcs-base-azure-runtime Runtime 12 Instalowanie tylko wstępnie utworzonych kół
dcs-base-azure-devel Devel 12 Kompilowanie rozszerzeń CUDA (wymaga nvcc)

Poniższy plik Dockerfile dodaje PyTorch do obrazu bazowego Databricks. Obrazy bazowe udostępniają Python w lokalizacji /opt/venv, zarządzanym przez uv. uv pip install domyślnie celuje w to środowisko. Aby użyć innego środowiska, stwórz i aktywuj wenv przed uruchomieniem uv pip install.

Aby dołączyć skrypt treningowy do obrazu, umieść train.py obok pliku Dockerfile. Dockerfile kopiuje go do /app/train.py. Jeśli przesyłasz kod aplikacji z code_source, pomiń instrukcję COPY i pozostaw train.py w katalogu źródłowym.

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

Buduj obraz lokalnie:

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

Następnie kliknij Rozpocznij się z Rejestrem Artefaktów , aby oznaczyć i przesunąć obraz do Rejestru Artefaktów. Użyj wynikowej <catalog>.<schema>.<image>:<tag> nazwy jako environment.unity_catalog_image w pliku YAML obciążenia.

Wskazówka

Alternatywnie, użyj poleceniadatabricks air images push pomocniczego w Databricks CLI i podążaj za interaktywnymi instrukcjami.

Limitations

  • Obrazy muszą być przechowywane w Rejestrze Artefaktów w przestrzeni roboczej, do której przesyłasz zadanie.
  • Rozmiar obrazu musi być mniejszy niż 20 GB.
  • WORKDIR nie jest uwzględniany podczas wykonywania. Użyj ścieżek bezwzględnych dla plików zawartych w obrazie. Na przykład użyj python /app/train.py, a nie python train.py.
  • Nie można używać environment.dependencies ani environment.version z environment.unity_catalog_image. Jeśli potrzebujesz dodatkowych pakietów poza tym, co znajduje się na obrazie, musisz dodać je do pliku Dockerfile.

Troubleshooting

W przypadku błędów uwierzytelniania związanych z rejestrem, błędów uprawnień, przesyłania obrazu lub wykrywania obrazów zobacz Rozwiązywanie problemów z Artifact Registry.

ssl. SSLError podczas ładowania zależności

Obraz niestandardowy może zakończyć się niepowodzeniem w czasie wykonywania z powodu błędu openSSL, gdy biblioteka próbuje utworzyć kontekst SSL, na przykład:

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

Błąd pojawia się podczas importowania bibliotek, które otwierają połączenia sieciowe, takie jak huggingface_hub, i uniemożliwiają ich ładowanie.

Dzieje się tak, ponieważ obciążenia AI Runtime działają na hostach z włączonymi Federal Information Processing Standards (FIPS). Jeśli biblioteki kryptograficzne obrazu nie są zgodne ze standardem FIPS, nie można zainicjować biblioteki OpenSSL w trybie FIPS, więc tworzenie kontekstu SSL kończy się niepowodzeniem.

Zalecane rozwiązanie:

Obciążenia przedsiębiorstwa, instytucji rządowych, opieki zdrowotnej i finansów często zależą od zgodności ze standardem FIPS 140-2 lub 140-3 dla inspekcji FedRAMP, CMMC lub HIPAA. Jeśli obciążenie musi pozostać zgodne ze standardem FIPS, skompiluj obraz przy użyciu bibliotek kryptograficznych zgodnych ze standardem FIPS.

Jeśli obciążenie nie wymaga zgodności ze standardem FIPS, możesz wyłączyć tryb FIPS, ustawiając zmienną OPENSSL_FORCE_FIPS_MODE środowiskową na 0wartość . Może to dyskretnie złamać wymagania dotyczące zgodności.

Aby wyłączyć tryb FIPS, ustaw to w pliku YAML obciążenia w sekcji env_variables:

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

Alternatywnie ustaw zmienną w pliku Dockerfile, aby była stosowana do każdego obciążenia, które używa obrazu:

ENV OPENSSL_FORCE_FIPS_MODE=0

Prześlij zadanie ponownie i upewnij się, że błąd SSL już się nie pojawia podczas ładowania zależności.

Dodatkowe zasoby

  • Ostatnia aktualizacja o