How We Productionise Data and AI on Databricks: Bundles, CI/CD, Testing, MLOps and Operations

A notebook that produces the right numbers once is a prototype. A production system produces the right numbers every day, after every change, at a known cost, and tells someone when it can't. Most of the work between the two is engineering discipline, not data science. This is the playbook we apply on Databricks engagements, for pipelines and for machine-learning and generative-AI systems.

What principles do we hold ourselves to?

1. Everything is code: transformations, jobs, pipelines, permissions, dashboards and alerts, all reviewed and versioned in Git.
2. Isolated environments: separate dev, staging and prod catalogs (and usually workspaces). Nothing is changed by hand in prod.
3. Promote what you tested: the commit and configuration that passed staging are what deploy to prod. Stricter teams build the wheel once and pass the same file between stages.
4. Tests at every layer: unit tests for logic, contracts for interfaces, expectations at runtime, integration tests in staging.
5. Idempotent and replayable: any run can be retried or backfilled without corrupting data.
6. Observable: freshness, volume, quality, cost and model behaviour are measured and alerted on.
7. Owned: every job has an owner, an on-call route and a runbook.

How do we structure the repository?

customer360/
├── databricks.yml              # bundle: targets, variables, artefacts
├── resources/
│   ├── jobs.yml                # job and task definitions
│   └── pipelines.yml           # Lakeflow declarative pipelines, if used
├── src/customer360/
│   ├── transforms/             # pure functions: DataFrame in -> DataFrame out
│   ├── io.py                   # reading and writing tables (thin, injectable)
│   └── entrypoints.py          # build_silver(), build_gold(): parse params, call transforms
├── tests/
│   ├── unit/                   # fast, local Spark, no workspace needed
│   └── integration/            # run as a job in staging against real tables
├── pyproject.toml
└── .github/workflows/ci.yml

The key design rule: business logic lives in pure functions, not in notebooks and not tangled with I/O. Notebooks are fine for exploration, and as thin entry points. That separation is what makes the logic testable in seconds.

# src/customer360/transforms/customers.py
from pyspark.sql import DataFrame, functions as F

def standardise_customers(df: DataFrame) -> DataFrame:
    return (
        df.withColumn("email", F.lower(F.trim("email")))
          .withColumn("country", F.upper(F.coalesce("country", F.lit("UNKNOWN"))))
          .dropDuplicates(["customer_id"])
    )

# tests/unit/test_customers.py
from pyspark.sql import SparkSession
from pyspark.testing import assertDataFrameEqual   # Spark 3.5+
from customer360.transforms.customers import standardise_customers

def test_standardise_customers(spark: SparkSession):
    raw = spark.createDataFrame(
        [(1, "  [email protected] ", None), (1, "[email protected]", "in"), (2, "[email protected]", "gb")],
        "customer_id INT, email STRING, country STRING",
    )
    out = standardise_customers(raw).orderBy("customer_id")
    assert out.count() == 2
    assert out.filter("customer_id = 1").first()["email"] == "[email protected]"
    assertDataFrameEqual(
        out.select("customer_id").distinct(),
        spark.createDataFrame([(1,), (2,)], "customer_id INT"),
    )

A session-scoped spark fixture in conftest.py creates a local SparkSession once. The whole unit suite runs in CI in a minute or two, with no cluster.

How do Databricks Asset Bundles fit in?

Databricks Asset Bundles describe a project's jobs, pipelines and other resources as YAML alongside the code, with targets for each environment. The Databricks CLI validates, deploys and runs them.

# databricks.yml
bundle:
  name: customer360

