Ray 會排程獨立的 Python 任務並追蹤進度。 在這台筆記本中，你可以在連接的 1xA10 運算器上啟動 Ray，提交八個 GPU 任務，且不等每個任務完成，並透過 Ray 儀表板觀看它們在可用 GPU 上一個一個執行。
Note
此範例需要 Databricks AI 環境版本 5 或以上。
要求
此筆記本需要使用 AI v5 環境的 AI Runtime。 這個逐步解說使用 1xA10 附加式 GPU 運算資源，讓你可以觀察 Ray 如何將工作排入佇列。
若要連接筆記型電腦：
- 在筆記本頂端選擇 「連接 」。
- 選擇 無伺服器 GPU。
- 在 環境 側面板，將 加速器 設為 1xA10。
- 選擇 AI v5 作為基礎環境。
- 選擇 套用，然後選擇 確認。
AI v5 包含 Ray 及支援 CUDA 的 PyTorch，因此此範例未安裝額外套件。 如果想同時執行這八個任務，你可以把筆記本連接到 8xH100 運算器並重新執行。
初始化 Ray
用 Ray 啟動筆記本會話 ray_init()。 此函式會顯示 Ray 上下文的相關資訊，並輸出可透過 Databricks 驅動程式 Proxy 存取的儀表板連結。
import ray
from serverless_gpu import ray_init
ray_init()
檢視 Ray 資源
在提交作品前，請先查看 Ray 發現的資源。 使用 1xA10 運算時，叢集應該會回報一張 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 任務，每次執行時會保留一顆 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 運算時，你會看到一個任務在執行，其他七個任務則等待 GPU 容量。 使用 8xH100 運算，八個任務可同時執行。
每次 .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 計算計算，結果顯示八個任務都連續使用相同的 GPU。 如果你使用 8xH100 運算資源重新執行筆記本，Ray 無須修改任何程式碼，就能將 8 個任務全部同時排程。