Build a Local ML Prediction API

Train, track, and serve an ML model as a REST API using MLflow and FastAPI.

Introduction

30 Second Summary

Every time you ask a voice assistant to identify a song or an app recommends what to watch next, a machine learning model is running behind an API, making predictions in real time.

In this project, you will build a complete local MLOps pipeline that trains a machine learning model, tracks your experiments, versions your best model, and serves predictions through a live REST API.

What You'll Build

You will train a classification model on the Iris dataset, track experiments with MLflow, register your best model in a model registry, and serve predictions through a FastAPI REST endpoint.

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

A trained scikit-learn model with logged parameters, metrics, and artifacts tracked in an MLflow experiment dashboard.
A versioned model stored in the MLflow Model Registry, ready for deployment.
A working FastAPI prediction API with a /predict endpoint and auto-generated interactive docs.
A production-ready API with a health check endpoint and input validation (Secret Mission).

Are there any prerequisites?

You just need basic Python knowledge (variables, functions, loops). No machine learning experience is required. This project keeps the ML simple so you can focus on the MLOps skills that employers are looking for.

Set Up Your MLOps Environment

Before you can train models and serve predictions, you need the right tools installed. This project uses four Python libraries that together cover the full MLOps lifecycle: training, tracking, versioning, and serving.

In This Step, Get Ready To:

Install all four Python libraries needed for this project.
Verify each library is installed correctly.
Create the project folder and file structure.

Install Dependencies

You need four libraries for this project. Each one handles a different part of the MLOps pipeline.

Open your terminal and run the following command to install all four packages at once.

pip install mlflow scikit-learn fastapi uvicorn

What do these libraries do?

mlflow is the experiment tracking and model registry tool. It logs your training runs and stores your models.
scikit-learn provides ML algorithms and built-in datasets. It is the simplest way to train a model in Python.
fastapi is the web framework you will use to build the prediction API. It auto-generates interactive documentation for your endpoints.
uvicorn is the ASGI server that runs your FastAPI application.

The installation may take a minute or two since MLflow has several sub-dependencies.

Verify Your Installations

A quick version check for each library confirms everything installed correctly. If any command fails, you will know exactly which package needs attention.

Run the following command to verify MLflow is installed.

python -c "import mlflow; print(mlflow.__version__)"

Run this command to verify scikit-learn is installed.

python -c "import sklearn; print(sklearn.__version__)"

Run this command to verify FastAPI is installed.

python -c "import fastapi; print(fastapi.__version__)"

I see version numbers

All three libraries are installed and ready to go.

I see an error

If you see a ModuleNotFoundError, the package did not install correctly.

Make sure you are using the same Python environment where you ran pip install.
Try running pip install mlflow scikit-learn fastapi uvicorn again.
If you have multiple Python versions installed, try using pip3 instead of pip.

Still seeing errors?

Paste the full error message and ask for help troubleshooting your installation.

Create the Project Folder Structure

Your project needs two Python files. One will handle model training and the other will run the prediction API. Setting up this structure now keeps your code organized from the start.

Create a new directory called mlops-prediction-api and navigate into it.

macOS

mkdir mlops-prediction-api
cd mlops-prediction-api

Windows

mkdir mlops-prediction-api
cd mlops-prediction-api

Create two empty Python files inside the project directory.

macOS

touch train.py app.py

Windows

New-Item train.py, app.py

What are these files for?

train.py is where you will write the model training script with MLflow experiment tracking. app.py is where you will build the FastAPI prediction API that serves your trained model.

Your folder structure should now look like this:

mlops-prediction-api/
  train.py
  app.py

Your environment is ready and your project is structured. Next up, you will write the training script and watch MLflow automatically track everything about your model training.

Train a Model and Track Experiments with MLflow

Your environment is set up and ready to go. Now it's time to train your first machine learning model and see MLflow in action.

Training a model is only useful if you can remember exactly how you trained it. What settings did you use? How accurate was it? Where is the saved model file? In real ML teams, losing track of this information is one of the biggest sources of wasted time. MLflow solves this by automatically logging every detail of your training run so you can reproduce and compare experiments later.

