AI ランタイム 上の Ray

注記

Databricks Runtime 機械学習を使用したDatabricksのクラシックコンピュートでRayを使用している場合は、Databricks上のRayを参照してください。

Rayは、Pythonワークロードをスケーリングするためのオープンソースフレームワークです。基盤となるGPUインフラストラクチャをプロビジョニングまたは管理することなく、AI ランタイムでRayを実行できます。

AIランタイムでRayを使用する理由

Ray は、モデルのトレーニング、バッチ推論、データ処理、ハイパーパラメーターチューニングなどの Python ワークロードを分散するためによく使用されます。AI ランタイム上では、Ray は Unity Catalog ボリューム内のデータにアクセスしたり、MLflow でエクスペリメントをトラッキングしたり、本番運用ワークフローの一部として実行したりできます。

AI ランタイム では、Ray Core、 Ray Data、 Ray Train、 Ray Tune などの一般的な Ray ライブラリを使用できます。既存の Ray ワークロードは、ライブラリの標準 APIs を引き続き使用できます。

AIランタイムでのRayの仕組み

すべての Ray アプリケーションは、1 つ以上のノードを含む可能性のある Ray クラスターで実行されます。AI ランタイムはノードをプロビジョニングし、Ray は利用可能な CPU および GPU リソース全体にわたってタスクとアクターをスケジュールします。Ray を起動する方法は、実行インターフェイスによって異なります。

インターフェース	ノード	Rayを起動する
ノートブック	ノートブックにアタッチされた単一のノード	AI ランタイム SDK (databricks.air) から ray_init() を呼び出す
Databricks CLI	1つ以上のノードの固定セット	再利用可能なブートストラップスクリプトを使用して、ノード0のヘッドを起動し、残りのノードをワーカーとして結合します

インターフェース	ノード	Rayを起動する
ノートブック	ノートブックにアタッチされた単一のノード	AI ランタイム SDK (databricks.air) から ray_init() を呼び出す
Databricks CLI	1つ以上のノードの固定セット	再利用可能なブートストラップスクリプトを使用して、ノード0のヘッドを起動し、残りのノードをワーカーとして結合します

ノートブックでRayを使用する

ノートブックがAI ランタイム GPUコンピュートに接続されている場合、Databricks AI ランタイム SDKの ray_init()を使用して、接続されたノードでRayを起動します。

Python

from databricks.air import ray_init

ray_init()

ray_init() 標準 ray.init() API を呼び出し、Databricks ドライバープロキシ経由でアクセスするように Ray ダッシュボードを構成し、ノートブックの出力にダッシュボード URL を出力します。ダッシュボードを使用して、コードのラン中に Ray ジョブ、タスク、アククタ、Logs、およびリソースの使用状況を検査します。

注記

ray_init() 環境バージョン5以降が必要です。バージョン5では、databricks.airからではなくserverless_gpuからインポートします。Databricks AI v6 には Ray が含まれています。Standard v6 を使用する場合は、ray_init()を呼び出す前に Ray をインストールしてください。「環境の設定」を参照してください。

ノートブックの例

例	説明
Ray Core ハローワールド	接続されたコンピュートで非同期 GPU タスクを送信し、Ray ダッシュボードでスケジューリングを検査し、結果を取得します。
Ray Tune を使用した CIFAR-10 ハイパーパラメータチューニング	PyTorch 画像分類器の並列分数 GPU 試行をランし、非同期逐次半減アルゴリズム (ASHA) を使用してパフォーマンスの低い構成を早期に停止します。
Ray DataとvLLMを使用したQwen2.5-32Bバッチ推論	8 つの H100 GPU 上で 8 つの永続的な vLLM レプリカを使用して多言語バッチ推論をランし、結果を Unity Catalog ボリュームに Parquet として保存します。

