Databricks AI Runtime GA: Gate the GPU Migration

A training script is only one part of a production GPU workload. Teams also need capacity, reproducible environments, controlled data access and evidence that each model change improves the task.

Managed GPU infrastructure can reduce infrastructure work. It does not replace the organisation's responsibility for data policy, evaluation and release approval.

The right starting point is a bounded pilot, not a platform-wide migration. GPU migration can weaken release controls; we recommend piloting one reproducible AI Runtime workload before expanding against quality, cost and governance gates.

What Changed

That pilot now has a new compute option: Databricks AI Runtime (AIR) is generally available.

The Databricks announcement introduces managed GPU infrastructure within Databricks. The verified AWS documentation describes fully managed serverless GPU compute for training, fine-tuning and broader deep learning workloads. Serverless means the provider manages the compute infrastructure rather than the customer operating GPU nodes.

The AWS documentation lists on-demand A10 and H100 GPUs, accessible through the Databricks CLI, notebooks, jobs and Ray. It also describes pre-built environments and dataloaders, or a user-provided Docker container. AIR includes MLflow experiment tracking, logging and debugging agents.

These capabilities create an opportunity to replace a separately operated GPU stack. They provide the technical basis for a pilot, not evidence that a particular workload will improve.

Why It Matters for Platform Teams

Replacing that stack changes the platform team's work, but the team still owns the release process.

We would focus on reproducibility: repeating a run from recorded code, data and environment inputs. Record the commit, dataset reference, container digest or environment definition, parameters and evaluation dataset reference together. This record supports failure investigation and model review.

Compare the complete cost of producing an acceptable model, including failed attempts, data preparation and engineering support. A lower training bill alone does not establish that migration was worthwhile. Serving remains a separate decision, covered in our inference cost guide.

Use Case: Governed Document Model Training

To test that value, consider a hypothetical logistics organisation with operational systems on-premises and an Azure data platform.

A small machine learning team trains a document classification model for back-office processing. The documents contain commercially sensitive information and occasional personal data. The team needs regular model updates without operating a separate GPU platform.

Its existing process exports training data to an independent GPU environment. Engineers maintain another access model, dependency set and collection of training logs. During approval, they must reconcile which data, code and environment produced each candidate.

For the proposed AIR pilot, the team selects one existing workload with a known evaluation baseline. It defines the processing purpose, data owner and retention before preparing a minimised dataset. Development and test use synthetic or anonymised documents.

The proposed flow starts with approved source data in bronze storage. Preparation creates a governed silver training dataset, and AIR trains the candidate model. MLflow records the experiment, while a fixed evaluation set determines whether the candidate meets the quality requirement.

Under our delivery standard, the team registers approved models in Unity Catalog. This is a proposed architecture, not a verified client deployment. Azure eligibility and security verification remain prerequisites.

The team records these measures before deciding whether to expand:

Implementation

These measures need a release contract and a controlled deployment process before the first pilot run.

1. Confirm the deployment boundary

Establish the service boundary before preparing sensitive inputs. Treat unresolved eligibility or security questions as blockers.

1. Verify the target cloud, EU region and workspace eligibility.
2. Record supported GPU options, prerequisites, pricing and service constraints in a design decision.
3. Define private network paths for storage, container registry, Key Vault and monitoring access.
4. Test DNS resolution and access from the training execution context.
5. Provision supporting infrastructure through version-pinned NexOps-managed Terraform modules.
6. Define Entra ID group access and privileged access through PIM.
7. Verify managed identity or workload identity federation support for each execution stage.
8. Test effective Unity Catalog permissions and workload isolation.

Keep secrets out of code and pipeline variables. Store unavoidable secrets in the environment's Key Vault.

2. Freeze the experiment inputs

Once the deployment boundary is confirmed, make the experiment repeatable. Keep training logic in a Python package rather than orchestration configuration.

1. Select a representative existing workload with an established baseline.
2. Record its code commit, dataset reference, dependency lock file and parameters.
3. For a custom container, record an immutable image digest.
4. Define checkpoint retention and output ownership.
5. Create a versioned evaluation set and keep it separate from training data.
6. Agree on the quality metric and permitted training spend.

3. Add an explicit acceptance check

