MLOps on Databricks

Databricks consolidates the MLOps stack around MLflow and Unity Catalog. Feature tables live in Unity Catalog as Delta tables. MLflow — fully managed — handles experiment tracking and the model registry. Lakehouse Monitoring runs drift detection on any Delta table. Model Serving deploys models as REST endpoints. The data engineer’s surface area is feature pipelines, monitoring setup, and batch inference jobs.

Feature Engineering Client

FeatureEngineeringClient is the interface to Databricks Feature Store. It wraps Delta tables with the metadata needed for point-in-time lookups and training-serving consistency.

from databricks.feature_engineering import FeatureEngineeringClient, FeatureLookup

fe = FeatureEngineeringClient()

Creating a feature table

fe.create_table(
    name="ml_catalog.churn.customer_features",
    primary_keys=["customer_id"],
    timestamp_keys=["feature_timestamp"],   # enables point-in-time lookups
    df=features_df,
    description="Customer behavioral features for churn prediction"
)

timestamp_keys is what enables point-in-time lookups during training — without it you get current values, which causes leakage on historical training data. It expects a timestamp column in the DataFrame.

Writing and updating

# Append new rows (for daily feature jobs)
fe.write_table(
    name="ml_catalog.churn.customer_features",
    df=new_features_df,
    mode="merge"   # upsert on primary_keys; use "overwrite" only to rebuild
)

Creating a training set

labels_df = spark.sql("""
    SELECT customer_id, observation_date, churned_within_30d
    FROM gold.churn_labels
    WHERE observation_date BETWEEN '2024-01-01' AND '2024-06-30'
""")

training_set = fe.create_training_set(
    df=labels_df,
    feature_lookups=[
        FeatureLookup(
            table_name="ml_catalog.churn.customer_features",
            lookup_key="customer_id",
            timestamp_lookup_key="observation_date"  # point-in-time: gets features AT OR BEFORE this date
        )
    ],
    label="churned_within_30d"
)

training_df = training_set.load_df()

The timestamp_lookup_key is the leakage prevention mechanism. For each row in labels_df, the feature store finds the most recent feature snapshot that existed at or before observation_date. The model trains on what it would have seen at prediction time.

Batch inference with score_batch

Use score_batch rather than loading features manually — it guarantees the same feature definitions used during training.

customers_to_score = spark.sql("""
    SELECT customer_id FROM silver.customers WHERE is_active = TRUE
""")

predictions = fe.score_batch(
    model_uri="models:/ml_models.production.churn_predictor@champion",
    df=customers_to_score   # only needs the lookup keys
)
# predictions: customer_id | prediction

Publishing to an online store

from databricks.feature_engineering.online_store_spec import AmazonDynamoDBSpec

fe.publish_table(
    name="ml_catalog.churn.customer_features",
    online_store_spec=AmazonDynamoDBSpec(
        region="eu-west-1",
        table_name="churn_features_online"
    ),
    mode="merge"
)

This syncs the latest feature values from the offline Delta table to DynamoDB. Run it at the end of your nightly feature job.

API reference

Method	Purpose
fe.create_table(name, primary_keys, timestamp_keys, df)	Create feature table
fe.write_table(name, df, mode)	Write/upsert features
fe.read_table(name)	Read feature table as DataFrame
fe.create_training_set(df, feature_lookups, label)	Point-in-time training data
fe.score_batch(model_uri, df)	Batch inference with consistent features
fe.publish_table(name, online_store_spec, mode)	Sync to online store

MLflow tracking

MLflow on Databricks is fully managed — no server to configure. Experiments are stored at workspace paths; runs within each experiment capture everything needed to reproduce training.

import mlflow
from mlflow.tracking import MlflowClient

mlflow.set_experiment("/Shared/churn/model_development")

with mlflow.start_run(run_name="xgboost_v3") as run:
    mlflow.log_params({
        "max_depth": 6,
        "learning_rate": 0.05,
        "n_estimators": 300,
        "subsample": 0.8
    })

    # ... train model ...

    mlflow.log_metrics({
        "auc": 0.91,
        "f1": 0.84,
        "precision": 0.80,
        "recall": 0.89
    })

    # Log data version — MLflow won't do this automatically
    mlflow.set_tag("training_data_version", "2024-06-30")
    mlflow.set_tag("git_commit", "abc1234")
    mlflow.set_tag("feature_table", "ml_catalog.churn.customer_features")

    mlflow.sklearn.log_model(model, "model")
    # or: mlflow.xgboost.log_model(model, "model")

