Ray on AI Runtime

Important

This feature is in Public Preview.

Note

This page covers Ray on AI Runtime. If you're using Ray on Databricks classic compute with Databricks Runtime ML, see Ray on Databricks.

Ray is an open source framework for scaling Python workloads. You can create Ray clusters and run Ray applications on AI Runtime, using serverless GPU compute that handles infrastructure provisioning automatically.

Use Ray in notebooks

When your notebook is connected to AI Runtime GPU compute, use ray_init() from the serverless_gpu package to start Ray on the attached compute:

from serverless_gpu import ray_init

ray_init()

ray_init() configures the Ray dashboard for access through the Databricks driver proxy and prints the dashboard URL in the notebook output. Use the dashboard to inspect Ray jobs, tasks, and resource usage while your code runs.

Note

ray_init() requires environment version 5 or later. Databricks AI v5 includes Ray. If you use Standard v5, install Ray before calling ray_init().

Notebook examples

Example	Description
Ray Core hello world	Submit asynchronous GPU tasks on attached compute, inspect scheduling in the Ray dashboard, and retrieve the results.
Qwen2.5-32B batch inference with Ray Data and vLLM	Run multilingual batch inference with eight persistent vLLM replicas on 8 H100 GPUs and save the results as Parquet in Unity Catalog.

Use Ray with the AI Runtime CLI

The AI Runtime CLI supports Ray in single-node and multi-node configurations. Include a bootstrap script under the workload's command, which starts the Ray cluster and coordinates head and worker nodes when the workload launches. The Ray hello world examples walk through this pattern.

AI Runtime supports Ray's core libraries, including Ray Core, Ray Data, Ray Train, and Ray Tune. With the Standard environment, add ray or the relevant extra (ray[data], ray[train], or ray[tune]) to your workload dependencies. The Databricks AI environment includes Ray.

CLI examples

Example	Description
Ray hello world examples	Minimal single-node and multi-node examples for Ray Core, Ray Train, Ray Data, and Ray Tune, including the cluster bootstrap pattern.
Distributed training with Ray Train	Fine-tune a large language model across multiple nodes using Ray Train and PyTorch.
Batch inference with Ray Data and vLLM	Run large-scale batch inference using Ray Data for distributed data loading and vLLM for efficient model serving.
Hyperparameter search with Ray Tune	Search LoRA fine-tuning hyperparameters for Qwen2.5 with Ray Tune, running one trial per GPU and stopping underperforming trials early with the ASHA scheduler.