Ray Core hello world em AI Runtime

O Ray agenda tarefas independentes em Python e acompanha o seu progresso para ti. Neste notebook, inicia o Ray no recurso de computação 1xA10 associado, submete oito tarefas de GPU sem esperar que cada uma termine e usa o painel do Ray para acompanhar a respetiva execução, uma de cada vez, na GPU disponível.

Requisitos

Este notebook requer o AI Runtime com o ambiente AI v5. O walkthrough utiliza 1xA10 de GPU de computação ligada para que possas observar o funcionamento da fila de Ray.

Para conectar o computador portátil:

1. Selecione Ligar no topo do caderno.
2. Selecione GPU Serverless.
3. No painel lateral Ambiente, defina Acelerador como 1xA10.
4. Seleciona AI v5 como ambiente base.
5. Selecione Candidatar-se e depois selecione Confirmar.

O AI v5 inclui Ray e PyTorch com CUDA, pelo que este exemplo não instala pacotes adicionais. Para ver as oito tarefas a serem executadas ao mesmo tempo, podes em vez disso ligar o portátil ao cálculo 8xH100 e voltar a executá-lo.

Inicializar Ray

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

import ray
from serverless_gpu import ray_init

ray_init()

Ver recursos do Ray

Antes de submeteres trabalhos, verifica os recursos que o Ray descobriu. Com capacidade de computação 1xA10, o cluster deverá indicar 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

Decorar uma função com @ray.remote(num_gpus=1) cria uma tarefa do Ray que reserva uma GPU sempre que é executada. Esta tarefa realiza um pequeno cálculo CUDA e devolve detalhes sobre a GPU que utilizou.

A chamada sleep mantém cada tarefa ativa durante tempo suficiente para a inspecionar no painel. Esta pausa serve apenas para a exploração de dashboards e não é para ser usada como referência.

@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 a ligação para o painel apresentada por ray_init(). Na página Tarefas, abre a tarefa em execução e acompanha a lista de tarefas depois da submissão. Com capacidade de computação 1xA10, deverás ver uma tarefa em execução enquanto as outras sete aguardam capacidade da GPU. Com o cálculo 8xH100, todas as oito tarefas podem ser executadas ao mesmo tempo.

Cada .remote() chamada submete uma tarefa sem esperar que esta termine e devolve uma referência de 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 dashboard, passe as referências ao objeto para ray.get(). Ray espera que as tarefas terminem e devolve os resultados por 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 reexecutares o portátil com computação 8xH100, o Ray pode agendar as oito tarefas ao mesmo tempo sem quaisquer alterações de código.

Bloco de notas de exemplo

Ray Core olá, mundo em AI Runtime

Obter caderno