, ,

Experiment Tracking and Model Registry for Infra Data (Infra to Data Science Series, Part 20)

Log every training run, register the good ones as immutable versions, and let a single champion alias decide what serves, so promotion and rollback each become one line. Built on the incident classifier, with the real failures that bite.

Infra to Data Science Series · Part 20 of 26

I once counted eleven files named model_final.joblib on a single training box, three of them dated the same afternoon, and not one of them recorded which CSV it trained on or what it scored. That pile is the natural end state of promoting models by hand, and it is exactly the mess an experiment tracker and a model registry exist to prevent.

Key takeaways: Log every training run to a tracking server so a model can be found again by its number, not its filename. Register the ones worth keeping as immutable versions, then move a single alias, champion, to name whichever version is live. Promotion becomes setting an alias and rollback becomes moving it back, one line each. Do not use the None, Staging, Production stages every tutorial still shows; they were deprecated in MLflow 2.9 and are being removed.
Who this is for: An infrastructure engineer or SRE who has trained, gated and served the incident classifier across Parts 13 to 19 and still keeps models as loose pickle files. Terms on first use: an experiment tracking server records the parameters, metrics and artifacts of each run so runs are comparable; a model registry is a versioned store where a trained model gets an immutable version number; an alias is a movable name, like champion, that points at one version and can be reassigned; a run is one execution of a training script.

Experiment Tracking as Logging for Model Runs

Where the project stands: last part we put a CI gate in front of promotion so a worse model could not ship. This part gives every run that passes through that gate a logged home and gives the promoted model a versioned, aliased slot you can roll back in one line. Tracking is logging for models. You already ship structured logs from every service and would never debug a production incident from memory, yet most teams train models with no record of which data, which parameters and which score produced the artifact now serving traffic.

A tracker records four things per run: parameters you chose, metrics you measured, artifacts you produced, and the code version that produced them. Framed that way it is observability pointed at training, the same reflex behind tracing and debugging for LLM systems. Because the Data Science Series already covers this ground for a pure data team, I spend one section on the mechanics here and link the deeper treatment of MLflow experiment tracking and the model registry for the parts I skip. Every command below was run against python 3.12, mlflow 3.11.1, scikit-learn 1.6.1 and pandas 2.2.3, and one detail bites anyone on an older tutorial: MLflow 3 changed the log_model call to take the estimator as sk_model and the artifact path as name, not the older artifact_path argument, so a copied snippet from 2.x will not run.

Logging a Training Run With MLflow

Take the same time aware training from Part 19 and wrap the fit in a run. Three log calls capture the parameters, the metric and the model itself, and the run gets an id you can quote in a review. Point the tracker at a shared server through an environment variable, never at the local disk, for a reason the next callout makes concrete.

# track.py  tested with python 3.12, mlflow 3.11.1, scikit-learn 1.6.1, pandas 2.2.3
import os, mlflow, mlflow.sklearn, pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import cross_val_score, TimeSeriesSplit

FEATURES = ['cpu_p95', 'mem_p95', 'load1_mean']
mlflow.set_tracking_uri(os.environ['MLFLOW_TRACKING_URI'])   # shared server, read from env
mlflow.set_experiment('incident-classifier')

df = pd.read_csv('features_train.csv').sort_values('window_start')
X, y = df[FEATURES], df['incident']
model = make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000))
auc = cross_val_score(model, X, y, cv=TimeSeriesSplit(5), scoring='roc_auc').mean()
model.fit(X, y)

with mlflow.start_run(run_name='logreg-baseline') as run:
    mlflow.log_param('model', 'logreg')
    mlflow.log_param('features', ','.join(FEATURES))
    mlflow.log_metric('cv_auc', auc)
    mlflow.sklearn.log_model(sk_model=model, name='incident-model')
    print('run_id', run.info.run_id[:12], 'cv_auc', round(auc, 3))
run_id 7f3c1a9d2b8e  cv_auc 0.874

That single run is now queryable: its cv_auc, its feature list and the exact model object live under one id on the server. Log the metric on a time aware split, not a random one, so the number you store is the honest number and not a leak inflated score, the mistake Part 14 dissected. One more habit pays off later, log sklearn.__version__ and a hash of the training CSV as parameters too, because a version you can find but cannot reproduce is only half an artifact. Everything sensitive stays out of the script: the tracking URI comes from MLFLOW_TRACKING_URI and any database credential from its own environment variable, so nothing lands in git.

You will read that mlflow.sklearn.autolog() captures all of this for free, and it does log parameters, metrics and the model on every fit without the explicit calls above. Two cautions keep it from misleading you. Autolog records the last fit it sees, so calling it before cross_val_score logs the metrics of the internal fold fits and can bury the score you actually care about under nested runs, and it captures whatever hyperparameters the estimator carries rather than the ones you reasoned about. Turn it on when you want a fast audit trail of a notebook, and log explicitly, as above, when a specific metric on a specific split is the thing a promotion decision will read.