The git_commit and training_data_version tags need to be set explicitly. Without them, the run is not reproducible.

Autologging

For supported frameworks (sklearn, XGBoost, LightGBM, PyTorch), autologging captures params, metrics, and the model artifact automatically:

mlflow.sklearn.autolog()
# or mlflow.xgboost.autolog()

with mlflow.start_run():
    model.fit(X_train, y_train)
    # params, metrics, model artifact logged automatically

Autologging logs what it knows — it still doesn’t know your data version or git commit. Set those tags manually inside the same start_run block:

import subprocess
from delta.tables import DeltaTable

mlflow.sklearn.autolog()

with mlflow.start_run():
    # Git commit — works when running from a Databricks Repo
    try:
        commit = subprocess.check_output(
            ["git", "rev-parse", "HEAD"], cwd="/Workspace/Repos/..."
        ).decode().strip()
    except Exception:
        # Fallback: Databricks job context has the repo info
        ctx = dbutils.notebook.entry_point.getDbutils().notebook().getContext()
        commit = ctx.tags().get("gitCommit").getOrElse("unknown")

    # Delta table version — pins the exact snapshot used for training
    feature_table = "ml_catalog.churn.customer_features"
    dt_version = (
        DeltaTable.forName(spark, feature_table)
        .history(1)
        .select("version")
        .first()[0]
    )

    mlflow.set_tags({
        "git_commit": commit,
        "feature_table": feature_table,
        "feature_table_version": str(dt_version),
    })

    model.fit(X_train, y_train)

The Delta table version is the most precise data pointer — it’s an integer that maps to an exact snapshot, queryable later with spark.read.format("delta").option("versionAsOf", version).table(feature_table). A date string works too but is ambiguous if the table is written to more than once per day.

If running in a Databricks Job (not a notebook), the job context exposes the git commit directly:

ctx = dbutils.notebook.entry_point.getDbutils().notebook().getContext()
commit = ctx.tags().get("gitCommit").getOrElse("unknown")
branch = ctx.tags().get("gitBranch").getOrElse("unknown")

This only works when the job is backed by a Git repo (Repos or a DAB with a git_source). For jobs running notebooks by path without a git source, use subprocess against the Repo working directory.

Querying runs

client = MlflowClient()

runs = client.search_runs(
    experiment_ids=[mlflow.get_experiment_by_name("/Shared/churn/model_development").experiment_id],
    filter_string="metrics.auc > 0.88 AND params.max_depth = '6'",
    order_by=["metrics.auc DESC"],
    max_results=10
)

best_run = runs[0]
print(best_run.info.run_id, best_run.data.metrics["auc"])

Tracking API reference

Method	Purpose
mlflow.set_experiment(path)	Set/create experiment
mlflow.start_run(run_name)	Start a run (use as context manager)
mlflow.log_param(key, value)	Log single parameter
mlflow.log_params(dict)	Log multiple parameters
mlflow.log_metric(key, value, step)	Log single metric
mlflow.log_metrics(dict)	Log multiple metrics
mlflow.log_artifact(path)	Log file
mlflow.set_tag(key, value)	Add metadata tag
mlflow.sklearn.log_model(model, name)	Log sklearn model
mlflow.sklearn.autolog()	Enable autologging
client.search_runs(experiment_ids, filter_string, order_by)	Query runs

Unity Catalog model registry

The Unity Catalog registry uses three-level naming (catalog.schema.model_name) and inherits UC permissions. Models registered here are governed the same way as tables.

mlflow.set_registry_uri("databricks-uc")
client = MlflowClient()

Registering a model

model_name = "ml_models.production.churn_predictor"

mlflow.register_model(
    model_uri=f"runs:/{run_id}/model",
    name=model_name
)
# Returns a ModelVersion object with version number

Aliases

Use aliases for production references, not version numbers. Aliases decouple deployment from versioning — moving an alias promotes the model without touching consumer code.

# Assign challenger alias to version 4
client.set_registered_model_alias(
    name=model_name,
    alias="challenger",
    version=4
)

# After validation, promote to champion (alias moves atomically)
client.set_registered_model_alias(
    name=model_name,
    alias="champion",
    version=4
)

# Remove alias from a version (optional cleanup)
client.delete_registered_model_alias(name=model_name, alias="challenger")

