メインコンテンツまでスキップ
に最終更新

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

備考

ベータ版

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

Docker Container Services では、airワークロードに独自のDockerコンテナイメージを持ち込むことができます。必要な場合は、カスタムイメージを使用します。

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

前提条件​

イメージを登録する​

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

パブリックイメージ​

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

Shell
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ログインの使用（インタラクティブな使用におすすめ）​

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

Shell
docker login

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

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

インタラクティブ認証を使用する​

Databricks シークレットスコープで資格情報を 1 つのステップで認証して保存する：

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

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

Databricksに事前に保存されたシークレットの使用（CI/スクリプトに推奨）​

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

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

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

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

YAML
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 にパッケージをインストールしてください。

ワークロードを投入してください：

Shell
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調整ポート（マルチノードのみ）。

例​

シングルノードA10​

YAML
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ベースイメージのいずれかに基づいてイメージを作成してください。

YAML
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）が事前に構成されたベースイメージを Docker Hub のdatabricksruntime/airで公開しています。

タグ

クラウド

バリアント

使用する場合

dcs-base-aws-runtime

AWS

ランタイム

ビルド済みホイールのみをインストールする

dcs-base-aws-devel

AWS

開発

CUDA拡張機能のコンパイル（nvccが必要）

dcs-base-azure-runtime

Azure

ランタイム

ビルド済みホイールのみをインストールする

dcs-base-azure-devel

Azure

開発

CUDA拡張機能のコンパイル（nvccが必要）

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

Databricks ベースイメージに PyTorch を追加する Dockerfile の例。ベースイメージは、uvによって管理される Python を /opt/venv で提供します。uv pip install はデフォルトでその環境をターゲットとします；別の環境を使用するには、uv pip install を実行する前に venv を作成してアクティブ化してください。

Dockerfile
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

ビルド、プッシュ、登録する。

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

要件​

  • イメージはDocker Hubでホストされている必要があります。Amazon ECR、Google GCR、およびGitHub GHCRはサポートされていません。
  • イメージサイズは20 GB未満である必要があります。
  • WORKDIR ランタイムでは認識されません。イメージに組み込まれたファイルには絶対パスを使用してください。たとえば、python train.pyではなくpython /app/train.pyを使用します。
  • environment.docker_image.url では environment.dependencies または environment.version を使用できません。イメージに含まれていない追加のパッケージが必要な場合は、Dockerfile に追加する必要があります。

関連項目​

このページの見出し