Serve an ML Model as a Production API

Train a heart disease classifier and serve it as a REST API with FastAPI.

Introduction

30 Second Summary

Every time you ask a chatbot a question or get a product recommendation, there's a machine learning model running behind an API somewhere. The hardest part isn't training the model. It's getting it into production where real users can actually reach it.

In this project, you will train a heart disease classifier, track the experiment with MLflow, and serve the model through a production-ready API built with FastAPI.

What You'll Build

A production-grade ML prediction API that accepts patient data as JSON, validates the input, and returns predictions with confidence scores.

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

  • A training script that trains a classifier and logs parameters, metrics, and artifacts to MLflow.
  • A FastAPI application that serves predictions via HTTP with input validation using Pydantic.
  • Production safety features including health checks, error handling, and model versioning.
  • Secret Mission: An A/B testing endpoint that routes requests between two model versions and tracks which performs better.

Are there any prerequisites?

You should have basic Python knowledge and some familiarity with training ML models (classification or regression). No prior experience with FastAPI or MLflow is needed.

Before We Start

Before We Start...

Set Up Your Python Environment

Before you can train a model or serve predictions, you need a clean Python environment with the right packages installed.

In this step, you will install Python, create an isolated virtual environment, and install the libraries that power your ML prediction API.

In this step, get ready to:
  • Install or verify Python 3.11+.
  • Create a virtual environment and install project dependencies.
  • Verify all packages are ready to use.
Verify your Python version

This project requires Python 3.11 or higher because scikit-learn 1.8.0 (the ML library you will use) only supports Python 3.11+.

  • In VS Code, click Terminal in the top menu bar.
  • Click New Terminal to open the integrated terminal.
  • Check your Python version by running:
python3 --version

✔️ I see version 3.11 or higher

You are all set. Your output shows something like Python 3.11.x or Python 3.12.x. Move on to the next substep.

ⓧ I see an older version

macOS often ships with an older Python that does not meet the 3.11 requirement. You need to install a newer version alongside it.

  • Go to python.org/downloads and download the latest Python 3.11+ installer for your operating system.
  • Run the installer. On Windows, check the Add Python to PATH box before clicking Install Now.
  • Close and reopen your terminal, then run python3 --version again to confirm you see 3.11 or higher.

macOS shadows new installs

If python3 --version still shows the old version after installing, try running python3.11 --version or python3.12 --version to confirm the new install exists. Use whichever versioned command works for the rest of this project in place of python3.

ⓧ Command not found

Python is not installed on your machine yet. You need to install it from the official site.

  • Go to python.org/downloads and download the latest Python 3.11+ installer for your operating system.
  • Run the installer. On Windows, check the Add Python to PATH box before clicking Install Now.
  • Close and reopen your terminal, then run python3 --version to confirm it shows 3.11 or higher.
  • On Windows, if python3 is not recognized, try python --version instead. Use whichever command works for the rest of this project.
Create your project folder and virtual environment

A virtual environment isolates your project's packages from the rest of your system. This prevents version conflicts between different projects.

  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir ml-prediction-api
cd ml-prediction-api

If both commands run without errors, your folder was created and you are now inside it.

  • Open this folder in VS Code by clicking File in the top menu bar.
  • Click Open Folder.
  • Navigate to Desktop, select the ml-prediction-api folder, then click Open.
  • Open a new terminal inside VS Code by clicking Terminal in the top menu bar, then New Terminal.

Your terminal prompt should now show ml-prediction-api as the current directory.

  • Create a virtual environment by running:
python3 -m venv .venv

What does this command do?

python3 -m venv .venv creates a new folder called .venv inside your project. This folder holds a self-contained Python installation just for this project.

The dot prefix makes it a hidden folder, keeping your project directory clean.

  • Activate the virtual environment by running the command for your operating system:

macOS

This command also works on Linux.

source .venv/bin/activate

Windows

.venv\Scripts\Activate.ps1

You should see (.venv) appear at the beginning of your terminal prompt. This confirms the virtual environment is active.

Don't see (.venv) in your prompt?

On Windows, you may get an error about execution policies. Run Set-ExecutionPolicy -Scope CurrentUser -ExecutionPolicy RemoteSigned in PowerShell, then try the activate command again.

On macOS, make sure you typed source before the path. Without it, the command will not activate the environment.

help me activate my Python virtual environment

Install and verify your dependencies

With your virtual environment active, you can now install the Python packages that power your ML prediction API.

  • Install all project dependencies by running:
pip install "fastapi[standard]" mlflow scikit-learn

What does this install?

  • fastapi[standard] installs FastAPI plus Uvicorn (the server that runs your API) and Pydantic (for input validation).
  • mlflow installs MLflow for experiment tracking and model management.
  • scikit-learn installs the ML library along with joblib (for saving models) and NumPy (for numerical operations).