Loading a model

import mlflow

# Load by alias — production code uses this form
model = mlflow.pyfunc.load_model(f"models:/{model_name}@champion")

# Load by version — for inspection or comparison
model_v3 = mlflow.pyfunc.load_model(f"models:/{model_name}/3")

Listing versions

versions = client.search_model_versions(f"name='{model_name}'")
for v in versions:
    print(v.version, v.aliases, v.tags)

Registry API reference

Method	Purpose
mlflow.register_model(uri, name)	Register a run’s model
client.set_registered_model_alias(name, alias, version)	Set/move alias
client.delete_registered_model_alias(name, alias)	Remove alias
mlflow.pyfunc.load_model("models:/name@alias")	Load by alias
mlflow.pyfunc.load_model("models:/name/version")	Load by version
client.update_model_version(name, version, description)	Update description
client.search_model_versions(filter_string)	List versions

Lakehouse Monitoring

Lakehouse Monitoring attaches to any Delta table and computes statistical profiles and drift metrics on a schedule. It writes results to two tables: {table}_profile_metrics and {table}_drift_metrics.

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.catalog import MonitorCronSchedule

w = WorkspaceClient()

Setting up a monitor

w.quality_monitors.create(
    table_name="ml_catalog.churn.customer_features",
    assets_dir="/Shared/monitoring/churn_features",
    output_schema_name="ml_catalog.monitoring",

    # Baseline: what "normal" looks like. Compare current data against this.
    baseline_table_name="ml_catalog.churn.customer_features_baseline",

    # Slice metrics by a categorical dimension (e.g., subscription tier)
    slicing_exprs=["subscription_tier"],

    schedule=MonitorCronSchedule(
        quartz_cron_expression="0 0 8 * * ?",   # daily at 08:00 UTC
        timezone_id="UTC"
    )
)

The baseline table should be a stable snapshot — a date range from when the model was trained performs well in production. Refresh it when you retrain.

What it tracks

Metric	Red flag threshold
Null rate per column	Sudden increase from baseline
Mean / median	> 2 standard deviations from baseline
Distinct count	Unexpected new categories
PSI (Population Stability Index)	> 0.2 = significant drift

Querying drift metrics

-- Features with significant drift in the latest window
SELECT
    column_name,
    drift_type,
    drift_score,
    psi_score,
    baseline_mean,
    current_mean
FROM ml_catalog.monitoring.customer_features_drift_metrics
WHERE window_end = (
    SELECT MAX(window_end)
    FROM ml_catalog.monitoring.customer_features_drift_metrics
)
AND (drift_score > 2.0 OR psi_score > 0.2)
ORDER BY psi_score DESC;

-- Null rate trend per column over the last 30 days
SELECT
    window_end,
    column_name,
    null_count / row_count AS null_rate
FROM ml_catalog.monitoring.customer_features_profile_metrics
WHERE window_end >= current_date() - INTERVAL 30 DAYS
ORDER BY column_name, window_end;

Monitoring API reference

Method	Purpose
w.quality_monitors.create(table_name, ...)	Create monitor
w.quality_monitors.get(table_name)	Get monitor status
w.quality_monitors.run_refresh(table_name)	Trigger immediate refresh
w.quality_monitors.delete(table_name)	Delete monitor

Model Serving

Databricks Model Serving deploys models from the registry as autoscaled REST endpoints. Provisioned throughput is available for LLMs; standard endpoints work for most ML models.

Calling an endpoint

import requests

endpoint_url = "https://{workspace}.cloud.databricks.com/serving-endpoints/churn-predictor/invocations"

response = requests.post(
    endpoint_url,
    headers={
        "Authorization": f"Bearer {databricks_token}",
        "Content-Type": "application/json"
    },
    json={
        "dataframe_records": [
            {
                "customer_id": "12345",
                "total_spend_30d": 150.0,
                "login_count_7d": 2,
                "support_tickets_30d": 3,
                "days_since_last_purchase": 15
            }
        ]
    }
)

result = response.json()
# {"predictions": [0.73]}

For endpoints backed by Feature Store models, pass only the lookup keys — the endpoint fetches features from the online store:

json={"dataframe_records": [{"customer_id": "12345"}]}

Inference Tables

Every endpoint can auto-log all requests and responses to a Delta table. Enable in the endpoint configuration. Schema:

