Ray Core Hola mundo en AI Runtime

Ray programa tareas independientes en Python y hace un seguimiento de su progreso por ti. En este cuaderno, inicias Ray con un recurso de cómputo 1xA10 asociado, envías ocho tareas de GPU sin esperar a que termine cada una de ellas y usas el panel de control de Ray para ver cómo se ejecutan una a una en la GPU disponible.

Requisitos

Este portátil requiere tiempo de ejecución de IA con el entorno de IA v5 . La guía utiliza un proceso de GPU conectado a 1xA10 para que pueda observar el funcionamiento de la cola de rayos.

Para conectar el cuaderno:

1. Selecciona Conectar en la parte superior del cuaderno.
2. Seleccione GPU sin servidor.
3. En el panel lateral Entorno, configura Acelerador en 1xA10.
4. Selecciona AI v5 como entorno base.
5. Selecciona Aplicar y luego selecciona Confirmar.

AI v5 incluye Ray y PyTorch habilitado para CUDA, por lo que este ejemplo no instala paquetes adicionales. Para ver que las ocho tareas se ejecutan a la vez, puedes conectar el cuaderno a un cálculo 8xH100 y volver a ejecutarlo.

Inicializar Ray

Inicia Ray para la sesión del cuaderno con ray_init(). La función muestra información sobre el contexto de Ray e imprime un enlace al panel que funciona a través del proxy del controlador de Databricks.

import ray
from serverless_gpu import ray_init

ray_init()

Examinar los recursos de Ray

Antes de entregar trabajos, revisa los recursos que Ray descubrió. Con el proceso 1xA10, el clúster debería mostrar una 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."
    )

Define una tarea de GPU

Decorar una función con @ray.remote(num_gpus=1) crea una tarea de Ray que reserva una GPU cada vez que se ejecuta. Esta tarea realiza un pequeño cálculo CUDA y devuelve detalles sobre la GPU que utilizó.

La sleep llamada mantiene cada tarea activa el tiempo suficiente para inspeccionarla en el panel de control. Esta pausa es solo para la exploración del panel y no está pensada para usarse como referencia.

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

Enviar tareas de forma asíncrona

Antes de ejecutar la siguiente celda, abre el enlace del panel impreso por ray_init(). Desde la página Jobs, abre la tarea en ejecución y observa la lista de tareas después de enviarla. Con un cálculo 1xA10, deberías ver una tarea ejecutándose mientras las otras siete esperan la capacidad de la GPU. Con el cálculo 8xH100, las ocho tareas pueden ejecutarse al mismo tiempo.

Cada .remote() llamada envía una tarea sin esperar a que termine y devuelve una referencia al objeto.

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

Recuperar resultados

Después de inspeccionar el panel, pasa las referencias a objetos a ray.get(). Ray espera a que terminen las tareas y devuelve sus resultados en orden de entrega.

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

Con el cálculo 1xA10, los resultados muestran que las ocho tareas usaron la misma GPU una tras otra. Si vuelves a ejecutar el cuaderno con capacidad de cómputo 8xH100, Ray puede planificar las ocho tareas al mismo tiempo sin ningún cambio en el código.

Cuaderno de ejemplo

Ray Core Hello World en AI Runtime

Obtención del cuaderno