Metadata-Version: 2.1
Name: databricks-sdk-air
Author: Databricks
Home-page: https://docs.databricks.com/aws/en/machine-learning/ai-runtime/
License: Databricks License
Description-Content-Type: text/markdown
Summary: Databricks AI Runtime Python package.
Project-URL: Documentation, https://docs.databricks.com/aws/en/machine-learning/ai-runtime/
Provides-Extra: all
Requires-Dist: cachetools==5.5.0; extra == 'all'
Requires-Dist: cloudpickle==3.1.2; extra == 'all'
Requires-Dist: databricks-sdk==0.101.0; extra == 'all'
Requires-Dist: ipywidgets==8.1.5; extra == 'all'
Requires-Dist: mlflow-skinny>=2.17,<4.0; extra == 'all'
Requires-Dist: packaging==25.0; extra == 'all'
Requires-Dist: psutil==6.1.0; extra == 'all'
Requires-Dist: requests==2.32.3; extra == 'all'
Requires-Dist: torch>=2,<3; extra == 'all'
Provides-Extra: data
Requires-Dist: torch>=2,<3; extra == 'data'
Requires-Dist: mlflow-skinny>=2.17,<4.0; extra == 'data'
Provides-Extra: distributed
Requires-Dist: cachetools==5.5.0; extra == 'distributed'
Requires-Dist: cloudpickle==3.1.2; extra == 'distributed'
Requires-Dist: databricks-sdk==0.101.0; extra == 'distributed'
Requires-Dist: ipywidgets==8.1.5; extra == 'distributed'
Requires-Dist: mlflow-skinny>=2.17,<4.0; extra == 'distributed'
Requires-Dist: packaging==25.0; extra == 'distributed'
Requires-Dist: psutil==6.1.0; extra == 'distributed'
Requires-Dist: requests==2.32.3; extra == 'distributed'
Version: 0.1.0

# 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:

```bash
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:

```bash
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:

```bash
pip install databricks-sdk-air[distributed]
```

```python
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/