例	説明
Ray Core ハローワールド	接続されたコンピュートで非同期 GPU タスクを送信し、Ray ダッシュボードでスケジューリングを検査し、結果を取得します。
Ray Tune を使用した CIFAR-10 ハイパーパラメータチューニング	PyTorch 画像分類器の並列分数 GPU 試行をランし、非同期逐次半減アルゴリズム (ASHA) を使用してパフォーマンスの低い構成を早期に停止します。
Ray DataとvLLMを使用したQwen2.5-32Bバッチ推論	8 つの H100 GPU 上で 8 つの永続的な vLLM レプリカを使用して多言語バッチ推論をランし、結果を Unity Catalog ボリュームに Parquet として保存します。

Ray と Databricks CLI を使用する

AIランタイムのDatabricks CLIコマンドは、シングルノードおよびマルチノードのRayワークロードをサポートします。ワークロード構成で、固定 accelerator_type と合計 num_accelerators を指定します。Ray Core の入門レシピには、再利用可能なブートストラップ スクリプトが含まれています。ワークロードの command で、 RAY_ENTRYPOINT を Python ファイルのパスに設定し、ブートストラップスクリプトをランします。ブートストラップスクリプトは次のことを行います。

• ノード 0 で Ray ヘッドを起動し、設定されたエントリポイントを実行します。
• 他のすべてのノードで、Ray ワーカーを起動し、エントリー ポイントの実行中に接続を維持します。
• エントリ ポイントが終了すると、Ray プロセスが停止します。

アプリケーションで、 ray.init(address="auto")を使用してクラスターに接続します。

標準環境では、ワークロードの依存関係に ray または ray[data]、ray[train]、ray[tune] などの必要な追加パッケージを追加します。Databricks AI v6にはRayが含まれています。

Rayの起動

次の最小限の例では、メインの Ray ライブラリを紹介し、同じクラスターブートストラップパターンを共有します。

例	説明
Ray Core	2つのA10 GPUでRayクラスターを起動し、GPUごとに1つのRay Coreタスクをスケジュールしてください。
Ray トレーニングする	GPU ごとに 1 つの Ray Train ワーカーを使用して、合成データで小さな多層パーセプトロンをトレーニングします。
Ray Data	Ray クラスター全体で合成データ処理パイプラインをランします。
Ray Tune	GPUあたり1つのトライアルで合成データ上で小さな多層パーセプトロンをチューニングします。

例	説明
Ray Core	2つのA10 GPUでRayクラスターを起動し、GPUごとに1つのRay Coreタスクをスケジュールしてください。
Ray トレーニングする	GPU ごとに 1 つの Ray Train ワーカーを使用して、合成データで小さな多層パーセプトロンをトレーニングします。
Ray Data	Ray クラスター全体で合成データ処理パイプラインをランします。
Ray Tune	GPUあたり1つのトライアルで合成データ上で小さな多層パーセプトロンをチューニングします。

追加の CLI 例

例	説明
Ray TrainによるResNet-50のファインチューニング	ResNet-50をファインチューニングするには、2つのノードで各A10 GPUに対して分散データ並列ワーカーを1つ使用します。
Ray Train を使用した Qwen2.5-3B のファインチューニング	8x H100 ノードで GPU あたり 1 つの分散データ並列ワーカーを使用して Qwen2.5-3B を微調整します。
Ray Data と vLLM を使用したバッチ推論	Qwen2.5-7B-Instructで応答を生成し、結果をParquetとしてUnity Catalogボリュームに書き込みます。
Ray Tune による LoRA ハイパーパラメータ検索	4 つの A10 GPU 全体で Qwen2.5-0.5B の LoRA 構成を検索し、ASHA スケジューラでパフォーマンスの低いトライアルを停止します。
Ray Tune を使用した XGBoost ハイパーパラメータ検索	4台のA10 GPUでXGBoostハイパーパラメータを検索し、最適な構成をLogし、Unity Catalogにモデルを登録します。

例	説明
Ray TrainによるResNet-50のファインチューニング	ResNet-50をファインチューニングするには、2つのノードで各A10 GPUに対して分散データ並列ワーカーを1つ使用します。
Ray Train を使用した Qwen2.5-3B のファインチューニング	8x H100 ノードで GPU あたり 1 つの分散データ並列ワーカーを使用して Qwen2.5-3B を微調整します。
Ray Data と vLLM を使用したバッチ推論	Qwen2.5-7B-Instructで応答を生成し、結果をParquetとしてUnity Catalogボリュームに書き込みます。
Ray Tune による LoRA ハイパーパラメータ検索	4 つの A10 GPU 全体で Qwen2.5-0.5B の LoRA 構成を検索し、ASHA スケジューラでパフォーマンスの低いトライアルを停止します。
Ray Tune を使用した XGBoost ハイパーパラメータ検索	4台のA10 GPUでXGBoostハイパーパラメータを検索し、最適な構成をLogし、Unity Catalogにモデルを登録します。

