Experiment Tracking

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

# 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

# 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:

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]

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

# 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:

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

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

Tests

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).