MLflow for Data Scientists: Tracking Experiments

MLflow for data scientists usually starts in a notebook. You try a model, change a parameter, try another, and after an hour you cannot remember which cell produced the best result. MLflow fixes that with a few habits rather than a platform project. This walkthrough uses them on the UCI Wine Quality data: predicting the quality score of red wines from their chemistry (UCI: Wine Quality).

import os
import shutil
import warnings

import mlflow
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_absolute_error
from sklearn.model_selection import train_test_split

warnings.filterwarnings("ignore")
shutil.rmtree("mlruns", ignore_errors=True)             # start clean so the example repeats exactly
if os.path.exists("mlflow.db"):
    os.remove("mlflow.db")
mlflow.set_tracking_uri("sqlite:///mlflow.db")          # MLflow 3.16 refuses the plain-folder store by default
mlflow.set_experiment("wine-quality-notebook")

wine = pd.read_csv("winequality-red.csv", sep=";")     # UCI Wine Quality, red wines
X, y = wine.drop(columns="quality"), wine.quality
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=42)

# 1. A quick baseline, logged by hand
with mlflow.start_run(run_name="baseline-ridge"):
    model = Ridge(alpha=1.0).fit(X_train, y_train)
    mlflow.log_param("model", "ridge")
    mlflow.log_metric("test_mae", mean_absolute_error(y_test, model.predict(X_test)))

# 2. A small sweep as nested runs under one parent, with autologging for parameters
mlflow.sklearn.autolog(log_models=False, log_datasets=False, silent=True)
with mlflow.start_run(run_name="forest-sweep"):
    for depth in (4, 8, None):
        for trees in (100, 300):
            with mlflow.start_run(run_name=f"depth={depth} trees={trees}", nested=True):
                rf = RandomForestRegressor(n_estimators=trees, max_depth=depth, random_state=42, n_jobs=1)
                rf.fit(X_train, y_train)
                mlflow.log_metric("test_mae", mean_absolute_error(y_test, rf.predict(X_test)))
                mlflow.set_tag("family", "random_forest")
mlflow.sklearn.autolog(disable=True)

# 3. Compare runs as a DataFrame, the way a notebook user would
runs = mlflow.search_runs(experiment_names=["wine-quality-notebook"], filter_string="metrics.test_mae > 0")
runs["params.max_depth"] = runs.get("params.max_depth", pd.Series(dtype=str)).fillna("-")
runs["params.n_estimators"] = runs.get("params.n_estimators", pd.Series(dtype=str)).fillna("-")
table = runs[["tags.mlflow.runName", "params.max_depth", "params.n_estimators", "metrics.test_mae"]].sort_values("metrics.test_mae")
print("M1 runs ranked by test MAE (quality points)")
for _, r in table.iterrows():
    print(f"   {r['tags.mlflow.runName']:<26} depth {r['params.max_depth']:<5} trees {r['params.n_estimators']:<4} MAE {r['metrics.test_mae']:.3f}")

every = mlflow.search_runs(experiment_names=["wine-quality-notebook"])
parent_id = every.loc[every["tags.mlflow.runName"].eq("forest-sweep"), "run_id"].iloc[0]
children = every[every["tags.mlflow.parentRunId"].eq(parent_id)]
print(f"\nM2 runs in the experiment: {len(every)} (1 baseline, 1 sweep parent, {len(children)} nested child runs)")
print(f"M3 parameter columns available for comparison, mostly from autologging: {sum(c.startswith('params.') for c in every.columns)}")

M1 runs ranked by test MAE (quality points)
   depth=None trees=300       depth None  trees 300  MAE 0.423
   depth=None trees=100       depth None  trees 100  MAE 0.423
   depth=8 trees=300          depth 8     trees 300  MAE 0.455
   depth=8 trees=100          depth 8     trees 100  MAE 0.456
   depth=4 trees=100          depth 4     trees 100  MAE 0.499
   depth=4 trees=300          depth 4     trees 300  MAE 0.500
   baseline-ridge             depth -     trees -    MAE 0.502

