fix: await

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(

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:

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
 

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