Ray Core Hello World op AI Runtime

Ray plant onafhankelijke Python-taken en volgt hun voortgang voor je. In dit notebook start je Ray op een aangesloten 1xA10-berekening, dien je acht GPU-taken in zonder te wachten tot ze klaar zijn, en gebruik je het Ray-dashboard om ze één voor één te zien draaien op de beschikbare GPU.

Requirements

Dit notebook vereist AI Runtime met de AI v5-omgeving . Deze walkthrough gebruikt 1xA10 gekoppelde GPU-rekenkracht, zodat je kunt zien hoe Ray werk in de wachtrij plaatst.

Om het notitieboekje te verbinden:

1. Selecteer Verbinden bovenaan het notitieboek.
2. Selecteer serverloze GPU.
3. Stel in het Environment-zijpaneel de Versneller in op 1xA10.
4. Selecteer AI v5 als basisomgeving.
5. Selecteer Toepassen en vervolgens Bevestigen.

AI v5 bevat Ray en CUDA-enabled PyTorch, dus dit voorbeeld installeert geen extra pakketten. Om alle acht taken tegelijk te zien draaien, kun je het notebook koppelen aan 8xH100 compute en het opnieuw uitvoeren.

Initialiseer Ray

Start Ray voor de notitieboeksessie met ray_init(). De functie toont informatie over de Ray-context en print een dashboardlink die werkt via de Databricks-driverproxy.

import ray
from serverless_gpu import ray_init

ray_init()

Bekijk Ray-bronnen

Voordat je werk indient, controleer de bronnen die Ray heeft ontdekt. Met 1xA10 compute zou de cluster één GPU moeten rapporteren.

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

Definieer een GPU-taak

Het decoreren van een functie met @ray.remote(num_gpus=1) creëert een Ray-taak die elke keer dat deze draait één GPU reserveert. Deze taak voert een kleine CUDA-berekening uit en geeft details terug over de gebruikte GPU.

De sleep aanroep houdt elke taak lang genoeg actief om deze in het dashboard te inspecteren. Deze pauze is alleen bedoeld voor dashboardonderzoek en is niet bedoeld als benchmark.

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

Dien taken asynchroon in

Voordat je de volgende cel uitvoert, open je de link naar het dashboard die door ray_init() wordt weergegeven. Open vanaf de Vacature-pagina de lopende opdracht en bekijk de takenlijst na het indienen. Met 1xA10 compute zou je één taak moeten zien draaien terwijl de andere zeven wachten op GPU-capaciteit. Met 8xH100 rekenkracht kunnen alle acht taken tegelijkertijd draaien.

Elke .remote() aanroep stuurt een taak in zonder te wachten tot deze klaar is en geeft een objectreferentie terug.

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

Resultaten ophalen

Nadat je het dashboard hebt geïnspecteerd, geef je de objectreferenties door aan ray.get(). Ray wacht tot de taken klaar zijn en geeft de resultaten in volgorde van indiening terug.

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

Met 1xA10 compute laten de resultaten zien dat alle acht taken één voor één dezelfde GPU gebruikten. Als je het notebook opnieuw uitvoert met 8xH100 rekenkracht, kan Ray alle acht taken tegelijk plannen zonder codewijzigingen.

Voorbeeld van notebook

Ray Core Hello World op AI Runtime

Notebook krijgen