---
title: "Experiment Tracking"
description: "MLflow-instrumented training with a hyperparameter sweep and a registry is the concrete, toolable skill most ML Engineer job descriptions list by name."
source: "https://confidentprep.com/paths/ml-engineering-on-aws/experiment-tracking/"
path: "ML Engineering on AWS"
repository: "https://github.com/confidentprep14-web/ai-engineering-projects/tree/main/experiment-tracking"
token_estimate: 3705
---

# Experiment Tracking

## Overview

MLflow-instrumented XGBoost training with a hyperparameter sweep, a model registry, and
inference against the registered model. This replaces project 4's manual folders with the
tool teams actually use — now every run is comparable, every model is registered, and
nothing depends on you remembering which folder had the best result.

## What to do

**Path:** 3 — ML Engineering on AWS
**Position:** 5 of 12
**Difficulty:** 🟢 (MLflow local server, no cloud)
**Time:** 3-4h
**AWS Cost:** None

---

### Agent Pickup Instructions

```bash
# Bootstrap
cd projects/confident-prep/p3-05-experiment-tracking
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

# Download dataset
mkdir -p data
curl -o data/adult.data \
  "https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.data"

# Start MLflow tracking server (leave running in a separate terminal)
python src/start_server.py

# In another terminal — run 5 experiments
python src/sweep.py

# Select best model and register it
python src/select_best.py

# Load from registry and infer
python src/load_and_infer.py

# Run tests
pytest tests/ -v
```

**Done when:**
- [ ] MLflow server starts at `http://localhost:5000` and UI is browsable
- [ ] `python src/sweep.py` creates 5 MLflow runs in experiment `"adult-income-sweep"`
- [ ] `python src/select_best.py` registers the best run's model with alias `"production"`
- [ ] `python src/load_and_infer.py` prints predictions for 5 sample rows without error
- [ ] `pytest tests/ -v` shows 4/4 passing
- [ ] MLflow UI at `http://localhost:5000` shows the experiment with 5 runs sortable by val_auc

---

### What this project is

You instrument the UCI Adult Income XGBoost training code with MLflow: logging hyperparameters, per-epoch metrics, and the final model. A sweep script runs 5 hyperparameter combinations as separate MLflow runs. A selection script queries the runs programmatically, identifies the best by val_AUC, registers it to the Model Registry with the alias "production", and a consumer script loads it by alias and runs inference. This replaces the folder-of-JSON-files approach from p3-04.

### What the learner achieves

"I can track ML experiments with MLflow, compare runs programmatically, register a model with a semantic alias, and load that model for inference — so I can reproduce any experiment and promote models without manual file management."

---

### Folder structure

```
p3-05-experiment-tracking/
├── data/
│   └── adult.data              # raw UCI dataset
├── mlruns/                     # created by MLflow automatically (do not commit)
├── src/
│   ├── train_with_mlflow.py   # single run with full MLflow instrumentation
│   ├── sweep.py               # 5 runs across hyperparameter grid
│   ├── select_best.py         # find best run by val_auc, register to Model Registry
│   ├── load_and_infer.py      # load "production" model from registry, predict
│   └── start_server.py        # launch mlflow server on port 5000
├── tests/
│   ├── test_tracking.py       # run created, correct experiment, best model selection
│   └── test_registry.py       # registry has "production" alias after select_best
├── .env.example
├── requirements.txt
└── README.md
```

---

### .env.example

```dotenv
# MLflow tracking URI — local file store (default, no server needed for tests)
MLFLOW_TRACKING_URI=http://localhost:5000

# MLflow experiment name
MLFLOW_EXPERIMENT_NAME=adult-income-sweep

# MLflow Model Registry name
MLFLOW_MODEL_NAME=adult-income-xgboost

# Local data path
LOCAL_DATA_DIR=data/
```

---

### requirements.txt

