AI Runtime 上での Ray Core ハローワールド

このページのノートブック バージョンを開く

Rayは独立したPythonタスクをスケジュールし、進捗を追跡してくれます。 このノートブックでは、接続された1xA10のコンピュートでRayを起動し、8つのGPUタスクを1つ一つ終わるのを待たずに提出し、Rayダッシュボードを使って利用可能なGPU上で1つずつ実行を観察できます。

Note

この例はDatabricks AI環境バージョン5以上が必要です。

Requirements

このノートブックは AI v5 環境でAIランタイムが必要です。 このウォークスルーでは、Ray のキューイング処理を確認できるように、1xA10 の接続された GPU コンピュートを使用しています。

ノートをつなげるには:

  1. ノートの上部で 「Connect 」を選択してください。
  2. [サーバーレス GPU] を選択します。
  3. 環境サイドパネルでアクセラレータを1xA10に設定してください。
  4. ベース環境として AI v5 を選択してください。
  5. 「 適用」を選択し、その後 「確認」を選択します。

AI v5にはRayとCUDA対応のPyTorchが含まれているため、この例では追加パッケージはインストールされません。 8つのタスクすべてを同時に実行したいなら、代わりにノートブックを 8xH100 コンピュートに接続して再実行できます。

レイの初期化

ray_init() を使用して、ノートブック セッション用に Ray を起動してください。 この関数はRayコンテキストの情報を表示し、Databricksドライバープロキシを介して動作するダッシュボードリンクを印刷します。

import ray
from serverless_gpu import ray_init

ray_init()

Rayリソースを検査する

作品を提出する前に、レイが見つけたリソースを確認してください。 1xA10のコンピュートがあれば、クラスターは1つのGPUを報告するはずです。

from pprint import pprint

cluster_resources = ray.cluster_resources()
available_resources = ray.available_resources()

pprint(
    {
        "cluster_resources": cluster_resources,
        "available_resources": available_resources,
    },
    sort_dicts=False,
)

if cluster_resources.get("GPU", 0) < 1:
    raise RuntimeError(
        "Ray did not detect a GPU. Attach the notebook to 1xA10 or 8xH100 compute, then rerun it."
    )

GPUタスクを定義する

関数を @ray.remote(num_gpus=1) で装飾すると、Ray タスクが生成され、実行するたびに1つのGPUを予約します。 このタスクは小さなCUDA計算を行い、使用したGPUの詳細を返します。

sleepコールは各タスクをダッシュボードで確認できるだけの時間を保持します。 この一時停止はダッシュボード探索のためであり、ベンチマークとして使うものではありません。

@ray.remote(num_gpus=1)
def run_gpu_task(task_id: int, inspection_seconds: int) -> dict:
    import os
    import time

    import ray
    import torch

    values = torch.arange(1, 5, dtype=torch.float32, device="cuda") + task_id
    computation_result = torch.square(values).sum().item()
    torch.cuda.synchronize()
    time.sleep(inspection_seconds)

    return {
        "task_id": task_id,
        "ray_gpu_ids": ray.get_gpu_ids(),
        "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
        "gpu_model": torch.cuda.get_device_name(0),
        "computation_result": computation_result,
    }

タスクを非同期で提出する

次のセルを実行する前に、 ray_init()が印刷したダッシュボードのリンクを開きます。 ジョブページから実行中のジョブを開き、提出後のタスクリストを確認してください。 1xA10の計算では、1つのタスクが実行され、他の7つのタスクがGPUの容量を待つ間に動作するはずです。 8xH100計算で8つのタスクすべてを同時に実行できます。

各 .remote() 呼び出しはタスクが完了するのを待たずにタスクを送信し、オブジェクト参照を返します。

task_refs = [run_gpu_task.remote(task_id, inspection_seconds=10) for task_id in range(8)]

print(f"Submitted {len(task_refs)} tasks.")
print(f"Each submission returned a {type(task_refs[0]).__name__}.")

結果の取得

ダッシュボードを検査した後、オブジェクト参照を ray.get()に渡します。 レイはタスクが終わるのを待ち、結果を提出順に返す。

results = ray.get(task_refs)
pprint(results, sort_dicts=False)

1xA10計算では、8つのタスクすべてが連続して同じGPUを使用していることが示されました。 8xH100のコンピュートでノートブックを再実行すれば、Rayはコード変更なしで8つのタスクすべてを同時にスケジューリングできます。

ノートブックの例

AI ランタイム上の Ray Core の「Hello World」

ノートブックを入手

  • 最終更新: