Train, Evaluate and Track with MLflow — ML Ops & Data Science in Production

What You'll Build

In this exercise you will train an XGBoost regression model that predicts IPL batsman strike rates from match-history features. You will run three experiment trials with different hyperparameter configurations, log parameters and evaluation metrics (MAE, RMSE, R²) to MLflow after each run, store the serialised model artefact, compare all three runs in the MLflow UI, and promote the best-performing run to the 'Staging' stage in the MLflow Model Registry. By the end you will have a reproducible, traceable training pipeline ready to feed into the FastAPI service in lesson 32.

Analogy🏏Cricket

🏏 Think of it like cricket: Before the IPL auction, a franchise's analytics team evaluates every player across dozens of trial matches — tracking Rohit Sharma's strike rate in power plays, Virat Kohli's average against pace, MS Dhoni's finishing rate in the death overs, Shubman Gill's consistency across pitches, and Jasprit Bumrah's economy in the middle overs. Each trial is recorded in a shared logbook so the selectors can compare and pick the best combination. MLflow is exactly that shared logbook for your ML experiments — every training run is a trial match, every metric is a scorecard entry, and the Model Registry is the final squad announcement. Keep the auction framing in mind throughout the exercise, because it fixes the discipline the steps teach: a franchise never signs a player off one good highlight reel, and you never register a model off one lucky run — you log every trial, compare them on identical conditions, and promote only with the full scorecard in front of you.

Prerequisites

Before starting this exercise ensure the following packages are installed: xgboost>=2.0, mlflow>=2.13, scikit-learn>=1.4, pandas>=2.2, numpy>=1.26. You should have MLflow's local tracking server accessible at http://localhost:5000 (start it with mlflow ui --port 5000). The exercise uses a synthetic IPL dataset that you will generate in the Setup step — no external data download is needed. Familiarity with scikit-learn's train_test_split and metrics modules is assumed from earlier lessons.

Analogy🏏Cricket

🏏 Think of it like cricket: the kit check before a net session. Just as a batter arriving for practice needs pads, gloves, and a bat they already know how to use — but doesn't need to have faced the new bowling machine before, because today's session is exactly where they'll learn it — this exercise expects you to arrive comfortable with Python, basic scikit-learn (fit, predict, train_test_split), and pandas, while MLflow itself is taught from scratch. Just as the coach insists on a properly prepared practice pitch — a clean, dedicated strip rather than the match square — you need a clean virtual environment or Conda environment where packages can be installed freely. And just as the only outside help needed is the equipment delivery van arriving once before practice, network access is required only for the initial pip install; after that everything runs locally, whether your 'net' is a laptop, a Docker container, or a cloud notebook. The payoff: checking your kit now means the session ahead is pure skill-building, with no stoppages for missing gear.

Setup

python

# setup.py — generate synthetic IPL dataset and configure MLflow
import numpy as np
import pandas as pd
import mlflow
from sklearn.model_selection import train_test_split

np.random.seed(42)

# Synthetic IPL player match records
ipl_players = ['Rohit Sharma', 'Virat Kohli', 'Shubman Gill', 'MS Dhoni', 'Jasprit Bumrah']
num_matches = 800  # total match-level records

batting_average   = np.random.normal(35, 10, num_matches).clip(10, 70)
innings_count     = np.random.randint(20, 250, num_matches)
boundary_pct      = np.random.beta(5, 5, num_matches)   # centred around 0.5
dot_ball_pct      = np.random.beta(3, 5, num_matches)   # centred around 0.375
opponent_bowling_quality = np.random.normal(7.5, 1.2, num_matches).clip(5, 10)
venue_encoded     = np.random.randint(0, 12, num_matches)
is_knockout       = np.random.randint(0, 2, num_matches)

# Target: strike rate has positive correlation with avg and boundary%, negative with dot%
strike_rate = (
    0.8 * batting_average
    + 60 * boundary_pct
    - 50 * dot_ball_pct
    + 2 * innings_count / 100
    - 3 * opponent_bowling_quality
    + 5 * is_knockout
    + np.random.normal(0, 8, num_matches)  # noise
).clip(80, 220)

ipl_match_df = pd.DataFrame({
    'batting_average': batting_average,
    'innings_count': innings_count,
    'boundary_pct': boundary_pct,
    'dot_ball_pct': dot_ball_pct,
    'opponent_bowling_quality': opponent_bowling_quality,
    'venue_encoded': venue_encoded,
    'is_knockout': is_knockout,
    'strike_rate': strike_rate
})

