Track ML Experiments with MLflow

Train models, compare runs visually, and register a champion model locally.

Introduction

30 Second Summary

Imagine tweaking a recipe dozens of times but never writing down what you changed. That is what training machine learning models without experiment tracking feels like. You lose your best results the moment you try something new.

In this project, you will build a complete ML experiment tracking system using MLflow and scikit-learn. You will train models to predict housing prices, track every experiment automatically, compare results visually, and register your best model for production use.

What You'll Build

By the end of this project, you will have a local MLflow dashboard open in your browser where you can compare training runs, spot the best-performing model at a glance, and load it for predictions in a single line of code.

By the end of this project, you'll have:

  • A trained machine learning model that predicts California housing prices from real census data, with results you can see printed in your terminal.
  • A visual experiment tracking dashboard where you compare multiple training runs side by side and identify which settings produced the best predictions.
  • A Model Registry entry where your best model is versioned and aliased as "champion", loadable for inference in one line of Python.
  • Secret Mission: Use GridSearchCV to automate hyperparameter search and watch MLflow create nested runs for every parameter combination.

Are there any prerequisites?

No prior machine learning experience is required. This project walks you through everything from installing the tools to registering a trained model. You just need a computer with VS Code and Python installed.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters.

Set Up Your ML Environment

Before you can train a model or track an experiment, you need the right Python packages installed in an isolated environment.

An isolated environment keeps this project's dependencies separate from everything else on your machine. In this step, you'll verify your Python version, create a virtual environment, and install MLflow alongside scikit-learn.

In this step, get ready to:
  • Check your Python version and create a virtual environment.
  • Install mlflow and scikit-learn packages.
  • Verify installations and start the MLflow tracking server.
Check your Python version and create a virtual environment

scikit-learn 1.9.0 requires Python 3.11 or higher. Let's confirm your installed version meets this minimum.

  • Check your Python version by running this command:

macOS

python3 --version

Windows

python --version

You should see a version number printed, such as Python 3.11.x or Python 3.12.x.

✔️ I see version 3.11 or higher

Your Python version meets the requirement. Continue to the next instruction below.

ⓧ I see an older version

You have Python installed, but it's below 3.11. scikit-learn 1.9.0 will not work with older versions.

  • Download Python 3.11 or newer from python.org/downloads.
  • Run the installer and check the box to add Python to your PATH.
  • Close and reopen your terminal, then rerun the version check above.

macOS users with multiple Python versions

macOS often ships with an older system Python. After installing a newer version, use python3.11 --version or python3.12 --version to confirm the new install is available. You may need to use the full binary name (python3.11) instead of python3 for the rest of this project.

ⓧ Command not found

Python is not in your system PATH.

  • Download Python 3.11 or newer from python.org/downloads.
  • Run the installer and check the box to add Python to your PATH.
  • Close and reopen your terminal, then rerun the version check above.
  • Open VS Code by pressing Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type VS Code and press Enter to open it.
  • Open the integrated terminal by pressing Ctrl+` (Windows) or Cmd+` (macOS).
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir mlflow-experiments
cd mlflow-experiments
  • Confirm you are inside the new folder by running:
pwd

You should see a path ending in mlflow-experiments.

  • Now open this folder in VS Code. Click File in the top menu bar.
  • Click Open Folder.
  • Navigate to your Desktop, select the mlflow-experiments folder, and click Open.

Why open the folder?

Opening the project folder in VS Code means the file explorer shows your project files and the integrated terminal starts in the correct directory. Every file you create in the next steps will appear in the sidebar automatically.

VS Code reloads with the folder open. Your previous terminal session closes when this happens.

  • Open a new terminal by pressing Ctrl+` (Windows) or Cmd+` (macOS).
  • Create a virtual environment by running this command:

macOS

python3 -m venv venv

Windows

python -m venv venv
  • Activate the virtual environment:

macOS

source venv/bin/activate

Windows

.\venv\Scripts\activate

Your terminal prompt now shows (venv) at the beginning. This confirms the virtual environment is active.

Don't see (venv) in your prompt?

  • Make sure you ran the activation command for your platform (macOS uses source, Windows uses .\venv\Scripts\activate).
  • On Windows, if you get a permissions error, run Set-ExecutionPolicy -Scope CurrentUser RemoteSigned in PowerShell, then try the activation command again.
  • Confirm the venv folder exists by running ls in your terminal.

Still stuck? Help me activate my Python virtual environment

Install mlflow and scikit-learn

With your virtual environment active, packages you install stay isolated inside this project.

  • Install the project dependencies by running:
pip install mlflow==3.14.0 scikit-learn==1.9.0 pandas

This downloads and installs all three packages plus their dependencies. The install takes a minute or two depending on your network speed.

What are these packages?

  • mlflow tracks experiments, logs parameters and metrics, and hosts the visual dashboard.
  • scikit-learn provides the machine learning algorithms you'll train.
  • pandas handles data loading and manipulation in table format.

When the install finishes, you should see Successfully installed followed by a list of package names near the bottom of the output.

Seeing a permissions error or install failure?

  • Make sure your prompt shows (venv). If it does not, rerun the activation command from the previous substep.
  • If you see a network timeout, check your internet connection and try running the install command again.

Need more help? Help me fix my pip install error

Verify installations and start the MLflow tracking server

Before you start building, let's confirm everything installed correctly and launch the tracking server you'll use for the rest of the project.

  • Verify MLflow is installed by running:
python -c "import mlflow; print(mlflow.__version__)"

Before you check the output, what version number do you expect to see based on what you just installed?

You should see 3.14.0 printed. This confirms MLflow is importable and the correct version is active inside your virtual environment.

Seeing ModuleNotFoundError?

  • Your virtual environment might not be active. Check that (venv) appears in your prompt and rerun the command.
  • If you see a different version than 3.14.0, run pip install mlflow==3.14.0 to pin the exact version.

Still stuck? Help me fix my MLflow import error

Now it's time to start the MLflow tracking server. This server stores your experiment data in a local SQLite database and serves the visual dashboard.

  • Start the MLflow server by running:
mlflow server --port 5000

What does this command do?

  • mlflow server launches the tracking server and web UI.
  • --port 5000 tells it to listen on port 5000.
  • The server automatically creates a mlflow.db SQLite file in your project folder to store all experiment data.

This command takes over your terminal

The server runs in the foreground and occupies this terminal session. You will not be able to type other commands here while it is running.

  • To open a second terminal for future commands, click the + icon in the terminal panel, or press Ctrl+Shift+` (Windows) or Cmd+Shift+` (macOS).
  • To stop the server later, press Ctrl+C in the terminal where it is running.

You should see server startup output including a line similar to Uvicorn running on http://127.0.0.1:5000.

  • Open your web browser and navigate to http://localhost:5000.

You should see the MLflow dashboard with an empty experiments list. This is the UI where all your tracked experiments will appear in the next steps.

Seeing a blank page or connection refused?

  • Make sure the server is still running in your first terminal tab. If you accidentally closed it, reactivate your venv and run mlflow server --port 5000 again.
  • Check that you are visiting http://localhost:5000 (not https).
  • If port 5000 is already in use by another application, stop that application first or choose a different port.

Still stuck? Help me get the MLflow server running

Your ML environment is locked and loaded. Next up, you'll train your first machine learning model and experience the exact problem that MLflow exists to solve.

Train a Model Without Tracking

Your ML environment is ready with scikit-learn and MLflow installed. Now it's time to train your first model.

Machine learning is an iterative process. You train a model, check how it performs, tweak settings called hyperparameters, and train again. But what happens to your previous results each time you retrain?

In this step, get ready to:
  • Load the California Housing dataset and explore its structure.
  • Train a RandomForestRegressor and evaluate its performance.
  • Change hyperparameters and discover that previous results are lost.
Load the California Housing dataset

The California Housing dataset contains real census data for 20,640 neighborhoods in California. Each row has 8 features (like median income, average rooms, and location) and a target value: the median house price for that neighborhood.

  • In VS Code, create a new file called train_no_tracking.py in your project folder.
  • Add the following code to load the dataset and print its structure:
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
import pandas as pd

# Load the California Housing dataset
housing = fetch_california_housing()
X = pd.DataFrame(housing.data, columns=housing.feature_names)
y = housing.target

print(f"Dataset shape: {X.shape}")
print(f"Features: {housing.feature_names}")
print(f"Target: Median house value (in $100,000s)")

