Ray Core hello world on AI Runtime

Open notebook version of this page

Ray schedules independent Python tasks and tracks their progress for you. In this notebook, you start Ray on attached 1xA10 compute, submit eight GPU tasks without waiting for each one to finish, and use the Ray dashboard to watch them run one at a time on the available GPU.

Note

This example requires the Databricks AI environment version 5 or above.

Requirements

This notebook requires AI Runtime with the AI v5 environment. The walkthrough uses 1xA10 attached GPU compute so that you can observe Ray queueing work.

To connect the notebook:

  1. Select Connect at the top of the notebook.
  2. Select Serverless GPU.
  3. In the Environment side panel, set Accelerator to 1xA10.
  4. Select AI v5 as the base environment.
  5. Select Apply, then select Confirm.

AI v5 includes Ray and CUDA-enabled PyTorch, so this example does not install additional packages. To see all eight tasks run at the same time, you can instead attach the notebook to 8xH100 compute and rerun it.

Initialize Ray

Start Ray for the notebook session with ray_init(). The function displays information about the Ray context and prints a dashboard link that works through the Databricks driver proxy.

import ray
from serverless_gpu import ray_init

ray_init()

Inspect Ray resources

Before you submit work, check the resources that Ray discovered. With 1xA10 compute, the cluster should report one 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 a GPU task

Decorating a function with @ray.remote(num_gpus=1) creates a Ray task that reserves one GPU each time it runs. This task performs a small CUDA calculation and returns details about the GPU it used.

The sleep call keeps each task active long enough to inspect it in the dashboard. This pause is only for dashboard exploration and is not meant to be used as a 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,
    }

Submit tasks asynchronously

Before you run the next cell, open the dashboard link printed by ray_init(). From the Jobs page, open the running job and watch the task list after submission. With 1xA10 compute, you should see one task running while the other seven wait for GPU capacity. With 8xH100 compute, all eight tasks can run at the same time.

Each .remote() call submits a task without waiting for it to finish and returns an object reference.

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

Retrieve results

After you have inspected the dashboard, pass the object references to ray.get(). Ray waits for the tasks to finish and returns their results in submission order.

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

With 1xA10 compute, the results show that all eight tasks used the same GPU one after another. If you rerun the notebook with 8xH100 compute, Ray can schedule all eight tasks at the same time without any code changes.

Example notebook

Ray Core hello world on AI Runtime

Get notebook

  • Last updated on