Ray Core hello world sur AI Runtime

Ouvrir dans Databricks

Ray planifie des tâches Python indépendantes et suit leur progression pour vous. Dans ce Notebook, vous start Ray sur un compute 1xA10 attaché, soumettez huit tâches GPU sans attendre que chacune se termine, et utilisez le tableau de bord Ray pour les regarder s’exécuter une par une sur le GPU disponible.

remarque

Cet exemple nécessite la version 5 ou une version ultérieure de l’environnement Databricks AI.

Configuration requise

Ce notebook nécessite AI Runtime avec l'environnement AI v5 . Le guide pratique utilise un compute GPU 1xA10 associé afin que vous puissiez observer le travail de mise en file d'attente de Ray.

Pour connecter le notebook :

1. Sélectionnez Connect en haut du Notebook.
2. Sélectionnez GPU serverless .
3. Dans le panneau latéral Environnement , réglez Accélérateur sur 1xA10 .
4. Sélectionnez AI v5 comme environnement de base.
5. Sélectionnez Apply , puis sélectionnez Confirm .

L'IA v5 inclut Ray et PyTorch compatible CUDA, cet exemple n'installe donc pas de packages supplémentaires. Pour voir les huit tâches s'exécuter simultanément, vous pouvez plutôt associer le notebook à un compute 8xH100 et le réexécuter.

Initialiser Ray

Start Ray pour la session de Notebook avec ray_init(). La fonction affiche des information sur le contexte Ray et imprime un Link de tableau de bord qui fonctionne via le proxy du Driver Databricks.

Python

import ray
from serverless_gpu import ray_init

ray_init()

Inspecter les ressources Ray

Avant de soumettre votre travail, vérifiez les ressources découvertes par Ray. Avec un compute 1xA10, le cluster doit signaler un GPU.

Python

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

Définissez une tâche GPU

Décorer une fonction avec @ray.remote(num_gpus=1) crée une tâche Ray qui réserve un GPU à chaque exécution. Cette tâche effectue un petit calcul CUDA et renvoie des détails sur le GPU utilisé.

L’appel sleep maintient chaque tâche active suffisamment longtemps pour l’inspecter dans le tableau de bord. Cette pause est uniquement destinée à l’exploration du tableau de bord et n’est pas conçue pour être utilisée comme un benchmark.

Python

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

Soumettre des tâches de manière asynchrone

Avant d’exécuter la cellule suivante, ouvrez le Link du tableau de bord imprimé par ray_init(). Depuis la page Jobs , ouvrez le job en cours d’exécution et surveillez la liste des tâches après la soumission. Avec un compute 1xA10, vous devriez voir une tâche en cours d’exécution tandis que les sept autres attendent de la capacité GPU. Avec un compute 8xH100, les huit tâches peuvent s’exécuter simultanément.

Chaque appel .remote() soumet une tâche sans attendre qu’elle se termine et renvoie une référence d’objet.

Python

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

Récupérer les résultats

Après avoir inspecté le tableau de bord, transmettez les références d'objet à ray.get(). Ray attend la fin des tâches et renvoie leurs résultats dans l'ordre de soumission.

Python

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

Avec un compute 1xA10, les résultats montrent que les huit tâches ont utilisé le même GPU les unes après les autres. Si vous réexécutez le notebook avec un compute 8xH100, Ray peut planifier les huit tâches en même temps sans aucune modification de code.

Exemple de Notebook

Ray Core hello world sur AI Runtime

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks