From 4ca525e0a5245d1be12e6231f8f68191fae8bd2a Mon Sep 17 00:00:00 2001 From: Trayan Azarov Date: Tue, 23 Jan 2024 22:56:34 +0200 Subject: [PATCH] [ENH]: FastAPI Shutdown hook (#1665) ## Description of changes *Summarize the changes made by this PR.* - Improvements & Bug fixes - On shutdown event system.stop() is called to properly stop all components ## Test plan *How are these changes tested?* - [x] Tests pass locally with `pytest` for python ## Documentation Changes N/A --- chromadb/server/fastapi/__init__.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/chromadb/server/fastapi/__init__.py b/chromadb/server/fastapi/__init__.py index f01f4137908..529606a6c36 100644 --- a/chromadb/server/fastapi/__init__.py +++ b/chromadb/server/fastapi/__init__.py @@ -141,6 +141,8 @@ def __init__(self, settings: Settings): allow_methods=["*"], ) + self._app.on_event("shutdown")(self.shutdown) + if settings.chroma_server_authz_provider: self._app.add_middleware( FastAPIChromaAuthzMiddlewareWrapper, @@ -280,6 +282,9 @@ def __init__(self, settings: Settings): use_route_names_as_operation_ids(self._app) instrument_fastapi(self._app) + def shutdown(self) -> None: + self._system.stop() + def app(self) -> fastapi.FastAPI: return self._app