AIランタイムでカスタムDockerイメージを使用する
プレビュー
この機能は パブリック プレビュー段階です。使用するには、ワークスペース管理者がワークスペースの プレビューページから Databricks アーティファクト Registry プレビューを有効にする必要があります。
AI ランタイムは、マネージド環境の代わりにカスタム Docker コンテナイメージでワークロードを実行できます。linux/amd64のイメージをビルドし、ワークスペースのアーティファクトレジストリにプッシュし、Databricks CLI で送信するワークロード YAML でその Unity Catalog 名を参照します。必要な場合は、カスタムイメージを使用します。
- システムライブラリ、または
environment.dependenciesでインストールできない複雑な依存関係。 - flash-attn やカスタムカーネルなど、ソースからコンパイルする CUDA 拡張機能。
- 開発、調査、本番運用環境全体で再現可能な環境。
- プラットフォームまたはセキュリティチームによって構築された、組織で承認済みのプライベートイメージ。
前提条件
- Databricks CLIバージョン1.19.0以降をインストールまたはアップデートします。Databricks CLI をインストールまたは更新するを参照してください。
- ワークスペース管理者に依頼して、 Databricks アーティファクト Registry プレビューを有効にしてもらってください。手順については、 ワークスペース レベルのプレビューを管理するを参照してください。
- 同じワークスペースに アーティファクトレジストリ を設定し、イメージをプッシュおよび読み取るための権限を取得します。「アーティファクトレジストリの使用開始」を参照してください。
- ローカルマシンにDockerをインストールして起動します。
アーティファクトレジストリにイメージをプッシュする
AI ランタイムでカスタムイメージを使用する前に、ワークロードを送信するために使用するのと同じワークスペースのアーティファクト Registry にイメージを保存します。最初にイメージをビルドするには、「独自のイメージをビルドする」を参照してください。
-
画像の Unity Catalog カタログとスキーマを作成または選択し、必要な アーティファクト Registry 特権を取得します。
-
ワークスペースのワークスペース OAuth プロファイルを作成または更新します。
Shelldatabricks auth login --profile my-databricks-profileDocker 認証情報ヘルパーには、
databricks auth loginによって作成されたプロファイルが必要です。パーソナルアクセストークンや OAuth マシン間 (M2M) 認証情報を使用するプロファイルはサポートしていません。 -
Docker を認証し、イメージにタグを付け、1 ステップでプッシュするには、
databricks air images pushヘルパーを実行します。Shelldatabricks air images push \--profile my-databricks-profile \--source my-training-image:v1 \--catalog main \--schema ml \--artifact training:v1別のレジストリからのイメージのコピーなど、その他のオプションについては、「Databricks CLI を使用したイメージのプッシュ」を参照してください。
代わりに
docker pushでプッシュするには、databricks auth docker configure --profile my-databricks-profileをランして Docker 認証情報ヘルパーを構成し、アーティファクトレジストリガイドの「イメージをプッシュする」に従います。 -
画像は次の形式で Unity Catalog 名に注意してください：
Text<catalog>.<schema>.<image>:<tag>たとえば、
main.ml.training:v1などです。ワークロード構成にレジストリホスト名を含めないでください。
ワークロードで Dockerイメージを使用する
ワークロードYAMLの environment.unity_catalog_image に画像のUnity Catalog名を指定します。
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
この例では、イメージ内の /app から train.py をランします。イメージを再構築せずにアプリケーションコードを個別にuploadするには、uploadされたアプリケーションコードのランを参照してください。
独自の Dockerイメージ を持ち込む場合、 environment.dependencies および environment.version はサポートされません。両方のフィールドに environment.unity_catalog_image を指定すると、エラーが Trigger されます。追加の依存関係がある場合は、代わりに Dockerfile にパッケージをインストールしてください。
ワークロードを送信します：
databricks air run -f workload.yaml -p my-databricks-profile
プロファイルは、イメージが保存されているワークスペースと同じワークスペースに対して認証する必要があります。
コンテナに挿入される環境変数
AIランタイムは、実行時にすべてのコンテナに次の環境変数を注入します。
-
CODE_SOURCE_PATH:code_sourceが構成されている場合の、uploadアプリケーションコードへのパス。 -
NUM_NODES：ノードの合計数。 -
LOCAL_WORLD_SIZE:ノードごとのGPU -
WORLD_SIZE：プロセスの総数。 -
POD_RANK：現在のノードランク（0始まり）。また、NODE_RANKとして注入されます。 -
MASTER_ADDR: ランク 0 コーディネーションアドレス（マルチノードのみ）。 -
MASTER_PORT:ランク0コーディネーションポート（マルチノードのみ）。
例
次の例は、upload アプリケーションコードとカスタムイメージを使用して分散トレーニングをランする方法を示しています。
uploadされたアプリケーションコードのラン
アプリケーション コードをイメージと別に upload するには、 code_source を使用します。イメージを再ビルドせずにコードを編集して再サブミットできます。コードの Python およびシステム依存関係をイメージにインストールします。
train.pyをworkload.yamlの隣のローカルsrcディレクトリに配置します。次の構成は、 src をuploadし、カスタムイメージ内でその train.py を実行します：
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 workload.yamlを基準として解決します。AIランタイムは、コンテナ内のuploadされたディレクトリのパスに$CODE_SOURCE_PATHを設定します。コードソースオプションについては、code_sourceを参照してください。
RDMA 対応のマルチノード H100
AWS p5 インスタンスで完全なネットワーク帯域幅を必要とするマルチノード H100 ジョブの場合は、NCCL および EFA が事前構成された Databricks ベースイメージのいずれかにイメージのベースを置きます：
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}" \
--local_addr="$(hostname -I | awk '{print $1}')" \
/app/train.py
独自のイメージを構築する
独自のイメージを構築する場合、Databricksでは、コーディングエージェントと一緒にdatabricks-ai-runtimeスキルを使用するか、Databricksベースイメージから開始することをお勧めします。
nvidia/cudaなど、別のベースイメージから起動することもできます。その場合、イメージには、ホストGPUドライバーがサポートするCUDAランタイムと、マルチGPU通信用のNCCLを含める必要があります。databricks-ai-runtimeスキルは、ゼロからのイメージの構築を対象としています。
マルチノードGPU_8xH100ワークロードの場合、別のベースイメージからビルドされたイメージにも EFA を含める必要があります。
コーディングエージェントの利用
databricks-AI-ランタイム Claude Code スキルをインストールすると、Dockerfile に関するステップバイステップのガイダンス（ゼロからの構築、CUDA/NCCL/EFA の互換性、一般的な問題、ビルド前のチェックリストなど）が得られます。このスキルには、Databricks CLI バージョン 1.0.0 以降が必要です。
databricks aitools install --skills databricks-ai-runtime
Databricksベース画像
DatabricksはDocker Hub（databricksruntime/air）でベースイメージを公開しており、CUDA、NCCL、クラウド固有のネットワーク（AWS EFAまたはAzure InfiniBand）があらかじめ構成されています。
タグ | バリアント | CUDA | 次の時に使用 |
|---|---|---|---|
| ランタイム | 12 | ビルド済みの wheel のインストールのみ |
| 開発用 | 12 | CUDA 拡張機能のコンパイル ( |
| ランタイム | 13 | CUDA 13でのビルド済みwheelのみのインストール |
| 開発用 | 13 | CUDA 13でのCUDA拡張機能のコンパイル（ |
次のDockerfileは、DatabricksベースイメージにPyTorchを追加します。ベースイメージは、 uvで管理される /opt/venvで Python を提供します。uv pip install は、defaultでその環境をターゲットにします。異なる環境を使用するには、 uv pip install を実行する前に venv を作成してアクティブ化します。
トレーニングスクリプトをイメージに含めるには、 train.py を Dockerfile の横に配置します。Dockerfile はそれを /app/train.pyにコピーします。code_sourceでアプリケーションコードをuploadする場合は、 COPY 命令を省略し、代わりに train.py をソースディレクトリに保持してください。
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.py /app/train.py
イメージはlinux/amd64プラットフォームを使用する必要があります。Apple シリコンを搭載した Mac では、Docker はdefaultで arm64イメージをビルドするため、常に --platform linux/amd64を渡す必要があります。イメージをローカルでビルドするには、次をランします。
docker build --platform linux/amd64 -t my-training-image:v1 .
次に、イメージを アーティファクト Registry にプッシュし、生成された <catalog>.<schema>.<image>:<tag> 名をワークロード YAML の environment.unity_catalog_image として使用します。
制限
- : イメージは、ワークロードを送信するワークスペースのアーティファクトレジストリに保存する必要があります。
- 画像サイズは20 GB未満である必要があります。
WORKDIR: は実行時に適用されません。: イメージに焼き込まれたファイルには絶対パスを使用します。: たとえば、python train.pyではなくpython /app/train.pyを使用します。environment.unity_catalog_imageでenvironment.dependenciesまたはenvironment.versionを使用することはできません。イメージに含まれている以外の追加パッケージが必要な場合は、Dockerfileに追加する必要があります。
トラブルシューティング
レジストリ関連の認証、許可、イメージプッシュ、またはイメージ探索のエラーについては、アーティファクトレジストリのトラブルシューティングを参照してください。
Docker認証情報ヘルパーには、databricks auth login によって作成されたプロファイルが必要です
次のようなエラーでイメージのプッシュが失敗することがあります。
Error: profile "my-profile" uses auth_type "". Docker credential helper requires a profile created by databricks auth login
このエラーは、プロファイルがパーソナルアクセストークンや、Docker 認証情報ヘルパーがサポートしていない別の認証方法を使用している場合に発生します。ワークスペース OAuth プロファイルを作成し、そのプロファイルでプッシュを再実行します。
databricks auth login --profile my-databricks-profile
複数のプロファイルがワークスペース ID と一致します
Docker認証の設定またはイメージのプッシュが、次のようなエラーで失敗する場合があります。
Error: multiple Databricks profiles match workspace ID 1111111111111111: profile-a and profile-b. Remove duplicate workspace_id entries before using Docker credential helper
このエラーは、設定ファイル（default では ~/.databrickscfg）の複数のプロファイルが同じワークスペースを指している場合に発生します。Dockerに使用しないプロファイルを削除し、残りのプロファイルでコマンドを再実行してください。プロファイルからログアウトして設定ファイルから削除するには、次をランします。
databricks auth logout --profile profile-a --delete
すべてのログアウトオプションについては、「auth コマンドグループ」を参照してください。
依存アイテムのロード中に ssl.SSLError が発生しました
ライブラリがSSLコンテキストを作成しようとしたときに、カスタムイメージが実行時にOpenSSLエラーで失敗することがあります。たとえば、次のようになります。
ssl.SSLError: [CRYPTO] unknown error (_ssl.c:3076)
huggingface_hubなど、ネットワーク接続を開くライブラリをインポート中にエラーが発生し、ロードが妨げられます。
これは、AI ランタイムのワークロードが、連邦情報処理標準 (FIPS) が有効になっているホストで実行されるために発生します。画像の暗号化ライブラリがFIPS準拠でない場合、OpenSSLはFIPSモードで初期化できず、SSLコンテキストの作成に失敗します。
推奨ソリューション：
企業、政府、ヘルスケア、および財務のワークロードは、FedRAMP、CMMC、または HIPPA の監査のために 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 エラーが発生しなくなったことを確認してください。