Flask
Machine Learning
Model Caching
Web Development
Python

How to cache a large machine learning model in Flask?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

The main rule is to load a large machine-learning model once per worker process, not once per request. In Flask, the most effective "cache" for the model itself is usually process memory, with a module-level singleton or app-factory initialization pattern. External caches such as Redis are great for prediction results and metadata, but they are usually the wrong place to store a large in-memory model object directly.

Why Per-Request Loading Is A Problem

If every HTTP request reloads the model from disk, startup cost dominates inference time and your latency becomes terrible. Large TensorFlow, PyTorch, or scikit-learn models may take hundreds of milliseconds or even seconds to deserialize.

That means the first optimization is not fancy caching infrastructure. It is simply moving model loading out of the request handler.

A Simple Flask Pattern

A common approach is a module-level lazy singleton.

python
1from flask import Flask, jsonify, request
2import joblib
3
4app = Flask(__name__)
5_model = None
6
7
8def get_model():
9    global _model
10    if _model is None:
11        _model = joblib.load("model.joblib")
12    return _model
13
14
15@app.route("/predict", methods=["POST"])
16def predict():
17    model = get_model()
18    payload = request.get_json()
19    features = [payload["x1"], payload["x2"]]
20    prediction = model.predict([features])[0]
21    return jsonify({"prediction": float(prediction)})

The first request pays the load cost. Later requests reuse the already loaded model inside that worker process.

Better Still: Load At Worker Startup

If you do not want the first request to pay the model-load penalty, initialize the model when the process starts.

python
1from flask import Flask
2import joblib
3
4app = Flask(__name__)
5model = joblib.load("model.joblib")

This works well when you know the model should always be present and the worker has enough memory to keep it resident.

The tradeoff is that each worker process gets its own copy. That is usually acceptable, but it matters for very large models.

Process Model Matters

With Gunicorn or uWSGI, Flask usually runs with multiple worker processes. Each process has its own memory space, so the model is cached per worker, not globally across the entire server.

That means:

  • one worker: one in-memory model instance
  • four workers: typically four in-memory model instances

This is often still the right tradeoff because inference becomes fast and simple. But you should size the worker count around the model's memory footprint rather than around generic web-server defaults.

Cache Results Separately From The Model

If many requests repeat the same inputs, cache prediction outputs separately. That is where Redis, Memcached, or Flask-Caching is useful.

python
1from flask import Flask, jsonify, request
2from flask_caching import Cache
3import joblib
4
5app = Flask(__name__)
6app.config["CACHE_TYPE"] = "SimpleCache"
7cache = Cache(app)
8model = joblib.load("model.joblib")
9
10
11@app.route("/predict", methods=["POST"])
12def predict():
13    payload = request.get_json()
14    key = f"pred:{payload['x1']}:{payload['x2']}"
15
16    cached = cache.get(key)
17    if cached is not None:
18        return jsonify({"prediction": cached, "cached": True})
19
20    prediction = float(model.predict([[payload["x1"], payload["x2"]]])[0])
21    cache.set(key, prediction, timeout=300)
22    return jsonify({"prediction": prediction, "cached": False})

This speeds up repeated requests without trying to serialize and store the model object itself in the cache backend.

When Flask Is Not Enough

If the model is very large, GPU-backed, or requires tight concurrency control, a dedicated model server may be cleaner than stuffing everything into a Flask app. Tools such as TorchServe, TensorFlow Serving, or a separate inference microservice are often better when model serving becomes the primary responsibility rather than a side feature.

Flask is fine for many cases, but it is not magic infrastructure for arbitrarily heavy inference workloads.

Common Pitfalls

  • Loading the model inside the request handler on every call.
  • Treating Redis or Memcached as the place to store the raw model object instead of keeping it in process memory.
  • Forgetting that each worker process holds its own copy of the model.
  • Choosing a worker count without accounting for the model's RAM or GPU footprint.
  • Mixing model caching and prediction-result caching as if they were the same problem.

Summary

  • Cache the model by loading it once per Flask worker process, not once per request.
  • A module-level singleton or startup-time load is usually enough.
  • Use external caches mainly for repeated prediction results, not for the model object itself.
  • Size your worker count around the model's memory cost.
  • If the serving problem becomes large enough, move to a dedicated model-serving architecture instead of stretching Flask too far.

Course illustration
Course illustration

All Rights Reserved.