Production gotcha: MLflow defaults its tracking store to a local ./mlruns folder. On an ephemeral CI runner that folder is deleted when the job ends, so every run you thought you logged vanishes with the container and the registry stays empty. I lost a week of nightly runs to exactly this before setting MLFLOW_TRACKING_URI to a shared server backed by Postgres and object storage. Set the tracking URI from an environment variable in the first two lines of the script, and confirm a run appears in the UI before you trust the pipeline.

A Registry in Place of a Folder of Pickles

A logged run is findable, but serving still needs a stable name to load. Registering promotes a run artifact into the model registry, where it becomes version 1, then version 2, and so on, each one immutable. Think of it in terms you already run: a registered model is a git repository, a version is an immutable build number, and an alias, coming next, is the load balancer pointer that decides which build takes traffic.

# after the run above, or from any run id
result = mlflow.register_model(
    f'runs:/{run.info.run_id}/incident-model', 'incident-classifier')
print('registered version', result.version)
Successfully registered model 'incident-classifier'.
Created version '1' of model 'incident-classifier'.
registered version 1

A version is immutable, and that matters more than it first sounds. You never edit version 3 to fix it; you register version 4 that supersedes it, and version 3 stays in the history as a record of what once served. That is the same guarantee a container image digest gives you, and it is what lets a rollback be exact rather than approximate. Log a model signature at the same time, the input columns and their types, so serving rejects a payload with the wrong schema instead of returning a confident wrong answer, and so the registry page shows precisely what shape of data each version expects.

Run the gated pipeline five times over a week and the registry accumulates five versions, each carrying the cv_auc it was logged with. That history is the point: you can see at a glance which version scored what, rather than guessing from a filename. Below are the five versions this project produced, with the score each one earned on the time aware split.

Registry conceptWhat you already runWhat it is for
Registered modela git repositoryone named line of a model over time
Model versionan immutable build numbera fixed artifact you can load and score
Alias, championa load balancer pointernames whichever version serves now
Version taga host label or annotationrecords gate status, owner, ticket
Runa build log entryparams, metrics and code for one fit
Five registered versions, one alias points at the bestcross validated AUC per registered version, champion sits on v40.850.870.890.830.900.88v1v2v3v4 championv5
Version 3 fell under the gate and never got the alias. Version 5 scored well but below v4, so champion stays on v4 until something beats it, not merely something newer.

Aliases for Promotion and One Line Rollback

A registry with five versions still has to answer one question for serving: which one is live. Aliases answer it. Set a champion alias on the winning version, load by that alias in serving, and promotion becomes moving the alias rather than editing a deploy script. Rollback is the same move in reverse, which is the property that turns a bad promotion from an incident into a one line correction.

from mlflow import MlflowClient
client = MlflowClient()

# promote, only after the Part 19 gate has passed on this version
client.set_registered_model_alias('incident-classifier', 'champion', 4)

champ = mlflow.sklearn.load_model('models:/incident-classifier@champion')
live = client.get_model_version_by_alias('incident-classifier', 'champion')
print('champion is version', live.version)

# rollback is one line: move the alias back to the last good version
client.set_registered_model_alias('incident-classifier', 'champion', 3)
print('rolled back to', client.get_model_version_by_alias('incident-classifier', 'champion').version)
champion is version 4
rolled back to 3

Serving loads from the URI models:/incident-classifier@champion, so it never names a version number and never needs a redeploy to switch models. Move the alias and the next load picks up the new version. More than one alias can sit on the registry at once, so a challenger alias on a candidate lets you shadow test it against champion before it takes traffic, which fixed stages could never do.

One operational detail decides whether rollback is actually fast. Most serving code loads the model once at process start and holds it in memory, which is right for latency but means moving the alias changes nothing until the process reloads. Decide the reload path before you need it: a scheduled reload every few minutes, a webhook that restarts workers on an alias change, or a health endpoint that reports the loaded version so you can tell whether a rollback has taken. I skipped this once and moved champion back during an incident, then watched the bad version keep serving for twenty more minutes because every worker still held it in memory. Aliases make rollback a one line decision, but the reload is what makes it a fast one.

flowchart LR
  R[Run logged] --> G{Gate passed}
  G -->|no| X[Stays a run, not registered]
  G -->|yes| V[Register new version]
  V --> A[Move champion alias]
  A --> S[Serving loads by alias]
  S --> B[Bad in prod, move alias back]
  B --> A
Promotion and rollback are the same edge in opposite directions. Everything hangs on one movable alias, so going back costs exactly what going forward did.