In This Step, Get Ready To:

Write a training script that trains a Logistic Regression model on the Iris dataset.
Use MLflow autologging to capture parameters, metrics, and model artifacts automatically.
Explore the MLflow UI to see everything that was logged about your training run.

Write the Training Script With MLflow Autologging

Your training script needs to do three things: load a dataset, train a model, and log everything to MLflow. The Iris dataset is a classic beginner ML dataset built into scikit-learn. It contains measurements of 150 flowers across three species: setosa, versicolor, and virginica. Each flower has four features: sepal length, sepal width, petal length, and petal width.

Open train.py in your code editor and paste in the following code:

import mlflow
import mlflow.sklearn
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

# Load the Iris dataset and split into training and test sets
X, y = datasets.load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Tell MLflow which experiment to log this run under
mlflow.set_experiment("iris-classification")

# Enable automatic logging for scikit-learn models
mlflow.sklearn.autolog()

# Define the hyperparameters for the model
params = {"solver": "lbfgs", "max_iter": 1000, "random_state": 8888}

# Start an MLflow run to track everything
with mlflow.start_run():
    # Train the model
    lr = LogisticRegression(**params)
    lr.fit(X_train, y_train)

    # Evaluate on the test set
    y_pred = lr.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)

    # Manually log the test accuracy metric
    mlflow.log_metric("test_accuracy", accuracy)

    # Save the trained model as an artifact
    mlflow.sklearn.log_model(sk_model=lr, name="iris_model")

    # Add a descriptive tag to this run
    mlflow.set_tag("Training Info", "Basic LR model for iris data")

    print(f"Test accuracy: {accuracy:.4f}")

What does this script do?

This script combines model training with experiment tracking. Here is what each section does:

load_iris(return_X_y=True) loads the Iris dataset directly from scikit-learn. No file downloads needed.
train_test_split splits the data into 80% for training and 20% for testing so you can measure how well the model performs on data it has never seen.
mlflow.sklearn.autolog() tells MLflow to automatically capture the model, hyperparameters, and training metrics without you writing extra logging code.
The with mlflow.start_run(): block groups everything into a single tracked experiment run.

Why both autolog and manual logging?

Autolog captures training metrics automatically, but it does not know about your test set evaluation. That is why you also call mlflow.log_metric() manually for the test accuracy. In real projects, you will often combine both approaches to get complete coverage.

Run the Training Script

With the script written, it is time to train your model. MLflow will create a local mlruns directory in your project folder to store all the experiment data.

Open your terminal and make sure you are inside your project directory.
Run the training script:

python train.py

You should see output that includes your test accuracy score and a run ID from MLflow confirming the experiment was logged.

I see a test accuracy score

Your model trained successfully and MLflow captured everything. The accuracy should be around 1.0 (100%) since the Iris dataset is a straightforward classification problem for Logistic Regression. You will also notice a new mlruns/ folder appeared in your project directory. This is where MLflow stores all your experiment data locally.

I see an error

Check these common issues:

If you see ModuleNotFoundError, one of your packages did not install correctly. Re-run the install command from Step 1.
Make sure you are running the command from inside your project directory where train.py is located.
If you see a syntax error, double-check that you copied the entire script correctly.

Still stuck? Describe the full error message you see when running python train.py.

Explore the MLflow UI

MLflow comes with a built-in web dashboard where you can visually explore everything that was logged during your training run. You need to start the MLflow tracking server to access it.

Open a separate terminal window. Keep your original terminal available for later.
Navigate to your project directory in this new terminal.
Start the MLflow tracking server:

mlflow server --port 5000

What does this command do?

This starts a local web server on port 5000 that reads the experiment data from your mlruns/ directory and displays it as an interactive dashboard. This terminal will stay busy running the server, which is why you need a separate terminal.

Open http://localhost:5000 in your browser.
Click on the iris-classification experiment in the left sidebar.
Click on your training run to open its details.

You should see a detailed view of your training run. Explore the different sections:

