Experiment Tracking and Model Registries, with Python

Data scientists try many versions of a model: different features, algorithms and settings. Without records, nobody can say which version is in production, how it was trained, or whether a new one is really better. Experiment tracking records every training run; a model registry stores the chosen versions with their metadata and marks which one is in production.

What to record for each run

• Settings: algorithm and hyperparameters.
• Metrics: on a fixed validation or test set.
• Data: a fingerprint or pointer to the exact training data.
• Code version: the Git commit.
• Artefacts: the model file, and plots or reports.

Tools such as MLflow, cloud ML platforms and experiment trackers do this with a few lines of code. The idea is simple enough to build by hand, which this example does.

A minimal registry

This example trains two versions of a model, stores each with its metadata in a registry folder, and promotes the new version only if it beats production by a clear margin:

import hashlib, json, tempfile, os
from datetime import datetime, timezone
import joblib
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score

rng = np.random.default_rng(3)
X = rng.normal(size=(2000, 4)); y = (X @ [0.8, -0.5, 0.3, 0.0] + rng.normal(0, 1, 2000) > 0).astype(int)
registry = tempfile.mkdtemp()            # stands in for a model registry or object storage

def train_and_register(X, y, params, version):
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.25, random_state=0)
    model = LogisticRegression(**params).fit(X_tr, y_tr)
    auc = roc_auc_score(y_te, model.predict_proba(X_te)[:, 1])
    folder = os.path.join(registry, f"churn-model/v{version}")
    os.makedirs(folder)
    joblib.dump(model, os.path.join(folder, "model.joblib"))
    meta = {"version": version, "params": params, "test_auc": round(auc, 4),
            "data_fingerprint": hashlib.sha256(X.tobytes() + y.tobytes()).hexdigest()[:12],
            "rows": len(y), "trained_at": datetime.now(timezone.utc).strftime("%Y-%m-%d")}
    json.dump(meta, open(os.path.join(folder, "metadata.json"), "w"))
    return meta

runs = [train_and_register(X, y, {"C": c, "max_iter": 1000}, v) for v, c in [(1, 0.001), (2, 1.0)]]
for m in runs:
    print({k: m[k] for k in ("version", "params", "test_auc", "data_fingerprint")})

# Promotion rule: the new version replaces production only if it is clearly better
production, candidate = runs[0], runs[1]
if candidate["test_auc"] >= production["test_auc"] + 0.005:
    production = candidate
print("in production: version", production["version"])
reloaded = joblib.load(os.path.join(registry, f"churn-model/v{production['version']}/model.joblib"))
print(f"reloaded v{production['version']} from the registry; churn probability for the first customer: "
      f"{reloaded.predict_proba(X[:1])[0, 1]:.2f}")

{'version': 1, 'params': {'C': 0.001, 'max_iter': 1000}, 'test_auc': 0.8282, 'data_fingerprint': '46db80c0b1f0'}
{'version': 2, 'params': {'C': 1.0, 'max_iter': 1000}, 'test_auc': 0.8299, 'data_fingerprint': '46db80c0b1f0'}
in production: version 1
reloaded v1 from the registry; churn probability for the first customer: 0.73

Version 2 scored slightly higher, but by less than the agreed margin of 0.005, so version 1 stays in production. That rule protects against switching models for differences that are just noise. Both runs share the same data fingerprint, which proves they were trained on identical data, so the comparison is fair.

Promotion rules

Decide in advance what a new model must do to replace the current one:

• Beat it by a minimum margin on the agreed metric and test data.
• Not get worse on guardrails, such as performance for important customer segments or fairness checks.
• Pass technical checks: prediction speed, input validation, size.

Then promote by changing the registry’s “production” pointer, which also makes rollback a one-line change.

Common mistakes

1. Comparing runs on different test data, so differences are meaningless.
2. Not recording data versions, so a model cannot be reproduced.
3. Promoting by gut feeling instead of an agreed rule.
4. Overwriting model files instead of keeping versions.

Key takeaway

Track every run’s settings, metrics, data fingerprint and code version, store chosen models as versions in a registry, and promote by explicit rules so rollback is easy. Next: getting predictions to users with model serving.