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.

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. Install the ray library or the relevant Ray extra (ray[data], ray[train], ray[tune]) as a dependency for your workload on the standard Databricks environment. If you use the Databricks AI environment, ray is preinstalled but extras must be installed separately.

Examples

Example	Description
Ray hello world examples (CLI)	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 (CLI)	Fine-tune a large language model across multiple nodes using Ray Train and PyTorch.
Batch inference with Ray Data and vLLM (CLI)	Run large-scale batch inference using Ray Data for distributed data loading and vLLM for efficient model serving.