Use the frozen inputs to evaluate the candidate independently of the training submission mechanism. This project-level Python script checks a higher-is-better quality metric and an approved training budget.

# src/training_gate.py
import argparse
import json
import math
from pathlib import Path
 
 
def read_number(path: Path, key: str) -> float:
    payload = json.loads(path.read_text(encoding="utf-8"))
    value = payload[key]
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ValueError(f"{key} must be numeric")
    if not math.isfinite(value):
        raise ValueError(f"{key} must be finite")
    return float(value)
 
 
def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--baseline", type=Path, required=True)
    parser.add_argument("--candidate", type=Path, required=True)
    parser.add_argument("--policy", type=Path, required=True)
    args = parser.parse_args()
 
    baseline_quality = read_number(args.baseline, "quality_score")
    candidate_quality = read_number(args.candidate, "quality_score")
    training_cost = read_number(args.candidate, "training_cost_eur")
    approved_budget = read_number(args.policy, "max_training_cost_eur")
 
    # Negative amounts are invalid for this budget check.
    if training_cost < 0 or approved_budget < 0:
        raise SystemExit("Training cost and budget must be non-negative")
 
    if candidate_quality < baseline_quality:
        raise SystemExit("Candidate failed the quality gate")
    if training_cost > approved_budget:
        raise SystemExit("Candidate exceeded the approved training budget")
 
 
if __name__ == "__main__":
    main()

The script consumes project-defined JSON files; these fields are not AIR API fields. The pipeline must generate them from recorded evaluation results and reconciled training costs, including failed attempts. The numeric checks cannot detect missing cost records.

1. Add unit tests for missing fields, invalid values, quality regression and budget failure.
2. Attach the MLflow run reference and input records to the release artefact.
3. Submit passing results for review.

A passing check provides review evidence, not automatic deployment permission. Use the same evaluation definition for baseline and candidate results.

4. Connect orchestration through reviewed configuration

The acceptance check then becomes one stage in the deployment process. Keep training submission, evaluation and evidence collection separate.

1. Deploy orchestration through Declarative Automation Bundles with shared templates, group permissions and a service principal execution identity.
2. Authenticate deployment pipelines through workload identity federation.
3. Run required tests, security scans, Terraform checks and bundle validation on every pull request.
4. Build the application artefact once.
5. Promote it through dev, test, acc and main, with main deploying to prd.
6. Require approvals for acceptance and production.
7. Run smoke tests after deployment.
8. If a smoke test fails, stop promotion.

The AWS scheduling guide defines GPU training through ai_runtime_task within a bundle job. It connects preprocessing through depends_on and defines recurring runs through schedule. The bundle identity documentation supports a service principal execution identity through run_as with service_principal_name.

Our bundle CI/CD guide explains the deployment structure. Give each environment separate identities, workspaces, catalogs, Key Vault and deployment state.

Pitfalls and When Not to Use It

Those promotion gates address the main risk: general availability does not establish suitability for every deployment.

Cloud scope is not interchangeable. The verified GPU and access details come from AWS documentation. They do not establish Azure regional availability or capacity guarantees.

Performance needs workload evidence. Measure cold and repeated launches, data loading, checkpoint writes and total experiment duration. Do not use an advertised launch time as an end-to-end training estimate.

Managed compute does not prove the security design. Test denied access to training data, checkpoints, logs and outputs. Confirm that execution does not depend on public network paths.

A container is not the complete experiment. Mutable data, uncontrolled downloads and unrecorded parameters can change model behaviour despite an unchanged image. Remove uncontrolled downloads and retain the complete input record.

Debugging assistance is not release evidence. An agent's suggested fix needs code review and a new evaluated run. A plausible diagnosis cannot replace a reproducible test.

Keep an existing GPU platform when it meets requirements and migration adds more operational work than it removes. Avoid custom training when prompting or retrieval already solves the task. Our fine-tuning decision guide explains that choice.

Key Takeaways

1. Select a bounded workload whose current results provide a useful comparison.
2. Resolve cloud, identity and network questions before introducing sensitive inputs.
3. Preserve the complete experiment record, not only the container image.
4. Gate promotion on evaluated quality, reconciled cost and reviewed evidence.
5. Expand only when the pilot supports both the model requirements and the operating process.