M2 runs in the experiment: 8 (1 baseline, 1 sweep parent, 6 nested child runs)
M3 parameter columns available for comparison, mostly from autologging: 19

Habit 1: use a local database, not a folder

Older tutorials point MLflow at a folder of files called mlruns. In MLflow 3.16, that file-based store is in maintenance mode, and the code above was refused when it tried to use it: MLflow asked for a database backend, or for the file store to be enabled explicitly. A local SQLite file works with no server, so the example uses sqlite:///mlflow.db. If you have old runs in a folder, MLflow's documentation describes a mlflow migrate-filestore command that moves them into SQLite (MLflow: migrate from file store).

Habit 2: log a baseline first

The first run is a plain ridge regression, logged by hand with one parameter and one metric. It scored a mean absolute error of 0.502 quality points. Every later model now has a reference: if a complex model cannot beat 0.502, it is not worth its cost.

Habit 3: group a sweep under one parent run

The random forest sweep tried three depths and two forest sizes. Each combination ran as a nested child run under a parent called forest-sweep, so the experiment holds eight runs in a tidy structure: one baseline, one parent, six children. In the MLflow UI the children fold under the parent, which keeps a busy experiment readable.

Habit 4: let autologging record the parameters

Scikit-learn autologging, turned on for the sweep, recorded the estimator's parameters automatically; the comparison table had 19 parameter columns without a single manual log call for them. Autologging can be enabled or disabled per library (MLflow: autologging). Here models were not logged, to keep runs small; turn model logging on when you want to reload a model later. The custom test metric was still logged by hand, because autologging records training metrics, and your own held-out evaluation is the one you will compare.

Habit 5: compare runs as a DataFrame

mlflow.search_runs returns a pandas DataFrame, which is where notebook users are comfortable. Sorted by test error, the results were clear:

• Unlimited depth gave the lowest error, 0.423, with 100 and 300 trees performing almost identically.
• Depth 8 reached 0.455, and depth 4 about 0.50, no better than the ridge baseline.

Two lessons came out of the table, not out of memory. Tripling the number of trees bought nothing measurable here, and shallow forests were not worth training. That is the point of tracking: decisions based on a table of results rather than on which cell you ran last.

Read the sweep honestly

A ranked table invites over-interpretation. The two best runs differ in the third decimal place, and with a single train-test split that difference is noise, not a finding. Before choosing a model, check whether the gap between candidates is larger than the variation you would see from a different split; cross-validation, logged as one run per fold or as a mean and standard deviation, gives you that. Also remember what the metric means: an error of 0.42 on a quality scale where most wines score 5 or 6 is useful for ranking wines, not for predicting a taster's exact score. Tracking makes these questions easier to ask, because every number has a run behind it.

Log more than metrics

Metrics tell you which run won; artifacts tell you why. Log a residual plot or a feature importance chart for promising runs, and a short note as a tag when a run tests a specific idea, such as "dropped sulphates to test leakage". Weeks later, those notes are what let you, or a colleague, understand the experiment without rerunning it.

Habits that keep results trustworthy

• Fix random seeds for splits and models, as here, so reruns give the same numbers.
• Log the data version, for example a file hash or a dataset tag, so you know what each run was trained on.
• Name runs meaningfully, as the sweep did with depth and tree count.
• Tag runs with the question they answer, such as "family: random_forest", to filter later.
• Keep the held-out set fixed across runs, or comparisons are meaningless.

When to move beyond the notebook

The local setup is enough for one person. When a team shares experiments, point MLflow at a shared tracking server; when a model goes to production, use the model registry. Those steps belong to MLOps and are covered in the MLflow tutorial, model versioning with MLflow and DVC and MLflow vs Weights & Biases.

Experiment tracking is part of the Data Science course.

Sources

• Automatic Logging with MLflow Tracking (MLflow documentation) — MLflow
• Migrate from File Store (MLflow documentation) — MLflow
• Wine Quality (UCI Machine Learning Repository) — UCI Machine Learning Repository