Parameters shows the hyperparameters you used (solver, max_iter, random_state).
Metrics shows your test_accuracy along with any metrics that autolog captured (like training_accuracy).
Artifacts contains the saved model file that MLflow stored for you.

What is experiment tracking?

Experiment tracking is the practice of recording every detail about a training run: what data you used, what settings you chose, how well the model performed, and where the model file is saved. Without tracking, ML teams end up with dozens of model files named things like model_final_v2_REAL.pkl and no way to know which one is actually the best.

You just trained a model and tracked the entire experiment with MLflow. Every parameter, metric, and artifact is saved and browsable. Next up, you will run more experiments with different settings and use the MLflow UI to compare them side by side.

Experiment with Hyperparameters and Compare Runs

In the previous step, you trained a model and saw how MLflow captures every detail of a training run automatically. But training one model isn't enough. Real ML work involves experimenting with different settings to find what works best.

In this step, you'll run multiple experiments with different hyperparameters, compare them side by side in the MLflow UI, and register the best model in the MLflow Model Registry so it's versioned and ready for deployment.

In This Step, Get Ready To:

Run multiple training experiments with different hyperparameter configurations.
Compare runs side by side in the MLflow UI to find the best model.
Register the best model in the MLflow Model Registry for versioning.

Modify Hyperparameters and Retrain

Hyperparameters are the settings you choose before training a model. Different combinations can produce different accuracy scores. The LogisticRegression model in scikit-learn supports several solvers, each using a different optimization algorithm to find the best model weights.

Open train.py in your code editor.
Change the hyperparameters dictionary to use a different solver. For your second run, try saga:

params = {"solver": "saga", "max_iter": 1000, "random_state": 8888}

Save the file and run the training script again.

python train.py

MLflow creates a brand new run entry for this experiment. Your first run is still there, untouched.

Now change the hyperparameters one more time for a third run. Try newton-cg as the solver with a lower max_iter:

params = {"solver": "newton-cg", "max_iter": 200, "random_state": 8888}

Save the file and run python train.py one more time.

Why are we keeping random_state the same?

The random_state parameter controls randomness in the training process. By keeping it at 8888 across all runs, you ensure that any difference in accuracy comes from the solver or max_iter change, not from random variation. This is how you run reproducible experiments.

You now have three runs logged in MLflow, each with a different solver configuration. Time to compare them.

Compare Runs in the MLflow UI

The MLflow UI lets you view all your experiment runs in one place and compare their parameters and metrics side by side. This is the core value of experiment tracking. Instead of keeping notes in a spreadsheet, MLflow captures everything automatically.

Make sure your MLflow tracking server is still running in a separate terminal. If it's not, start it again:

mlflow server --port 5000

Open http://localhost:5000 in your browser.
Click on the iris-classification experiment in the left sidebar.

You should see three runs listed, each with different parameter values logged.

Select all three runs using the checkboxes next to each run.
Click the Compare button to see a side-by-side comparison.
Look at the test_accuracy metric for each run. Identify which solver produced the highest accuracy.

What should I see?

The comparison view shows a table with each run as a column. You can see the parameters (solver, max_iter, random_state) and metrics (test_accuracy, training_accuracy, and others captured by autolog) lined up for easy comparison. One of your runs will have the highest test_accuracy. That's the model you want to deploy.

Why This Approach Instead of a Spreadsheet?

You might be tempted to track experiments manually in a spreadsheet or notebook. MLflow captures parameters, metrics, and the actual model artifact together in one place. If you need to reproduce a result from three weeks ago, you can load the exact model and see exactly what settings produced it. A spreadsheet can't do that.

Register the Best Model

The MLflow Model Registry is where you promote a model from "just an experiment" to "ready for deployment." Registering a model gives it a name and version number, making it easy to load in your API later.

In the MLflow UI, click on the run with the best test_accuracy.
Scroll down to the Artifacts section and find the logged model.
Click Register Model.
Create a new registered model and name it iris-classifier.

What is the Model Registry?

