Databricks AI Runtime SDK
=========================

The Databricks AI Runtime SDK (``databricks-sdk-air``) runs multi-GPU workloads on
`Databricks AI Runtime <https://docs.databricks.com/aws/en/machine-learning/ai-runtime>`_
directly from `Databricks notebooks <https://docs.databricks.com/aws/en/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 <https://api-docs.databricks.com/python/serverless_gpu/index.html>`_
(``serverless_gpu``). To migrate, see
`Migrate from serverless_gpu <https://docs.databricks.com/aws/en/machine-learning/ai-runtime/distributed-training#migrate-serverless-gpu>`_.

For concepts, usage guides, and examples, see the `AI Runtime documentation <https://docs.databricks.com/aws/en/machine-learning/ai-runtime>`_.

.. toctree::
   :maxdepth: 1

   api/reference
   api/data
   api/environment_variables

Indices and tables
==================

* :ref:`genindex`
* :ref:`modindex`
* :ref:`search`