Custom Docker images for Databricks AI Runtime

Build a custom Docker image when your training workload needs a CUDA extension (flash-attn, apex, xformers, custom kernels), a non-standard framework, or any system-level dependency that doesn't fit a plain pip install. Reference the resulting image from the air YAML via environment.docker_image.url and register it once with air register image (see the last section).

Two paths:

1. Databricks base images (recommended) — CUDA, NCCL, and the correct AWS/Azure networking stack are pre-configured. You only add your ML framework and code.
2. From scratch — full control, but you install CUDA, NCCL, and the cloud-specific RDMA layer yourself.