Think of the Model Registry as a catalog of your production-ready models. Each registered model gets a version number. When you update and re-register a model later, it becomes version 2, version 3, and so on. Your API can always load the latest version without you changing any code.

Pro tip: register programmatically

You can also register models from your training script by adding this line after the model is logged:

mlflow.register_model(model_uri, "iris-classifier")

The model_uri is the URI from the run, which you can get from the MLflow UI's artifact details. This is useful when you want to automate model registration in a CI/CD pipeline.

Don't see the Register Model button?

Make sure you started the MLflow server with mlflow server --port 5000 and are accessing the UI at http://localhost:5000. The Register Model button appears in the Artifacts section of an individual run page, not on the experiment overview. Help me find the Register Model button in the MLflow UI.

Your best model is now registered, versioned, and ready to be loaded by an API. In the next step, you'll build a FastAPI prediction endpoint that pulls this model directly from the registry and serves real-time predictions.

Build a Prediction API with FastAPI

You have a trained model registered in MLflow. That is great for tracking experiments, but a model sitting in a registry does not help anyone make predictions.

To make your model useful, you need to wrap it in an API that other applications can call. In this step, you will build a REST API using FastAPI that loads your registered model and serves real-time predictions through a /predict endpoint.

In This Step, Get Ready To:

Write a FastAPI application that loads your model from the MLflow Model Registry.
Start the API server and test predictions through the auto-generated Swagger UI.

Write the FastAPI Application

FastAPI is a modern Python web framework that auto-generates interactive API docs and validates incoming data using Pydantic models. You will define the shape of the input data, load the model from MLflow, and create an endpoint that returns predictions.

Open app.py in your code editor.
Paste the following code into the file.

from fastapi import FastAPI
from pydantic import BaseModel
import mlflow

# Define the expected shape of incoming prediction requests
class IrisInput(BaseModel):
    sepal_length: float
    sepal_width: float
    petal_length: float
    petal_width: float

# Create the FastAPI app instance
app = FastAPI(title="Iris Prediction API")

# Connect to the MLflow tracking server and load the registered model
mlflow.set_tracking_uri("http://localhost:5000")
model = mlflow.pyfunc.load_model("models:/iris-classifier/latest")

# Map numeric predictions to species names
SPECIES = ["setosa", "versicolor", "virginica"]

@app.post("/predict")
def predict(input_data: IrisInput):
    # Convert the input into the format the model expects
    features = [[input_data.sepal_length, input_data.sepal_width,
                 input_data.petal_length, input_data.petal_width]]
    prediction = model.predict(features)
    species = SPECIES[int(prediction[0])]
    return {"prediction": int(prediction[0]), "species": species}

What does this code do?

The IrisInput class defines the four measurements the model needs to make a prediction. FastAPI uses this class to automatically validate incoming requests.

The mlflow.set_tracking_uri("http://localhost:5000") line tells MLflow where to find the tracking server. Then mlflow.pyfunc.load_model("models:/iris-classifier/latest") loads the latest version of your registered model.

The /predict endpoint takes the four input values, formats them as a list of lists (the shape scikit-learn expects), runs the prediction, and maps the numeric result to a species name.

Why FastAPI Instead of Flask?

FastAPI auto-generates interactive API documentation at the /docs endpoint and validates incoming data through Pydantic. Flask would require you to add these features manually with extra libraries.

FastAPI is the web framework used in production ML serving across the industry. Learning it here gives you a skill that transfers directly to real MLOps work.

Start the API Server

Your API needs two things running at the same time: the MLflow tracking server (so the model can be loaded) and the FastAPI server itself. You will need two separate terminal windows for this.

Confirm that the MLflow tracking server is still running on port 5000 in your first terminal. If it is not, start it again.

mlflow server --port 5000

Open a second terminal window and navigate to your project directory.
Run the following command to start the FastAPI server.

uvicorn app:app --reload --port 8000

What does this command do?

The uvicorn command starts the ASGI server that runs your FastAPI app. app:app tells uvicorn to look in the file app.py for the FastAPI instance named app. The --reload flag auto-restarts the server when you change code. The --port 8000 flag sets the port so it does not conflict with MLflow on port 5000.

