Ray Core: „Witaj, świecie!” w środowisku AI Runtime

Otwórz wersję notebooka tej strony

Ray planuje niezależne zadania Python i śledzi ich postępy za ciebie. W tym notatniku uruchamiasz Ray na dołączonym środowisku obliczeniowym 1xA10, przesyłasz osiem zadań GPU bez czekania na zakończenie każdego z nich i używasz panelu Ray, aby obserwować, jak są uruchamiane pojedynczo na dostępnym GPU.

Note

Ten przykład wymaga środowiska AI Databricks w wersji 5 lub wyższej.

Wymagania

Ten notebook wymaga AI Runtime w środowisku AI v5 . Opis krok po kroku wykorzystuje zasoby obliczeniowe GPU z dołączonym 1xA10, dzięki czemu można obserwować kolejkowanie zadań w Ray.

Aby podłączyć laptopa:

  1. Wybierz Połącz na górze notesu.
  2. Wybierz pozycję Bezserwerowy procesor GPU.
  3. W panelu bocznym Środowisko ustaw opcję Akcelerator na 1xA10.
  4. Wybierz AI v5 jako środowisko bazowe.
  5. Wybierz Złóż aplikację, a następnie wybierz Potwierdź.

AI v5 zawiera PyTorch z obsługą Ray i CUDA, więc ten przykład nie instaluje dodatkowych pakietów. Aby zobaczyć wszystkie osiem zadań uruchamianych jednocześnie, możesz zamiast tego podłączyć notatnik do obliczeń 8xH100 i uruchomić go ponownie.

Inicjalizuj promieni

Uruchom Ray dla sesji notatnika za pomocą ray_init(). Funkcja wyświetla informacje o kontekście Ray i wyświetla link do panelu zarządzającego, który działa przez sterownik Databricks.

import ray
from serverless_gpu import ray_init

ray_init()

Sprawdź zasoby Ray

Zanim złożysz pracę, sprawdź zasoby, które odkrył Ray. Przy obliczeniach 1xA10 klaster powinien zgłaszać jedną kartę graficzną.

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

Zdefiniuj zadanie GPU

Opatrzenie funkcji dekoratorem @ray.remote(num_gpus=1) powoduje utworzenie zadania Ray, które przy każdym uruchomieniu rezerwuje jedną jednostkę GPU. To zadanie wykonuje niewielkie obliczenia CUDA i zwraca dane dotyczące użytej karty graficznej.

Wywołanie sleep utrzymuje każde zadanie aktywne wystarczająco długo, aby można było je sprawdzić w panelu. Ta pauza służy tylko do eksploracji dashboardu i nie ma być używana jako punkt odniesienia.

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

Wysyłaj zadania asynchronicznie

Zanim uruchomisz kolejną komórkę, otwórz link do pulpitu nawigacyjnego wyświetlony przez ray_init(). Na stronie Jobs otwórz działające zadanie i obserwuj listę zadań po przesłaniu. Przy zasobie obliczeniowym 1xA10 powinieneś zobaczyć jedno uruchomione zadanie, podczas gdy pozostałe siedem czeka na dostępne zasoby GPU. Dzięki obliczeniom 8xH100 wszystkie osiem zadań może działać jednocześnie.

Każde .remote() wywołanie wysyła zadanie bez oczekiwania na jego zakończenie i zwraca referencję do obiektu.

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

Pobierz wyniki

Po sprawdzeniu panelu przekaż odwołania do obiektów do ray.get(). Ray czeka na zakończenie zadań i zwraca wyniki w kolejności złożenia.

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

Dzięki obliczeniom 1xA10 wyniki pokazują, że wszystkie osiem zadań korzystało z tej samej karty graficznej jedno po drugim. Jeśli ponownie uruchomisz notebook z obliczeniami 8xH100, Ray może zaplanować wszystkie osiem zadań jednocześnie bez żadnych zmian w kodzie.

Przykładowy notatnik

Ray Core witaj świat na AI Runtime

Pobierz laptopa

  • Last updated on