What does this code do?

The imports bring in everything you need: the dataset loader, train/test splitting, the model, evaluation metrics, and pandas for data handling.

fetch_california_housing() downloads the dataset directly from scikit-learn's built-in collection. Wrapping the data in a DataFrame gives it column names, making it easier to work with.

  • Save train_no_tracking.py.
  • Run the script in your terminal (the one that has your virtual environment active) by running:
python train_no_tracking.py

You should see output showing the dataset has 20,640 rows, 8 features, and a target measured in $100,000s.

Seeing a ModuleNotFoundError?

Make sure your virtual environment is active. Your terminal prompt should show (venv) at the beginning. If it doesn't, activate it again with .\venv\Scripts\activate (Windows) or source venv/bin/activate (macOS/Linux).

Help me troubleshoot this ModuleNotFoundError

Train a RandomForestRegressor and evaluate its performance

A Random Forest is an ensemble of decision trees that vote together to make predictions. Its behavior changes based on hyperparameters like how many trees to use and how deep each tree can grow.

  • Add the following code to the bottom of train_no_tracking.py, below the print statements:
# Split into training and test sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Train a RandomForestRegressor
model = RandomForestRegressor(n_estimators=50, max_depth=10, random_state=42)
model.fit(X_train, y_train)

# Evaluate the model
y_pred = model.predict(X_test)
r2 = r2_score(y_test, y_pred)
mse = mean_squared_error(y_test, y_pred)

print(f"\nResults:")
print(f"R-squared: {r2:.4f}")
print(f"Mean Squared Error: {mse:.4f}")

What does this code do?

train_test_split holds back 20% of the data for testing, so you can measure performance on data the model has never seen.

The model is configured with n_estimators=50 (50 trees) and max_depth=10 (each tree can be 10 levels deep). These are the hyperparameters you will change later.

R-squared measures how well predictions match reality (1.0 = perfect). Mean Squared Error measures the average squared difference between predictions and actual values (lower = better).

  • Save train_no_tracking.py and run it again:
python train_no_tracking.py

You should see the dataset info followed by an R-squared score and MSE value. Take a mental note of that R-squared number.

Seeing a NameError?

Make sure you added the new code below the existing print statements, not above them. The training code uses X and y which are defined in the first section of the file.

Help me fix this NameError

✔️ Awesome, I've got everything!

Great. Make sure you saved your file.

ⓧ I'd like to double check the full code

Scroll up and confirm your file has both chunks in order:

  • Chunk 1: imports, dataset loading, and three print statements showing dataset shape, features, and target.
  • Chunk 2: train_test_split, model training with n_estimators=50 and max_depth=10, evaluation, and results printing.
Change hyperparameters and observe what happens

You just got a score from your model. Now imagine you want to try a different configuration to see if you can do better. This is the core loop of machine learning: try, measure, compare, repeat.

  • In train_no_tracking.py, find this line:
model = RandomForestRegressor(n_estimators=50, max_depth=10, random_state=42)
  • Change n_estimators=50 to n_estimators=100 and max_depth=10 to max_depth=20, so the line looks like this:
model = RandomForestRegressor(n_estimators=100, max_depth=20, random_state=42)
  • Save the file.

Before you run this, think: what was your R-squared score from the last run? Can you remember the exact number?

  • Run the script again:
python train_no_tracking.py

You see a new R-squared score and MSE. But the previous results from your first run are gone from your terminal.

Why is this a problem?

You just lost your first experiment. You might remember roughly what the R-squared was, but you cannot precisely compare the two runs.

In real ML work, you might try dozens or hundreds of configurations across days or weeks. Without a tracking system, you would need to manually copy scores into a spreadsheet, hope you recorded the exact parameters, and trust that nothing got mixed up. This is the problem MLflow solves.

  • Now set the values back to the originals. Change n_estimators=100 to n_estimators=50 and max_depth=20 to max_depth=10:
model = RandomForestRegressor(n_estimators=50, max_depth=10, random_state=42)
  • Save train_no_tracking.py.

Why reset the values?

This keeps your baseline configuration in the file for reference. In the next step, you will build a new script that uses MLflow to track every configuration automatically, so you never lose a result again.

