Skip to main content

Module 3 — Experiment tracking with MLflow

The churn notebook is now reproducible: given the same code, data, and environment, it produces the same model. Reproducibility solves "how do I get this run back?". It does not solve "which of the fifty runs I did last week was actually the best, and why?". That is the job of experiment tracking, and MLflow is what we will use throughout the course.

What a run is, and what to log​

An experiment is a project — for us, telecom-churn. A run is one execution of the training code with specific inputs. For each run we log four things:

  • Parameters: hyperparameters and any configuration that affects the outcome. Learning rate, tree depth, feature list, data version, seed.
  • Metrics: values that measure quality, potentially over time. ROC AUC, F1, precision at a fixed recall, calibration error. Metrics can be single values or time series (loss per epoch).
  • Artifacts: files produced by the run. The trained model, the confusion matrix as a PNG, the SHAP summary, the feature importance table.
  • Metadata: git commit, user, start time, duration, host name. MLflow captures most of these automatically.

The rule: anything you would need to answer "why did this run beat the previous one?" belongs in the tracking system, not in memory or in a spreadsheet.

Instrumenting the churn training​

The instrumentation is deliberately intrusive: it wraps the training in a with mlflow.start_run() block and logs at the points where each value is produced.

import mlflow
import mlflow.sklearn
from sklearn.metrics import roc_auc_score, f1_score

mlflow.set_tracking_uri("http://mlflow.internal:5000")
mlflow.set_experiment("telecom-churn")

with mlflow.start_run(run_name="gbm-baseline"):
    mlflow.log_params({
        "model": "GradientBoosting",
        "n_estimators": 300,
        "max_depth": 5,
        "learning_rate": 0.05,
        "data_version": "v2026-08-31",
        "seed": 42,
    })

    model.fit(X_train, y_train)
    proba = model.predict_proba(X_val)[:, 1]

    mlflow.log_metric("roc_auc", roc_auc_score(y_val, proba))
    mlflow.log_metric("f1", f1_score(y_val, proba > 0.5))

    mlflow.sklearn.log_model(model, artifact_path="model")
    mlflow.log_artifact("reports/confusion_matrix.png")

Two subtleties. log_metric accepts a step argument for time series: log the loss with step=epoch and MLflow builds the curve automatically. log_model writes not just the pickle but a signature (input and output schemas) and the exact conda or pip environment — enough to redeploy the model anywhere.

Autologging: quick to enable, wrong to trust blindly​

mlflow.sklearn.autolog() (also mlflow.pytorch, mlflow.tensorflow, and others) hooks into the training call and logs parameters, metrics and the model automatically. One line before fit() and everything appears.

The trap is that autologging captures what the library exposes, not what your project needs. It does not know that data_version matters, nor which threshold you use to compute F1, nor that the split was stratified on tenure buckets. Autologging is a good default for prototyping and a bad substitute for explicit logging in a project that will go to production. Use it, then add explicit log_param and log_metric calls for the values it misses.

Comparing runs​

The MLflow UI lets you filter runs (params.model = "GradientBoosting" and metrics.roc_auc > 0.85), sort them, and select several to view their parallel coordinates plot. Two lessons from using this daily:

  • A metric improvement of 0.002 with an unchanged random seed is noise, not a gain. Vary the seed across three runs and look at the mean and standard deviation before celebrating.
  • The best run on validation is often not the best in production, because validation reflects a snapshot. Track the top three candidates through modules 4 and 5, not just the leader.

Local file store versus shared server​

mlflow.set_tracking_uri("file:./mlruns") writes runs into a local folder. That is fine for solo prototyping and useless for a team: runs are invisible to others, artifacts live on your laptop, and there is no audit trail.

A shared MLflow tracking server (a small FastAPI-style service) with a database backend (PostgreSQL is typical) and an object store for artifacts (S3, GCS, MinIO) fixes this. Every team member points MLFLOW_TRACKING_URI at the same server, and every run — from a laptop, from CI, from a scheduled retraining — appears in the same UI. The registry we will use in module 5 lives on this same server.

Log the data version, not just the model

The single most valuable parameter to log is the version of the dataset used. Without it, a run's metrics have no denominator: you cannot tell whether a gain came from a better model or from a lucky data snapshot. Module 4 makes data versioning concrete with DVC; the tag it produces goes straight into log_params.

Summary​

  • A run logs parameters, metrics, artifacts and metadata; the criterion is "would I need this to explain a difference between runs?".
  • Instrument explicitly with log_param, log_metric and log_model; autologging helps but misses project-specific values.
  • Compare runs in the UI with mean and standard deviation across seeds — a two-thousandth gain on one seed is noise.
  • A shared tracking server with a database and an object store turns individual runs into a team-visible history that CI and the registry will build on.

Next module: DVC — versioning the CSVs and the model files themselves so that a run's data_version parameter points at bytes you can recover.