```
mlflow==2.19.0
xgboost==2.1.3
scikit-learn==1.6.0
pandas==2.2.3
numpy==2.2.1
python-dotenv==1.0.1
pytest==8.3.4
```

---

### src/ — what to implement

#### `src/train_with_mlflow.py`

**CLI arguments:**
- `--max-depth` (int, default 5)
- `--learning-rate` (float, default 0.1)
- `--n-estimators` (int, default 100)
- `--data` (str, default `data/adult.data`)
- `--experiment-name` (str, default from `MLFLOW_EXPERIMENT_NAME` env var or `"adult-income-sweep"`)
- `--run-name` (str, optional)

**`load_and_split(data_path: str) -> tuple`**
- Same as p3-01/p3-04: assign column names, strip whitespace, replace `?` with NaN, drop NaN rows, pd.get_dummies, encode target
- Returns `(X_train, X_test, y_train, y_test)`, 80/20 split, `random_state=42`

**`train_and_log(X_train, X_test, y_train, y_test, params: dict, experiment_name: str, run_name: str | None) -> str`**

Sets up MLflow run and logs:

```python
mlflow.set_tracking_uri(os.getenv("MLFLOW_TRACKING_URI", "http://localhost:5000"))
mlflow.set_experiment(experiment_name)

with mlflow.start_run(run_name=run_name) as run:
    # Log hyperparameters
    mlflow.log_param("max_depth", params["max_depth"])
    mlflow.log_param("learning_rate", params["learning_rate"])
    mlflow.log_param("n_estimators", params["n_estimators"])

    # Train
    model = XGBClassifier(use_label_encoder=False, eval_metric="logloss", **params)
    model.fit(X_train, y_train,
              eval_set=[(X_test, y_test)],
              verbose=False)

    # Log per-round validation logloss
    results = model.evals_result()
    for i, loss in enumerate(results["validation_0"]["logloss"]):
        mlflow.log_metric("val_logloss", loss, step=i)

    # Final metrics
    val_auc = roc_auc_score(y_test, model.predict_proba(X_test)[:, 1])
    val_acc = accuracy_score(y_test, model.predict(X_test))
    mlflow.log_metric("val_auc", val_auc)
    mlflow.log_metric("val_accuracy", val_acc)

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

    return run.info.run_id
```

Returns the MLflow run_id.

**`main()`**
- Parse args → `load_and_split` → `train_and_log`
- Print: `Run {run_id} complete. val_auc={val_auc:.4f}`

---

#### `src/sweep.py`

**`SWEEP_CONFIGS: list[dict]`**
```python
SWEEP_CONFIGS = [
    {"max_depth": 3, "learning_rate": 0.1,  "n_estimators": 100},
    {"max_depth": 5, "learning_rate": 0.1,  "n_estimators": 100},
    {"max_depth": 7, "learning_rate": 0.1,  "n_estimators": 100},
    {"max_depth": 5, "learning_rate": 0.05, "n_estimators": 200},
    {"max_depth": 5, "learning_rate": 0.01, "n_estimators": 200},
]
```

**`run_sweep(configs: list[dict], data_path: str, experiment_name: str) -> list[str]`**
- Iterates through configs
- For each: calls `train_with_mlflow.train_and_log(...)` with `run_name=f"sweep-{i+1}"`
- Prints `Running sweep {i+1}/5: {config}`
- Returns list of run_ids

**`main()`**
- Reads experiment name from env or default
- Calls `run_sweep`
- Prints: `Sweep complete. {n} runs in experiment "{experiment_name}"`
- Prints: `View at: http://localhost:5000`

---

#### `src/select_best.py`

**CLI:** `python src/select_best.py [--experiment-name adult-income-sweep] [--model-name adult-income-xgboost]`

**`find_best_run(experiment_name: str, metric: str = "val_auc") -> pd.DataFrame`**
- Calls `mlflow.search_runs(experiment_names=[experiment_name], order_by=[f"metrics.{metric} DESC"])`
- Returns the top run as a single-row DataFrame
- Raises `ValueError("No runs found")` if experiment has no runs

