diff --git a/src/fastserve/batching.py b/src/fastserve/batching.py index b3e7b31..9ade6e2 100644 --- a/src/fastserve/batching.py +++ b/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: diff --git a/src/fastserve/fastserve.py b/src/fastserve/fastserve.py index 2e3acf8..6e8de65 100644 --- a/src/fastserve/fastserve.py +++ b/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, ):