Databricks details AI Runtime workflow for agent-led model research

Databricks said its AI Research team uses AI Runtime to provision and schedule model training. The work ranges from small experiments to runs occupying hundreds of GPUs for days or weeks. According to the company, those workloads include supervised fine-tuning, offline reinforcement learning and online reinforcement learning on models ranging from a few billion parameters to nearly one trillion.

Researchers submit their GPU requirements programmatically through the AI Runtime command-line interface. AI Runtime then handles provisioning and scheduling, including high-bandwidth networking for distributed training across clouds. Databricks said workload priorities and preemption protect time-sensitive experiments, while lower-priority exploratory jobs can use spare capacity.

Before a job starts, the runtime checks the GPU nodes. During execution, it monitors for communication failures, thermal or power throttling, and driver problems. If one of those faults causes a job to fail, Databricks said the system cordons the unhealthy node so later jobs are not placed there. A separate account of its GPU reliability process describes bootstrap checks, continuous monitoring and multi-node fabric tests, followed by reset and retesting for quarantined machines.

A CLI submission produces a Databricks job run for execution details and an MLflow run for every attempt, according to the company’s observability documentation. The MLflow records can retain parameters, metrics, GPU system measurements, logs and artifacts, allowing researchers to compare attempts rather than treat a retry as the same run. Teams can also use Declarative Automation Bundles to schedule GPU work and connect CPU preprocessing with GPU training in a dependency graph.

In the internal loop Databricks described, software agents propose model or training changes, submit several experiments in parallel, evaluate training metrics and use the results to plan another round. The company claims the number of training jobs across its AI Functions research workloads grew eightfold as it adopted agent-driven research. Databricks did not disclose a baseline job count, measurement period, inclusion rules or methodology, and the increase has not been independently verified.

AI Runtime and its CLI commands remain in public preview, according to Databricks documentation. A workspace administrator must enable the features before users can submit workloads.