diff --git a/private_gpt/arq/worker.py b/private_gpt/arq/worker.py index 9117966e..c248154a 100644 --- a/private_gpt/arq/worker.py +++ b/private_gpt/arq/worker.py @@ -33,7 +33,7 @@ async def startup(ctx: dict[Any, Any]) -> None: async def shutdown(ctx: dict[Any, Any]) -> None: del ctx - clean_global_injector() + await clean_global_injector() async def run_chat_job( diff --git a/private_gpt/celery/base.py b/private_gpt/celery/base.py index 4643c9a4..ae513874 100644 --- a/private_gpt/celery/base.py +++ b/private_gpt/celery/base.py @@ -271,7 +271,11 @@ class StatelessBackgroundTask(_BackgroundTask): raise e finally: if loop: - clean_global_injector(loop) + if not loop.is_closed(): + try: + asyncio.run_coroutine_threadsafe(clean_global_injector(loop), loop).result(timeout=10) + except Exception: + pass with contextlib.suppress(RuntimeError): loop.call_soon_threadsafe(loop.stop) if thr: @@ -369,7 +373,7 @@ class StatefulBackgroundTask(_BackgroundTask): @classmethod async def _shutdown_async(cls) -> None: - clean_global_injector() + await clean_global_injector() @classmethod def shutdown_runtime(cls) -> None: diff --git a/private_gpt/di.py b/private_gpt/di.py index 8fd75e87..b8ceed8a 100644 --- a/private_gpt/di.py +++ b/private_gpt/di.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import inspect import logging import threading from asyncio import AbstractEventLoop @@ -97,7 +98,7 @@ def set_injector(injector: Injector) -> None: _global_injector = injector -def clean_global_injector(loop: AbstractEventLoop | None = None) -> None: +async def clean_global_injector(loop: AbstractEventLoop | None = None) -> None: try: loop = loop or asyncio.get_running_loop() if not hasattr(loop, _INJECTOR_KEY): @@ -111,9 +112,12 @@ def clean_global_injector(loop: AbstractEventLoop | None = None) -> None: with contextlib.suppress(Exception): impl: Any = injector.get(interface) if hasattr(impl, "close"): - impl.close() + res = impl.close() + if inspect.isawaitable(res): + await res with _loop_injector_lock: - delattr(loop, _INJECTOR_KEY) + if hasattr(loop, _INJECTOR_KEY): + delattr(loop, _INJECTOR_KEY) except RuntimeError: pass diff --git a/tests/test_di.py b/tests/test_di.py index ff617074..0fd6456f 100644 --- a/tests/test_di.py +++ b/tests/test_di.py @@ -159,7 +159,7 @@ def test_clean_global_injector() -> None: old_global_injector.get(NodeStoreComponent) running_loop = asyncio.get_running_loop() - clean_global_injector(running_loop) + await clean_global_injector(running_loop) new_injector = get_global_injector() assert new_injector is not None