Databricks AI Runtime SDK

The Databricks AI Runtime SDK (databricks-sdk-air) runs multi-GPU workloads on Databricks AI Runtime directly from Databricks notebooks. Import it as databricks.air.

It provides the same functionality as running torchrun train.py or accelerate launch train.py from the command line, but integrates seamlessly into the Databricks notebook environment.

The SDK is preinstalled in GPU environment version 6 and above, and replaces the deprecated Serverless GPU Python API (serverless_gpu). To migrate, see Migrate from serverless_gpu.

For concepts, usage guides, and examples, see the AI Runtime documentation.

• API Reference
• Data Loading and Checkpointing
• Environment Variables

Indices and tables

• Index

• Module Index

• Search Page