refactor

src/fastserve/batching.py

@@ -6,7 +6,9 @@
 from threading import Event, Thread
 from typing import Any, Callable, Dict, List
 
-from loguru import logger
+from logging import Logger, INFO
+
+logger = Logger(__name__, level=INFO)
 
 
 class BatchedQueue:

src/fastserve/fastserve.py

@@ -1,3 +1,4 @@
+from contextlib import asynccontextmanager
 from typing import Any, List
 
 from fastapi import FastAPI
@@ -12,15 +13,18 @@ class BaseRequest(BaseModel):
 
 class FastServe:
     def __init__(self, batch_size=2, timeout=0.5) -> None:
-        self.batch_processing = BatchProcessor(
-            func=self.handle, bs=batch_size, timeout=timeout
-        )
-        self._app = FastAPI()
-
-        @self._app.on_event("shutdown")
-        def shutdown_event():
+        self.batch_processing: BatchProcessor = None
+
+        @asynccontextmanager
+        async def lifespan(app: FastAPI):
+            self.batch_processing = BatchProcessor(
+                func=self.handle, bs=batch_size, timeout=timeout
+            )
+            yield
             self.batch_processing.cancel()
 
+        self._app = FastAPI(lifespan=lifespan)
+
     def serve(
         self,
     ):