Here is where the tutorial default is wrong, and it is worth saying plainly. Almost every guide still teaches the None, Staging and Production stages, and MLflow deprecated them in version 2.9 with removal planned. Build on stages today and you build on an API scheduled to disappear, so reach for aliases and version tags instead. A tag such as validation_status set to approved records why a version earned its alias, and unlike a stage you can attach as many as the audit needs.

One failure will find you the first time serving starts against a fresh registry. Load by an alias before anything set it, and MLflow refuses rather than guessing.

# serving boots against a registry where champion was never set
champ = mlflow.sklearn.load_model('models:/incident-classifier@champion')
mlflow.exceptions.RestException: RESOURCE_DOES_NOT_EXIST: Registered model
alias 'champion' not found for model 'incident-classifier'

Do not paper over this by loading models:/incident-classifier/latest instead, which is the reflex fix and the wrong one, because latest returns the highest version number whether or not it passed a gate. Set the champion alias as the final step of promotion, after the gate, so the alias existing is itself the signal that a version is fit to serve. Then serving reads the alias and fails loudly on a cold registry rather than quietly serving an untested build.

War story: For a month I served the classifier by loading models:/incident-classifier/latest, reasoning that the newest registered version was the best. Then a teammate registered an experimental version 7 straight from a notebook to test a feature idea, never ran it through the gate, and left for the day. Serving picked it up on the next reload because it was latest, and cv_auc on live traffic fell from 0.90 to 0.71 for six hours until an on call engineer noticed incident alerts had gone quiet. The fix was two changes: serving loads by the champion alias only, and the alias is set exclusively by the promotion step after the gate. Latest means newest, never best, and I had shipped newest to production.

Comparing Candidate Runs Before Promotion

Before the alias moves, someone has to decide that a new version beats the current champion, and the tracker makes that a query rather than an argument. Pull the two runs, put their metrics side by side, and promote only on a real gain, not a rounding difference. This is the same discipline behind building an eval set before an LLM feature, a fixed comparison the change has to clear.

champ_v = client.get_model_version_by_alias('incident-classifier', 'champion')
champ_auc = client.get_run(champ_v.run_id).data.metrics['cv_auc']
cand_auc = client.get_run(run.info.run_id).data.metrics['cv_auc']
print(f'champion {champ_auc:.3f}  candidate {cand_auc:.3f}  delta {cand_auc - champ_auc:+.3f}')
promote = cand_auc - champ_auc > 0.01   # a real gain, not noise
print('promote' if promote else 'keep champion')
champion 0.900  candidate 0.884  delta -0.016
keep champion

Version 5 scored 0.884 against a champion at 0.900, a loss dressed as a fresh model, so the comparison keeps the alias where it is. Set a minimum gain, here 0.01, so you do not churn the champion on noise between runs that differ only by a lucky fold. That threshold belongs in the repo next to the gate value from Part 19, chosen from what a real improvement is worth on your systems rather than from a default. This comparison recipe, two run ids to a promote decision, is the reference artifact of this part; keep it beside the promotion script.

QuestionDeprecated stagesAliases and tags
How many can a version holdone stage onlymany aliases and tags
Shadow a challengernot possiblea second alias on the candidate
Roll backtransition stage againmove the alias back, one line
Supported aheadremoval planned since 2.9the current path

Register Broadly, Alias Only the Winner

None of this is a new muscle for you. A registry is version control for a binary, an alias is a pointer to the current release, and a rollback is moving that pointer back, all moves you make every deploy day. What changes is that the artifact under version control is a model whose quality lives in a number, so the promotion decision reads a metric from the tracker instead of a green check from a test. Register every gated run, keep the history, and let a single alias carry the weight of which one is live.

Do this on Monday: point MLFLOW_TRACKING_URI at a shared server, wrap one training run of your own telemetry in mlflow.start_run with three log calls, and register the result. Then set a champion alias by hand and change serving to load models:/your-model@champion. By the afternoon you will have replaced a folder of nameless pickles with a versioned model you can compare, promote and roll back on purpose. Next part turns this project into a real infra data build, anomaly detection on metrics and time series, using the tracked and registered model as its spine.

Infra to Data Science Series · Part 20 of 26
« Previous: Part 19  |  Guide  |  Next: Part 21 »

References

About The Author


Discover more from Journal of Intelligent Infrastructure

Subscribe to get the latest posts sent to your email.

Leave a Reply

Your email address will not be published. Required fields are marked *

Architect’s Toolkit

About the Author

Dr. Pranay Jha is a Cloud and AI Consultant with 18+ years of experience in hybrid cloud, virtualization, and enterprise infrastructure transformation. He specializes in VMware technologies, multi-cloud strategy, and Generative AI solutions. He holds a PhD in Computer Applications with research focused on Cloud and AI, has published multiple research papers, and has been a VMware vExpert since 2016 and a VMUG Community Leader.

Discover more from Journal of Intelligent Infrastructure

Subscribe now to keep reading and get access to the full archive.

Continue reading