The install may take a minute or two. Wait until you see your terminal prompt return.

Seeing a permissions error or install failure?

Make sure you see (.venv) in your terminal prompt. If you do not, your virtual environment is not active. Go back and run the activation command for your OS.

If you see errors about pip being outdated, run pip install --upgrade pip first, then try the install command again.

help me fix my pip install error

  • Verify all packages installed correctly by running:
python -c "import fastapi; import mlflow; import sklearn; print(f'FastAPI: {fastapi.__version__}'); print(f'MLflow: {mlflow.__version__}'); print(f'scikit-learn: {sklearn.__version__}')"

You should see version numbers printed for all three packages with no errors.

Seeing a ModuleNotFoundError?

If Python cannot find one of the packages, your virtual environment may not be active. Check for (.venv) in your prompt, then rerun the pip install command.

help me fix my Python import errors

Your Python environment is locked and loaded. Next up, you will train a machine learning model and track the entire experiment with MLflow.

Train a Model and Track It with MLflow

Your environment is ready. Now it is time to train a model and track the entire experiment so nothing gets lost.

In production ML, teams need a record of every training run: what data went in, which hyperparameters were used, and how well the model performed. MLflow handles this automatically.

In this step, get ready to:
  • Load and split the breast cancer dataset for training.
  • Train a GradientBoostingClassifier with MLflow tracking enabled.
  • View the experiment in the MLflow dashboard.
Load the dataset and prepare training data

Before training a model, you need data split into a training set and a test set. scikit-learn includes a built-in breast cancer classification dataset with 30 numeric features, perfect for binary classification.

  • In VS Code, press Cmd+N (macOS) or Ctrl+N (Windows) to create a new file.
  • Save it immediately as train.py by pressing Cmd+S (macOS) or Ctrl+S (Windows). Make sure it saves inside your project folder.
  • Add the following imports, MLflow setup, and data loading code to train.py:
import mlflow
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import accuracy_score, precision_score, recall_score
from joblib import dump
import json
from datetime import datetime, timezone

# Enable MLflow experiment tracking
mlflow.set_experiment("heart-disease-prediction")
mlflow.sklearn.autolog()

# Load dataset (using breast cancer as a binary classification health dataset)
data = load_breast_cancer()
X_train, X_test, y_train, y_test = train_test_split(
    data.data, data.target, test_size=0.2, random_state=42
)

What does this code do?

The imports pull in everything you need: MLflow for tracking, scikit-learn for the model and data, and joblib/json/datetime for saving artifacts later.

  • mlflow.set_experiment("heart-disease-prediction") creates (or connects to) a named experiment. All runs get grouped under this name.
  • mlflow.sklearn.autolog() tells MLflow to automatically capture parameters, metrics, and the trained model for any scikit-learn training that happens next.
  • train_test_split splits the 569 samples into 80% training and 20% test data. The random_state=42 ensures reproducible splits.

You will not see any output yet. The next chunk adds the training logic that produces results.

Train the model inside an MLflow run

MLflow groups everything that happens inside a "run" together. You will wrap your training and evaluation code in a run context so that every parameter and metric gets logged automatically.

  • Add the following training code below the data loading section in train.py:
with mlflow.start_run():
    # Train model
    model = GradientBoostingClassifier(
        n_estimators=100, max_depth=3, learning_rate=0.1, random_state=42
    )
    model.fit(X_train, y_train)

    # Evaluate
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    precision = precision_score(y_test, y_pred)
    recall = recall_score(y_test, y_pred)

    print(f"Accuracy: {accuracy:.3f}")
    print(f"Precision: {precision:.3f}")
    print(f"Recall: {recall:.3f}")

What does this code do?

Everything inside with mlflow.start_run(): gets tracked as a single experiment run.

  • GradientBoostingClassifier is an ensemble method that builds decision trees sequentially, each correcting the errors of the previous one. It performs well on tabular data without much tuning.
  • model.fit(X_train, y_train) trains the model. Because autolog is enabled, MLflow captures all the hyperparameters (n_estimators, max_depth, learning_rate) automatically.
  • The evaluation section calculates accuracy, precision, and recall on the held-out test set. These metrics tell you how well the model generalises to unseen data.
  • Save train.py.
  • Run the training script by executing this command in your terminal:
python3 train.py

You should see three metrics printed to your terminal, similar to this:

Accuracy: 0.956
Precision: 0.963
Recall: 0.972

What should I see?

Your exact numbers may vary slightly, but accuracy should be above 0.90. The script also creates an mlruns folder in your project directory where MLflow stores all experiment data locally.

Seeing an import error?