Seeing a Connection Error When the Server Starts?

If you see an error about connecting to the MLflow tracking server, make sure the MLflow server is running in your other terminal window. The FastAPI app tries to load the model from MLflow at startup, so the tracking server must be available first.

If you see an error about the model not being found, make sure you registered your model as iris-classifier in the previous step. Help me debug my model loading error.

Test the Prediction API

FastAPI automatically generates interactive documentation for your API. You can use this to send test requests without writing any client code.

Open http://localhost:8000/docs in your browser.

You should see the Swagger UI with your POST /predict endpoint listed. This interactive documentation is one of the key advantages of FastAPI.

Click on the POST /predict endpoint to expand it.
Click Try it out.
Replace the request body with the following sample input and click Execute.

{
  "sepal_length": 5.1,
  "sepal_width": 3.5,
  "petal_length": 1.4,
  "petal_width": 0.2
}

What should I see?

The response should return a prediction of "setosa". These measurements describe a flower with a short, narrow petal, which is characteristic of the setosa species.

Now test with a different set of values to confirm the model can distinguish between species.

Click Try it out again and enter the following values.

{
  "sepal_length": 6.7,
  "sepal_width": 3.0,
  "petal_length": 5.2,
  "petal_width": 2.3
}

This time the response should predict "virginica". The longer petal measurements are a strong signal for the virginica species.

Your model is now live as a prediction API. Any application, whether it is a web app, a mobile app, or another service, can send a POST request to your /predict endpoint and get a species prediction back. This is exactly how ML models are served in production.

Secret mission

Make Your API Production-Ready

Your prediction API works, but it is not ready for the real world yet. Production APIs need health checks so monitoring tools can verify the service is alive. They also need input validation to reject bad data before it reaches your model. Let's add both to your API.

Before you go

Clean Up Your Resources

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

Resources you used:

MLflow tracking server running on port 5000.
FastAPI server running on port 8000.
Project directory mlops-prediction-api/ containing train.py and app.py.
MLflow experiment data stored in the mlruns/ directory.
Python packages installed via pip: mlflow, scikit-learn, fastapi, uvicorn.

Keep everything running

No action needed. Choose this if you are still actively building or want to keep experimenting with your MLOps pipeline.

Your MLflow server will continue running on port 5000.
Your FastAPI server will continue running on port 8000.
All experiment data and registered models remain available in the MLflow UI.

Pause - I'll come back to this later

Shut down the running servers to free up your terminal and ports. Your project files and experiment data will stay on disk for when you come back.

Go to the terminal running the MLflow server and press Ctrl+C to stop it.
Go to the terminal running the FastAPI server and press Ctrl+C to stop it.
To resume later, navigate to your project directory and restart the servers with mlflow server --port 5000 and uvicorn app:app --reload --port 8000.

Delete - I don't want to use this again

Remove all project resources and start fresh. This will delete your experiment data, registered models, and project files.

Stop both running servers by pressing Ctrl+C in each terminal.
Delete the project directory by running rm -rf mlops-prediction-api/ (macOS/Linux) or removing the folder manually on Windows.
Optionally, uninstall the Python packages by running pip uninstall mlflow scikit-learn fastapi uvicorn -y.

Mission accomplished

That's a Wrap!

Nice work! You just built a complete local MLOps pipeline from scratch. You trained a model, tracked experiments, registered your best model, and served predictions through a live API.

Here's what you learned:

Train a scikit-learn classification model and track experiments with MLflow, logging parameters, metrics, and model artifacts automatically.
Compare multiple training runs in the MLflow UI and register the best model in the MLflow Model Registry for versioning.
Build a FastAPI prediction API that loads a registered model and serves real-time predictions through a /predict endpoint.
Make your API production-ready with a /health endpoint, Pydantic input validation, and error handling.

This is the core MLOps loop that employers look for: train, track, version, deploy. You now have hands-on proof that you can do all four.

Ready to quiz yourself?