Použití vlastních imagí Dockeru

Důležité

Vlastní obrazy Dockeru pro úlohy AI Runtime CLI jsou ve fázi beta.

Docker Container Services (DCS) umožňuje přenést do úloh vlastní image air kontejneru Dockeru. Pokud potřebujete, použijte vlastní image:

  • Konkrétní verze systémových knihoven.
  • Složité závislosti, které se nevejdou čistě do environment.dependencies.
  • Přesné prostředí pro reprodukování výsledků výzkumu.
  • Standardní obrazy vytvořené platformním nebo bezpečnostním týmem vaší organizace.

Předpoklady

Zaregistrovat obraz

Před spuštěním úlohy s vlastním image jej zaregistrujte pomocí air register image. Registrace načítá a ukládá image do mezipaměti na platformě Databricks. Každý uživatel musí zaregistrovat image pouze jednou pro každý tag image. Znovu se zaregistrujte pouze když nahrajete nový tag nebo obměníte přihlašovací údaje. Registrace trvá 2 až 6 minut a blokuje, dokud image není připravená.

Veřejné obrazy

Zaregistrujte veřejné image zadáním URL adresy image Dockeru a svého profilu Databricks:

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

Funguje také zkrácený tvar odkazu na image. Například: library/ubuntu:latest.

Soukromé obrazy Docker Hub

Pokud chcete zaregistrovat privátní Docker Hub image, vygenerujte nejprve osobní přístupový token. V nastavení účtu Docker Hub klikněte na osobní přístupové tokenyVygenerovat nový token. Stačí přístup jen pro čtení.

Zvolte jednu z následujících metod ověřování:

Přihlaste se k Docker Hub v terminálu. Zobrazí se výzva k zadání uživatelského jména Docker Hub a osobního přístupového tokenu:

docker login

Tím se vaše přihlašovací údaje uloží do ~/.docker/config.json. Pak zaregistrujte image – air přečte přihlašovací údaje automaticky:

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

Použití interaktivního ověřování

Ověřte a uložte přihlašovací údaje v Databricks secret scope v jednom kroku:

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

Budete vyzváni k zadání uživatelského jména služby Docker Hub a osobního přístupového tokenu. Přihlašovací údaje se ukládají do rozsahu tajných údajů ve vašem pracovním prostoru pro budoucí registrace.

Uložte přihlašovací údaje do tajného kódu Databricks a odkazujte na něj přímo:

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

Použít obraz Dockeru v úloze

V YAML úlohy zadejte image Dockeru v části 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

Při použití vlastního obrazu Dockeru nejsou environment.dependencies a environment.version podporovány. Zadání environment.docker_image.url do kteréhokoli z obou polí způsobí chybu. Pokud máte další závislosti, nainstalujte balíčky do souboru Dockerfile.

Odešlete úlohu:

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

Proměnné prostředí vložené do vašeho kontejneru

AI Runtime za běhu vkládá do každého kontejneru následující proměnné prostředí:

  • NUM_NODES — celkový počet uzlů.
  • LOCAL_WORLD_SIZE — GPU na uzel.
  • WORLD_SIZE — celkový počet procesů.
  • POD_RANK — pořadí aktuálního uzlu (číslováno od 0). Injektuje se také jako NODE_RANK.
  • LOCAL_ADDR — IP adresa místního uzlu (pouze s více uzly).
  • MASTER_ADDR — koordinační adresa rank 0 (pouze pro víceuzlové konfigurace).
  • MASTER_PORT — koordinační port rank-0 (pouze pro víceuzlové konfigurace).

Příklady

Jednouzlový A10

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

Víceuzlový H100 s RDMA

Pro víceuzlové úlohy H100, které na instancích AWS p5 vyžadují plnou šířku síťového pásma, vytvořte image na základě jedné ze základních imagí Databricks s předem nakonfigurovanými technologiemi NCCL a EFA:

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

Vytvoření vlastní image

Základní image Databricks

Databricks zveřejňuje základní obrazy na Docker Hubu na adrese databricksruntime/air, s předem nakonfigurovanými technologiemi CUDA a NCCL a cloudově specifickým síťovým nastavením (AWS EFA nebo Azure InfiniBand).

Tag Cloud Variant Použít, když
dcs-base-aws-runtime AWS Runtime Instalace pouze předdefinovaných kol
dcs-base-aws-devel AWS Devel Kompilace rozšíření CUDA (vyžaduje nvcc)
dcs-base-azure-runtime Azure Runtime Instalace pouze předdefinovaných kol
dcs-base-azure-devel Azure Devel Kompilace rozšíření CUDA (vyžaduje nvcc)

Použijte variantu runtime, pokud v souboru Dockerfile nekompilujete rozšíření CUDA, jako jsou flash-attn, apex nebo vlastní jádra.

Příklad souboru Dockerfile, který přidává PyTorch do základní image Databricks. Základní image poskytují Python ve /opt/venvsprávě uv. uv pip install cílí na toto prostředí ve výchozím nastavení; pokud chcete použít jiné prostředí, vytvořte a aktivujte venv před spuštěním 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

Sestavte, nahrajte a zaregistrujte:

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

Requirements

  • Image musí být hostované na Docker Hub. Amazon ECR, Google GCR a GitHub GHCR se nepodporují.
  • Velikost obrázku musí být pod 20 GB.
  • WORKDIR není dodrženo za běhu. Pro soubory upečené do obrázku použijte absolutní cesty. Například použijte python /app/train.py, nikoli python train.py.
  • Nelze použít environment.dependencies ani environment.version s environment.docker_image.url. Pokud potřebujete další balíčky, které nejsou součástí obrazu, musíte je přidat do souboru Dockerfile.

Viz také