データとDatabricks機能の操作

Rayワークロードは、Unity Catalogボリューム内のデータにアクセスし、MLflowでランを記録し、宣言型オートメーションバンドルで定義されたLakeflow Jobsタスクとして実行できます。

Unity Catalogのデータにアクセスする

Unity Catalogテーブル内の表形式データに対しては、 ray.data.read_databricks_tables() を使用してテーブルを読み取るか、 Databricks SQL warehouseを通じてSQLクエリーを実行します。AI RuntimeはローカルSparkセッションを提供しないため、warehouse ID、カタログ、およびスキーマを指定します。Rayが結果を処理する前に、SQLを使用して結合、集計、およびフィルタを実行できます。

ファイルベースおよび非構造化データの場合は、 Unity Catalog ボリュームを使用します。Ray ワーカーは /Volumes/<catalog>/<schema>/<volume>/... パスを使用してファイルにアクセスできます。たとえば、 Ray データは ray.data.read_parquet() を使用して Parquet ファイルを読み取り、 Dataset.write_parquet() を使用して結果を書き込むことができます。

ワークロードの追跡とオーケストレーション

MLflowを使用して、Rayワークロードからパラメーター、メトリクス、アーティファクトを記録します。エクスペリメントの追跡と観察可能性を参照してください。

ワークロードをスケジュールしたり、CPUデータ準備タスクで構成したりするには、Lakeflow Jobsと宣言型オートメーションバンドルを使用します。GPUワークロードをスケジュールし、タスクを構成するを参照してください。

既存のRayワークロードを移行する

自己管理型 Ray ワークロードを移行するには、インフラストラクチャの構成とサービス統合を AI ランタイム 実行モデルに置き換える必要があります。

ワークロードの評価

移行の前に、 制限事項を確認してください。ワークロードが1つのアクセラレータタイプで固定サイズのクラスターでランできることを確認します。

ワークロードの移行

セルフマネージド型Rayワークロードを移行するには：

1. 対話型のシングルノード開発にはノートブックを、シングルノードまたはマルチノードのジョブにはDatabricks CLIを選択してください。
2. 既存のクラスターリソースを AI ランタイム アクセラレータタイプとアクセラレータ数にマッピングします。ハードウェアのオプションを参照してください。
3. ノートブックの既存のクラスターStartup、または Databricks CLI ワークロードのドキュメント化されたヘッドおよびワーカーブートストラップを ray_init() に置き換えます。
4. ワークロードの Python 依存関係を宣言し、そのデータ、チェックポイント、モデル、出力パスを更新します。
5. 最小の適用可能な構成でワークロードを検証します。意図した固定トポロジをテストする前に、そのデータアクセス、リソース要求、および出力を確認してください。

制限

AI ランタイム上の Ray には次の制限があります：

• Ray クラスターは、ワークロードの実行中は固定ノード数になります。Ray クラスターのオートスケールはサポートされていません。
• ワークロード内のすべてのノードで同じ accelerator_type が使用されます。Ray クラスターには、個別の CPU 専用および GPU ワーカーグループを含めることはできません。また、CPU および GPU の容量を個別にスケーリングすることもできません。
• ノートブックは、アタッチされた単一ノードで Ray を起動します。マルチノード Ray クラスターには Databricks CLI を使用してください。
• Databricks ドライバープロキシ経由の Ray ダッシュボードアクセスは、現在ノートブックワークロードでのみ利用できます。

ヒント

CPU の準備と GPU の処理を個別のステージとして実行できる場合は、マルチタスクの LakeFlow Job を使用し、Unity Catalog ボリュームを使用してタスク間でデータを渡します。