カスタム Docker イメージを使用する

Important

AI ランタイム CLI ワークロード用のカスタム Docker イメージはベータ版です。

Docker Container Services (DCS) を使用すると、独自の Docker コンテナーイメージを air ワークロードに取り込みます。必要な場合は、カスタムイメージを使用します。

特定のシステムライブラリバージョン。
environment.dependenciesに完全には収まらない複雑な依存関係。
研究成果を再現するための正確な環境。
組織のプラットフォームまたはセキュリティチームによって構築された標準イメージ。

前提条件

AI ランタイム CLI をインストールします。
プライベートイメージの場合、イメージにアクセスできるDocker Hub アカウント。

イメージを登録する

カスタムイメージを使用してワークロードを実行する前に、 air register imageに登録します。登録により、Databricks プラットフォームでイメージがプルされ、キャッシュされます。各ユーザーは、イメージタグごとに 1 回イメージを登録する必要があります。新しいタグをプッシュするか、資格情報をローテーションする場合にのみ再登録します。登録は2～6分かかり、イメージの準備が完了するまで待機します。

パブリックイメージ

Docker イメージの URL と Databricks プロファイルを指定して、パブリックイメージを登録します。

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

短縮形式のイメージ参照も使用できます。たとえば、「 library/ubuntu:latest 」のように入力します。

プライベート Docker Hub イメージ

プライベート Docker Hub イメージを登録するには、まず個人用アクセストークンを生成します。 Docker Hub アカウントの設定で、[個人用アクセストークン] → [新しいトークンの生成] をクリックします。読み取り専用アクセスで十分です。

次のいずれかの認証方法を選択します。

ターミナルでDocker Hubにログインします。 Docker Hubユーザー名と個人用アクセストークンの入力を求められます。

docker login

これにより、資格情報が ~/.docker/config.jsonに格納されます。次に、イメージを登録します。 air は資格情報を自動的に読み取ります。

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

対話型認証の使用

1 つの手順で資格情報を認証し、Databricks シークレットスコープに格納します。

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

Docker Hubユーザー名と個人用アクセストークンの入力を求められます。資格情報は、今後の登録のためにワークスペースシークレットスコープに格納されます。

事前に格納された Databricks シークレットの使用 (CI/スクリプトに推奨)

資格情報を Databricks シークレットに格納し、それを直接参照します。

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

ワークロードで Docker イメージを使用する

environment.docker_image.urlでワークロード YAML で Docker イメージを指定します。

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

独自の Docker イメージを取り込む場合、 environment.dependencies と environment.version はサポートされません。いずれかのフィールドで environment.docker_image.url を指定すると、エラーがトリガーされます。追加の依存関係がある場合は、代わりに Dockerfile にパッケージをインストールします。

ワークロードを送信します。

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

コンテナーに挿入された環境変数

AI ランタイムは、実行時にすべてのコンテナーに次の環境変数を挿入します。

NUM_NODES — ノードの合計数。
LOCAL_WORLD_SIZE — ノードあたりの GPU。
WORLD_SIZE — プロセスの合計数。
POD_RANK — 現在のノードランク (0 インデックス付き)。また、 NODE_RANKとして注入されます。
LOCAL_ADDR — ローカルノード IP (マルチノードのみ)。
MASTER_ADDR — ランク 0 の調整アドレス (マルチノードのみ)。
MASTER_PORT — rank-0 調整ポート (マルチノードのみ)。

Examples

単一ノード 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

RDMA を使用したマルチノード H100

AWS p5 インスタンスで完全なネットワーク帯域幅を必要とするマルチノード H100 ジョブの場合は、NCCL と EFA が事前に構成された Databricks 基本イメージのいずれかにイメージを基にします。

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

独自のイメージをビルドする

Databricks の基本イメージ

Databricks は、CUDA、NCCL、クラウド固有のネットワーク (AWS EFA または Azure InfiniBand) が事前に構成されたdatabricksruntime/airで、Docker Hubにベースイメージを発行します。

タグ	雲	Variant	次の場合に使用します。
`dcs-base-aws-runtime`	AWS	Runtime	事前に構築されたホイールのみをインストールする
`dcs-base-aws-devel`	AWS	Devel	CUDA 拡張機能のコンパイル ( `nvcc`が必要)
`dcs-base-azure-runtime`	Azure	Runtime	事前に構築されたホイールのみをインストールする
`dcs-base-azure-devel`	Azure	開発	CUDA 拡張機能のコンパイル ( `nvcc`が必要)

Dockerfile が flash-attn、apex、カスタムカーネルなどの CUDA 拡張機能をコンパイルしない限り、 ランタイム バリアントを使用します。

Databricks 基本イメージに PyTorch を追加する Dockerfile の例。基本イメージは、/opt/venvによって管理されるuvでPythonを提供します。 uv pip install 既定では、その環境を対象とします。別の環境を使用するには、 uv pip installを実行する前に venv を作成してアクティブ化します。

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

ビルド、プッシュ、および登録:

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

Requirements

イメージはDocker Hubでホストする必要があります。 Amazon ECR、Google GCR、および GitHub GHCR はサポートされていません。
イメージサイズは 20 GB 以下である必要があります。
WORKDIR は実行時に受け入れられない。イメージにベイクされたファイルには絶対パスを使用します。たとえば、python /app/train.pyではなく python train.pyを使用します。
environment.dependenciesでenvironment.versionまたはenvironment.docker_image.urlを使用することはできません。イメージ内のパッケージ以外の追加パッケージが必要な場合は、それらを Dockerfile に追加する必要があります。

こちらも参照ください

フィードバック

このページはお役に立ちましたか?

Last updated on 2026-06-29