MLflow Tutorial: Tracking Experiments and Models

Every data scientist has had this moment: a model from three weeks ago scored better than anything since, and nobody can say exactly which data, settings and code produced it. Experiment tracking removes that problem by recording every run as it happens. MLflow is the most widely used open-source tool for it, and it also handles the next step: keeping track of which model is in production.

This tutorial goes from training to a production-style model reference on a laptop, with real data.

The pieces

MLflow Tracking is an API and UI for logging parameters, code versions, metrics and output files, and experiments group the runs (MLflow: tracking). By default runs go to a local mlruns folder. The script uses a local SQLite file instead, which the registry needs, and a team would point the same code at a shared tracking server.

The tutorial

The data is the UCI Bike Sharing dataset: daily rental counts in Washington, D.C., with weather and calendar information, licensed under CC BY 4.0 (UCI: Bike Sharing). The models train on January 2011 to June 2012 and are tested on the second half of 2012, so they are judged on the future, as a production model would be.

import mlflow
import pandas as pd
from mlflow import MlflowClient
from mlflow.models import infer_signature
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error

mlflow.set_tracking_uri("sqlite:///mlflow.db")          # a local database file instead of a server
mlflow.set_experiment("bike-demand")

# UCI Bike Sharing (CC BY 4.0): daily rentals in Washington, D.C., 2011-2012
day = pd.read_csv("data/day.csv")
features = ["season", "yr", "mnth", "holiday", "weekday", "workingday", "weathersit", "temp", "atemp", "hum", "windspeed"]
train, test = day[day.dteday < "2012-07-01"], day[day.dteday >= "2012-07-01"]   # train on the past, test on the future
X, y, X_test, y_test = train[features], train.cnt, test[features], test.cnt

# 1. Three experiments, each logged as a run: parameters, a metric and the model itself
model_uris = {}
for n_trees, depth in [(50, 4), (200, 8), (200, None)]:
    with mlflow.start_run(run_name=f"forest-{n_trees}-trees-depth-{depth}") as run:
        model = RandomForestRegressor(n_estimators=n_trees, max_depth=depth, random_state=42).fit(X, y)
        mae = mean_absolute_error(y_test, model.predict(X_test))
        mlflow.log_params({"n_estimators": n_trees, "max_depth": depth, "train_rows": len(X)})
        mlflow.log_metric("test_mae", mae)
        # MLflow saves scikit-learn models in the safe skops format and asks which types to trust;
        # this forest was trained here, so its tree type is trusted explicitly
        info = mlflow.sklearn.log_model(model, name="model", input_example=X_test.head(2),
                                        signature=infer_signature(X_test, model.predict(X_test)),
                                        skops_trusted_types=["sklearn.tree._tree.Tree"])
        model_uris[run.info.run_id] = info.model_uri

# 2. Compare runs by querying the tracking store, not by remembering
runs = mlflow.search_runs(order_by=["metrics.test_mae ASC"])
print(f"{'run':<32} {'trees':>6} {'depth':>6} {'test MAE (rentals a day)':>26}")
for _, r in runs.iterrows():
    print(f"{r['tags.mlflow.runName']:<32} {r['params.n_estimators']:>6} {r['params.max_depth']:>6} {r['metrics.test_mae']:>26.0f}")

# 3. Register the best model and point an alias at it; applications load the alias, not a file path
best = runs.iloc[0]
version = mlflow.register_model(model_uris[best.run_id], "bike-demand")
client = MlflowClient()
client.set_registered_model_alias("bike-demand", "champion", version.version)
print(f"\nregistered bike-demand version {version.version}, alias champion -> {best['tags.mlflow.runName']}")

champion = mlflow.pyfunc.load_model("models:/bike-demand@champion")
first_days = test.head(3)
for date, actual, predicted in zip(first_days.dteday, first_days.cnt, champion.predict(first_days[features])):
    print(f"{date}: predicted {predicted:,.0f} rentals, actual {actual:,}")
print("champion run parameters:", client.get_run(best.run_id).data.params)

run                               trees  depth   test MAE (rentals a day)
forest-200-trees-depth-None         200   None                        885
forest-200-trees-depth-8            200      8                        906
forest-50-trees-depth-4              50      4                        971

registered bike-demand version 1, alias champion -> forest-200-trees-depth-None
2012-07-01: predicted 6,142 rentals, actual 5,531
2012-07-02: predicted 6,372 rentals, actual 6,227
2012-07-03: predicted 6,288 rentals, actual 6,660
champion run parameters: {'n_estimators': '200', 'max_depth': 'None', 'train_rows': '547'}

What the output shows

• Every run was recorded without extra effort: its parameters, its test error and the model itself.
• Comparison is a query, not a memory. search_runs returned the runs sorted by test error. The unrestricted forest won, with an average error of 885 rentals a day against 971 for the small one.
• The registry gave the winner a name and a version: bike-demand version 1.
• The alias is the interface. The application loaded models:/bike-demand@champion and never needed a run ID or a file path. The registry documentation describes exactly this: point an alias at a version, and change what production loads by reassigning it (MLflow: model registry).
• Lineage survives. From the champion, the script recovered the run's parameters, including how many rows trained it.
• The predictions are honest about difficulty. On the first test days the model was off by between 145 and 611 rentals, for days it had never seen.

Saving models safely

One line in the script deserves attention. Recent MLflow versions save scikit-learn models in the skops format by default, which the API reference describes as guaranteeing safe deserialisation, and ask you to name any types you trust (MLflow: mlflow.sklearn). The tree structure inside a random forest is one of them. Without skops_trusted_types, logging the model fails.

That friction is deliberate. The older pickle format can run arbitrary code when a file is loaded, and Python's documentation warns never to unpickle data from a source you do not trust (Python: pickle). Trust types only for models you trained yourself or received from a source you trust. See model packaging from pickle to production.

What to log in every run

Parameters and one metric are the minimum. A run you can trust months later also records:

• The data version: a dataset hash, a DVC version or at least the date range and row count, as the script logs train_rows.
• The code version: the Git commit. MLflow records it automatically when the script runs from a Git repository.
• Several metrics, including one that matches the business goal, not only the one you optimised.
• The evaluation split, so nobody later compares a random split with a time-based one.
• Plots and reports as artefacts, such as residuals by month.
• The environment: package versions, which MLflow stores with logged models.

From laptop to team

The same code scales up by changing where it logs:

1. Local files or a local database, as here, for one person.
2. A shared tracking server with a database and artefact storage, so a team sees every run.
3. Registry permissions and reviews, so only checked models get the champion alias.

The habits matter more than the setup: log every run, including the bad ones, record the data version, and let applications load aliases rather than paths. For the registry's role across the lifecycle, see the model registry. For versioning the data itself, see data versioning with DVC.

Common mistakes

• Logging only the best run, so nobody knows what was tried.
• Not recording the data version alongside the parameters.
• Random train and test splits for time-ordered data, which flatter the score.
• Hard-coding model file paths in applications instead of loading an alias.
• Trusting every type to make a model load, or loading pickles from unknown sources.
• A tracking store on one laptop for a team project.

MLflow basics, experiment tracking, the model registry and model versioning form a module of the MLOps course.

Sources

• ML Experiment Tracking (MLflow documentation) — MLflow
• Model Registry (MLflow documentation) — MLflow
• mlflow.sklearn (MLflow API reference) — MLflow
• pickle: Python object serialization (Python documentation) — Python Software Foundation
• Bike Sharing dataset (UCI Machine Learning Repository) — UCI Machine Learning Repository