Tens of thousands of training and evaluation runs sit behind the document-parsing models in Databricks AI Functions such as ai_parse_document and ai_extract. The company’s AI Research team described this week how it gets that work done on Databricks AI Runtime (AIR), the GPU platform it has used since the product’s first year.
The team’s research covers Genie, Smart Routing and AI Functions. It has recently published work on making search three times faster, new techniques for offline reinforcement learning and complex document extraction. At any moment there are hundreds of research workstreams in flight.
The scale of the work
Models range from a few billion parameters to nearly a trillion. Depending on the job, the team trains them with supervised fine-tuning, offline reinforcement learning or online reinforcement learning, on NVIDIA H100 or B300 GPUs. The largest runs need hundreds of GPUs working together for days or weeks, linked by high-bandwidth networking.
Hardware that size fails. AIR runs preflight checks on assigned GPU nodes to confirm they are healthy and can talk to each other before training starts. During a run it watches for NCCL communication failures, thermal or power throttling and GPU driver problems. If a bad node kills a job, AIR cordons it off so later jobs do not land on the same machine. The team said this matters most for projects such as KARL, where post-training models of more than 100 billion parameters can take thousands of GPU hours.
Day-to-day tooling
Researchers change configurations and code constantly. With the AIR CLI they launch workloads programmatically, name the GPUs they need and leave provisioning and scheduling to the platform, including distributed training across clouds. The same custom Docker images used in local development run in training, and Databricks Artifact Registry shares them across experiments. Faster launch times mean jobs start in seconds.
AIR creates and manages MLflow experiments automatically, so the team can follow job status, stream logs and pull metrics. All research data lives in Unity Catalog, which AIR exposes through FUSE mounting as if it were a local filesystem, even across workspaces.
Capacity is shared by priority. A debugging run may need a few GPUs while a big training job needs hundreds. AIR’s prioritization and preemption let urgent experiments go first, and spare capacity can feed exploratory runs over a weekend.
Agents run the experiments
Inspired by the autoresearch idea popularized by Andrej Karpathy, the team increasingly lets agents propose changes, launch experiments, judge results and decide what to try next. Across AI Functions research, the number of training jobs has grown eightfold. When agents launch hundreds of jobs, a slow start or a vague error message becomes a bottleneck, so the team wants fast launches, programmatic interfaces and clear errors. Through the Databricks CLI and databricks-air-skills, its agents start concurrent workloads, check status and read MLflow metrics.
AIR entered public preview in March 2026 as a serverless GPU service for every Databricks customer.
Sources
- How our AI Research team uses AI Runtime
- Introducing AI Runtime: Scalable, serverless NVIDIA GPUs on Databricks for training and finetuning | Qi Zheng
