Loading Scikit-learn Models in FastAPI Startup

Why Load a Model on Startup?

Imagine your API receives 100 prediction requests per second. If the model takes 200 milliseconds to load from the disk, loading it for every request would be a disaster. The API would be sluggish and waste enormous amounts of CPU and I/O resources reloading the same data over and over.

The solution is to load the model into the application's memory a single time, right when the server process begins. This way, the model is always ready, and prediction requests can be handled with minimal latency.

The official FastAPI documentation has a great section that introduces this exact use case.

The FastAPI Way: Lifespan Events

To manage tasks that should run at startup and shutdown, FastAPI provides a powerful and elegant feature called lifespan events. This allows you to define a block of code that executes before the application starts accepting requests (for setup) and another block that executes after the application has finished (for cleanup).

For our purposes, the "setup" phase is the perfect place to load our ML model.

The recommended way to implement this is with Python's asynccontextmanager. While the name might seem complex, the concept is straightforward: it creates a managed context for our application's life.

1. Code before yield: This is our startup logic. It runs once, before any requests are processed. We will load our model here.
2. yield: This keyword passes control to the application, which then runs and serves requests for its entire life.
3. Code after yield: This is our shutdown logic. It runs once, after the server is stopped. We can use this to clean up resources.

Implementing a Lifespan Function to Load a Model

Let's see how this looks in practice. We will create a lifespan function that loads a pre-trained model saved with a library like pickle or joblib.

The official documentation provides the definitive guide to implementing this.

Hands-On Practice

Now it's your turn to apply this. Let's assume we have a pre-trained scikit-learn model saved in a file named model.pkl. Your task is to write the code to load it on application startup.

For this exercise, you don't need a real .pkl file. We'll simulate the loading process, just as the documentation did with fake_answer_to_everything_ml_model.

Here is a starting main.py:

# main.py

from fastapi import FastAPI
from contextlib import asynccontextmanager
import pickle




# Let's pretend this is our complex, pre-trained model object
# In a real scenario, this would be loaded from a file.
class MyMLModel:
    def predict(self, input_data: float) -> float:



        # A very sophisticated AI model
        return input_data * 2.0




# 1. Create a dictionary to hold your model.
# This makes it accessible to your endpoints later.
ml_models = {}




# 2. Create the lifespan context manager.
@asynccontextmanager
async def lifespan(app: FastAPI):



    # --- Code to run on startup ---
    # Load the "model" into the ml_models dictionary.
    # We will simulate loading from a file.
    



    # In a real project, you would do something like:
    # with open("model.pkl", "rb") as f:
    #     ml_models["my_model"] = pickle.load(f)
    



    # For this exercise, just instantiate our fake model:
    ml_models["my_model"] = MyMLModel()
    print("AI Model loaded!")
    
    yield
    



    # --- Code to run on shutdown ---
    # Clean up the models and release the resources
    ml_models.clear()
    print("AI Model cleared.")





# 3. Create the FastAPI app instance, passing in the lifespan manager.
app = FastAPI(lifespan=lifespan)




# 4. A root endpoint to confirm the app is running.
@app.get("/")
def read_root():
    return {"Hello": "World"}




# In the next lesson, we will create an endpoint that USES the model.
# For example:
# @app.post("/predict")
# def predict(data: float):
#     prediction = ml_models["my_model"].predict(data)
#     return {"prediction": prediction}

Copy this code into your main.py file. Run your application with Uvicorn:

uvicorn main:app --reload

When the server starts, you should see the AI Model loaded! message printed to your console. This confirms that your startup logic ran correctly. When you stop the server (with Ctrl+C), you should see AI Model cleared., confirming your shutdown logic also ran.

A Note on Other Methods

As you explore other tutorials, you might see different ways to load models. It's useful to recognize them:

1. Loading at the module level: Some code simply runs model = pickle.load(...) at the top of the main.py file. This works, but it's less flexible. The code runs as soon as the file is imported, which can slow down things like running simple unit tests that don't even need the model. The video from AssemblyAI (Deploy ML models with FastAPI...) demonstrates this simpler approach.
2. Deprecated on_event handlers: Older versions of FastAPI used decorators like @app.on_event("startup"). The lifespan context manager is now the recommended approach because it neatly groups setup and cleanup logic together. You can see this mentioned in the "Alternative Events (deprecated)" section of the Lifespan Events documentation you read earlier.

While these methods work, using the lifespan context manager is the modern, robust, and officially recommended best practice.

Conclusion

Fantastic work! You have successfully implemented one of the most important performance optimizations for serving machine learning models. By loading your model on startup, you ensure your API is fast, efficient, and ready to handle prediction requests.

Key Takeaways:

• Loading ML models or other heavy resources for every request is highly inefficient.
• FastAPI's lifespan event handler is the recommended way to run code once on application startup and shutdown.
• Using an asynccontextmanager provides a clean structure for defining startup logic (before yield) and cleanup logic (after yield).
• Store the loaded model in a globally accessible object (like a dictionary) so your endpoint functions can use it.

In our next lesson, we will build on this foundation. We will define Pydantic schemas for our model's input features and its prediction output, and then create the /predict endpoint that uses the model you just loaded into memory.