The Project | About Us | Contribute | Donations | License
HOME My courses Completed
◄ Previous
◄

Experiment tracking with MLflow

Next ►
►

Module 2 left CineClick with a reproducible project and an open question: comparing one run against another with dvc metrics diff works, but exploring dozens of hyperparameter combinations and remembering what was tried three weeks ago calls for a different tool. In this lesson we instrument the churn pipeline with MLflow Tracking: every training run will be recorded with its parameters, metrics, code, and data, queryable through a web interface with a table, filters, and comparisons. With that team memory in place, Laura will launch a small hyperparameter search that finds a clearly better candidate for the business — and we'll finally be able to solve the mystery of "v2" that has been dragging along since module 1.

Contents

  1. The problem: comparing dozens of runs (and why dvc metrics diff falls short)
  2. MLflow concepts: tracking server, experiment, run, params, metrics, tags, and artifacts
  3. Instrumenting CineClick's train.py and evaluate.py
  4. The MLflow interface: run table, filters, and comparison
  5. A hyperparameter search: 12 runs and a new candidate
  6. Where tracking lives: local mlruns vs. a shared server
  7. DVC and MLflow together, and the mystery of "v2" solved

The problem: comparing dozens of runs

In module 2, when Laura tried n_estimators: 300, the flow was: edit params.yaml, run dvc repro, and compare with dvc metrics diff. Perfect for answering "does this change improve on what's in Git?". But DVC's diff compares the current state against one revision: it's a snapshot of two points, not a historical memory.

The questions Laura cannot answer with that flow:

  • Which combinations have we already tried? If she tries 12 configurations, she'd need 12 commits (or jot them down by hand in a spreadsheet, which is what she used to do with the notebook).
  • What was the best recall across all of last month's experiments, and with which parameters? git log stores commits, not a sortable table of metrics.
  • Was that experiment launched by Laura or Marc, and with which version of the data? There's nowhere to look it up.

Experiment tracking solves exactly this: a centralized log where every training run is automatically recorded with its full context. It doesn't replace DVC (we'll see shortly how they coexist): it adds the dimension DVC lacks — large-scale exploration and its history.

