DatabricksEthan Tang, Ivan Zhou, Irene Dea, Steven Chen, Bruce Fontaine, Ben HansenProblem & solution
Databricks’ AI Research team uses AI Runtime (AIR) to manage hundreds of parallel research workstreams for products like Genie, Smart Routing, and AI Functions, training models from a few billion to nearly a trillion parameters. AIR performs preflight checks on GPU nodes to ensure health and communication before training, and monitors node health during execution to detect issues like NCCL failures or thermal throttling. The AIR CLI and custom Docker containers enable fast experimentation by launching jobs in seconds without rebuilding environments, supporting both supervised fine-tuning and reinforcement learning. MLflow integration automatically tracks experiments, while Unity Catalog FUSE mounting gives governed access to research data across workspaces without changing training code. The team has scaled to agent-led training, where AI agents propose changes, launch trials, and evaluate results, increasing training job volume by 8x across AI Functions research workloads.
Databricks’ AI Research team develops models for products like Genie, Smart Routing, and AI Functions. They face challenges with resiliency for long training runs, fast experimentation, and GPU fleet management. They use AI Runtime (AIR) to address these challenges, leveraging built-in GPU resiliency, fast job launches with CLI and custom containers, MLflow integration, and Unity Catalog data access. They have also adopted agent-led training, increasing experimentation volume by 8x.
The team adopted Databricks AI Runtime (AIR) as their GPU infrastructure platform. AIR provides built-in GPU resiliency with preflight checks and node health monitoring. The AIR CLI allows programmatic job launches with specified GPUs, and custom Docker containers ensure environment consistency. MLflow integration automatically tracks experiments, and Unity Catalog FUSE mounting provides governed data access. Workload prioritization and preemption help manage GPU capacity across projects. These features support both traditional research workflows and agent-led training.
AIR performs preflight checks on GPU nodes and monitors health during execution to detect issues like NCCL failures or thermal throttling, cordoning off unhealthy nodes to prevent future scheduling on problematic infrastructure.
The team uses the AIR CLI and custom Docker containers to launch training jobs in seconds, enabling fast iteration without rebuilding environments, and supports both SFT and RL workloads.
Agent-led training has increased training job volume by 8x, and AIR supports this by providing programmatic interfaces, fast job launches, clear errors, reliable execution, and accessible training results.
AIR performs preflight checks on assigned GPU nodes to ensure they are healthy and can communicate before training starts. It monitors node health during execution and detects issues like NCCL communication failures or thermal throttling. If a node becomes unhealthy and causes a job to fail, AIR cordons off the affected node so subsequent jobs aren't scheduled on it.
AIR provides workload prioritization and preemption capabilities. This allows the team to allocate resources across projects, prioritize time-sensitive experiments, and use excess capacity for exploratory work without manual coordination.
Agents interact with AIR programmatically through the Databricks CLI and databricks-air-skills. They can launch concurrent workloads, monitor job status, retrieve MLflow metrics, and iterate based on results. AIR's fast job launches, built-in reliability, and scalability are well suited for this automated experimentation cycle.
Summaries are generated from the original by a language model. Check the original before relying on it.
See the idea
AI-DLC is an AI-driven development lifecycle that runs delivery through Inception, Design, and Construction, with specs, steering files, and hooks keeping agents aligned to agreed intent.
Keep reading
The Good Stuff
Paid picks we use ourselves. Links carry a referral code — a subscription through one earns us a commission at no extra cost to you.
One useful email. Unsubscribe anytime.