You just felt the core pain of untracked experiments: every re-run overwrites your previous results. In the next step, you will add MLflow tracking so every run is automatically saved and comparable.

Add MLflow Experiment Tracking

In the previous step, you trained a model twice with different hyperparameters and lost the first result forever. That is the exact problem MLflow solves.

MLflow automatically captures every parameter, metric, and model artifact each time you train. In this step, you will add MLflow tracking to your training script so every run is permanently recorded and browsable in a dashboard.

In this step, get ready to:
  • Add MLflow autologging to a new training script.
  • Run the script and observe the automatic logging output.
  • Open the MLflow UI and explore your logged run with parameters, metrics, and a model artifact.
Add MLflow autologging to your training script

MLflow's autolog feature watches your scikit-learn training calls and automatically records every parameter, metric, and model artifact. You will create a new file that wraps the same training logic inside MLflow's tracking context.

  • Create a new file called train_with_tracking.py in your project folder and add the following imports and MLflow setup:
import mlflow
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
import pandas as pd

# Enable MLflow autologging for scikit-learn
mlflow.sklearn.autolog()

# Set the experiment name
mlflow.set_experiment("California Housing")

# Load the California Housing dataset
housing = fetch_california_housing()
X = pd.DataFrame(housing.data, columns=housing.feature_names)
y = housing.target

# Split into training and test sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

What does this code do?

  • mlflow.sklearn.autolog() tells MLflow to watch every scikit-learn model you train and automatically log its parameters, metrics, and the fitted model artifact.
  • mlflow.set_experiment("California Housing") groups all your runs under a named experiment so they are easy to find in the UI.
  • The data loading and splitting is identical to your previous script. The difference is what happens around the training step next.
  • Below the train/test split, add the MLflow-tracked training block:
# Train with MLflow tracking
with mlflow.start_run(run_name="random_forest_baseline"):
    model = RandomForestRegressor(n_estimators=50, max_depth=10, random_state=42)
    model.fit(X_train, y_train)

    # Evaluate and log custom metrics
    y_pred = model.predict(X_test)
    r2 = r2_score(y_test, y_pred)
    mse = mean_squared_error(y_test, y_pred)
    mlflow.log_metric("test_r2_score", r2)
    mlflow.log_metric("test_mse", mse)

    print(f"Run logged successfully!")
    print(f"R-squared: {r2:.4f}")
    print(f"Mean Squared Error: {mse:.4f}")

What does this code do?

  • with mlflow.start_run(run_name="random_forest_baseline") creates a named tracking context. Everything inside this block gets recorded as one experiment run.
  • Autologging captures n_estimators, max_depth, and random_state automatically. You do not need to log them manually.
  • mlflow.log_metric() adds custom metrics beyond what autolog captures. Here you log the test set R-squared and MSE so you can compare models on held-out data later.

Why log custom metrics on top of autologging?

Autologging records training metrics (how the model scores on its own training data). These can be misleading because a model can memorize training data and still fail on new inputs.

By explicitly logging test set metrics (test_r2_score and test_mse), you measure how the model performs on data it has never seen. This is what matters for real-world use.

  • Save train_with_tracking.py.

✔️ Awesome, I've got everything!

Great. Double check you have saved the file before moving on.

ⓧ I'd like to double check the full code

Your train_with_tracking.py should contain both sections you just wrote, in this order:

  • The imports and MLflow setup chunk (imports, mlflow.sklearn.autolog(), mlflow.set_experiment(), data loading, and the train/test split).
  • The training block chunk (with mlflow.start_run(...), model creation, evaluation, custom metric logging, and print statements).

Scroll up and confirm each chunk is present in your file, in order.

Run the script and observe the automatic logging output

Your script is ready. Time to see MLflow in action.

  • Before you run this, think: what do you expect MLflow to capture? Will the terminal output look any different from your untracked script?
  • Run the tracked training script by entering this command in your terminal:
python train_with_tracking.py

You should see output similar to this:

Run logged successfully!
R-squared: 0.8066
Mean Squared Error: 0.2531

The terminal output looks almost identical to before. The magic happened silently in the background. MLflow recorded every parameter, metric, and the trained model artifact to its database.

Seeing a connection error?