**`register_best_model(best_run: pd.Series, model_name: str, alias: str = "production") -> str`**
- Gets the run_id and constructs the model URI: `f"runs:/{run_id}/model"`
- Calls `mlflow.register_model(model_uri, model_name)` — returns `ModelVersion`
- Sets alias using `mlflow.MlflowClient().set_registered_model_alias(model_name, alias, version)`
- Returns the registered model version number as string

**`main()`**
- Find best run → register → print:
  ```
  Best run: {run_id}
  val_auc: {val_auc:.4f}
  Registered as: {model_name} version {version} with alias "production"
  ```

---

#### `src/load_and_infer.py`

**CLI:** `python src/load_and_infer.py [--model-name adult-income-xgboost] [--alias production] [--data data/adult.data]`

**`load_production_model(model_name: str, alias: str) -> mlflow.pyfunc.PyFuncModel`**
- Constructs URI: `f"models:/{model_name}@{alias}"`
- Calls `mlflow.pyfunc.load_model(uri)`
- Returns the loaded model

**`run_inference(model, data_path: str, n_samples: int = 5) -> pd.DataFrame`**
- Loads `n_samples` rows from the dataset (rows 0 to n_samples-1)
- Preprocesses exactly as in training (column names, dummies — must use same columns as training)
- Calls `model.predict(X_sample)`
- Returns DataFrame with columns `["row", "prediction", "label"]` where label is the original income value

**`main()`**
- Load model → run inference
- Print predictions table
- Print: `Model loaded from alias "{alias}". {n_samples} predictions made.`

**Edge cases:**
- If model registry does not have the alias, catch `MlflowException` and print: `No model with alias "{alias}". Run select_best.py first.`

---

#### `src/start_server.py`

**`main()`**
- Runs: `mlflow server --host 0.0.0.0 --port 5000 --backend-store-uri ./mlruns --default-artifact-root ./mlruns`
- Uses `subprocess.run()` so it can be `Ctrl+C`'d
- Prints before starting: `Starting MLflow server at http://localhost:5000`

---

### tests/ — what to test

#### `tests/test_tracking.py`

**`test_mlflow_run_created_with_correct_experiment`**
- Set `MLFLOW_TRACKING_URI` to a temp directory (`file:///tmp/mlruns_test`)
- Call `train_with_mlflow.train_and_log(...)` with `experiment_name="test-experiment"`
- Assert the returned run_id is a non-empty string
- Assert `mlflow.get_experiment_by_name("test-experiment")` is not None

**`test_best_model_selection_picks_highest_auc`**
- Create 3 mock MLflow runs in a temp experiment with val_auc: 0.85, 0.91, 0.87
- Call `select_best.find_best_run("test-experiment")`
- Assert the returned run has `metrics.val_auc == 0.91`

#### `tests/test_registry.py`

**`test_production_alias_set_after_select_best`**
- Run `sweep.run_sweep(SWEEP_CONFIGS[:2], data, experiment_name)` to create 2 runs
- Run `select_best.register_best_model(best_run, model_name="test-model", alias="production")`
- Assert `mlflow.MlflowClient().get_model_version_by_alias("test-model", "production")` returns a result

**`test_load_and_infer_produces_predictions_without_error`**
- After registry test above, call `load_and_infer.load_production_model("test-model", "production")`
- Call `load_and_infer.run_inference(model, data_path, n_samples=3)`
- Assert result DataFrame has 3 rows and `prediction` column

---

### README.md content

```markdown
# P3-05 — Experiment Tracking

MLflow-instrumented XGBoost training with hyperparameter sweep, model registry, and inference.

## Prerequisites

- Python 3.11+
- No AWS credentials needed

## Quick start

**Terminal 1 — start MLflow server:**
```bash
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
python src/start_server.py
```

**Terminal 2 — run experiments:**
```bash
source .venv/bin/activate
mkdir -p data && curl -o data/adult.data \
  "https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.data"
