Analysing and serving · Chapter 23

MLflow and the Model Lifecycle

MLflow is the least glamorous component in this part of the manual and the one that most often decides whether a model survives contact with a regulator. This chapter covers what it does, and why the Databricks version of it is different from a model registry bolted on beside your data.

Why the Databricks version is different #

MLflow is open source and originated at Databricks, which means the API you learn here is the API you would use anywhere. The managed version adds one thing that changes the shape of the problem, which is that the model registry lives inside Unity Catalog. A registered model is a catalog object with a three-level name, the same grants, the same lineage graph, and the same audit trail as a table.

On the other clouds the registry is a separate service with its own permission model. That is workable, and teams work it every day, but it means the question "who can deploy this model" and the question "who can read the data it was trained on" are answered by two different systems that nobody has reconciled. The gap between those two answers is where most of the uncomfortable audit findings we have seen actually live.

The four things MLflow does #

Tracking records runs. Every training execution logs its parameters, its metrics, the code version behind it, and any artefacts it produced, so that a result found interesting on a Thursday can be reconstructed the following month. Reproducibility is the real product here rather than the leaderboard of metrics, because the leaderboard tells you which experiment won while reproducibility tells you how to build that winner again.

Models is a packaging format. A trained model is saved with its dependencies and a declared signature describing its inputs and outputs, wrapped in what MLflow calls a flavour, meaning the framework-specific representation such as scikit-learn, PyTorch, or a generic Python function. The flavour concept is what lets a model trained in one place be served in another without the serving team needing to know how it was built.

Model Registry in Unity Catalog holds versions and aliases. A version is immutable and numbered. An alias is a movable label such as champion or challenger that applications point at instead of pinning a number, so promoting a new model is a governed metadata change rather than a redeployment. Promotion across environments becomes an action with a permission attached to it and an entry in the audit log.

Evaluation runs a model against a held-out dataset and records the results as part of the run. mlflow.evaluate handles conventional metrics for classification and regression, while generative systems have a harness of their own in mlflow.genai.evaluate, which scores the recorded execution of a request rather than only its final text and is the subject of a later section in this chapter. As Agents on Databricks argues at length, the evaluation set should be written by the people who will use the system rather than the people building it, and that holds for a credit risk model as much as for an agent.

import mlflow
from mlflow import MlflowClient

mlflow.set_registry_uri("databricks-uc")

with mlflow.start_run(run_name="risk_score_gbm") as run:
    mlflow.log_params({"max_depth": 6, "learning_rate": 0.05})
    mlflow.log_metric("auc_validation", 0.834)
    mlflow.sklearn.log_model(
        sk_model=model,
        name="model",
        input_example=X_train.head(),
        registered_model_name="originations_prd.models.application_risk_score",
    )

# Aliases move; version numbers do not.
MlflowClient().set_registered_model_alias(
    name="originations_prd.models.application_risk_score",
    alias="challenger",
    version=7,
)

Lineage is what makes an audit survivable #

Because training reads governed tables, the platform records which table version fed which model version without anybody being asked to document it. That link is the one a regulator actually asks about, and it is the one hand-maintained documentation always loses first.

In the auto finance work this manual keeps returning to, a model influencing a lending decision needs a defensible answer to three questions. What data trained this, who approved it, and what did it do in production. The first is lineage, the second is the grant and audit record on the registered model, and the third is inference logging. All three are properties of the platform rather than artefacts of a team's diligence, which matters because diligence is the thing that lapses when a deadline arrives.

Serving and closing the loop #

Model Serving puts a registered version behind an authenticated endpoint that scales with traffic. Inference tables capture request and response traffic into governed Delta tables, which is the part that closes the loop, because production traffic that lands in a table you can query is production traffic you can turn into an evaluation dataset. The questions your model handled badly last quarter are sitting there waiting to be labelled. For a generative endpoint the equivalent record needs to be richer than a request and a response, because the part worth reading happened between the two, which is what the tracing section below is about.

Drift monitoring should watch inputs and not only outputs. Output drift tells you something has already gone wrong, whereas a shift in the distribution of applicant income or vehicle age tells you the world moved before the model's errors accumulate enough to be visible in aggregate. Monitoring belongs with the rest of your platform telemetry, covered in Observability.

These are capable products and we would not pretend otherwise. SageMaker has excellent training and tuning tooling, and Vertex has a genuinely good pipelines story that some teams prefer to ours. The argument is narrower than a feature comparison, and it is that on Databricks the governance surface is shared and the lineage back to the training data is a property of the platform, whereas elsewhere both are things your team builds and then has to defend.

Advanced: what production actually demands #

Everything above gets a governed model into production. What follows is where large organisations discover the work they had not planned for.