If you see an error about failing to connect to the tracking server, your MLflow server from earlier may have stopped. Switch back to the terminal where you started it and run mlflow server --port 5000 again. Then re-run your training script in the other terminal.

Help me fix my MLflow tracking server connection error.

Open the MLflow UI and explore the logged run

The real payoff is seeing everything captured in a visual dashboard. Your MLflow server is already serving the UI at the same address.

  • Open your browser and go to http://localhost:5000.
  • In the Experiments list, click the California Housing experiment.
  • Click the run named random_forest_baseline to open its detail page.

On the run detail page, you will see three sections that prove everything was captured:

  • Parameters shows n_estimators, max_depth, random_state, and more. These were captured automatically by autologging.
  • Metrics shows test_r2_score and test_mse (your custom metrics) alongside autologged training metrics.
  • Artifacts contains the serialized model file that MLflow saved. You could load this model later without retraining.

Why is this better than printing to the terminal?

When you ran the untracked script twice, the first result vanished the moment the second run finished. Now every run is permanently stored with a timestamp, its exact parameters, and the model itself.

Next time you change a hyperparameter, the old result stays right there in the dashboard for comparison.

Your first experiment is tracked and browsable. Next, you will train multiple models with different hyperparameters and use the MLflow comparison view to find the best one.

Compare Multiple Experiment Runs

In the last step, you tracked a single training run and saw how MLflow captures parameters, metrics, and artifacts automatically. Now you have proof that tracking works.

The real power of experiment tracking is comparing runs side by side to find the best hyperparameters. Right now you have one run in your experiment. If you changed hyperparameters and trained again, you could compare them manually, but that does not scale. In this step, you will train four models with different configurations in one script and use the MLflow UI to visually compare them.

In this step, get ready to:
  • Create a script that trains four models with different hyperparameter configurations.
  • Run the script and observe multiple runs logged in the MLflow UI.
  • Use the Chart view to visually compare runs and identify the best model.
Create a comparison script with multiple configurations

Instead of editing and rerunning a script four times, you will define a list of hyperparameter configurations and loop through them. Each iteration starts a separate MLflow run, so every configuration is tracked independently.

  • Create a new file called compare_experiments.py in your project folder.
  • Add the imports, autologging setup, and data preparation by pasting this code:
import mlflow
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
import pandas as pd

# Enable MLflow autologging for scikit-learn
mlflow.sklearn.autolog()

# Set the experiment name
mlflow.set_experiment("California Housing")

# Load the California Housing dataset
housing = fetch_california_housing()
X = pd.DataFrame(housing.data, columns=housing.feature_names)
y = housing.target

# Split into training and test sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

What does this code do?

This is the same foundation you used in train_with_tracking.py. It enables autologging, sets the experiment name to "California Housing" (so new runs appear alongside your baseline), and prepares the train/test split.

  • Now add the hyperparameter configurations and the training loop below the code you just pasted:
# Define different hyperparameter configurations to compare
configs = [
    {"n_estimators": 50, "max_depth": 5, "min_samples_split": 2},
    {"n_estimators": 100, "max_depth": 10, "min_samples_split": 2},
    {"n_estimators": 100, "max_depth": 15, "min_samples_split": 5},
    {"n_estimators": 200, "max_depth": 20, "min_samples_split": 10},
]

for i, params in enumerate(configs):
    with mlflow.start_run(run_name=f"rf_config_{i+1}"):
        model = RandomForestRegressor(random_state=42, **params)
        model.fit(X_train, y_train)

        y_pred = model.predict(X_test)
        r2 = r2_score(y_test, y_pred)
        mse = mean_squared_error(y_test, y_pred)
        mlflow.log_metric("test_r2_score", r2)
        mlflow.log_metric("test_mse", mse)

        print(f"Config {i+1}: n_estimators={params['n_estimators']}, "
              f"max_depth={params['max_depth']} -> R2={r2:.4f}, MSE={mse:.4f}")

print("\nAll runs logged! Open http://localhost:5000 to compare.")

What does this code do?

The configs list holds four dictionaries, each with a different combination of n_estimators, max_depth, and min_samples_split.

The for loop iterates through each configuration. Inside each iteration, mlflow.start_run() opens a new tracked run, the model trains with those specific parameters, and custom metrics are logged. Each run gets a unique name like rf_config_1 so you can identify them in the UI.

  • Save compare_experiments.py.