Make sure your virtual environment is activated. You should see its name in parentheses at the start of your terminal prompt. If not, activate it and try again.

If you see a module not found error for a specific package, install it with pip install <package-name> and re-run the script.

help me debug an import error when running train.py

✔️ Awesome, I've got everything!

Great. Double check you have saved train.py and you can see the accuracy, precision, and recall output in your terminal.

ⓧ I'd like to double check the full code

Here is the complete train.py file as it should look right now:

import mlflow
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import accuracy_score, precision_score, recall_score
from joblib import dump
import json
from datetime import datetime, timezone

# Enable MLflow experiment tracking
mlflow.set_experiment("heart-disease-prediction")
mlflow.sklearn.autolog()

# Load dataset (using breast cancer as a binary classification health dataset)
data = load_breast_cancer()
X_train, X_test, y_train, y_test = train_test_split(
    data.data, data.target, test_size=0.2, random_state=42
)

with mlflow.start_run():
    # Train model
    model = GradientBoostingClassifier(
        n_estimators=100, max_depth=3, learning_rate=0.1, random_state=42
    )
    model.fit(X_train, y_train)

    # Evaluate
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    precision = precision_score(y_test, y_pred)
    recall = recall_score(y_test, y_pred)

    print(f"Accuracy: {accuracy:.3f}")
    print(f"Precision: {precision:.3f}")
    print(f"Recall: {recall:.3f}")
View your experiment in the MLflow UI

MLflow includes a web dashboard that visualises everything your training run captured. You can browse parameters, compare metrics, and inspect model artifacts.

  • Open a new terminal window (keep it separate from the one you used to run the script).
  • Make sure you are in the same project folder where you ran train.py.
  • Start the MLflow dashboard by running:
mlflow ui

You should see output similar to:

INFO:     Uvicorn running on http://127.0.0.1:5000 (Press CTRL+C to quit)
  • Open your browser and go to http://127.0.0.1:5000.
  • Click on the heart-disease-prediction experiment in the left sidebar.
  • Click on the run listed in the table to view its details.

What should I see?

The run detail page shows the hyperparameters MLflow captured automatically (n_estimators, max_depth, learning_rate), the training metrics, and a logged model artifact.

This is what autolog gives you for free. No manual logging code required.

Dashboard not loading?

Make sure you opened a new terminal for the mlflow ui command. It needs to keep running while you view the dashboard.

If port 5000 is already in use, stop whatever is using it or run mlflow ui --port 5001 and visit http://127.0.0.1:5001 instead.

help me troubleshoot the MLflow UI not loading

Your model is trained and tracked. Next up, you will save this model to disk and create a metadata file so your API can load it in production.

Save and Version Your Model

In the previous step, you trained a model and watched MLflow automatically capture every detail of your experiment. Your model exists in memory, but the moment you close your terminal, it vanishes.

In this step, you will save your trained model to a file so your API can load it later. You will also create a metadata file that tracks the model's version, training date, and performance metrics.

In this step, get ready to:
  • Export your trained model to a file using joblib.
  • Create a model metadata file that tracks version and performance.
  • Understand why model versioning matters in production systems.
Export your trained model using joblib

Your trained model currently lives only in your Python process's memory. To serve predictions through an API, you need the model saved as a file that the API can load on startup.

joblib is a Python library that serializes objects to disk. scikit-learn's official documentation recommends joblib over Python's built-in pickle because it is more efficient when working with large NumPy arrays.

Why joblib over pickle?

ML models often contain large NumPy arrays (the weights learned during training). joblib uses memory-mapping techniques that make saving and loading these arrays faster and more memory-efficient than standard pickle.

Since joblib is already a dependency of scikit-learn, you don't need to install anything extra. The from joblib import dump line is already at the top of your train.py file.

Create a model metadata file

