httpx aclose patch

Author: binary-husky

ajet/backbone/warm_up.py

@@ -6,6 +6,8 @@
 import asyncio
 import logging
 import os
+from ajet.utils.async_utils import apply_httpx_aclose_patch
+apply_httpx_aclose_patch()
 
 
 def init_parallel_rollout_logger(experiment_name):

ajet/utils/async_utils.py

@@ -2,7 +2,6 @@
 import concurrent.futures
 from typing import Any
 
-
 def run_async_coroutine_with_timeout(coro, timeout: int = 3600) -> Any:
     """
     Run an async coroutine with a timeout, supporting both inside and outside event loops.
@@ -32,3 +31,40 @@ def run_async_coroutine_with_timeout(coro, timeout: int = 3600) -> Any:
             except Exception:
                 raise
     return final_res
+
+
+def apply_httpx_aclose_patch():
+    try:
+        from openai._base_client import AsyncHttpxClientWrapper
+
+        _original_init = AsyncHttpxClientWrapper.__init__
+
+        def _patched_init(self, *args, **kwargs):
+            try:
+                self._created_loop = asyncio.get_running_loop()
+            except RuntimeError:
+                self._created_loop = None
+            _original_init(self, *args, **kwargs)
+
+        def _patched_del(self) -> None:
+            if self.is_closed:
+                return
+
+            try:
+                current_loop = asyncio.get_running_loop()
+            except RuntimeError:
+                return
+
+            if getattr(self, "_created_loop", None) is not None and current_loop is not self._created_loop:
+                return
+
+            try:
+                current_loop.create_task(self.aclose())
+            except Exception:
+                pass
+
+        AsyncHttpxClientWrapper.__init__ = _patched_init
+        AsyncHttpxClientWrapper.__del__ = _patched_del
+        print("Applied httpx aclose patch.")
+    except ImportError:
+        pass