✔️ Awesome, I've got everything!

Great. Double check you have saved compare_experiments.py and you are ready to run it.

ⓧ I'd like to double check the full code

Scroll up and confirm your file has both chunks in order:

  • Chunk 1: Imports, autologging, experiment setup, data loading, and train/test split.
  • Chunk 2: The configs list, the for loop with mlflow.start_run(), and the final print statement.
Run the script and observe multiple runs in the MLflow UI

Before you run this, what do you think will happen? You have four configurations, so predict how many new runs will appear in the MLflow UI.

  • Run the comparison script by executing this command in your terminal:
python compare_experiments.py

You should see four lines of output, one per configuration, each showing the hyperparameters used and the resulting R2 and MSE scores. The final line confirms all runs were logged.

Seeing a connection error?

Make sure your MLflow tracking server is still running in a separate terminal. If it stopped, restart it with mlflow server --port 5000 and run the script again.

help me fix a connection error when running my MLflow comparison script

  • Switch to your browser and refresh the MLflow UI at http://localhost:5000.
  • Click on the California Housing experiment in the left sidebar.

You should now see five runs in the table: your original random_forest_baseline from the last step, plus rf_config_1 through rf_config_4.

Why do all five runs appear together?

All runs share the same experiment name: "California Housing". MLflow groups runs by experiment, so everything you log under that name appears in one place. This is how real ML teams organize work. One experiment per problem, many runs per experiment.

Use the comparison view to identify the best model

The Table view shows all runs, but you need to compare them visually to spot which configuration performed best.

  • In the Table view, click the Columns button above the runs table.
  • Add the test_r2_score metric column and the test_mse metric column so you can see scores for each run.
  • Click the test_r2_score column header to sort runs by R2 score descending.

The run with the highest R2 score is your best-performing configuration. You can see exactly which hyperparameters produced that score without scrolling through terminal history.

  • Switch to Chart view using the view toggle above the runs table.
  • Select the Parallel coordinates chart type.
  • Configure it to show n_estimators, max_depth, and test_mse as coordinates.

Each line in the chart represents one run. Lines shading toward lower MSE values indicate better-performing configurations. You can hover over individual lines to see each run's details.

What am I looking at?

The Parallel coordinates chart draws one vertical axis per parameter or metric. Each run becomes a line connecting its values across axes. This makes it easy to spot patterns: for example, you can see that deeper trees with more estimators tend to produce lower MSE.

This is exactly what was impossible when you ran train_no_tracking.py in Step 2. Back then, you lost each result the moment you changed hyperparameters. Now every run is preserved and comparable.

Not seeing the Chart view option?

Make sure you are on the experiment page showing all runs, not inside a single run's detail page. Click California Housing in the sidebar to return to the full experiment view.

help me find the chart comparison view in the MLflow UI

You now have a visual experiment comparison system. Next up, you will take the best-performing model and register it in the MLflow Model Registry so it can be versioned and deployed.

Register Your Best Model

You have five tracked experiment runs and you know which hyperparameter combination produced the best R-squared score. But that model is still just an artifact buried inside a single run.

The Model Registry solves this by giving a model a name, a version number, and an alias so anyone on a team can load it reliably without hunting through experiment runs. In this step, you will register your best model, tag it as the "champion", and prove you can load it for inference in a single line of code.

In this step, get ready to:
  • Register the best-performing model into the MLflow Model Registry.
  • Assign the "champion" alias to the registered model version.
  • Load the model using its registry alias and make predictions on new data.
Register the best-performing model into the Model Registry

From your comparison in Step 4, the configuration with n_estimators=100, max_depth=15, and min_samples_split=5 delivered the strongest results. You will train this configuration one more time and register the resulting model in a single step using mlflow.sklearn.log_model().

  • Create a new file called register_model.py in your project folder.
  • Add the imports and data setup by pasting this code:
import mlflow
from mlflow import MlflowClient
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
import pandas as pd

# Enable MLflow autologging for scikit-learn
mlflow.sklearn.autolog()

# Set the experiment name
mlflow.set_experiment("California Housing")

# Load the California Housing dataset
housing = fetch_california_housing()
X = pd.DataFrame(housing.data, columns=housing.feature_names)
y = housing.target