include:
  - resources/*.yml

artifacts:
  default:
    type: whl
    build: python -m build --wheel
    path: .

variables:
  catalog:
    description: Unity Catalog catalog the jobs read and write

targets:
  dev:
    mode: development            # prefixes resource names per developer, pauses schedules
    default: true
    workspace:
      host: https://adb-1111111111111111.11.azuredatabricks.net
    variables:
      catalog: dev

  staging:
    workspace:
      host: https://adb-2222222222222222.12.azuredatabricks.net
    variables:
      catalog: staging
    run_as:
      service_principal_name: 11111111-2222-3333-4444-555555555555

  prod:
    mode: production             # enforces production safeguards
    workspace:
      host: https://adb-3333333333333333.13.azuredatabricks.net
    variables:
      catalog: prod
    run_as:
      service_principal_name: 66666666-7777-8888-9999-000000000000

# resources/jobs.yml
resources:
  jobs:
    customer360_daily:
      name: customer360-daily
      max_concurrent_runs: 1
      timeout_seconds: 7200
      schedule:
        quartz_cron_expression: "0 30 2 * * ?"
        timezone_id: Asia/Kolkata
      email_notifications:
        on_failure: [[email protected]]
      parameters:
        - name: catalog
          default: ${var.catalog}
        - name: run_date
          default: ""                       # empty = yesterday; set explicitly for backfills
      job_clusters:
        - job_cluster_key: main
          new_cluster:
            spark_version: 16.4.x-scala2.12     # pin an LTS runtime
            node_type_id: Standard_D8ds_v5
            autoscale: { min_workers: 2, max_workers: 8 }
            data_security_mode: SINGLE_USER
      tasks:
        - task_key: build_silver
          job_cluster_key: main
          max_retries: 2
          min_retry_interval_millis: 300000
          python_wheel_task:
            package_name: customer360
            entry_point: build_silver
          libraries:
            - whl: ./dist/*.whl
        - task_key: build_gold
          depends_on: [{ task_key: build_silver }]
          job_cluster_key: main
          python_wheel_task:
            package_name: customer360
            entry_point: build_gold
          libraries:
            - whl: ./dist/*.whl

entry_point refers to an entry point declared in pyproject.toml (for example under [project.scripts]), not just a function name.

Many teams now run jobs on serverless compute instead of job clusters, declaring dependencies through job environments. That removes cluster sizing and start-up time. The bundle structure is the same either way.

What does the CI/CD pipeline look like?

# .github/workflows/ci.yml
name: ci
on:
  pull_request:
  push:
    branches: [main]

jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with: { python-version: "3.12" }
      - run: pip install -e ".[dev]"
      - run: ruff check . && pytest tests/unit -q

  deploy-staging:
    needs: test
    if: github.ref == 'refs/heads/main'
    runs-on: ubuntu-latest
    environment: staging
    env:
      DATABRICKS_HOST: ${{ vars.DATABRICKS_HOST }}
      DATABRICKS_CLIENT_ID: ${{ secrets.DATABRICKS_CLIENT_ID }}        # service principal (OAuth M2M)
      DATABRICKS_CLIENT_SECRET: ${{ secrets.DATABRICKS_CLIENT_SECRET }}
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with: { python-version: "3.12" }
      - run: pip install build                      # used by the bundle's wheel artefact build
      - uses: databricks/setup-cli@main             # pin to a release tag in your repo
      - run: databricks bundle validate -t staging
      - run: databricks bundle deploy -t staging
      - run: databricks bundle run -t staging integration_tests     # an integration-test job defined in resources/

  deploy-prod:
    needs: deploy-staging
    runs-on: ubuntu-latest
    environment: production          # GitHub environment with required reviewers
    env:
      DATABRICKS_HOST: ${{ vars.DATABRICKS_HOST }}
      DATABRICKS_CLIENT_ID: ${{ secrets.DATABRICKS_CLIENT_ID }}
      DATABRICKS_CLIENT_SECRET: ${{ secrets.DATABRICKS_CLIENT_SECRET }}
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with: { python-version: "3.12" }
      - run: pip install build
      - uses: databricks/setup-cli@main
      - run: databricks bundle deploy -t prod

Pull requests run linting and unit tests. Merging to main deploys to staging and runs integration tests. Production deploys need an approval and use a separate service principal. Where the platform supports it, prefer OIDC workload identity federation from GitHub to Databricks over long-lived client secrets. People have read access to production data, but only the deployment identity can change production jobs and tables.

How do we test data, not just code?

Unit tests prove the logic. They do not prove that today's data is sane. We layer the checks:

Layer	What it catches	How
Unit tests	Logic bugs	pytest + local Spark, small hand-built DataFrames
Contracts	Breaking schema changes between producers and consumers	Schema assertions in tests; schema registry for events
Runtime expectations	Bad records, nulls, invalid ranges	Pipeline expectations or explicit checks that drop, quarantine or fail
Reconciliation	Missing or duplicated data	Row counts and checksums vs source, per day and key range
Integration tests	Wiring, permissions, performance	Full job in staging on representative data
Anomaly monitors	Silent drift in volume and distributions	Table monitoring and alerts on metrics history

Quarantine, don't discard. Records that fail validation go to a _quarantine table with the failure reason, so they can be fixed and replayed. Silently dropping data is how numbers drift without anyone noticing.

How do we make runs idempotent and backfills boring?

Every job takes its processing window as a parameter (run_date, or a start and end), and writes in a way that makes re-running the same window produce the same result:

def write_daily(df, table: str, run_date: str) -> None:
    (df.write
       .mode("overwrite")
       .option("replaceWhere", f"order_date = '{run_date}'")   # atomically replace one day only
       .saveAsTable(table))

• Batch facts: overwrite the window with replaceWhere, as above, or MERGE on natural keys.
• Incremental streams: checkpoints plus idempotent writes. See our CDC post for sequence-guarded merges and txnAppId / txnVersion.
• Backfills: the same job, run with explicit window parameters, often fanned out with a for-each task and bounded concurrency. No special backfill code means no special backfill bugs.

How do we take models to production?

Tracking and registration. Every training run logs parameters, metrics, the code version, the training-data version and the model itself to MLflow. Models are registered in Unity Catalog as catalog.schema.model, which gives them the same permissions, lineage and audit trail as tables.

import mlflow
from mlflow import MlflowClient

mlflow.set_registry_uri("databricks-uc")

class ChurnProbability(mlflow.pyfunc.PythonModel):
    """Serve the positive-class probability, not the 0/1 label that predict() returns."""
    def __init__(self, model):
        self.model = model

    def predict(self, context, model_input, params=None):
        return self.model.predict_proba(model_input)[:, 1]

with mlflow.start_run(run_name="churn-gbdt") as run:
    mlflow.log_params(params)
    mlflow.log_metrics({"val_aucpr": val_aucpr, "val_logloss": val_logloss})
    info = mlflow.pyfunc.log_model(
        name="model",                               # MLflow 3 signature
        python_model=ChurnProbability(model),
        input_example=X_valid.head(5),
        registered_model_name="prod.ml.churn",
    )

client = MlflowClient()
client.set_registered_model_alias("prod.ml.churn", "challenger", info.registered_model_version)

Promotion with a gate. A deployment job evaluates @challenger against @champion on the same recent hold-out data: business metric, calibration, fairness slices and latency. It moves the champion alias only if the challenger wins on agreed thresholds. Rollback means pointing the alias back.

Serving. For batch scoring we load the alias inside Spark. For real-time scoring we use a Model Serving endpoint.

score = mlflow.pyfunc.spark_udf(spark, "models:/prod.ml.churn@champion", result_type="double")
scored = features.withColumn("churn_score", score(*feature_columns))

Because jobs reference the alias, not a version number, promotion and rollback need no code change or redeploy.

Monitoring. We watch input feature drift, prediction distribution drift and, as labels arrive, real performance by segment. Retraining is triggered by schedule, by drift thresholds or by performance decay. It always flows through the same evaluation gate.

Generative AI systems get the same discipline, adapted. Prompts and retrieval configurations are versioned. Evaluation datasets mix real questions with expected answers or rubrics. LLM-judge scorers are calibrated against human ratings. Request tracing captures retrieval, tool calls and token cost, and production traces are sampled back into the evaluation set. Our posts on LLM fine-tuning and RAG in production go deeper.

What do we monitor and alert on?

We define service-level objectives for data products, not just for jobs:

SLO	Example target
Freshness	gold.daily_revenue updated by 06:00 IST on 99% of days
Completeness	Daily row count within ±3σ of the trailing 28-day pattern
Validity	< 0.1% of records quarantined
Model quality	Weekly AUC-PR on labelled outcomes ≥ agreed floor
Cost	Pipeline cost per day within budget, with a 20% alert threshold

Databricks system tables make much of this queryable without extra tooling. For example, failed job runs in the last day:

SELECT job_id, run_id, result_state, period_start_time, period_end_time
FROM system.lakeflow.job_run_timeline
WHERE result_state = 'FAILED'
  AND period_start_time >= current_timestamp() - INTERVAL 1 DAY;

Similar queries over system.billing.usage attribute cost to jobs and tags, and system.access.audit tracks who accessed what. Alerts route to the owning team's on-call channel, and every alert links to a runbook.

How do we control cost and keep it secure?

• Compute policies cap cluster sizes and enforce tags and auto-termination. Serverless and budget policies attribute spend to teams and projects.
• Right-size by evidence: use query profiles and job metrics, not guesswork. Most savings come from data layout (fewer files scanned) and incremental processing, not bigger clusters.
• Secrets live in secret scopes or Unity Catalog connections, never in code or notebook outputs.
• Least privilege: groups, not individuals, hold grants; PII is protected by row filters and column masks; production writes are limited to service principals.

What is our definition of done?

A data or AI product is production-ready when:

• Code, jobs and permissions are deployed from Git by CI/CD to a prod target
• Unit tests, contracts and runtime expectations exist, and integration tests pass in staging
• Runs are parameterised, idempotent and backfillable
• Freshness, volume, quality and cost SLOs are defined and alerted
• Models are registered in Unity Catalog with an evaluation gate, alias-based promotion and monitoring
• Lineage is visible, and PII is governed with filters and masks
• An owner, on-call route and runbook exist, and a failure drill has been run once

None of this is glamorous, and all of it is what makes data and AI trustworthy enough for a business to act on.