Need dvc metrics diff Tracking (MLflow)
Compare current state vs. a commit Excellent Possible but indirect
A table of 50 runs sorted by recall No Yes, natively
Filter "all runs with recall > 0.5" No Yes (search syntax)
Know who launched what and when Via commits (if everything is committed) Automatic on every run
Reproduce a run exactly Excellent (it's its reason for being) Only if you record commit and data

Look at the last row: tracking observes; it doesn't guarantee reproducibility on its own. That's why the two tools complement each other instead of competing.

MLflow concepts

MLflow is an open source platform for managing the ML lifecycle. It has several components (Tracking, Model Registry, Models, Projects); in this lesson we use Tracking, and in the next one the registry. Its vocabulary:

  • Tracking server: the service that receives and stores the records. It can be simply a local folder (mlruns/) or a shared server with a database — we cover this in section 6.
  • Experiment: a named container that groups related runs. For us: churn-cineclick. A project usually has one or a few experiments.
  • Run: one concrete training execution. It's the central unit: every dvc repro that trains a model, or every iteration of a search loop, becomes a run with its unique identifier (run_id).
  • Params: the hyperparameters and configuration of that run (n_estimators=100, test_size=0.2). They are immutable key-value pairs: written once.
  • Metrics: numeric result values (recall=0.43). They can be logged multiple times per run (for example, one metric per epoch in deep learning), although in our case it will be a final value.
  • Tags: free-form metadata for organizing and searching (author=laura, dvc_data=af12bc...). MLflow adds some automatically, such as the Git commit.
  • Artifacts: files attached to the run — the serialized model, plots, the params.yaml that was used. Whatever doesn't fit in a key-value pair goes here.
graph TD
    E["Experiment: churn-cineclick"] --> R1["Run a3f9... (baseline)"]
    E --> R2["Run 7c21... (300 trees)"]
    E --> R3["Run f0d4... (balanced)"]
    R3 --> P["Params: n_estimators, max_depth, class_weight"]
    R3 --> M["Metrics: accuracy, precision, recall, f1"]
    R3 --> T["Tags: git commit, DVC data, author"]
    R3 --> A["Artifacts: model, signature, input example"]

Instrumenting CineClick's train.py and evaluate.py

First, MLflow enters the project's dependencies like any other library (module 2: nothing gets installed "by hand"):

# We add mlflow to pyproject.toml (dependencies section) and recompile the lock
pip-compile pyproject.toml -o requirements.lock
pip-sync requirements.lock

Our pipeline has one peculiarity: training and evaluating are separate stages in dvc.yaml, but conceptually they are one and the same run (a model and its metrics). The standard solution: train.py opens the run, saves its run_id to a file, and evaluate.py reopens that same run to add the metrics.

src/cineclick_churn/train.py, instrumented (the relevant fragment):

import json
from pathlib import Path

import mlflow
import pandas as pd
import yaml
from mlflow.models import infer_signature
from sklearn.ensemble import RandomForestClassifier

def train(features_path: str, model_path: str, params_path: str = "params.yaml"):
    # 1. Load the parameters from params.yaml, the single source of truth
    params = yaml.safe_load(Path(params_path).read_text())["train"]

    df = pd.read_csv(features_path)
    X, y = df.drop(columns=["churned"]), df["churned"]

    # 2. Select (or create) the project's experiment
    mlflow.set_experiment("churn-cineclick")

    # 3. Open a run: everything logged inside the "with" block is attached to it
    with mlflow.start_run() as run:
        # 4. Log ALL the params from the 'train' section in one go
        mlflow.log_params(params)

        model = RandomForestClassifier(
            n_estimators=params["n_estimators"],
            random_state=params["random_state"],
        )
        model.fit(X, y)

        # 5. The signature describes the model's input/output schema;
        #    the input_example stores 5 real rows as documented sample input.
        signature = infer_signature(X, model.predict(X))
        mlflow.sklearn.log_model(
            model,
            artifact_path="model",
            signature=signature,
            input_example=X.head(5),
        )

        # 6. Save the run_id so the 'evaluate' stage can reopen this same run
        Path("models/mlflow_run_id.txt").write_text(run.info.run_id)

    # The .pkl is still produced: it's the output DVC tracks (for now)
    import joblib
    joblib.dump(model, model_path)

Points worth understanding well:

  • mlflow.set_experiment("churn-cineclick") creates the experiment if it doesn't exist and activates it. Without this line, runs would fall into the default experiment (Default), a junk drawer best avoided.
  • mlflow.log_params(params) receives the entire dictionary: we don't copy values by hand, so what's logged always matches what was executed. This detail kills a whole class of errors ("the spreadsheet said 200 trees but it was 100").
  • The signature and the input_example look like bureaucracy, but they're gold: they document which columns, with which types, the model expects. When we have to serve it in module 4, that will be the specification of the input contract, and MLflow will validate against it.

And evaluate.py, which reopens the run and adds the metrics (the same ones it keeps writing to metrics.json for DVC):

import json
from pathlib import Path

import mlflow

def evaluate(model_path: str, test_path: str, metrics_path: str = "metrics.json"):
    metrics = _compute_metrics(model_path, test_path)  # accuracy, precision, recall, f1

    # metrics.json still exists: it's the file dvc.yaml declares as a metric
    Path(metrics_path).write_text(json.dumps(metrics, indent=2))

    # Reopen the run that train.py opened (and closed), passing its id
    run_id = Path("models/mlflow_run_id.txt").read_text().strip()
    with mlflow.start_run(run_id=run_id):
        mlflow.log_metrics(metrics)

Now an ordinary dvc repro leaves, in addition to the usual outputs, a complete run in MLflow. Nothing from module 2's flow has broken: we've only added observation.

The MLflow interface: table, filters, and comparison

To explore what's been logged, we start the local web interface:

# From the repo root (where the mlruns/ folder that MLflow created lives)
mlflow ui --port 5000
# Open http://127.0.0.1:5000 in the browser

The essentials of the UI:

  • Run table: one row per run; configurable columns with params and metrics. You can sort by any metric — "give me the table sorted by recall, descending" is one click.
  • Search with filters: the search box accepts a SQL-like syntax. Useful examples:
metrics.recall > 0.5
metrics.recall > 0.5 and params.class_weight = 'balanced'
tags.mlflow.source.git.commit = 'a1b2c3d'
  • Run comparison: select several rows, click Compare, and you get params side by side (with differences highlighted) and metric-vs-parameter scatter plots. It's the visual, multi-run version of dvc metrics diff.
  • Detail view: inside a run you find its browsable artifacts — the model, its signature, the input_example.

A hyperparameter search: 12 runs and a new candidate

Let's recall the business problem: the baseline model has recall 0.43 — out of every 100 customers who leave, it only catches 43. For a retention campaign that lets more than half of the savable customers slip away. Laura suspects the class imbalance (~15% churn) is the culprit: the forest learns that "almost nobody leaves" and calls it a day.

The exploration: a loop over three hyperparameters. This is exploration, not the blessed pipeline, so it goes in a separate script (scripts/hyperparameter_search.py) and doesn't touch dvc.yaml:

import itertools

import mlflow
import pandas as pd
from mlflow.models import infer_signature
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
from sklearn.model_selection import train_test_split

df = pd.read_csv("data/processed/features.csv")
X, y = df.drop(columns=["churned"]), df["churned"]
# Same split as the pipeline: test_size=0.2, random_state=42, stratified
X_tr, X_te, y_tr, y_te = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

mlflow.set_experiment("churn-cineclick")

# 2 x 3 x 2 = 12 combinations
grid = itertools.product(
    [100, 300],            # n_estimators
    [10, 20, None],        # max_depth
    [None, "balanced"],    # class_weight
)

for n_est, depth, weight in grid:
    with mlflow.start_run(run_name=f"rf-{n_est}-{depth}-{weight}"):
        mlflow.log_params({
            "n_estimators": n_est,
            "max_depth": depth,
            "class_weight": weight,
        })
        mlflow.set_tag("type", "hyperparameter-search")

        model = RandomForestClassifier(
            n_estimators=n_est, max_depth=depth,
            class_weight=weight, random_state=42,
        )
        model.fit(X_tr, y_tr)
        pred = model.predict(X_te)

        mlflow.log_metrics({
            "accuracy": accuracy_score(y_te, pred),
            "precision": precision_score(y_te, pred),
            "recall": recall_score(y_te, pred),
            "f1": f1_score(y_te, pred),
        })
        signature = infer_signature(X_tr, model.predict(X_tr))
        mlflow.sklearn.log_model(model, "model", signature=signature,
                                 input_example=X_tr.head(5))

Twelve runs in a few minutes. In the UI, Laura sorts by recall and the pattern jumps out — summarized:

run n_estimators max_depth class_weight accuracy precision recall F1
rf-300-20-balanced 300 20 balanced 0.842 0.55 0.68 0.61
rf-100-20-balanced 100 20 balanced 0.839 0.54 0.67 0.60
rf-300-None-balanced 300 None balanced 0.848 0.57 0.62 0.59
rf-300-None-None 300 None None 0.861 0.62 0.44 0.51
rf-100-None-None (≈ baseline) 100 None None 0.860 0.61 0.43 0.50

The reading, which is what matters:

  • class_weight="balanced" is the decisive factor: it tells the forest that each customer of the minority class (churned=1) weighs more during training, compensating for the 15% imbalance. Recall jumps from 0.43 to 0.68.
  • The price: precision drops from 0.61 to 0.55 (more false alarms) and accuracy falls from 0.86 to 0.842. Is it a good trade? For CineClick, yes: a retention call to a customer who wasn't going to leave costs little; an undetected churner is a lost subscription. Catching 68 out of every 100 instead of 43 is a real business improvement. F1 rises from 0.50 to 0.61, confirming that the trade-off is favorable overall.
  • Compare with module 2: bumping to 300 trees without touching the class weights gave a marginal improvement (recall 0.44). The systematic search found in one afternoon what one-off tweaks couldn't.

The new candidate: RandomForest with n_estimators=300, max_depth=20, class_weight="balanced" — accuracy 0.842, precision 0.55, recall 0.68, F1 0.61. To consecrate it, module 2's flow still applies: add max_depth and class_weight to the train: section of params.yaml (and as arguments in train.py), run dvc repro, and commit. Exploration with MLflow, consecration with DVC + Git. What to do with this "winning" model — versioning it and marking it formally — is exactly the topic of the next lesson.

Where tracking lives: local mlruns vs. a shared server

So far everything has been stored in an mlruns/ folder next to the repo (worth adding to .gitignore: it's tracking data, not code). That works for one person, but the goal was team memory: if Marc launches runs on his machine, Laura doesn't see them.

MLflow's architecture separates two stores:

  • Backend store: where params, metrics, and tags go — structured, queryable data. For real use: a database (PostgreSQL, MySQL, or SQLite to start with).
  • Artifact store: where the large files go (models, plots). For real use: object storage (S3, Azure Blob, GCS...) — the same kind of place where DVC's storage remote lives.
graph LR
    L["Laura<br/>train.py"] -->|"log_params, log_metrics"| S["Tracking server<br/>(mlflow server)"]
    M["Marc<br/>hyperparameter_search.py"] -->|"log_model"| S
    S --> B[("Backend store<br/>PostgreSQL: params, metrics, tags")]
    S --> A[("Artifact store<br/>S3/Blob: models, files")]
    U["Team's browser<br/>MLflow UI"] --> S

Standing up a shared server (here in its minimal local version, with SQLite):

mlflow server \
  --backend-store-uri sqlite:///mlflow.db \
  --artifacts-destination ./mlartifacts \
  --host 0.0.0.0 --port 5000

And the scripts point at it with an environment variable — without changing a single line of code:

export MLFLOW_TRACKING_URI=http://mlflow-server.cineclick.internal:5000

In real production, that server is just another internal service (a container with its managed database and its bucket), but the mechanics you've learned are identical: only the URI changes.

DVC and MLflow together, and the mystery of "v2" solved

The division of labor CineClick will use for the rest of the course:

Responsibility Tool Why
Version datasets and their changes DVC It's its native function (module 2)
Define and run the reproducible pipeline DVC (dvc.yaml + dvc repro) Selective invalidation, tied to Git
Log and compare experiments at scale MLflow Tracking Table, filters, team history
Store the model with signature and context MLflow Artifacts + signature + registry (next lesson)

DVC governs (decides what runs and guarantees it's reproducible); MLflow observes and remembers (records what happened and with what result). For that observation to carry full traceability, every run must be tied to code and data:

  • Code: MLflow automatically records the Git commit in the mlflow.source.git.commit tag (if you launch the script from a repo with the work committed — one more reason to commit before training).
  • Data: we tie it ourselves with a tag, reading the hash DVC already computes:
import yaml
lock = yaml.safe_load(open("dvc.lock"))
data_md5 = lock["stages"]["prepare_data"]["deps"][0]["md5"]
mlflow.set_tag("dvc_data_md5", data_md5)

And with this, let's settle the outstanding account. In module 1, the anecdote: a churn_model_v2.pkl existed that someone had trained, nobody knew with which parameters or which data, and the eternal doubt of whether "v2 was worse". The honest answer is that it was unanswerable: the information needed to answer it was never stored anywhere. It wasn't a mystery, it was an absence of record. With this lesson's system, the question can't even be asked: any future "v2" will be a run with its params, its metrics, its code commit, and its data hash — and comparing it with v1 will cost two clicks in the Compare tab. That is the moral of tracking: it doesn't make better models, it makes it impossible not to know where each one came from.

Common Mistakes and Tips

  • Logging metrics by hand instead of from variables: if you write mlflow.log_metric("recall", 0.68) with the number copy-pasted, the tracking lies as soon as anything changes. Always log the computed result.
  • Forgetting set_experiment and filling up Default: you'll end up with 200 runs from three projects mixed together. Set the experiment at the top of the script, always.
  • Runs orphaned of context: a run with no associated Git commit (because you launched with uncommitted changes) or no reference to the data is right back at the "v2" problem. Team convention: commit before launching "serious" experiments.
  • Putting the hyperparameter search inside dvc.yaml: the DVC pipeline is the blessed, deterministic path; exploration is divergent by nature. Keep them separate: exploration scripts log to MLflow; whatever wins gets consecrated in params.yaml and dvc repro.
  • Committing mlruns/: it's tracking data (potentially gigabytes of models). Into .gitignore. Shared memory is achieved with a tracking server, not with Git.
  • Confusing params and metrics: params are input (you choose them), metrics are output (the evaluation produces them). If you log test_size as a metric, the UI's filters and plots stop making sense.

Exercises

  1. UI filters. Write the MLflow search expressions for: (a) runs with F1 above 0.55; (b) runs from the hyperparameter search (tag type) that use class_weight balanced; (c) runs with recall greater than 0.6 and precision greater than 0.5.

  2. Extending the grid. Modify scripts/hyperparameter_search.py to add the hyperparameter min_samples_leaf with values [1, 5] to the grid. How many runs will the loop generate now? Also add a batch tag with value "grid-2" so you can tell this batch apart from the first one in the UI.

  3. Full traceability. Marc finds an interesting run from two weeks ago in the UI and wants to reproduce it exactly on his machine. List the steps, indicating which field of the run each piece of information comes from (assuming the run has the mlflow.source.git.commit tag and the dvc_data_md5 tag).

Solutions

  1. The three expressions:
metrics.f1 > 0.55
tags.type = 'hyperparameter-search' and params.class_weight = 'balanced'
metrics.recall > 0.6 and metrics.precision > 0.5

Note: param and tag values are compared as strings (single quotes); metrics, as numbers.

  1. The grid becomes 2 × 3 × 2 × 2 = 24 runs. Changes in the script:
grid = itertools.product(
    [100, 300], [10, 20, None], [None, "balanced"],
    [1, 5],  # min_samples_leaf
)

for n_est, depth, weight, leaf in grid:
    with mlflow.start_run(run_name=f"rf-{n_est}-{depth}-{weight}-{leaf}"):
        mlflow.log_params({..., "min_samples_leaf": leaf})
        mlflow.set_tag("type", "hyperparameter-search")
        mlflow.set_tag("batch", "grid-2")
        model = RandomForestClassifier(..., min_samples_leaf=leaf, random_state=42)

Tip: grids grow multiplicatively; before adding a fourth hyperparameter, ask yourself whether the results of the previous grid suggest it's worth it.

  1. Marc's steps: (1) from the run, copy the commit from the mlflow.source.git.commit tag and run git checkout <commit> in his clone — that recovers the code, params.yaml, dvc.yaml, and the .dvc metafiles; (2) run dvc checkout (with dvc pull if data is missing from his cache) — since the commit's metafiles point to the data of that day, the run's dvc_data_md5 serves as verification that he has the right version; (3) recreate the environment with that commit's lock (pip-sync requirements.lock); (4) run dvc repro. If the run was exploratory (a script outside the pipeline), run the script with the params the run shows. All the necessary context was in the run: that's exactly what "v2" never had.

Conclusion

CineClick now has memory: every training run is recorded in the churn-cineclick experiment with its parameters, metrics, model signature, code commit, and data hash, queryable and comparable in the MLflow UI. The division of labor is clear — DVC governs data and pipeline, MLflow observes and remembers — and the first systematic hyperparameter search has borne fruit: with class_weight="balanced", n_estimators=300, and max_depth=20, recall rises from 0.43 to 0.68 (precision 0.55, F1 0.61), a trade-off clearly favorable for the retention business. But that winning model is, right now, just "run f0d4... that turned out well" and a churn_model.pkl that someone will have to copy to wherever it needs to go. How does a winning run become the official model, with a version, a status, and a name production can ask for it by? That's the job of the model registry, the piece that manages the "blessed" model — and the topic of the next lesson.

◄ Previous
◄
Next ►
►
The Project | About Us | Contribute | Donations | License
© Copyright 2026. All rights reserved