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.
68
83%
Does it follow best practices?
Run evals on this skill
Adds up to 20 points to the overall score
View guide
Passed
No findings from the security scan
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
- Repository
- databricks/databricks-agent-skills
- Commit
691deaa
- Last updated
- First committed
Is this your skill?
If you maintain this skill, you can claim it as your own. Once claimed, you can manage eval scenarios, bundle related skills, attach documentation or rules, and ensure cross-agent compatibility.