AI Runtime 上での Ray Core の hello world
Ray は独立した Python タスクをスケジュールし、その進捗状況を追跡します。このノートブックでは、アタッチされた 1xA10 コンピュート上で Ray を起動し、各タスクの完了を待たずに 8 つの GPU タスクを送信します。また、Ray ダッシュボードを使用して、利用可能な GPU 上でそれらが一度に 1 つずつ実行される様子を監視します。
この例では、Databricks AI 環境のバージョン 5 以降が必要です。
要件
このノートブックには、 AI v5 環境を備えた AI Runtime が必要です。このウォークスルーでは 1xA10 アタッチ済み GPU コンピュートを使用するため、Ray が作業をキューイングする様子を観察できます。
ノートブックを接続するには：
- ノートブックの上部にある [接続] を選択します。
- Serverless GPU を選択します。
- [環境] サイドパネルで、 [アクセラレータ] を 1xA10 に設定します。
- 基本環境として AI v5 を選択します。
- [適用] を選択し、次に [確定] を選択します。
AI v5 には Ray と CUDA 対応の PyTorch が含まれているため、この例では追加のパッケージをインストールしません。8つのタスクを同時に実行する様子を確認するには、代わりにノートブックを 8xH100 コンピュートにアタッチして再実行します。
Ray を初期化する
ray_init() を使用して、ノートブックセッション用に Ray を起動します。この関数は Ray コンテキストに関する情報を表示し、Databricks ドライバープロキシ経由で機能するダッシュボードLinkを出力します。
import ray
from serverless_gpu import ray_init
ray_init()
Ray リソースの調査
作業を送信する前に、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) でデコレートすると、実行されるたびに 1 つの GPU を予約する Ray タスクが作成されます。このタスクは小規模な 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() によって出力されたダッシュボードLinkを開いてください。 ジョブ ページから、実行中のジョブを開き、送信後のタスクリストを確認します。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() に渡します。Ray はタスクが終了するのを待機し、送信順に結果を返します。
results = ray.get(task_refs)
pprint(results, sort_dicts=False)
1xA10 コンピュートを使用した場合、8 つのタスクすべてが同じ GPU を順番に使用したことが結果からわかります。8xH100 コンピュートでノートブックを再実行すると、コードを変更することなく、Ray が 8 つのタスクすべてを同時にスケジュールできます。