AI ランタイム上のレイ

RayはPythonワークロードのスケーリングを目的としたオープンソースフレームワークです。 AIランタイム上でRayを動かすなら、基盤となるGPUインフラのプロビジョニングや管理をしなくても可能です。

なぜAIランタイムでRayを使うのか

Rayは、モデルトレーニング、バッチ推論、データ処理、ハイパーパラメータチューニングなどのPythonワークロードの配布に一般的に使われています。 AI Runtimeでは、RayはUnity Catalogボリュームのデータにアクセスし、MLflowで実験を追跡し、本番ワークフローの一部として実行できます。

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

レイがAIランタイムでどのように動作するか

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

ノートブックでRayを使う

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

from databricks.air import ray_init

ray_init()

ray_init() 標準 ray.init() APIを呼び出し、Databricksドライバープロキシ経由でRayダッシュボードへのアクセスを設定し、ノートブック出力にダッシュボードURLを印刷します。 ダッシュボードを使って、コードが実行されている間にRayジョブ、タスク、アクター、ログ、リソース使用状況をチェックできます。

ノートブックの例

DatabricksのCLIでRayを使いましょう

AIランタイム用のDatabricksのCLIコマンドは、シングルノードのRayワークロードをサポートします。 ワークロード構成では、固定された accelerator_type と num_acceleratorsを指定します。 Qwen2.5-3B Ray Trainのレシピには再利用可能なブートストラップスクリプトが含まれています。 ワークロードのcommandで、RAY_ENTRYPOINTをPythonファイルのパスに設定し、ブートストラップスクリプトを実行します。 ブートストラップスクリプトは以下のことをします:

• ノード上のレイヘッドを起動し、設定されたエントリポイントを実行します。
• エントリーポイントが終わると、レイプロセスは停止します。

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

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

CLI の例

データと Databricks の機能を使用する

RayワークロードはUnity Catalogボリュームのデータにアクセスし、MLflowで実行を記録し、宣言的自動化バンドルで定義されたLakeflowジョブタスクとして実行できます。

Unity カタログにおけるデータへのアクセス

Unityカタログテーブル内の表形式データについては、 ray.data.read_databricks_tables() を使ってテーブルを読み込んだり、DatabricksのSQLウェアハウスでSQLクエリを実行したりします。 倉庫ID、カタログ、スキーマを指定してください。AI RuntimeはローカルのSparkセッションを提供しません。 SQLを使って、Rayが結果を処理する前に結合、集約、フィルターを実行できます。

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

ワークロードを追跡して統合管理

MLflowを使って、Rayワークロードのパラメータ、指標、アーティファクトを記録しましょう。 「 実験の追跡と観測可能性」を参照してください。

ワークロードをスケジュール化したり、CPUのデータ準備タスクで組み合わせる場合は、Lakeflow JobsとDeclarative Automation Bundlesをご利用ください。 「GPU ワークロードのスケジュール設定とタスクの作成」を参照してください。

既存のRayワークロードの移行

セルフマネージドのRayワークロードの移行には、インフラ構成やサービス統合をAIランタイム実行モデルに置き換えることが含まれます。

ワークロードを評価する

移行前に制限を確認してください。 ワークロードが1種類のアクセラレータで固定サイズのクラスタで動作可能か確認してください。

ワークロードを移行する

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

1. 対話型開発にはノートブックを、ジョブの送信には Databricks CLI を選択してください。
2. 既存のクラスターリソースをAIランタイムのアクセラレータタイプとアクセラレータ数にマッピングします。 詳細は、ハードウェア を参照してください。
3. ノートブック内の既存のクラスター起動処理を ray_init() に置き換えるか、Databricks CLI ワークロードに記載されているブート ストラップ スクリプトを使用してください。
4. ワークロードのPython依存関係を宣言し、データ、チェックポイント、モデル、出力パスを更新します。
5. 最小の適用構成でワークロードを検証します。 意図した固定トポロジーをテストする前に、データアクセス、リソースリクエスト、出力を確認してください。

Limitations

AIランタイム上のRayには以下の制限があります:

• Rayクラスターは、ワークロード期間中固定されたノード数を持っています。 レイクラスターのオートスケーリングはサポートされていません。
• ワークロード内のすべてのノードは同じ accelerator_typeを使用します。 RayクラスターはCPU専用とGPUワーカーグループを別々に含めたり、CPUとGPUの容量を独立してスケールさせることもできません。
• Azure Databricksドライバープロキシを通じたRayダッシュボードへのアクセスはノートブックワークロードのみに利用可能です。

• RayはノートブックとDatabricksのCLIワークロードの両方で単一のノード上で動作します。