Ray Core hello world auf AI Runtime

Notizbuchversion dieser Seite öffnen

Ray plant unabhängige Python-Aufgaben und verfolgt deren Fortschritt für dich. In diesem Notizbuch startest du Ray mit angeschlossener 1xA10-Compute, gibst acht GPU-Aufgaben ab, ohne auf jede einzelne zu warten, und nutzt das Ray-Dashboard, um sie einzeln auf der verfügbaren GPU ausführen zu sehen.

Note

Dieses Beispiel benötigt die Databricks KI-Umgebung Version 5 oder höher.

Anforderungen

Dieses Notebook benötigt KI-Laufzeit mit der KI-v5-Umgebung . Die Vorgehensweise verwendet angefügte 1xA10-GPU-Computeressourcen, sodass Sie die Ray-Warteschlangearbeit beobachten können.

So schließen Sie das Notebook an:

1. Wähle "Verbinden " am oberen Rand des Notizbuchs.
2. Wählen Sie serverlose GPU aus.
3. Im Environment-Seitenpanel stellen Sie den Beschleuniger auf 1xA10 ein.
4. Wähle AI v5 als Basisumgebung aus.
5. Wähle Anwenden und dann Bestätigen.

AI v5 enthält Ray und CUDA-fähiges PyTorch, daher installiert dieses Beispiel keine zusätzlichen Pakete. Um alle acht Aufgaben gleichzeitig ausführen zu sehen, kannst du das Notizbuch stattdessen an 8xH100 Compute anschließen und erneut ausführen.

Initialisieren von Ray

Starte Ray für die Notebook-Sitzung mit ray_init(). Die Funktion zeigt Informationen über den Ray-Kontext an und druckt einen Dashboard-Link, der über den Databricks-Treiber-Proxy funktioniert.

import ray
from serverless_gpu import ray_init

ray_init()

Ray-Ressourcen prüfen

Bevor du Arbeiten einreichst, schau dir die Ressourcen an, die Ray entdeckt hat. Mit 1xA10-Berechnung sollte der Cluster eine GPU angeben.

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

Definiere eine GPU-Aufgabe

Das Dekorieren einer Funktion mit @ray.remote(num_gpus=1) erzeugt eine Ray-Aufgabe, die jedes Mal eine GPU reserviert, wenn sie ausgeführt wird. Diese Aufgabe führt eine kleine CUDA-Berechnung durch und liefert Details zur verwendeten GPU.

Der Aufruf sleep hält jede Aufgabe lange genug aktiv, um sie im Dashboard zu überprüfen. Diese Pause dient nur der Erkundung von Dashboards und dient nicht 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,
    }

Aufgaben asynchron einreichen

Bevor du die nächste Zelle ausführst, öffne den Dashboard-Link, der von ray_init() ausgegeben wird. Öffne auf der Jobs-Seite den laufenden Job und sieh dir nach der Abgabe die Aufgabenliste an. Mit 1xA10-Berechnung solltest du eine Aufgabe laufen sehen, während die anderen sieben auf GPU-Kapazität warten. Mit 8xH100-Berechnung können alle acht Aufgaben gleichzeitig ausgeführt werden.

Jeder .remote() Aufruf sendet eine Aufgabe, ohne auf deren Abschluss zu warten, und gibt eine Objektreferenz zurück.

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

Ergebnisse abrufen

Nachdem Sie das Dashboard überprüft haben, übergeben Sie die Objektreferenzen an ray.get(). Ray wartet, bis die Aufgaben fertig sind, und gibt die Ergebnisse in der Abgabereihenfolge zurück.

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

Mit 1xA10-Berechnung zeigen die Ergebnisse, dass alle acht Aufgaben nacheinander dieselbe GPU verwendeten. Wenn du das Notebook mit 8xH100-Berechnung erneut ausführst, kann Ray alle acht Aufgaben gleichzeitig planen, ohne dass der Code geändert wird.

Beispiel-Notebook

Ray Core Hello World auf AI Runtime

Notebook abrufen