import mlflow
import mlflow.sklearn
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import roc_auc_score

mlflow.set_experiment("/Users/jane@acme.com/fraud-detection")

with mlflow.start_run(run_name="gb-baseline") as run:
    params = {"n_estimators": 200, "max_depth": 4, "learning_rate": 0.05}
    mlflow.log_params(params)

    clf = GradientBoostingClassifier(**params)
    clf.fit(X_train, y_train)

    auc = roc_auc_score(y_val, clf.predict_proba(X_val)[:, 1])
    mlflow.log_metric("val_auc", auc)
    mlflow.log_metric("train_size", len(X_train))

    # Sklearn flavor handles env capture, signature, input example
    mlflow.sklearn.log_model(
        sk_model=clf,
        artifact_path="model",
        signature=mlflow.models.infer_signature(X_train, clf.predict(X_train)),
        input_example=X_train.head(),
    )

    print(f"Run {run.info.run_id}: AUC={auc:.4f}")

Tracking URI Patterns

The MLflow tracking URI tells the client where to send runs. On Databricks the patterns to know:

import mlflow
mlflow.set_tracking_uri("databricks")
mlflow.set_registry_uri("databricks-uc")  # register models into Unity Catalog

The split between tracking URI and registry URI lets runs flow to the workspace tracking server while models register into Unity Catalog. This is the modern setup; do not register new models into the legacy workspace registry.

Unity Catalog Model Registry

Models in Unity Catalog have fully qualified names (catalog.schema.model) and follow the same GRANT / REVOKE access model as tables.

# Register a logged model into Unity Catalog
result = mlflow.register_model(
    model_uri=f"runs:/{run.info.run_id}/model",
    name="prod.ml.fraud_score",
)
print(f"Version {result.version} registered as prod.ml.fraud_score")

Stage management uses aliases (e.g., champion, challenger, archived) rather than the legacy fixed stages (None / Staging / Production). Aliases are semantic and assignable, which makes blue-green and shadow patterns cleaner:

from mlflow import MlflowClient
client = MlflowClient()

# Promote version 7 to champion
client.set_registered_model_alias(
    name="prod.ml.fraud_score",
    alias="champion",
    version=7,
)

# Production code references the alias, not a fixed version
import mlflow.pyfunc
model = mlflow.pyfunc.load_model("models:/prod.ml.fraud_score@champion")

Model Serving: From Registry to REST Endpoint

Databricks Model Serving spins up a managed endpoint from a registered model in roughly 5-10 minutes. The endpoint scales 0-N replicas, supports A/B traffic split between model versions, and integrates with Unity Catalog for access control.

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import (
    EndpointCoreConfigInput, ServedEntityInput,
)

w = WorkspaceClient()
w.serving_endpoints.create(
    name="fraud-score",
    config=EndpointCoreConfigInput(
        served_entities=[
            ServedEntityInput(
                entity_name="prod.ml.fraud_score",
                entity_version="7",
                workload_size="Small",  # 4 concurrent requests
                scale_to_zero_enabled=True,
            ),
        ],
    ),
)

Scale-to-zero is the cost-saver: idle endpoints drop to zero replicas after the configured timeout and cold-start in roughly 10-30 seconds on the next request. Latency-sensitive endpoints serving real-time inference should disable scale-to-zero and run a warm minimum replica count.

Tracking Inference: Inference Tables and Lakehouse Monitoring

Production ML needs the loop closed: predictions and the features that produced them must land in a Delta table so you can detect drift, debug bad scores, and retrain on real data. Databricks does this automatically via inference tables:

# When creating the endpoint
auto_capture_config = AutoCaptureConfigInput(
    enabled=True,
    catalog_name="prod",
    schema_name="ml",
    table_name_prefix="fraud_score_inference",
)

Each inference request and response lands as a row in prod.ml.fraud_score_inference_payload. Pair it with Lakehouse Monitoring to detect feature drift, label drift, and prediction distribution drift on a schedule. Without this loop, your "production model" is essentially flying blind.

Operating Patterns Across Customer Engagements

Five patterns separate ML platforms that work from those that just exist:

1. Catalog hierarchy mirrors environments — dev.ml, staging.ml, prod.ml as separate UC catalogs with hard ACL boundaries. Dev models do not get loaded into prod by accident.
2. Aliases over versions in code — production code references @champion, never a fixed version number. Promotion is one alias swap, not a redeploy.
3. Inference tables on by default — every prod endpoint logs to a Delta table from day one, even before monitoring is wired up. You cannot retrofit observability after a customer incident.
4. Jobs cluster training, not all-purpose — training runs scheduled via Workflows on jobs clusters or serverless GPU. The DBU cost difference for a single weekly training run is meaningful at scale.
5. Feature engineering in Delta tables, not in notebooks — feature pipelines are jobs producing curated Delta tables. Models read from those tables. This makes feature reuse, lineage, and drift detection tractable.

Gotchas

• Mismatched signature and input — endpoint returns 400 because the JSON column order doesn't match the signature. Always log a representative input_example and test with the same JSON shape
• Library version drift — model trained on scikit-learn 1.4 fails to load on 1.5. MLflow captures the env, but the serving runtime must match; pin the runtime image explicitly
• Run pollution — thousands of throwaway runs from interactive experimentation make the experiment unusable. Use separate experiments for ad-hoc vs. tracked work, or use mlflow.delete_run as part of the cleanup loop
• Scale-to-zero on latency-critical endpoints — cold start of 10-30 seconds breaks SLAs. Disable scale-to-zero and run min-replicas above zero for these endpoints

How Opsio Helps

Opsio's mlops services practice designs and operates MLflow-on-Databricks platforms for industrial, automotive, and financial-services customers. We integrate MLflow with databricks implementation services workspaces, build the catalog hierarchy and CI/CD around the model registry, wire inference tables into Lakehouse Monitoring, and operate the resulting platform under SLAs. For broader use cases like demand forecasting and anomaly detection, see also our predictive analytics delivery.