Ray Core olá, mundo no AI Runtime

Ray agenda tarefas independentes em Python e acompanha o progresso delas para você. Neste notebook, você inicia o Ray no ambiente de computação 1xA10 anexado, envia oito tarefas de GPU sem esperar que cada uma termine e usa o painel do Ray para observá-las em execução, uma de cada vez, na GPU disponível.

Requirements

Este notebook requer tempo de execução de IA com o ambiente AI v5 . O passo a passo usa computação de GPU anexada 1xA10 para que você possa observar o Ray enfileirando tarefas.

Para conectar o notebook:

1. Selecione Conectar no topo do caderno.
2. Selecione GPU sem servidor.
3. No painel lateral Ambiente, defina Acelerador para 1xA10.
4. Selecione AI v5 como ambiente base.
5. Selecione Aplicar, depois selecione Confirmar.

A AI v5 inclui Ray e PyTorch habilitado para CUDA, então este exemplo não instala pacotes adicionais. Para ver as oito tarefas rodando ao mesmo tempo, você pode conectar o notebook ao cálculo 8xH100 e rodá-lo novamente.

Inicializar Ray

Inicie o Ray para a sessão do notebook com ray_init(). A função exibe informações sobre o contexto Ray e imprime um link de painel que funciona através do proxy do driver Databricks.

import ray
from serverless_gpu import ray_init

ray_init()

Inspecionar os recursos do Ray

Antes de enviar trabalhos, confira os recursos que Ray encontrou. Com capacidade computacional 1xA10, o cluster deve informar uma 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."
    )

Defina uma tarefa de GPU

Aplicar o decorador @ray.remote(num_gpus=1) a uma função cria uma tarefa do Ray que reserva uma GPU cada vez que é executada. Essa tarefa realiza um pequeno cálculo CUDA e retorna detalhes sobre a GPU que utilizou.

A chamada sleep mantém cada tarefa ativa por tempo suficiente para ser examinada no painel. Essa pausa é apenas para exploração de dashboards e não é para ser usada como 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,
    }

Submeter tarefas de forma assíncrona

Antes de executar a próxima célula, abra o link do painel impresso por ray_init(). Na página Jobs, abra o job em execução e acompanhe a lista de tarefas após o envio. Com capacidade de processamento 1xA10, você deve ver uma tarefa em execução enquanto as outras sete aguardam disponibilidade de capacidade da GPU. Com o cálculo 8xH100, todas as oito tarefas podem ser executadas ao mesmo tempo.

Cada .remote() chamada envia uma tarefa sem esperar que ela termine e retorna uma referência ao 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

Depois de inspecionar o painel, passe as referências aos objetos para ray.get(). Ray espera as tarefas terminarem e devolve os resultados em ordem de submissão.

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

Com o cálculo 1xA10, os resultados mostram que todas as oito tarefas usaram a mesma GPU uma após a outra. Se você reexecutar o notebook com capacidade computacional 8xH100, o Ray poderá agendar todas as oito tarefas ao mesmo tempo sem precisar alterar o código.

Notebook de exemplo

Ray Core: Olá, mundo no AI Runtime

Obter laptop