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

AI ランタイムでのカスタム Dockerイメージの使用

備考

この機能はベータ版です。これを使用するには、ワークスペース管理者がワークスペースのプレビューページ から AI Runtime Beta Features および Databricks Artifact Registry のプレビューを有効にする必要があります。

AI ランタイム は、アーティファクト Registryに保存されているカスタム Docker コンテナーイメージを実行できます。必要な場合は、カスタム画像を使用します。

  • environment.dependenciesではインストールできないシステムライブラリまたは複雑な依存関係。
  • 開発、リサーチ、本番運用を通じて再現可能な環境。
  • プラットフォームまたはセキュリティチームによって構築された、プライベートで組織承認済みのイメージ。

前提条件​

  • Databricks CLI バージョン 1.19.0 以降をインストールまたはアップデートします。Databricks CLI のインストールまたは更新を参照してください。
  • ワークスペース管理者に依頼して、 AI ランタイム ベータ版機能（AI Runtime Beta Features） のプレビューを有効にしてもらってください。手順については、ワークスペースレベルのプレビューを管理するを参照してください。
  • 同じワークスペースで アーティファクト Registry を設定し、イメージのプッシュおよび読み取りの権限を取得します。Artifact Registry の使用を開始するを参照してください。
  • ローカルマシンにDockerをインストールして起動します。

アーティファクト Registryに画像をプッシュする​

AI Runtimeでカスタムイメージを使用する前に、ワークロードの送信に使用するワークスペースと同じワークスペース内のArtifact Registryに保存します。

  1. イメージ用の Unity Catalog カタログとスキーマを作成または選択します。

  2. Docker認証の設定、必要な権限の付与、およびイメージのプッシュを行うには、アーティファクト Registry の起動に従ってください。

  3. 次の形式で画像の 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 名を指定します。

YAML
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するには、Run uploaded application codeを参照してください。

独自の Docker イメージを使用する場合、environment.dependencies および environment.version はサポートされていません。 environment.unity_catalog_image をどちらかのフィールドで指定すると、エラーが発生します。追加の依存関係がある場合は、代わりに Dockerfile にパッケージをインストールしてください。

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

Shell
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としてもインジェクトされます。
  • LOCAL_ADDR：ローカルノードIP（マルチノードのみ）。
  • MASTER_ADDR：ランク 0 コーディネーションアドレス（マルチノードのみ）。
  • MASTER_PORT：ランク 0 コーディネーションポート（マルチノードのみ）。

例​

次の例は、uploadアプリケーションコードとカスタムイメージを使用した分散トレーニングをランする方法を示しています。

アップロードされたアプリケーションコードのラン​

code_sourceを使用して、イメージとは別にアプリケーションコードをuploadします。イメージを再ビルドせずに、コードを編集して再送信できます。イメージ内のコードの Python およびシステム依存関係をインストールします。

workload.yamlの隣のローカルのsrcディレクトリにtrain.pyを配置します。次の構成では、srcをuploadし、カスタムイメージ内でそのtrain.pyをランします。

YAML
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に設定します。コードソースのオプションについては、コードソースを参照してください。

RDMAを備えたマルチノードH100​

AWS p5インスタンスでフルネットワーク帯域幅を必要とするマルチノードH100ジョブの場合、NCCLおよびEFAが事前に構成されたDatabricksベースイメージのいずれかに基づいてイメージを作成してください。

YAML
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 ベースイメージから開始することを推奨しています。

コーディングエージェントを使用してください​

ステップバイステップの Dockerfile ガイダンス（ゼロからの構築、CUDA/NCCL/EFA の互換性、一般的な問題、および事前構築チェックリストを含む）については、databricks-ai-runtime Claude Code スキルをインストールしてください。このスキルには、Databricks CLI バージョン 1.0.0 以降が必要です。

Shell
databricks aitools install --skills databricks-ai-runtime

Databricksベースイメージ​

Databricks は、CUDA、NCCL、およびクラウド固有のネットワーク（AWS EFA または Azure InfiniBand）が事前に構成されたベースイメージを Docker Hub のdatabricksruntime/airで公開しています。

タグ

バリアント

CUDA

使用する場合

dcs-base-aws-runtime

ランタイム

12

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

dcs-base-aws-devel

開発

12

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

dcs-base-aws-runtime-cu13

ランタイム

13

CUDA 13 上でのビルド済み wheel のみのインストール

dcs-base-aws-devel-cu13

開発

13

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

タグ

バリアント

CUDA

使用する場合

dcs-base-aws-runtime

ランタイム

12

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

dcs-base-aws-devel

開発

12

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

dcs-base-aws-runtime-cu13

ランタイム

13

CUDA 13 上でのビルド済み wheel のみのインストール

dcs-base-aws-devel-cu13

開発

13

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

次のDockerfileは、DatabricksのベースイメージにPyTorchを追加します。ベースイメージはuvによって管理されるPythonを/opt/venvに提供します。uv pip installはdefaultでその環境をターゲットにします。別の環境を使用するには、uv pip installを実行する前にvenvを作成して有効化してください。

トレーニングスクリプトをイメージに含めるには、train.pyをDockerfileの隣に配置します。Dockerfileは、それを/app/train.pyにコピーします。code_sourceを使用してアプリケーションコードをuploadする場合は、COPYの命令を省略し、代わりにtrain.pyをソースディレクトリに保持します。

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.py /app/train.py

イメージをローカルでビルドします：

Shell
docker build --platform linux/amd64 -t my-training-image:v1 .

続いてArtifact Registry の使用を開始するに従ってイメージにタグ付けし、Artifact Registry にプッシュします。生成された <catalog>.<schema>.<image>:<tag> 名をワークロード YAML の environment.unity_catalog_image として使用します。

ヒント

または、Databricks CLI で databricks air images push ヘルパーコマンドを使用し、インタラクティブなプロンプトに従います。

制限事項​

  • イメージは、ワークロードを送信するワークスペース内のArtifact Registryに保存する必要があります。
  • イメージサイズは20 GB未満である必要があります。
  • WORKDIR ランタイムでは認識されません。イメージに組み込まれたファイルには絶対パスを使用してください。たとえば、python train.pyではなくpython /app/train.pyを使用します。
  • environment.unity_catalog_image では environment.dependencies または environment.version を使用できません。イメージに含まれていない追加のパッケージが必要な場合は、Dockerfile に追加する必要があります。

トラブルシューティング​

レジストリ関連の認証、権限、イメージのプッシュ、またはイメージの検出に関するエラーについては、アーティファクトレジストリのトラブルシューティングを参照してください。

依存関係のロード時の ssl.SSLError​

カスタムイメージは、ライブラリがSSLコンテキストを作成しようとしたときに、ランタイムにOpenSSLエラーで失敗する可能性があります（例：）。

Text
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モードを無効にするには、ワークロードYAMLのenv_variablesの下で設定します。

YAML
env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

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

Dockerfile
ENV OPENSSL_FORCE_FIPS_MODE=0

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

その他のリソース​

このページの見出し