Champion and challenger, and shadow deployment. A challenger version receives a copy of live traffic and is scored without its predictions being acted on. This is the only honest way to learn how a model behaves on real inputs before it makes real decisions, and it costs the compute of running two models rather than one.

Training and serving skew. A model trained on a feature computed one way in a batch job and served a feature computed another way in application code will quietly degrade, and the degradation is invisible to every test either team runs. A feature store fixes this by making the feature definition a single governed object that both training and serving read, which is less about convenience than about removing an entire class of silent failure.

Promotion as code rather than clicks. Alias moves, endpoint configuration, and monitoring should be described in the same deployment bundles that carry the rest of the estate, as set out in CI/CD and Environments. A promotion performed in a browser is one nobody can review beforehand or reconstruct afterwards.

Serving cost is an unpredictable line item. Endpoints bill on provisioned capacity and uptime rather than the query patterns your finance team has learned to forecast, and a rarely used endpoint left running is pure waste. Scale to zero where latency permits, and read Cost and Performance before committing to always-on capacity.

Deprecating a model with live consumers is harder than deploying one. Inference tables tell you who is still calling an endpoint, which is the information that turns a deprecation from a negotiation into a schedule.

MLflow 3 and tracing for generative systems #

MLflow 3 divides the product along the line that generative work had already drawn in practice. The classical side keeps everything described above and adds logged models, an entity that follows a model across runs and environments rather than being pinned to the run that produced it, together with deployment jobs that carry evaluation, approval and deployment as one governed workflow. The generative side gains something earlier versions did not have at all, which is tracing.

A trace is the recorded execution of a single request through your application, held as a tree of spans. A span is one step, meaning a model call, a tool execution or a retrieval, and it records that step's inputs and outputs, when it started and finished, whether it succeeded or failed, and which span it hangs beneath. Spans carry declared types including CHAT_MODEL, TOOL, RETRIEVER, AGENT and RERANKER, so a reader can tell without guessing whether the slow half of a nine second answer was the vector search or the third model call. Traces arrive either automatically for a supported library, or by decorating the functions whose logic is your own.

import mlflow
from mlflow.genai.scorers import Guidelines, Safety

mlflow.openai.autolog()          # every model call becomes a span

@mlflow.trace                     # your own steps join the same tree
def answer(question: str) -> str:
    chunks = retrieve_policy(question, k=5)
    return generate(question, chunks)

mlflow.genai.evaluate(
    data=eval_set,
    predict_fn=answer,
    scorers=[
        Safety(),
        Guidelines(guidelines="Cite the policy clause relied on.", name="cited"),
    ],
)

Evaluation in MLflow 3 attaches to those traces rather than sitting beside them. Scorers, whether the built-in judges, judges you define against criteria you have written down, or ordinary code returning a pass and a reason, run over traces during development, and the same scorers can be run over production traces so that the definition of quality does not change when a system leaves the laboratory. Governance is the part we would emphasise, because Unity Catalog covers prompts, applications and traces on the same terms it already covers your tables, and MLflow 3 gives prompts a registry of their own. Traces written to the catalog land as Delta tables in OpenTelemetry format, which means the record of what an agent did last Tuesday is queryable in SQL, grantable to a named group, and subject to the retention rules you run for everything else.

Why this changes debugging. Before tracing, a wrong answer from a generative system was a mystery you investigated by running the request again and hoping it went wrong the same way. With a trace it becomes a specific step you can point at. Retrieval returned three passages and none of them mentioned the policy, or retrieval was fine and the model ignored the passage that did, or the tool was called with the loan number in the wrong format, returned nothing, and the model apologised instead of saying so. Those are three different faults with three different fixes, and from outside they are indistinguishable. Instrument from the first prototype rather than after the first incident, because the run you most want to read is one that has already happened.

Traces are data and behave like data. A busy assistant produces a trace per request, each carrying the full prompt, the retrieved passages and every intermediate output, so the volume is a multiple of your traffic rather than a rounding error against it. The documented ingestion ceilings are generous and there is no per-experiment limit on how many traces you keep, so the constraint you meet first is your own storage bill rather than the platform's. Decide retention on the day you switch tracing on. What we usually recommend is full traces for a short recent window where debugging actually happens, a reduced set kept longer where the value is trend and evaluation rather than forensics, and the same restrictive grants you would put on any table holding whatever a customer typed, because that is precisely what a trace holds.

Prompts and agents are versioned artefacts too #

An agent registered as a model in Unity Catalog inherits everything described in this chapter, meaning the same versions, aliases, approval path, lineage, and inference logging. A prompt revision is a production change and belongs under the same versioning as the weights, which is what the prompt registry is for, so that an incident months later can name the prompt version, the model version, the trace of the request itself, and the data behind the answer. The agent-specific material sits in Agents on Databricks, and the governance underneath it is the same governance you already run for a credit risk model.

Table of Contents