A saved model file alone isn't enough for production. You also need a metadata file that tells the API which version of the model is loaded, when it was trained, and how well it performed. This is what lets you track changes and roll back if a new model underperforms.

  • In train.py, find the line that says print(f"Recall: {recall:.3f}").
  • Below that line, add the following code to save your model and its metadata:
    # Save model for the API
    dump(model, "model.joblib")

    # Save model metadata
    metadata = {
        "model_name": "GradientBoostingClassifier",
        "version": "1.0.0",
        "trained_at": datetime.now(timezone.utc).isoformat(),
        "metrics": {
            "accuracy": round(accuracy, 4),
            "precision": round(precision, 4),
            "recall": round(recall, 4),
        },
        "features": list(data.feature_names),
        "n_features": len(data.feature_names),
    }
    with open("model_metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)

    print("\nModel saved to model.joblib")
    print("Metadata saved to model_metadata.json")
    print("\nRun 'mlflow ui' to view the experiment dashboard")

What does this code do?

This code does two things: saves the model, and saves a record of the model's identity.

  • dump(model, "model.joblib") serializes your trained GradientBoostingClassifier to a binary file. The API will load this file to make predictions.
  • The metadata dictionary records the model's name, version (1.0.0), UTC training timestamp, accuracy/precision/recall metrics, feature names, and feature count.
  • json.dump(metadata, f, indent=2) writes the metadata to model_metadata.json in a human-readable format. Your API will read this file to report model info.
  • Save train.py.
  • Re-run the training script by running this command:
python3 train.py
  • You should see the accuracy, precision, and recall metrics printed, followed by confirmation that model.joblib and model_metadata.json were saved.
  • Confirm the files exist by running:
ls
  • You should see model.joblib and model_metadata.json listed alongside train.py in the output.

Not seeing the new files?

  • Make sure your new code is indented inside the with mlflow.start_run(): block. If the indentation is off, the code runs outside the MLflow context or raises an IndentationError.
  • Check that you added the code below the existing print statements, not above or outside the with block.
  • If you see a NameError for dump or json, verify that from joblib import dump and import json are at the top of your file.

Why isn't my train.py saving model.joblib and model_metadata.json?

✔️ Awesome, I've got everything!

Great. Double check you've saved train.py and confirmed both output files exist.

ⓧ I'd like to double check the full code

Here is the complete train.py file. Compare yours against it:

import mlflow
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import accuracy_score, precision_score, recall_score
from joblib import dump
import json
from datetime import datetime, timezone

# Enable MLflow experiment tracking
mlflow.set_experiment("heart-disease-prediction")
mlflow.sklearn.autolog()

# Load dataset (using breast cancer as a binary classification health dataset)
data = load_breast_cancer()
X_train, X_test, y_train, y_test = train_test_split(
    data.data, data.target, test_size=0.2, random_state=42
)

with mlflow.start_run():
    # Train model
    model = GradientBoostingClassifier(
        n_estimators=100, max_depth=3, learning_rate=0.1, random_state=42
    )
    model.fit(X_train, y_train)

    # Evaluate
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    precision = precision_score(y_test, y_pred)
    recall = recall_score(y_test, y_pred)

    print(f"Accuracy: {accuracy:.3f}")
    print(f"Precision: {precision:.3f}")
    print(f"Recall: {recall:.3f}")

    # Save model for the API
    dump(model, "model.joblib")

    # Save model metadata
    metadata = {
        "model_name": "GradientBoostingClassifier",
        "version": "1.0.0",
        "trained_at": datetime.now(timezone.utc).isoformat(),
        "metrics": {
            "accuracy": round(accuracy, 4),
            "precision": round(precision, 4),
            "recall": round(recall, 4),
        },
        "features": list(data.feature_names),
        "n_features": len(data.feature_names),
    }
    with open("model_metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)

    print("\nModel saved to model.joblib")
    print("Metadata saved to model_metadata.json")
    print("\nRun 'mlflow ui' to view the experiment dashboard")
Understand why model versioning matters

You just created a versioned model artifact. The metadata file records version 1.0.0, the exact timestamp the model was trained, and the metrics it achieved. This is the foundation of production ML operations.

Why does model versioning matter in production?

In production, models get retrained regularly as new data arrives. Without versioning, you have no way to know which model is currently serving predictions or how it compares to earlier versions.

Model versioning enables two critical workflows. First, rollbacks: if a new model underperforms, you can instantly switch back to the previous version because you still have its file and metadata. Second, A/B testing: you can serve two versions side-by-side and compare their real-world performance before fully committing to a new model.

Your model is trained, saved, and versioned. Next up, you will build the FastAPI application that loads this model and serves predictions over HTTP.

Build the Prediction API with FastAPI

You have a trained model and a metadata file sitting in your project folder. Right now, the only way to get a prediction is to open a Python shell and manually call the model.

In this step, you will wrap the model in a FastAPI application that accepts HTTP requests, validates the input data, and returns predictions as JSON. This is how ML models are served in production environments.

In this step, get ready to:
  • Define Pydantic validation schemas for API requests and responses.
  • Build a prediction endpoint that loads the model on startup and returns predictions with confidence scores.
  • Add health check and model info endpoints for production monitoring.
Create Pydantic models for request validation and response format

Before writing any endpoint logic, you need to define what data your API accepts and what it returns. Pydantic models enforce strict validation on incoming requests. If someone sends invalid data, the API rejects it automatically with a clear error message.

  • Create a new file called app.py in your project folder (the same folder where train.py lives).
  • Add the imports and the input validation model by pasting this code:
import json
from contextlib import asynccontextmanager

import numpy as np
from fastapi import FastAPI, HTTPException
from joblib import load
from pydantic import BaseModel, Field


# --- Pydantic Models ---
class PredictionRequest(BaseModel):
    features: list[float] = Field(
        ...,
        min_length=30,
        max_length=30,
        description="List of 30 numeric features for prediction",
    )

What does this code do?

  • The imports bring in FastAPI for the web framework, Pydantic for data validation, joblib to load your saved model, and numpy to reshape feature arrays.
  • PredictionRequest defines what the API expects: exactly 30 floating-point numbers. min_length=30 and max_length=30 ensure the caller sends the right number of features.
  • The ... (ellipsis) in Field(...) means the field is required. Requests without it are rejected automatically.
  • Now add the response models below PredictionRequest in app.py:
class PredictionResponse(BaseModel):
    prediction: int
    prediction_label: str
    confidence: float
    model_version: str


class HealthResponse(BaseModel):
    status: str
    model_loaded: bool


class ModelInfoResponse(BaseModel):
    model_name: str
    version: str
    trained_at: str
    metrics: dict
    n_features: int

What do these response models do?

Each class defines the exact shape of a JSON response. PredictionResponse includes the numeric prediction, a human-readable label, confidence score, and model version.

HealthResponse tells monitoring tools whether the service is up and the model is loaded. ModelInfoResponse exposes metadata about the currently running model.

  • Save app.py.

Seeing a red underline on list[float]?

If your editor shows an error on list[float], make sure you are using Python 3.11 or higher. The lowercase list type hint syntax requires Python 3.9+.

help me fix a type hint error in my FastAPI Pydantic model

Build the /predict endpoint that loads the model and returns predictions

Your API needs to load model.joblib into memory when the server starts, not on every request. FastAPI's lifespan context manager handles this startup logic cleanly.

  • Add the application state, lifespan, and FastAPI app creation below your Pydantic models in app.py:
# --- Application State ---
ml_model = None
model_metadata = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load model and metadata on startup."""
    global ml_model, model_metadata
    try:
        ml_model = load("model.joblib")
        with open("model_metadata.json", "r") as f:
            model_metadata = json.load(f)
        print("Model loaded successfully")
    except Exception as e:
        print(f"Warning: Could not load model: {e}")
    yield


# --- FastAPI App ---
app = FastAPI(
    title="ML Prediction API",
    description="Production ML model serving with FastAPI",
    version="1.0.0",
    lifespan=lifespan,
)

What does this code do?

  • ml_model and model_metadata are global variables that hold your loaded model and its metadata in memory.
  • The lifespan function runs once when the server starts. It loads model.joblib and model_metadata.json so every request can use them without re-reading from disk.
  • The yield statement marks the boundary between startup and shutdown. Code before yield runs on startup. Code after it (if any) would run on shutdown.
  • Now add the /predict endpoint below the app = FastAPI(...) declaration:
@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    """Make a prediction using the loaded model."""
    if ml_model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    try:
        features_array = np.array(request.features).reshape(1, -1)
        prediction = int(ml_model.predict(features_array)[0])
        probabilities = ml_model.predict_proba(features_array)[0]
        confidence = float(max(probabilities))
        label = "malignant" if prediction == 0 else "benign"

        return PredictionResponse(
            prediction=prediction,
            prediction_label=label,
            confidence=round(confidence, 4),
            model_version=model_metadata["version"] if model_metadata else "unknown",
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

What does the /predict endpoint do?

  • It first checks that the model is actually loaded. If not, it returns a 503 (Service Unavailable) response.
  • features_array = np.array(request.features).reshape(1, -1) converts the input list into a 2D numpy array that scikit-learn expects.
  • predict_proba returns probability scores for each class. The highest probability becomes the confidence score.
  • The label maps the numeric prediction (0 or 1) to a human-readable string: "malignant" or "benign".
  • Save app.py.

Seeing an import error for numpy?

If your editor cannot find numpy, make sure your virtual environment is activated and that you installed all dependencies in Step 1. Numpy is installed as a dependency of scikit-learn.

help me resolve a missing numpy import in my FastAPI project

Add a /health endpoint and a /model/info endpoint

Production APIs need endpoints that monitoring tools can ping to check service health. They also need a way for consumers to inspect which model version is running.

  • In app.py, find the line @app.post("/predict", response_model=PredictionResponse).
  • Add the following /health endpoint above that line:
@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint for monitoring."""
    return HealthResponse(status="healthy", model_loaded=ml_model is not None)

Why a health check endpoint?

Load balancers and container orchestrators (like Kubernetes) continuously poll /health to decide if your service can receive traffic. If the model fails to load, model_loaded returns false and the orchestrator routes traffic elsewhere.

  • Now add the /model/info endpoint directly below the /health endpoint (still above /predict):
@app.get("/model/info", response_model=ModelInfoResponse)
async def model_info():
    """Return model metadata and performance metrics."""
    if model_metadata is None:
        raise HTTPException(status_code=503, detail="Model metadata not available")
    return ModelInfoResponse(
        model_name=model_metadata["model_name"],
        version=model_metadata["version"],
        trained_at=model_metadata["trained_at"],
        metrics=model_metadata["metrics"],
        n_features=model_metadata["n_features"],
    )

What does /model/info return?

This endpoint reads from the model_metadata.json you created during training. It exposes the model name, version, training timestamp, and accuracy metrics. This helps API consumers understand exactly which model is serving their predictions.

  • Save app.py.
  • Start the FastAPI development server by running this command:
fastapi dev app.py
  • Open your browser and navigate to http://127.0.0.1:8000/docs.
  • You should see the Swagger UI interactive documentation page listing your three endpoints: /health, /model/info, and /predict.

Server not starting?

  • Make sure you are running the command from the same folder where app.py, model.joblib, and model_metadata.json all exist.
  • Check that your virtual environment is activated. If you see a "command not found" error for fastapi, re-activate your environment and confirm the package is installed.
  • If you see a syntax error, double check that all parentheses and brackets are properly closed in app.py.

help me debug why my FastAPI server won't start

✔️ Awesome, I've got everything!

Great. Double check you have saved app.py and your server is running with all three endpoints visible at http://127.0.0.1:8000/docs.

ⓧ I'd like to double check the full code

import json
from contextlib import asynccontextmanager

import numpy as np
from fastapi import FastAPI, HTTPException
from joblib import load
from pydantic import BaseModel, Field


# --- Pydantic Models ---
class PredictionRequest(BaseModel):
    features: list[float] = Field(
        ...,
        min_length=30,
        max_length=30,
        description="List of 30 numeric features for prediction",
    )


class PredictionResponse(BaseModel):
    prediction: int
    prediction_label: str
    confidence: float
    model_version: str


class HealthResponse(BaseModel):
    status: str
    model_loaded: bool


class ModelInfoResponse(BaseModel):
    model_name: str
    version: str
    trained_at: str
    metrics: dict
    n_features: int


# --- Application State ---
ml_model = None
model_metadata = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load model and metadata on startup."""
    global ml_model, model_metadata
    try:
        ml_model = load("model.joblib")
        with open("model_metadata.json", "r") as f:
            model_metadata = json.load(f)
        print("Model loaded successfully")
    except Exception as e:
        print(f"Warning: Could not load model: {e}")
    yield


# --- FastAPI App ---
app = FastAPI(
    title="ML Prediction API",
    description="Production ML model serving with FastAPI",
    version="1.0.0",
    lifespan=lifespan,
)


@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint for monitoring."""
    return HealthResponse(status="healthy", model_loaded=ml_model is not None)


@app.get("/model/info", response_model=ModelInfoResponse)
async def model_info():
    """Return model metadata and performance metrics."""
    if model_metadata is None:
        raise HTTPException(status_code=503, detail="Model metadata not available")
    return ModelInfoResponse(
        model_name=model_metadata["model_name"],
        version=model_metadata["version"],
        trained_at=model_metadata["trained_at"],
        metrics=model_metadata["metrics"],
        n_features=model_metadata["n_features"],
    )


@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    """Make a prediction using the loaded model."""
    if ml_model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    try:
        features_array = np.array(request.features).reshape(1, -1)
        prediction = int(ml_model.predict(features_array)[0])
        probabilities = ml_model.predict_proba(features_array)[0]
        confidence = float(max(probabilities))
        label = "malignant" if prediction == 0 else "benign"

        return PredictionResponse(
            prediction=prediction,
            prediction_label=label,
            confidence=round(confidence, 4),
            model_version=model_metadata["version"] if model_metadata else "unknown",
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

Your prediction API is live. Next up, you will test it with real data, intentionally break it with invalid inputs, and add the production safety features that make it deployment-ready.

Add Production Safety Features

Your API is built, but a production system needs to handle bad data gracefully. If someone sends 5 features instead of 30, or the model file goes missing, your API should return a clear error instead of crashing.

In this step, you will review the error handling already built into your code, test it live through the auto-generated Swagger docs, and run the full system with both the MLflow dashboard and the FastAPI server running simultaneously.

In this step, get ready to:
  • Review the three layers of production safety in your API.
  • Test the API with valid data, invalid data, and edge cases using Swagger UI.
  • Run the full system with MLflow and FastAPI side by side.
Review the error handling in your API

Your app.py already has three layers of production safety built in from the previous step. Let's walk through each one.

  • Open app.py in VS Code and find the PredictionRequest class near the top of the file:
class PredictionRequest(BaseModel):
    features: list[float] = Field(
        ...,
        min_length=30,
        max_length=30,
        description="List of 30 numeric features for prediction",
    )

What are the three safety layers?

Layer 1: Pydantic input validation. The min_length=30 and max_length=30 constraints reject any request that does not contain exactly 30 numeric features. Bad data never reaches the model.

Layer 2: Service availability checks. The /predict endpoint raises a 503 if ml_model is None. The /model/info endpoint raises a 503 if model_metadata is None. This protects callers when the model fails to load at startup.

Layer 3: Catch-all error handling. The try/except block inside /predict catches unexpected errors during inference and returns a 500 with a human-readable detail message. The server stays running even if one prediction fails.

Test the API with Swagger UI

FastAPI auto-generates interactive API documentation called Swagger UI. You can use it to send real requests to your endpoints without writing any client code.

  • Start the FastAPI development server by running this command:
fastapi dev app.py

What does this command do?

The fastapi dev command starts a development server with auto-reload enabled. It serves your API on http://127.0.0.1:8000 and auto-generates interactive docs at http://127.0.0.1:8000/docs.

  • Open your browser and navigate to http://127.0.0.1:8000/docs.
  • You should see the Swagger UI page listing all three endpoints: /health, /model/info, and /predict.

Don't see the Swagger page?

Make sure your virtual environment is activated and the terminal is in your project folder. Check the terminal output for error messages about missing modules.

help me debug why my FastAPI server won't start

First, test a successful prediction with valid data.

  • Click on the POST /predict endpoint to expand it.
  • Click Try it out.
  • Paste the following JSON into the request body field:
{
  "features": [17.99, 10.38, 122.8, 1001.0, 0.1184, 0.2776, 0.3001, 0.1471, 0.2419, 0.07871, 1.095, 0.9053, 8.589, 153.4, 0.006399, 0.04904, 0.05373, 0.01587, 0.03003, 0.006193, 25.38, 17.33, 184.6, 2019.0, 0.1622, 0.6656, 0.7119, 0.2654, 0.4601, 0.1189]
}
  • Click Execute.
  • You should see a 200 response with a JSON body containing prediction, prediction_label, confidence, and model_version fields.

What should I see?

The response should look something like this: a prediction of 0 or 1, a prediction_label of "malignant" or "benign", a confidence score between 0 and 1, and a model_version of "1.0.0".

Now test what happens with invalid data.

  • Click Try it out again on the POST /predict endpoint.
  • Replace the request body with this invalid payload (only 3 features instead of 30):
{
  "features": [1.0, 2.0, 3.0]
}
  • Click Execute.
  • You should see a 422 Validation Error response. Pydantic rejected the request because the features list has fewer than 30 items.

Why does this matter in production?

Without input validation, bad data would reach the model and either crash the server or produce garbage predictions silently. The 422 response tells the caller exactly what went wrong so they can fix their request.

  • Test the GET /health endpoint by clicking on it, then clicking Try it out and Execute.
  • You should see a 200 response with "status": "healthy" and "model_loaded": true.
Run the full system with MLflow and FastAPI

In a real production environment, you would have both the experiment tracking dashboard and the prediction API running at the same time. Let's replicate that setup locally.

  • Leave the FastAPI server running in its current terminal.
  • Open a second terminal window. Press Cmd+Shift+P (macOS) or Ctrl+Shift+P in VS Code to open the Command Palette.
  • Type Terminal: Create New Terminal and press Enter.
  • Activate your virtual environment in the new terminal by running:
source .venv/bin/activate
  • Start the MLflow UI by running:
mlflow ui

What does this command do?

The mlflow ui command starts a web dashboard at http://127.0.0.1:5000 where you can view your experiment runs, compare metrics, and browse model artifacts. It reads from the mlruns/ directory that was created when you ran train.py.

  • Open http://127.0.0.1:5000 in a new browser tab.
  • You should see your experiment run listed under the heart-disease-prediction experiment name, with logged metrics for accuracy, precision, and recall.

You now have both services running: your prediction API at port 8000 and the MLflow dashboard at port 5000.

Port already in use?

If you see an error about the port being in use, a previous server process may still be running. Press Ctrl+C in the terminal to stop it, then re-run the command.

help me fix a port conflict between MLflow and FastAPI

✔️ Awesome, I've got everything!

Great. Double check that both services are running and you can access Swagger UI at port 8000 and MLflow at port 5000.

ⓧ I'd like to double check the full code

Here is the complete app.py file:

import json
from contextlib import asynccontextmanager

import numpy as np
from fastapi import FastAPI, HTTPException
from joblib import load
from pydantic import BaseModel, Field


# --- Pydantic Models ---
class PredictionRequest(BaseModel):
    features: list[float] = Field(
        ...,
        min_length=30,
        max_length=30,
        description="List of 30 numeric features for prediction",
    )


class PredictionResponse(BaseModel):
    prediction: int
    prediction_label: str
    confidence: float
    model_version: str


class HealthResponse(BaseModel):
    status: str
    model_loaded: bool


class ModelInfoResponse(BaseModel):
    model_name: str
    version: str
    trained_at: str
    metrics: dict
    n_features: int


# --- Application State ---
ml_model = None
model_metadata = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load model and metadata on startup."""
    global ml_model, model_metadata
    try:
        ml_model = load("model.joblib")
        with open("model_metadata.json", "r") as f:
            model_metadata = json.load(f)
        print("Model loaded successfully")
    except Exception as e:
        print(f"Warning: Could not load model: {e}")
    yield


# --- FastAPI App ---
app = FastAPI(
    title="ML Prediction API",
    description="Production ML model serving with FastAPI",
    version="1.0.0",
    lifespan=lifespan,
)


@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint for monitoring."""
    return HealthResponse(status="healthy", model_loaded=ml_model is not None)


@app.get("/model/info", response_model=ModelInfoResponse)
async def model_info():
    """Return model metadata and performance metrics."""
    if model_metadata is None:
        raise HTTPException(status_code=503, detail="Model metadata not available")
    return ModelInfoResponse(
        model_name=model_metadata["model_name"],
        version=model_metadata["version"],
        trained_at=model_metadata["trained_at"],
        metrics=model_metadata["metrics"],
        n_features=model_metadata["n_features"],
    )


@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    """Make a prediction using the loaded model."""
    if ml_model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    try:
        features_array = np.array(request.features).reshape(1, -1)
        prediction = int(ml_model.predict(features_array)[0])
        probabilities = ml_model.predict_proba(features_array)[0]
        confidence = float(max(probabilities))
        label = "malignant" if prediction == 0 else "benign"

        return PredictionResponse(
            prediction=prediction,
            prediction_label=label,
            confidence=round(confidence, 4),
            model_version=model_metadata["version"] if model_metadata else "unknown",
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

Your ML serving system is complete. You have a trained model tracked in MLflow, served through a production-safe API that validates inputs, handles errors, and reports its own health status.

Secret mission

Model A/B Testing

In production, you never just swap one model for another and hope for the best. Real ML teams run both models simultaneously, split traffic between them, and let the data decide which one wins. Build an A/B testing system that does exactly this.

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:

  • Running processes: FastAPI server on port 8000 and MLflow UI on port 5000.
  • Project files: train.py, app.py, model_a.joblib, model_b.joblib, and model_metadata.json.
  • Python virtual environment with installed packages (FastAPI, MLflow, scikit-learn).
  • MLflow experiment data stored in the mlruns/ directory.

Keep everything running

No action needed. Choose this if you're still actively building or want to continue experimenting with the API.

  • Your FastAPI server and MLflow UI will keep running until you close the terminal windows or press Ctrl+C.
  • Your project files, trained models, and experiment data remain on your Desktop.
  • Your virtual environment stays available for future use.

Pause - I'll come back to this later

Shut down running processes to free up memory, but keep all your files intact for later.

  • In each terminal running a server, press Ctrl+C to stop the process.
  • Deactivate your virtual environment by running:
deactivate
  • Your project folder, trained models, and MLflow experiment data stay on disk for when you return.
  • To resume later, navigate back to your project folder, reactivate the virtual environment, and start the servers again.

Delete - I don't want to use this again

Remove all project resources and start fresh.

  • In each terminal running a server, press Ctrl+C to stop the process.
  • Deactivate your virtual environment by running:
deactivate
  • Navigate to your Desktop and delete the entire project folder. On macOS or Linux, run:
cd ~/Desktop
rm -rf ml-prediction-api

On Windows, run:

cd ~\Desktop
rmdir /s /q ml-prediction-api
  • This removes all project files, trained models, the virtual environment, and the mlruns/ experiment data.
  • Python itself remains installed on your system. You can keep it for future projects.

That's a Wrap!

That's a Wrap!

Nice work! You just built a production-grade ML serving system that mirrors how real companies deploy machine learning models.

You've learned how to:

  • Train an ML model and track the entire experiment (parameters, metrics, artifacts) using MLflow.
  • Serve a model as a production API with FastAPI, complete with Pydantic input validation and error handling.
  • Implement production patterns including health checks, model metadata endpoints, and model versioning.
  • Build a model A/B testing system to compare model versions in real-time.

Ready to quiz yourself?