gabrielmoreira Databricks AI Runtime (air) CLI
Databricks AI Runtime (air) is a CLI tool for submitting GPU training workloads to Databricks serverless compute. It manages environment setup, distributed training configuration, and workload lifecycle — without requiring you to manage clusters or infrastructure.
A typical workload YAML looks like:
experiment_name: my-training-job
compute:
num_accelerators: 1
accelerator_type: GPU_1xA10
environment:
dependencies:
- mlflow
version: "AI5"
command: echo "Hello World"
Submit with air run --file workload.yaml -p <databricks_config_profile>.
Bring your own custom Docker images
Use a custom Docker image instead of environment.version when your workload needs specific system libraries, CUDA extensions (flash-attn, apex, custom kernels), or dependencies that don't fit environment.dependencies.
Read docker-images.md for step-by-step guidance on:
- Using Databricks-provided base images
- Dockerfile patterns
- Pre-build compatibility checklist (CUDA/driver, PyTorch, NCCL, EFA/RDMA)
- Registering images with
air register image