FEATURE_COLS = [
    'batting_average', 'innings_count', 'boundary_pct',
    'dot_ball_pct', 'opponent_bowling_quality', 'venue_encoded', 'is_knockout'
]

X = ipl_match_df[FEATURE_COLS].values
y = ipl_match_df['strike_rate'].values

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Point MLflow at a local tracking server
mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("ipl-strike-rate-experiment")

print(f"Dataset shape: {ipl_match_df.shape}")
print(f"Train samples: {X_train.shape[0]}, Test samples: {X_test.shape[0]}")
print("MLflow experiment set: ipl-strike-rate-experiment")

Step 1: Train Three XGBoost Runs with Different Hyperparameters

XGBoost's key hyperparameters for regression are the number of boosting rounds (n_estimators), the learning rate (which controls how much each new tree corrects the previous ensemble), the maximum tree depth (which controls model complexity), and the column subsampling ratio (which adds regularisation by training each tree on a random subset of features). We will try three distinct configurations to simulate a real hyperparameter search: a conservative low-learning-rate configuration, a deeper high-capacity configuration, and a regularised balanced configuration.

Analogy🏏Cricket

🏏 Think of it like cricket: Rohit Sharma adapts his batting tempo across powerplay, middle-overs, and death overs — three distinct strategies for three contexts. Hyperparameter configurations are the same: a conservative approach (low learning rate, shallow trees) is like Rohit rotating the strike patiently; a high-capacity approach (deep trees, more rounds) is like hitting a six every over; a regularised approach balances risk and reward. MLflow records all three 'innings' so you can pick the highest scorer.

python

import xgboost as xgb
import mlflow
import mlflow.xgboost
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import numpy as np

# Three hyperparameter configurations to compare
xgb_experiment_configs = [
    {
        "run_name": "conservative-low-lr",
        "n_estimators": 200,
        "learning_rate": 0.03,
        "max_depth": 4,
        "subsample": 0.8,
        "colsample_bytree": 0.8,
    },
    {
        "run_name": "high-capacity-deep",
        "n_estimators": 300,
        "learning_rate": 0.1,
        "max_depth": 8,
        "subsample": 0.9,
        "colsample_bytree": 0.9,
    },
    {
        "run_name": "regularised-balanced",
        "n_estimators": 150,
        "learning_rate": 0.05,
        "max_depth": 5,
        "subsample": 0.75,
        "colsample_bytree": 0.75,
    },
]

run_results = []

for cfg in xgb_experiment_configs:
    with mlflow.start_run(run_name=cfg["run_name"]) as run:
        params = {k: v for k, v in cfg.items() if k != "run_name"}

        # Log all hyperparameters
        mlflow.log_params(params)
        mlflow.log_param("feature_cols", FEATURE_COLS)
        mlflow.log_param("train_size", X_train.shape[0])

        # Train XGBoost model
        ipl_strike_rate_model = xgb.XGBRegressor(
            objective="reg:squarederror",
            random_state=42,
            **params
        )
        ipl_strike_rate_model.fit(X_train, y_train)

        # Evaluate on held-out test set
        batting_predictions = ipl_strike_rate_model.predict(X_test)
        mae  = mean_absolute_error(y_test, batting_predictions)
        rmse = np.sqrt(mean_squared_error(y_test, batting_predictions))
        r2   = r2_score(y_test, batting_predictions)

        # Log evaluation metrics
        mlflow.log_metric("mae",  mae)
        mlflow.log_metric("rmse", rmse)
        mlflow.log_metric("r2",   r2)

        # Log model artefact
        mlflow.xgboost.log_model(ipl_strike_rate_model, artifact_path="model")

        run_results.append({
            "run_id":   run.info.run_id,
            "run_name": cfg["run_name"],
            "mae":  mae,
            "rmse": rmse,
            "r2":   r2,
        })
        print(f"[{cfg['run_name']}] MAE={mae:.3f}  RMSE={rmse:.3f}  R²={r2:.4f}")

print("\nAll 3 runs completed and logged to MLflow.")

Step 2: Compare Runs and Select the Best

With all three runs logged, MLflow's comparison view lets you sort columns by any metric. Programmatically, the MlflowClient.search_runs API returns runs as a list of Run objects sorted by any metric column. We will use this to identify the run with the lowest RMSE — our primary selection criterion — then secondary-check that its R² is above 0.75 to confirm the model is genuinely capturing signal rather than just predicting the mean. This two-gate criterion mirrors real model-selection policies used in production ML pipelines.

