Ray Core Hallå världen på AI Runtime

Ray schemalägger oberoende Python-uppgifter och följer deras framsteg åt dig. I denna anteckningsbok startar du Ray på ansluten 1xA10-beräkning, skickar in åtta GPU-uppgifter utan att vänta på att varje är klar, och använder Rays dashboard för att se dem köras en i taget på det tillgängliga grafikkortet.

Requirements

Denna anteckningsbok kräver AI Runtime med AI v5-miljön. Genomgången använder 1xA10-ansluten GPU-beräkning så att du kan se Ray queue-arbetet.

För att koppla ihop anteckningsboken:

1. Välj Ansluta högst upp i anteckningsboken.
2. Välj Serverlös GPU.
3. I sidopanelen Miljö ställer du in Accelerator till 1xA10.
4. Välj AI v5 som basmiljö.
5. Välj Apply, välj sedan Bekräfta.

AI v5 inkluderar Ray och CUDA-aktiverad PyTorch, så detta exempel installerar inga ytterligare paket. För att se alla åtta uppgifter köras samtidigt kan du istället koppla notebooken till 8xH100 compute och köra om den.

Initiera Ray

Starta Ray för anteckningsbokssessionen med ray_init(). Funktionen visar information om Ray-kontext och skriver ut en länk till instrumentpanelen som fungerar via Databricks-drivrutinsproxyn.

import ray
from serverless_gpu import ray_init

ray_init()

Inspect Ray-resurser

Innan du skickar in arbete, kontrollera de resurser som Ray upptäckte. Med 1xA10 beräkning bör klustret rapportera ett 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."
    )

Definiera en GPU-uppgift

Att dekorera en funktion med @ray.remote(num_gpus=1) skapar en Ray-uppgift som reserverar ett GPU varje gång den körs. Denna uppgift utför en liten CUDA-beräkning och returnerar detaljer om GPU:n den använde.

Anropet till sleep håller varje uppgift aktiv tillräckligt länge för att granska den i instrumentpanelen. Denna paus är endast för utforskning av instrumentpanelen och är inte avsedd att användas som en referenspunkt.

@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,
    }

Lämna in uppgifter asynkront

Innan du kör nästa cell, öppna dashboard-länken som skrivs ut av ray_init(). Från sidan Jobb, öppna det pågående jobbet och övervaka uppgiftslistan efter att du har skickat in jobbet. Med 1xA10-beräkning bör du se en uppgift köras medan de andra sju väntar på GPU-kapacitet. Med 8xH100 beräkning kan alla åtta uppgifter köras samtidigt.

Varje .remote() anrop skickar in en uppgift utan att vänta på att den ska bli klar och returnerar en objektreferens.

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__}.")

Hämta resultat

Efter att du har inspekterat dashboarden, skicka objektreferenserna till ray.get(). Ray väntar på att uppgifterna ska bli klara och returnerar resultaten i inlämningsordning.

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

Med 1xA10-beräkning visar resultaten att alla åtta uppgifter använde samma GPU en efter en. Om du kör anteckningsboken med 8xH100 beräkning kan Ray schemalägga alla åtta uppgifter samtidigt utan några kodändringar.

Exempelanteckningsbok

Ray Core Hej världen på AI Runtime

Hämta anteckningsbok