# Databricks AI Runtime

> Databricks AI Runtime (`air`) CLI — the command-line tool for submitting and managing GPU training workloads on Databricks serverless compute. Use for: running `air` workloads, custom Docker image setup, environment configuration, and troubleshooting `air` jobs.

- Skill: `gabrielmoreira/databricks-ai-runtime` (Agent Skill, multi-file: 2 files)
- Install (CLI): `npx skillmds@latest add gabrielmoreira/databricks-ai-runtime`
- Raw SKILL.md: https://api.skillmd.com/api/skills/gabrielmoreira/databricks-ai-runtime/raw
- Safety review: pending (external: skill-scanner PASS, skillspector PASS)
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: DevOps & Infra
- Author: gabrielmoreira (https://skillmd.com/u/gabrielmoreira)
- Updated: 2026-09-09
- Page: https://skillmd.com/skills/gabrielmoreira/databricks-ai-runtime

---


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

```yaml
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](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`