Analogy🏏Cricket

🏏 Think of it like cricket: When selectors pick the Indian squad, they do not just look at batting average — they also consider strike rate, recent form, and fielding. A batsman with a high average but poor strike rate may not fit the T20 format. Our two-gate criterion (RMSE < threshold AND R² > 0.75) is exactly that: we want the model that is both accurate in absolute terms and explains a meaningful portion of variance, not a model that games one metric at the expense of the other.

python

from mlflow import MlflowClient
import pandas as pd

client = MlflowClient(tracking_uri="http://localhost:5000")
experiment = client.get_experiment_by_name("ipl-strike-rate-experiment")

# Retrieve all runs sorted by RMSE ascending
all_runs = client.search_runs(
    experiment_ids=[experiment.experiment_id],
    order_by=["metrics.rmse ASC"]
)

# Build comparison table
comparison_rows = []
for r in all_runs:
    comparison_rows.append({
        "run_name": r.data.tags.get("mlflow.runName", r.info.run_id[:8]),
        "run_id":   r.info.run_id,
        "mae":      round(r.data.metrics.get("mae",  0), 4),
        "rmse":     round(r.data.metrics.get("rmse", 0), 4),
        "r2":       round(r.data.metrics.get("r2",   0), 4),
    })

comparison_df = pd.DataFrame(comparison_rows)
print("=== MLflow Run Comparison ===")
print(comparison_df.to_string(index=False))

# Select best run: lowest RMSE with R² > 0.75
eligible_runs = [r for r in all_runs if r.data.metrics.get("r2", 0) > 0.75]
if not eligible_runs:
    raise RuntimeError("No run meets the R² > 0.75 quality gate.")

best_run = eligible_runs[0]  # already sorted by RMSE ASC
best_run_id   = best_run.info.run_id
best_rmse     = best_run.data.metrics["rmse"]
best_r2       = best_run.data.metrics["r2"]
best_run_name = best_run.data.tags.get("mlflow.runName", "unknown")

print(f"\nBest run selected: {best_run_name}")
print(f"  Run ID : {best_run_id}")
print(f"  RMSE   : {best_rmse:.4f}")
print(f"  R²     : {best_r2:.4f}")

Step 3: Register and Promote the Best Model to Staging

The MLflow Model Registry provides a versioned catalogue of your trained models. Each time you register a run's artefact under a model name, a new version is created. Version stages — None, Staging, Production, Archived — control which version downstream consumers (the FastAPI service) load at startup. Transitioning the best run to Staging and archiving all previous Staging versions ensures the API always loads the latest validated model without requiring a code change or redeploy.

Analogy🏏Cricket

🏏 Think of it like cricket: The BCCI maintains a central squad list — probables, main squad, playing XI. Promoting a model to Staging is like moving a player from 'probables' to the 'travelling squad': they have been validated but are not yet in the playing XI (Production). MS Dhoni was famously recalled to the squad at Staging (IPL comeback) before being confirmed in the playing XI. The registry gives you that same controlled, auditable transition path.

python

import mlflow
from mlflow import MlflowClient

client = MlflowClient(tracking_uri="http://localhost:5000")
MODEL_NAME = "ipl-strike-rate-predictor"

# Register the best run's model artefact
model_uri = f"runs:/{best_run_id}/model"
print(f"Registering model from: {model_uri}")

registered_model_version = mlflow.register_model(
    model_uri=model_uri,
    name=MODEL_NAME
)
print(f"Registered as version: {registered_model_version.version}")

# Transition to Staging, archiving any existing Staging version
client.transition_model_version_stage(
    name=MODEL_NAME,
    version=registered_model_version.version,
    stage="Staging",
    archive_existing_versions=True
)
print(f"Version {registered_model_version.version} promoted to Staging")

# Verify the promotion
staging_models = client.get_latest_versions(MODEL_NAME, stages=["Staging"])
for mv in staging_models:
    print(f"\nStaging model details:")
    print(f"  Name    : {mv.name}")
    print(f"  Version : {mv.version}")
    print(f"  Stage   : {mv.current_stage}")
    print(f"  Run ID  : {mv.run_id}")

# Confirm the model loads correctly
loaded_model = mlflow.pyfunc.load_model(f"models:/{MODEL_NAME}/Staging")
test_prediction = loaded_model.predict(X_test[:3])
print(f"\nSample predictions (first 3 test rows): {test_prediction.round(2)}")
print(f"Actual strike rates               : {y_test[:3].round(2)}")

