Model Deployment with Flask: Build and Serve ML Models as REST APIs (2026)

Training a machine learning model is only half the job. A model that lives only in a Jupyter notebook provides zero business value. Deployment — making your model available to applications, users, and other services — is what turns research into product. Flask is the most beginner-friendly way to serve a Python ML model as a REST API, and understanding this workflow is a critical skill for any data scientist who wants their work to actually be used.

Training and Serialising Your Model

Before you can serve a model, you need to save it to disk. Joblib is the recommended tool for scikit-learn models; pickle works for anything; ONNX and TorchScript are standard for neural networks. Always save the preprocessing pipeline alongside the model — the API must apply the same transformations to incoming data that were applied during training:

import joblib
import pandas as pd
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.model_selection import train_test_split

# Train a model (example: churn prediction)
df = pd.read_csv('customer_data.csv')
X = df.drop('churned', axis=1)
y = df['churned']

# Build pipeline that includes preprocessing + model
pipeline = Pipeline([
    ('scaler', StandardScaler()),
    ('model',  GradientBoostingClassifier(n_estimators=200, random_state=42))
])

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
pipeline.fit(X_train, y_train)
print(f"Accuracy: {pipeline.score(X_test, y_test):.3f}")

# Save the entire pipeline — preprocessing + model together
joblib.dump(pipeline, 'churn_model.pkl')

# Also save feature names for input validation
feature_names = list(X.columns)
joblib.dump(feature_names, 'feature_names.pkl')
print("Model saved.")

Building the Flask API

Flask lets you expose your model as an HTTP endpoint in under 50 lines. The API receives JSON, validates input, transforms it, predicts, and returns JSON:

from flask import Flask, request, jsonify
import joblib
import pandas as pd
import numpy as np
from datetime import datetime

app = Flask(__name__)

# Load model once at startup — not on every request
pipeline     = joblib.load('churn_model.pkl')
feature_names = joblib.load('feature_names.pkl')

@app.route('/health', methods=['GET'])
def health():
    # Health check endpoint for load balancers
    return jsonify({'status': 'ok', 'timestamp': datetime.utcnow().isoformat()})

@app.route('/predict', methods=['POST'])
def predict():
    # Predict churn probability.
    # Input:  JSON object with feature values
    # Output: prediction (0 or 1), probability (float), label (str)
    # Parse incoming JSON
    data = request.get_json(force=True)
    if data is None:
        return jsonify({'error': 'Request body must be valid JSON'}), 400

    # Validate that all required features are present
    missing = [f for f in feature_names if f not in data]
    if missing:
        return jsonify({'error': f'Missing features: {missing}'}), 400

    # Build DataFrame with correct column order
    X = pd.DataFrame([{f: data[f] for f in feature_names}])

    # Predict
    try:
        prediction   = int(pipeline.predict(X)[0])
        probabilities = pipeline.predict_proba(X)[0]
        churn_prob   = float(probabilities[1])
    except Exception as e:
        return jsonify({'error': f'Prediction failed: {str(e)}'}), 500

    return jsonify({
        'prediction':   prediction,
        'probability':  round(churn_prob, 4),
        'label':        'churn' if prediction == 1 else 'no churn',
        'confidence':   round(max(probabilities), 4)
    })

@app.route('/predict/batch', methods=['POST'])
def predict_batch():
    # Accept a list of records and predict for all
    data = request.get_json(force=True)
    if not isinstance(data, list):
        return jsonify({'error': 'Expected a JSON array of records'}), 400

    records = []
    for i, record in enumerate(data):
        missing = [f for f in feature_names if f not in record]
        if missing:
            return jsonify({'error': f'Record {i} missing features: {missing}'}), 400
        records.append({f: record[f] for f in feature_names})

    X = pd.DataFrame(records)
    predictions  = pipeline.predict(X).tolist()
    probabilities = pipeline.predict_proba(X)[:, 1].tolist()

    return jsonify([
        {'prediction': int(p), 'probability': round(prob, 4)}
        for p, prob in zip(predictions, probabilities)
    ])

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=False)

Testing Your API

Test the API locally before deployment. Use Python’s requests library or curl:

import requests

url = 'http://localhost:5000/predict'

# Single prediction
payload = {
    'tenure_months': 24,
    'monthly_charges': 65.5,
    'total_charges': 1572.0,
    'contract_type': 1,   # 0=month-to-month, 1=one year, 2=two year
    'tech_support': 0,
    'internet_service': 1
}

response = requests.post(url, json=payload)
print(response.json())
# {'prediction': 0, 'probability': 0.1234, 'label': 'no churn', 'confidence': 0.8766}

Production Considerations

Flask’s built-in development server is not suitable for production — it’s single-threaded and not designed for high concurrency. For production, run Flask behind a WSGI server like Gunicorn: gunicorn -w 4 -b 0.0.0.0:5000 app:app. The -w 4 flag starts 4 worker processes, handling 4 requests simultaneously. For I/O-bound workloads, consider asynchronous frameworks like FastAPI (which is now often preferred over Flask for new ML API projects due to native async support and automatic documentation generation via Swagger). Add request logging, error tracking (Sentry), and rate limiting for any public-facing API.

Frequently Asked Questions

Flask vs FastAPI — which should I use for ML deployment?
FastAPI is the modern choice for new projects — it has automatic OpenAPI documentation, native async support, data validation via Pydantic, and is generally faster than Flask. Flask is still excellent for simple models and is widely understood. If you’re familiar with Flask, it’s a fine choice; if starting fresh, FastAPI is worth learning.

How do I handle model versioning?
Include the model version in your API URL (e.g., /v1/predict, /v2/predict) so clients can pin to a specific version. Store models in a model registry (MLflow, Weights & Biases, or even S3 with versioning). Log every prediction with its input, output, and model version — this is essential for debugging and drift detection.

How do I keep the model updated without downtime?
Load the new model into memory, then atomically swap the reference. Most production systems use a rolling deployment — spin up new instances with the new model, route a small percentage of traffic to them, verify performance, then fully switch over. Container orchestration (Kubernetes) makes this straightforward.