AI RuntimeでカスタムDockerイメージを使いましょう

Important

この機能は ベータ版です。 これを使用するには、ワークスペース管理者がワークスペースプレビューページからAIランタイムベータ機能とDatabricksアーティファクトレジストリのプレビューを有効にしなければなりません。

AIランタイムはArtifact Registryに保存されたカスタムDockerコンテナイメージを実行できます。 必要な場合は、カスタム イメージを使用します。

  • システムライブラリや environment.dependenciesではインストールできない複雑な依存関係。
  • 開発、研究、生産にまたがる再現可能な環境。
  • 貴社のプラットフォーム チームまたはセキュリティ チームが構築した、組織で承認された非公開のイメージ。

前提条件

画像をアーティファクトレジストリにプッシュする

AI Runtimeでカスタムイメージを使う前に、ワークロードを提出したのと同じワークスペース内のArtifact Registryに保存してください。

  1. 画像用のUnityカタログカタログとスキーマを作成または選択してください。

  2. 「 Artifact Registryで始めて 」に従い、Docker認証を設定し、必要な権限を付与し、イメージをプッシュします。

  3. 画像のUnity Catalog名を以下の形式で確認してください:

    <catalog>.<schema>.<image>:<tag>

    たとえば、「 main.ml.training:v1 」のように入力します。 ワークロード設定にはレジストリホスト名を含めないでください。

Tip

または、Databricks CLIの databricks air images push ヘルパーコマンド を使うこともできます。

ワークロードでDockerイメージを使う

ワークロードのYAMLで画像のUnity Catalog名を environment.unity_catalog_image以下に指定してください:

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

この例では、画像内でtrain.pyから/appを実行します。 イメージを再構築せずにアプリケーションコードを別途アップロードするには、「 アップロードされたアプリケーションコードを実行する」をご覧ください。

独自の Docker イメージを取り込む場合、 environment.dependencies と environment.version はサポートされません。 いずれかのフィールドで environment.unity_catalog_image を指定すると、エラーがトリガーされます。 追加の依存関係がある場合は、代わりに Dockerfile にパッケージをインストールします。

ワークロードを送信します。

databricks air run -f workload.yaml -p my-databricks-profile

プロファイルは画像が保存されている同じワークスペースに認証しなければなりません。

コンテナーに挿入された環境変数

AI ランタイムは、実行時にすべてのコンテナーに次の環境変数を挿入します。

  • CODE_SOURCE_PATH: code_source が設定されたときにアップロードされたアプリケーションコードへのパス。
  • NUM_NODES:ノードの総数。
  • LOCAL_WORLD_SIZE:ノードあたりのGPU数。
  • WORLD_SIZE:プロセスの総数。
  • POD_RANK: 現在のノードランク(0インデックス)。 また、 NODE_RANKとして注入されます。
  • LOCAL_ADDR:ローカルノードIP(マルチノードのみ)。
  • MASTER_ADDR: ランク0の調整アドレス(マルチノードのみ)。
  • MASTER_PORT: ランク0のコーディネーションポート(マルチノードのみ)。

Examples

以下の例は、アップロードされたアプリケーションコードとカスタムイメージで分散トレーニングを実行する方法を示しています。

アップロードしたアプリケーションコードを実行する

code_sourceを使って画像とは別にアプリケーションコードをアップロードしてください。 画像を再構築せずにコードを編集して再提出できます。 コードのPythonおよびシステム依存関係を画像にインストールしてください。

train.py workload.yamlの隣にあるローカルsrcディレクトリに置いてください。 以下の構成は src をアップロードし、カスタムイメージ内で 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 Runtimeはアップロードされたディレクトリのコンテナ内のパスに $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 基本イメージから開始することをお勧めします。

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

ゼロからのビルド、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にベース イメージを発行します。

タグ Variant CUDA 次の場合に使用します。
dcs-base-azure-runtime Runtime 12 事前に構築されたホイールのみをインストールする
dcs-base-azure-devel Devel 12 CUDA 拡張機能のコンパイル ( nvccが必要)

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

画像にトレーニングスクリプトを含めるには、Dockerファイルの横に train.py を配置してください。 Dockerfileはそれを /app/train.pyにコピーします。 code_sourceでアプリケーションコードをアップロードする場合は、COPY指示を省略し、ソースディレクトリにtrain.pyを残してください。

FROM databricksruntime/air:dcs-base-azure-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 .

次に「 Artifact Registryで始めて 」に従い、画像をタグ付けしてArtifact Registryにプッシュしてください。 その結果得られた <catalog>.<schema>.<image>:<tag> 名をワークロードYAMLの environment.unity_catalog_image として使います。

Tip

または、Databricks CLIの databricks air images push ヘルパーコマンド を使い、インタラクティブなプロンプトに従ってください。

Limitations

  • 画像は、作業を提出するワークスペースのArtifact Registryに保存する必要があります。
  • イメージ サイズは 20 GB 以下である必要があります。
  • WORKDIR は実行時に受け入れられない。 イメージにベイクされたファイルには絶対パスを使用します。 たとえば、python /app/train.pyではなく python train.pyを使用します。
  • environment.dependenciesでenvironment.versionまたはenvironment.unity_catalog_imageを使用することはできません。 イメージ内のパッケージ以外の追加パッケージが必要な場合は、それらを Dockerfile に追加する必要があります。

Troubleshooting

レジストリ関連の認証、許可、イメージプッシュ、または画像発見エラーについては 、「アーティファクトレジストリのトラブルシューティング」を参照してください。

依存関係の読み込み時の ssl.SSLError

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

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

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

これは、AIランタイムワークロードが連邦情報処理基準(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 エラーが表示されなくなったのを確認します。

その他のリソース

  • 最終更新: