Ray Core hello world no AI Runtime

Open in Databricks

O Ray programa tarefas independentes do Python e rastreia o progresso delas para você. Neste notebook, você começa o Ray em um compute 1xA10 anexado, envia oito tarefas de GPU sem esperar que cada uma termine e usa o painel do Ray para observá-las sendo executadas uma de cada vez na GPU disponível.

nota

Este exemplo requer o ambiente Databricks AI versão 5 ou acima.

Requisitos

Este notebook requer o AI Runtime com o ambiente AI v5 . O passo a passo usa compute de GPU 1xA10 anexado para que você possa observar o trabalho de enfileiramento do Ray.

Para conectar o notebook:

1. Selecione Conectar na parte superior do notebook.
2. Selecione GPU serverless .
3. No painel lateral Ambiente , defina Acelerador como 1xA10 .
4. Selecione AI v5 como o ambiente base.
5. Selecione Aplicar e, em seguida, selecione Confirmar .

O AI v5 inclui Ray e PyTorch habilitado para CUDA, portanto, este exemplo não instala pacotes adicionais. Para ver todas as oito tarefas serem executadas ao mesmo tempo, você pode, em vez disso, anexar o notebook ao compute 8xH100 e executá-lo novamente.

Inicializar o Ray

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

Python

import ray
from serverless_gpu import ray_init

ray_init()

Inspecionar recursos do Ray

Antes de enviar o trabalho, verifique os recursos que o Ray descobriu. Com compute 1xA10, o cluster deve relatar uma 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."
    )

Definir uma tarefa de GPU

Decorar uma função com @ray.remote(num_gpus=1) cria uma tarefa do Ray que reserva uma GPU cada vez que é executada. Esta 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 inspecioná-la no dashboard. Esta pausa é apenas para exploração do dashboard e não deve ser usada como 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,
    }

Enviar tarefas de forma assíncrona

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

Cada chamada .remote() envia uma tarefa sem aguardar a conclusão e retorna uma referência de objeto.

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

Recuperar resultados

Após inspecionar o dashboard, passe as referências de objeto para ray.get(). O Ray aguarda a conclusão das tarefas e retorna seus resultados na ordem de envio.

Python

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

Com compute 1xA10, os resultados mostram que todas as oito tarefas usaram a mesma GPU uma após a outra. Se você executar novamente o notebook com compute 8xH100, o Ray poderá programar todas as oito tarefas ao mesmo tempo sem nenhuma alteração no código.

Notebook de exemplo

Ray Core hello world no AI Runtime

Abrir notebook em uma nova abaOpen in Databricks