litserve 的callback 以及中间件简单说明

对于中间件实际上就是fastapi 内部或者starlette的一些能力，使用比较简单，callback 比较有用

中间件使用

• 参考

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)

callback 使用

对于callback 目前提供了不少

• 当前提供的

on_server_start: 服务器启动时触发。
on_serve_end: 服务器停止服务时触发。
on_before_setup: 在设置过程开始之前触发。
on_after_setup: 在设置过程完成后触发。
on_before_decode_request: 在请求解码之前触发。
on_after_decode_request: 在请求解码后触发。
on_before_predict: 在预测操作开始之前触发。
on_after_predict: 在预测操作完成后触发。
on_before_encode_response: 在响应编码之前触发。
on_after_encode_response: 在响应编码后触发。
on_request: 在端点函数接收到请求时触发。
on_response: 在推理工作者创建响应并且在返回响应给客户端之前触发。

• 参考使用

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()

说明

callback 以及中间件还是比较有用的，可以提升系统的灵活性以及实现一些请求记录

参考资料

https://lightning.ai/docs/litserve/features/callbacks-middleware