Databricks SDK AIR

databricks-sdk-air provides the databricks.air Python APIs for Databricks AI Runtime (AIR): utilities for loading training data from Unity Catalog volumes and for running distributed GPU training workloads.

databricks.air is in Beta; the API is subject to change.

The base package installs the databricks.air APIs with no extra dependencies. Add the extras below for the features you use, or [all] for everything:

pip install databricks-sdk-air          # base, no extra dependencies
pip install "databricks-sdk-air[all]"   # data + distributed

databricks.air.data

PyTorch data utilities for streaming Unity Catalog volumes into GPU training pipelines: UCVolumeDataset, a checkpoint-aware DataLoader, distributed-checkpoint UCVolumeReader/UCVolumeWriter, and the Checkpointable protocol. Install with the data extra:

pip install databricks-sdk-air[data]

databricks.air.distributed

APIs for launching and monitoring single-node multi-GPU training workloads. Install with the distributed extra:

pip install databricks-sdk-air[distributed]

import databricks.air


@databricks.air.distributed(num_accelerators=8, accelerator_type="GPU_8xH100")
def train_model():
    pass

from databricks.air import distributed, ray_init, ray_launch is also supported.

Documentation

• AI Runtime: https://docs.databricks.com/aws/en/machine-learning/ai-runtime/