Testing & Verification

Run the verification script below to confirm that all three experiment runs were logged correctly, the model registry contains a version in Staging, and the loaded Staging model produces sensible predictions. The script checks for the presence of required metrics and that RMSE is within an acceptable range for the synthetic dataset. If any assertion fails, re-run the corresponding step above after checking that the MLflow tracking server is still running on port 5000.

Analogy🏏Cricket

🏏 Think of it like cricket: the verification script is the match referee's pre-game inspection, and its value is exactly that it checks the *system*, not the players. Before any international fixture, the referee walks a fixed checklist — pitch hardness, boundary rope distance, floodlight levels, sightscreen operation — because a brilliant team on a defective ground still produces an invalid match. Your script does the referee's walk over the tracking setup: can it reach the tracking store, do the expected six runs exist, does each run carry its parameters and metrics, is exactly one model version sitting in Production? Each check is boring in isolation and decisive in combination. The deeper habit this builds is the MLOps one: never certify a pipeline by eyeballing one pretty dashboard screenshot, the way a referee never certifies a ground by glancing at it from the pavilion. Write the checklist as executable code, run it after every change, and let a green result — like the referee's signed pre-match form — be the only thing that clears the ground for play. When the same script runs in CI, every future teammate inherits the inspection for free.

python

# verify_mlflow.py — automated verification checks
from mlflow import MlflowClient
import mlflow
import numpy as np

client = MlflowClient(tracking_uri="http://localhost:5000")
MODEL_NAME = "ipl-strike-rate-predictor"

# Check 1: Experiment exists
experiment = client.get_experiment_by_name("ipl-strike-rate-experiment")
assert experiment is not None, "FAIL: Experiment not found"
print("[PASS] Experiment 'ipl-strike-rate-experiment' exists")

# Check 2: At least 3 runs exist
all_runs = client.search_runs([experiment.experiment_id])
assert len(all_runs) >= 3, f"FAIL: Expected >= 3 runs, got {len(all_runs)}"
print(f"[PASS] {len(all_runs)} runs found in experiment")

# Check 3: All runs have required metrics
for r in all_runs:
    for metric in ["mae", "rmse", "r2"]:
        assert metric in r.data.metrics, f"FAIL: Run {r.info.run_id} missing metric '{metric}'"
print("[PASS] All runs have mae, rmse, r2 metrics")

# Check 4: Staging model exists in registry
staging_versions = client.get_latest_versions(MODEL_NAME, stages=["Staging"])
assert len(staging_versions) == 1, "FAIL: Expected exactly 1 Staging model version"
print(f"[PASS] Staging model version {staging_versions[0].version} found")

# Check 5: RMSE of staging model < 20 (sanity range for synthetic data)
staging_run = client.get_run(staging_versions[0].run_id)
staging_rmse = staging_run.data.metrics["rmse"]
assert staging_rmse < 20, f"FAIL: Staging RMSE {staging_rmse:.4f} exceeds threshold 20"
print(f"[PASS] Staging RMSE = {staging_rmse:.4f} (below threshold 20)")

print("\n=== All verification checks PASSED ===")

Common Mistake: Forgetting to call archive_existing_versions=True when transitioning a new version to Staging. If you leave this out, the Model Registry ends up with multiple Staging versions — and the FastAPI service (which calls get_latest_versions with stages=['Staging']) will return whichever version was registered last, not necessarily the best one. Always archive existing versions when promoting to avoid ambiguity.

Pro Tip

Open http://localhost:5000 in your browser after running all three steps. Navigate to 'Experiments → ipl-strike-rate-experiment' and click 'Chart view'. MLflow will render a parallel coordinates plot showing how each hyperparameter configuration maps to RMSE — an extremely fast way to visually identify which parameter drove the most improvement. This view is more informative than a table when you have 10+ experiment runs.

XGBoost hyperparameters n_estimators, learning_rate, max_depth, and colsample_bytree all interact — always evaluate combinations, not individual parameters in isolation.
Use mlflow.log_params for the full hyperparameter dict and mlflow.log_metric for each evaluation metric after fitting on the test set.
The two-gate promotion policy (RMSE below threshold AND R² above 0.75) prevents a model that games one metric from reaching the Staging registry.
archive_existing_versions=True is mandatory when transitioning to Staging to prevent multiple concurrent Staging versions from confusing downstream consumers.
The mlflow.pyfunc.load_model API with a models:/<name>/Staging URI decouples the FastAPI service from specific run IDs, enabling model updates without code changes.