# Split into training and test sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

What does this code do?

  • from mlflow import MlflowClient imports the client class you will use later to assign an alias to the registered model.
  • The rest is the same data setup pattern from your earlier scripts: enable autologging, set the experiment name, load and split the California Housing data.
  • Save register_model.py.
  • Confirm the file appears in the VS Code Explorer sidebar.
  • Now add the training and registration block below the data split by pasting this code at the bottom of the file:
# Train the best model configuration and register it
with mlflow.start_run(run_name="best_model_for_registry"):
    model = RandomForestRegressor(
        n_estimators=100, max_depth=15, min_samples_split=5, random_state=42
    )
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)
    r2 = r2_score(y_test, y_pred)
    mse = mean_squared_error(y_test, y_pred)
    mlflow.log_metric("test_r2_score", r2)
    mlflow.log_metric("test_mse", mse)

    # Log and register the model
    model_info = mlflow.sklearn.log_model(
        sk_model=model,
        name="CaliforniaHousingModel",
        registered_model_name="CaliforniaHousingModel",
    )

    print(f"Model registered! R2={r2:.4f}, MSE={mse:.4f}")
    print(f"Model URI: {model_info.model_uri}")

What does this code do?

  • mlflow.start_run(run_name="best_model_for_registry") opens a new tracked run using the winning hyperparameters.
  • mlflow.sklearn.log_model() saves the trained model artifact to the run. The registered_model_name parameter simultaneously creates an entry in the Model Registry and assigns it version 1.
  • The name parameter sets the artifact path within the run, while registered_model_name creates the registry entry.
  • Save register_model.py.
  • Run the script in your terminal:
python register_model.py

You should see output showing the R-squared score, MSE, and a line reading Model registered! followed by the model URI.

Seeing a connection error?

  • Make sure your MLflow tracking server is still running in its own terminal at http://localhost:5000. If you closed it, restart with mlflow server --port 5000.
  • Confirm your virtual environment is activated in the terminal where you run the script.

Help me debug my MLflow model registration script.

Assign the "champion" alias to the registered model version

Your model now lives in the registry as CaliforniaHousingModel version 1. An alias like "champion" acts as a stable pointer that downstream systems use to load the currently best model, regardless of which version number it is.

  • Add this code below the with mlflow.start_run(...) block (outside and after it) in register_model.py:
# Assign the "champion" alias to version 1
client = MlflowClient()
client.set_registered_model_alias(
    name="CaliforniaHousingModel", alias="champion", version="1"
)
print("Alias 'champion' assigned to version 1")

What does this code do?

  • MlflowClient() creates a client that talks directly to the MLflow tracking server.
  • set_registered_model_alias() attaches the label "champion" to version 1. Anyone can now load this model using the alias without knowing the version number.
Load the model using its registry alias and make predictions

The whole point of the registry is that consumers never hard-code a run ID or artifact path. They load by name and alias.

  • Add this code at the very bottom of register_model.py:
# Load the model using the alias and make a prediction
loaded_model = mlflow.pyfunc.load_model("models:/CaliforniaHousingModel@champion")
sample = X_test.iloc[:5]
predictions = loaded_model.predict(sample)
print(f"\nSample predictions (in $100,000s): {predictions}")
print(f"Actual values: {y_test[:5]}")

What does this code do?

  • mlflow.pyfunc.load_model("models:/CaliforniaHousingModel@champion") fetches the model artifact using the registry alias. No run IDs, no file paths.
  • The loaded model exposes a .predict() method that works on any pandas DataFrame with the right columns.

Before you run this, what do you think the predictions will represent?

  • Save register_model.py and run it:
python register_model.py

You should see the model register, the alias assigned, and then sample predictions printed in units of $100,000s alongside the actual values.

What should I see?

The output includes registration confirmation, then Alias 'champion' assigned to version 1, followed by an array of predicted house values and the corresponding actual values. The predictions and actuals should be close, confirming the model works.

Now verify that the model is permanently loadable from the registry outside of this script. What do you expect to see when you load the model in a one-liner?

  • Run this command in your terminal:
python -c "import mlflow; model = mlflow.pyfunc.load_model('models:/CaliforniaHousingModel@champion'); print(model)"

