カスタム 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 ログインの使用 (対話型使用に推奨)

ターミナルで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-ai-runtime スキルを使用するか、Databricks 基本イメージから開始することをお勧めします。

コーディング エージェントを使用する

ゼロからのビルド、CUDA/NCCL/EFA の互換性、一般的な問題、ビルド前チェックリストなど、Dockerfile のステップ バイ ステップ ガイダンス用 の databricks-ai-runtime Claude Code スキルをインストールします。 このスキルには、Databricks CLI バージョン 1.0.0 以降が必要です。

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

Databricks の基本イメージ

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

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 に追加する必要があります。

Troubleshooting

ssl.SSLError: [CRYPTO] 依存関係の読み込み時に不明なエラー (_ssl.c)

ライブラリが SSL コンテキストを作成しようとすると、実行時に OpenSSL エラーが発生してカスタム イメージが失敗する場合があります。次に例を示します。

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

このエラーは、 huggingface_hubなどのネットワーク接続を開き、読み込みを妨げるライブラリのインポート中に表示されます。

これは、 air ワークロードが FIPS 対応ホストで実行されるために発生します。 イメージの暗号化ライブラリが FIPS に準拠していない場合、OpenSSL は FIPS モードで初期化に失敗するため、SSL コンテキストの作成は失敗します。

推奨される解決策:

エンタープライズ、政府機関、医療、および財務のワークロードは、多くの場合、FedRAMP、CMMC、または HIPAA 監査の FIPS 140-2 または 140-3 コンプライアンスに依存します。 ワークロードが FIPS に準拠している必要がある場合は、FIPS 準拠の暗号化ライブラリを使用してイメージを構築します。

ワークロードで FIPS 準拠が必要ない場合は、 OPENSSL_FORCE_FIPS_MODE 環境変数を 0 に設定して FIPS モードを無効にすることができます。 これを行うと、コンプライアンス要件が自動的に損なわれます。

FIPS モードを無効にするには、 env_variablesのワークロード YAML で設定します。

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

または、イメージを使用するすべてのワークロードに適用されるように、Dockerfile に変数を設定します。

ENV OPENSSL_FORCE_FIPS_MODE=0

ワークロードを再送信し、依存関係の読み込み時に SSL エラーが表示されなくなったのを確認します。

こちらも参照ください

• AI ランタイム CLI のクイック スタート
• ワークロード YAML リファレンス
• AI ランタイム CLI コマンド リファレンス