AIランタイムでカスタムDockerイメージを使用する
この機能はベータ版です。使用するには、ワークスペース管理者がワークスペースの プレビューページから AI ランタイム ベータ機能 および Databricks アーティファクト Registry プレビューを有効にする必要があります。
AI ランタイムは、アーティファクトレジストリに保存されているカスタムDockerコンテナイメージを実行できます。カスタムイメージが必要な場合は、それを使用します。
- システムライブラリ、または
environment.dependenciesでインストールできない複雑な依存関係。 - 開発、調査、本番運用環境全体で再現可能な環境。
- プラットフォームまたはセキュリティチームによって構築された、組織で承認済みのプライベートイメージ。
前提条件
- Databricks CLIバージョン1.19.0以降をインストールまたはアップデートします。Databricks CLI をインストールまたは更新するを参照してください。
- ワークスペース管理者に依頼して、 AIランタイム ベータ版機能 のプレビューを有効にしてもらってください。手順については、「 ワークスペース レベルのプレビューを管理する」を参照してください。
- 同じワークスペースに アーティファクトレジストリ を設定し、イメージをプッシュおよび読み取るための権限を取得します。「アーティファクトレジストリの使用開始」を参照してください。
- ローカルマシンにDockerをインストールして起動します。
アーティファクトレジストリにイメージをプッシュする
AI ランタイムでカスタムイメージを使用する前に、ワークロードの送信に使用するのと同じワークスペース内のアーティファクトレジストリに保存してください。
-
画像用のUnity Catalogカタログとスキーマを作成または選択します。
-
アーティファクトレジストリの入門に従って、Docker 認証を構成し、必要な権限を付与し、イメージをプッシュします。
-
画像は次の形式で Unity Catalog 名に注意してください：
Text<catalog>.<schema>.<image>:<tag>たとえば、
main.ml.training:v1などです。ワークロード構成にレジストリホスト名を含めないでください。
代わりに、Databricks CLI の databricks air images push ヘルパーコマンドを使用します。
ワークロードで 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:ノードごとのGPUWORLD_SIZE：プロセスの総数。POD_RANK：現在のノードランク（0始まり）。また、NODE_RANKとして注入されます。LOCAL_ADDR: ローカルノードIP (マルチノードのみ)。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}" \
/app/train.py
独自のイメージを構築する
独自のイメージを構築する場合、Databricksでは、コーディングエージェントと一緒にdatabricks-ai-runtimeスキルを使用するか、Databricksベースイメージから開始することをお勧めします。
コーディングエージェントの利用
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
画像をローカルでビルドします。
docker build --platform linux/amd64 -t my-training-image:v1 .
次に、アーティファクトレジストリの起動を参照して、イメージにタグを付け、アーティファクト Registryにプッシュします。生成された<catalog>.<schema>.<image>:<tag>の名前を、ワークロードYAML内でenvironment.unity_catalog_imageとして使用します。
代わりに、Databricks CLIでdatabricks air images pushヘルパーコマンドを使用して、対話型プロンプトに従うことができます。
制限
- : イメージは、ワークロードを送信するワークスペースのアーティファクトレジストリに保存する必要があります。
- 画像サイズは20 GB未満である必要があります。
WORKDIR: は実行時に適用されません。: イメージに焼き込まれたファイルには絶対パスを使用します。: たとえば、python train.pyではなくpython /app/train.pyを使用します。environment.unity_catalog_imageでenvironment.dependenciesまたはenvironment.versionを使用することはできません。イメージに含まれている以外の追加パッケージが必要な場合は、Dockerfileに追加する必要があります。
トラブルシューティング
レジストリ関連の認証、許可、イメージプッシュ、またはイメージ探索のエラーについては、アーティファクトレジストリのトラブルシューティングを参照してください。
依存アイテムのロード中に 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 エラーが発生しなくなったことを確認してください。