Column	Description
timestamp	Request time
request_id	Unique identifier
model_name	Endpoint name
model_version	Version that responded
input_features	JSON of input payload
predictions	JSON of output
latency_ms	Response time

Query for monitoring:

-- Prediction distribution over the last 24 hours
SELECT
    DATE_TRUNC('hour', timestamp) AS hour,
    AVG(predictions[0]) AS avg_churn_score,
    PERCENTILE(predictions[0], 0.9) AS p90_churn_score,
    COUNT(*) AS request_count
FROM ml_catalog.monitoring.churn_predictor_inference_logs
WHERE timestamp > current_timestamp() - INTERVAL 24 HOURS
GROUP BY 1
ORDER BY 1;

Use inference tables as the input to a downstream job that joins predictions to ground truth when outcomes arrive.

Workflows

A standard ML batch pipeline in Databricks Workflows:

# Databricks Asset Bundle (databricks.yml) or API representation
workflow = {
    "name": "churn_prediction_daily",
    "schedule": {
        "quartz_cron_expression": "0 0 2 * * ?",   # 02:00 UTC daily
        "timezone_id": "UTC"
    },
    "tasks": [
        {
            "task_key": "refresh_features",
            "notebook_task": {
                "notebook_path": "/Repos/prod/churn/01_refresh_features"
            }
        },
        {
            "task_key": "run_inference",
            "depends_on": [{"task_key": "refresh_features"}],
            "notebook_task": {
                "notebook_path": "/Repos/prod/churn/02_batch_inference"
            }
        },
        {
            "task_key": "export_to_crm",
            "depends_on": [{"task_key": "run_inference"}],
            "notebook_task": {
                "notebook_path": "/Repos/prod/churn/03_export_predictions"
            }
        }
    ]
}

Task refresh_features should end with fe.write_table(..., mode="merge") and fe.publish_table(...). Task run_inference uses fe.score_batch() and writes to the predictions table and the inference log. Keep these as separate tasks so failures are isolated and retries don’t recompute everything.

For CI/CD, use Databricks Asset Bundles (DABs). Bundle targets map to environments (dev/staging/prod), and databricks bundle deploy applies the configuration.

DE responsibilities

Feature pipelines

Build:

• Feature computation jobs (Spark/SQL) writing to feature tables via fe.write_table()
• Feature table schema with primary_keys and timestamp_keys
• Online store sync at end of each feature job
• Feature freshness monitoring (alert if last write > SLA)

Maintain:

• Pipeline scheduling and dependencies in Workflows
• Data quality checks in feature computation (null rates, row counts, value distributions)
• Feature table optimization (OPTIMIZE + ZORDER, VACUUM)
• Baseline table refreshes after model retraining

Support:

• Debug feature computation issues (wrong values, unexpected nulls, missing customers)
• Provide lineage when data scientists ask where a feature comes from
• Help onboard new features (schema, point-in-time setup)

MLflow and registry

Build:

• Experiment folder structure in the workspace
• Artifact storage configuration (S3/ADLS path)
• Model registry permissions via Unity Catalog grants

Maintain:

• Experiment cleanup policy (archive runs older than N days)
• Artifact storage costs (large models + many runs accumulate fast)
• Ensure production models have git_commit and training_data_version tags

Support:

• Rollback: move @champion alias back to previous version
• Debug model loading failures (library version mismatches, missing artifacts)

Red flags to escalate

Category	Signal	Why it matters
Feature pipeline	Feature computation returns all nulls	Model will produce garbage predictions
Feature pipeline	Row count drops > 20% from previous run	Missing data = missing predictions
Feature pipeline	Unexpected new values in categorical columns	Model may not handle unseen categories
Feature pipeline	Feature freshness > SLA	Predictions based on stale data
Registry	Production model has no training_data_version tag	Cannot reproduce or audit
Registry	@champion alias points to a version with no owner tag	No accountability in production
Registry	Model promoted without comparison metrics logged	Cannot verify improvement
Drift	PSI > 0.2 sustained > 1 week	Significant distribution shift, retraining likely needed
Drift	Prediction distribution shifts > 20% from baseline	Model behavior changed unexpectedly
Accuracy	Accuracy drops > 5% week-over-week	Model degradation

---

This document was produced by Claude based on David Morel’s instructions and questions, drawing on the Databricks and MLflow documentation. Treat it as a starting point, not a canonical reference — verify against official docs before relying on specific API signatures.