Callbacks, middleware

Callbacks, middleware

Middleware and callbacks let developers extend LitServe with custom, self-contained code that runs at various points in the request lifecycle, without modifying the core LitAPI.

Background

Middleware

In ASGI (Asynchronous Server Gateway Interface) frameworks like FastAPI, Uvicorn, and Starlette, middleware refers to a layer of components that process requests and responses before they reach the application or after they leave it. LitServe is based on FastAPI and also supports middleware.

Some of the use cases of middlewares are:

• Authentication: Validate user credentials before allowing access to the API.
• Rate limiting: Limit the number of requests a client can make within a certain timeframe.
• Input validation: Check the format and content of incoming requests to ensure they meet certain criteria.
• Response transformation: Modify the response data format before sending it back to the client.

Callbacks

A LitServe callback, (aka: hook), is a way to inject custom behavior into the server lifecycle at various stages, such as at the start or end of setup, predict, etc. Callbacks extend the functionality of the server without modifying the core logic of the LitAPI.

Using middleware in LitServe

Middleware in LitServe enables you to customize the request-response cycle by adding functionality such as logging, authentication, and request transformation. By integrating middleware, you can manage cross-cutting concerns like security and performance without complicating your core application logic.

To integrate middleware with LitServe, pass a list of middlewares as argument to the LitServer class. Here’s an example of integrating a middleware that adds a custom X-Request-Id header to every response:

In this example:

• RequestIdMiddleware is a middleware class that adds a custom header X-Request-Id to every response.
• The middleware is instantiated with a length parameter, which determines the length of the X-Request-Id value.
• When the server is run, each response will include the X-Request-Id header with a value of "00000".

The following client code sends a request to the server and checks the response:

import requests

response = requests.post("http://localhost:8000/predict", json={"input": 4.0})
print(response.json())
print(response.headers["X-Request-Id"])  # This will print "00000"

Example: CORS middleware

from starlette.middleware.cors import CORSMiddleware
import litserve as ls

if __name__ == "__main__":
    # Define your LitAPI
    api = ls.test_examples.SimpleLitAPI()

    # Define the CORS settings
    cors_middleware = (
        CORSMiddleware,
        {
            "allow_origins": ["*"],  # Allows all origins
            "allow_methods": ["GET", "POST"],  # Allows GET and POST methods
            "allow_headers": ["*"],  # Allows all headers
        }
    )

    # Create the LitServer and pass the CORS middleware
    server = ls.LitServer(api, middlewares=[cors_middleware])

    # Run the server on port 8000
    server.run(port=8000)

You can further customize the CORS settings by modifying the allow_origins, allow_methods, allow_headers, allow_credentials, expose_headers, and max_age parameters in the CORSMiddleware.

For instance, if you only want to allow requests from a specific domain, you could set: "allow_origins": ["https://example.com"]

Callbacks with LitServe

LitServe trigger callbacks in the following events:

• on_server_start: Triggered when the server starts.
• on_serve_end: Triggered when the server stops serving.
• on_before_setup: Triggered before the setup process begins.
• on_after_setup: Triggered after the setup process completes.
• on_before_decode_request: Triggered before a request is decoded.
• on_after_decode_request: Triggered after a request has been decoded.
• on_before_predict: Triggered before a prediction operation starts.
• on_after_predict: Triggered after a prediction operation completes.
• on_before_encode_response: Triggered before a response is encoded.
• on_after_encode_response: Triggered after a response has been encoded.
• on_request: Triggered on receiving a request in the endpoint function.
• on_response: Triggered once the response is created from the inference worker and just before returning the response to the client.

(The programming equivalent of these events are called hooks.)

To create a callback, LitServe provides a Callback class that includes the necessary hooks to be triggered. For example, to print the model's prediction time, you can implement the on_before_predict and on_after_predict hooks and display the measured time as follows:

import time
import litserve as ls

class PredictionTimeLogger(ls.Callback):
    def on_before_predict(self, lit_api):
        t0 = time.perf_counter()
        self._start_time = t0

    def on_after_predict(self, lit_api):
        t1 = time.perf_counter()
        elapsed = t1 - self._start_time
        # log/print or store in a database
        print(f"Prediction took {elapsed:.2f} seconds", flush=True)

To use this callback with the server, pass an instance of the callback to the LitServer class as follows:

import time
import litserve as ls

class PredictionTimeLogger(ls.Callback):
    def on_before_predict(self, lit_api):
        t0 = time.perf_counter()
        self._start_time = t0

    def on_after_predict(self, lit_api):
        t1 = time.perf_counter()
        elapsed = t1 - self._start_time
        print(f"Prediction took {elapsed:.2f} seconds", flush=True)

if __name__ == '__main__':
    lit_api = ls.test_examples.SimpleLitAPI()
    server = ls.LitServer(lit_api, callbacks=[PredictionTimeLogger()])
    server.run()

Conclusion

Middleware and callbacks are powerful tools in LitServe that enable developers to fine-tune the behavior of their servers. By understanding and utilizing these features, you can implement robust, flexible, and scalable APIs tailored to your specific needs. Whether it's modifying requests on the fly, handling errors gracefully, or performing custom actions at server startup, middleware and callbacks provide the hooks you need to build advanced server-side functionality in LitServe.