You should see the model object printed, confirming it loads successfully from the registry using only its name and alias.

Seeing a model not found error?

  • Confirm you ran register_model.py successfully before this command. The model must be registered before it can be loaded.
  • Check that your MLflow tracking server is still running at http://localhost:5000.
  • Verify you spelled CaliforniaHousingModel exactly as shown, including capitalization.

Help me debug loading a model from the MLflow Model Registry.

✔️ Awesome, I've got everything!

Great. Double check that you have saved register_model.py and confirmed the model loads from the registry.

ⓧ I'd like to double check the full code

Your register_model.py should contain these sections in order. Scroll up and confirm each is present:

  • Imports (mlflow, MlflowClient, sklearn, pandas).
  • Autologging and experiment setup (mlflow.sklearn.autolog() and mlflow.set_experiment("California Housing")).
  • Data loading and train/test split.
  • Training block inside with mlflow.start_run(run_name="best_model_for_registry"): that trains, evaluates, and registers the model with registered_model_name="CaliforniaHousingModel".
  • Alias assignment using client.set_registered_model_alias() outside the with block.
  • Model loading with mlflow.pyfunc.load_model("models:/CaliforniaHousingModel@champion") and prediction printing.

Your model is now registered, versioned, and aliased. You have a complete ML lifecycle: train, track, compare, and deploy. The Secret Mission takes this further by automating the hyperparameter search itself.

Secret mission

Automated Hyperparameter Search with Nested Runs

In the main project, you manually chose hyperparameter combinations one by one. What if you could search across dozens of combinations automatically and have MLflow track every single one? In this secret mission, you will use scikit-learn's GridSearchCV to automate that search and observe how MLflow creates a nested run structure to organize the results.

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely locally, so there are no ongoing costs.

Resources you used:

  • MLflow tracking server process (running at http://localhost:5000).
  • MLflow tracking data stored in mlflow.db (SQLite database) and an mlartifacts directory.
  • Python virtual environment with mlflow, scikit-learn, and pandas packages.
  • Python scripts: train_no_tracking.py, train_with_tracking.py, compare_experiments.py, register_model.py, and grid_search_tracking.py.

Keep everything running

No action needed. Choose this if you are still actively building or want to explore MLflow further.

  • Your MLflow data persists in the local mlflow.db file and mlartifacts directory. No storage costs apply.
  • If you closed the terminal running the MLflow server, restart it anytime by activating your virtual environment and running:
mlflow server --port 5000
  • Your registered model and all experiment runs will still be there when you come back.

Pause - I'll come back to this later

Shut down running processes to free up your terminal. Your files and data stay on disk for next time.

  • Switch to the terminal running the MLflow server and press Ctrl+C to stop it.
  • Deactivate your virtual environment by running:
deactivate
  • Your project folder, mlflow.db, and all scripts remain on your Desktop. Resume anytime by activating the virtual environment and restarting the server.

Delete - I don't want to use this again

Remove all project resources and start fresh. This permanently deletes your experiment data, trained models, and scripts.

  • Switch to the terminal running the MLflow server and press Ctrl+C to stop it.
  • If your virtual environment is still active, deactivate it by running:
deactivate
  • Navigate to your Desktop and delete the entire project folder. On macOS/Linux, run:
cd ~/Desktop
rm -rf mlflow-tracking
  • On Windows, run:
cd ~\Desktop
rmdir /s /q mlflow-tracking
  • This removes all Python scripts, the virtual environment, mlflow.db, the mlartifacts directory, and any other generated files.

Nice Work!

Nice Work!

Nice work! You just built a complete ML experiment tracking system from scratch. You trained a model, tracked every experiment run, compared results visually, and registered your best model for production use.

You've learned how to:

  • Track ML experiments with MLflow autologging, capturing parameters, metrics, and model artifacts automatically every time you train.
  • Compare multiple experiment runs visually in the MLflow UI to identify the best-performing hyperparameter configuration.
  • Register a versioned model in the MLflow Model Registry with a champion alias, then load it for inference in a single line of code.
  • Secret Mission: Used GridSearchCV with MLflow to automate hyperparameter search across a parameter grid, producing nested parent and child runs in the UI.

Ready to quiz yourself?