python src/sweep.py
python src/select_best.py
python src/load_and_infer.py
```

## View experiments

Open `http://localhost:5000` in your browser.

## Run a single experiment

```bash
python src/train_with_mlflow.py --max-depth 5 --learning-rate 0.05
```

## Tests

```bash
pytest tests/ -v
```

Note: tests use a temp MLflow tracking directory, not the local server.
```

---

### GUIDE.md content

```markdown
# Guide — P3-05 Experiment Tracking

## What MLflow gives you vs the JSON folder

| Capability | p3-04 (JSON files) | p3-05 (MLflow) |
|---|---|---|
| Search by metric | Scan all JSON files manually | `mlflow.search_runs(order_by=["metrics.val_auc DESC"])` |
| Compare runs visually | Not possible | MLflow UI with parallel coordinates |
| Link metric to code version | Not possible | MLflow auto-logs git SHA |
| Load model for inference | Hardcode the model path | `mlflow.pyfunc.load_model("models:/name@production")` |
| Promote model to production | Copy file, update a config | Set alias in registry |

## What `mlflow.log_metric(..., step=i)` does

Logging with a step creates a time series in MLflow. When you log `val_logloss` at each
boosting round, the MLflow UI shows you the learning curve — you can see if the model
is still improving or has plateaued.

## The Model Registry vs artifact storage

The model artifact (the XGBoost `.xgb` file) lives in the artifact store (your `./mlruns`
directory or S3 in production). The Model Registry is a metadata layer on top: it tracks
versions, aliases (`production`, `staging`, `archived`), and descriptions.

When you `load_model("models:/adult-income-xgboost@production")`, MLflow looks up the
registry to find which artifact version has the `production` alias, then retrieves it.
You never hard-code a file path.

## The alias pattern

Using aliases (`production`, `staging`) instead of version numbers decouples your inference
code from the registry history. When you promote a new model to production, you move the
alias — the inference code does not change.

## Interview framing

"I use MLflow for experiment tracking — every run logs its hyperparameters, per-epoch metrics,
and model artifact. I can compare 5 runs in one query, register the winner with a semantic alias,
and load it by alias at inference time. That's the workflow I'd use on a team."
```

---

### Failure modes to handle

| Failure | Where | How to handle |
|---|---|---|
| MLflow server not running | `train_with_mlflow.py` | If `MLFLOW_TRACKING_URI` points to a server and it's unreachable, MLflow raises `MlflowException` — catch and print "Is the MLflow server running? Run: python src/start_server.py" |
| Experiment has no runs when selecting best | `select_best.py: find_best_run` | Raise `ValueError("No runs found in experiment. Run sweep.py first.")` |
| `production` alias not set when loading | `load_and_infer.py` | Catch `MlflowException`, print actionable message |
| Column mismatch between training and inference | `load_and_infer.py: run_inference` | `pd.get_dummies` must produce same columns — use `reindex(columns=train_columns, fill_value=0)` |
| `mlruns/` directory permission error | `start_server.py` | Print the error and suggest `chmod 755 mlruns/` |

---

### The metric this project measures

**What:** Validation AUC across all 5 sweep runs; spread between best and worst

**Format:** Per-run: `val_auc={val:.4f}` logged to MLflow and printed. Summary: `Best: {max_auc:.4f}, Worst: {min_auc:.4f}, Spread: {spread:.4f}`

**Target:** Best run AUC > 0.88 on UCI Adult. Spread between best and worst should be > 0.01 (demonstrating that hyperparameter choice matters).

## Source code

A reference implementation of Experiment Tracking — README, guide, source and tests — is at https://github.com/confidentprep14-web/ai-engineering-projects/tree/main/experiment-tracking.