Move _run_coro to base class, prevent LiteLLM logging warnings by introducing the AsyncLoopRunner

Author: Simba

adt_eval/mlflow_base.py

@@ -2,16 +2,66 @@
 
 from __future__ import annotations
 
+import asyncio
 import inspect
+import threading
 from abc import abstractmethod
 from datetime import datetime
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 import mlflow
 
 from adt_eval.base import BaseEvaluator
 
 
+class AsyncLoopRunner:
+    """Run coroutines on a long-lived event loop in a background thread."""
+
+    def __init__(self) -> None:
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._thread: Optional[threading.Thread] = None
+        self._ready = threading.Event()
+
+    def start(self) -> None:
+        if self._thread and self._thread.is_alive():
+            return
+
+        def _runner() -> None:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            self._loop = loop
+            self._ready.set()
+            try:
+                loop.run_forever()
+            finally:
+                try:
+                    pending = asyncio.all_tasks(loop)
+                    for task in pending:
+                        task.cancel()
+                    if pending:
+                        loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+                finally:
+                    loop.close()
+
+        self._thread = threading.Thread(target=_runner, daemon=True)
+        self._thread.start()
+        self._ready.wait()
+
+    def submit(self, coro):
+        if not self._loop:
+            raise RuntimeError("AsyncLoopRunner not started")
+        future = asyncio.run_coroutine_threadsafe(coro, self._loop)
+        return future.result()
+
+    def close(self) -> None:
+        if not self._loop or not self._thread:
+            return
+        self._loop.call_soon_threadsafe(self._loop.stop)
+        self._thread.join()
+        self._loop = None
+        self._thread = None
+
+
 class MLflowEvaluatorBase(BaseEvaluator):
     """Base evaluator that wraps the core run in an MLflow run."""
 
@@ -77,6 +127,11 @@ def log_run_metrics(self, metrics: Dict[str, Any]) -> None:
         if score is not None:
             mlflow.log_metric("score", score)
 
+    def _run_coro(self, coro):
+        if not hasattr(self, "_loop_runner") or self._loop_runner is None:
+            raise RuntimeError("Async loop runner not initialized")
+        return self._loop_runner.submit(coro)
+
     def get_report_results_and_metrics(self, eval_results) -> tuple[List[Dict[str, Any]], Dict[str, Any]]:
         """Return report-ready results and metrics from mlflow.genai.evaluate output."""
         metrics = {}
@@ -96,21 +151,26 @@ async def run(self):
         self.configure_mlflow()
         run_name = self.get_run_name()
         nested = mlflow.active_run() is not None
-
-        with mlflow.start_run(run_name=run_name, nested=nested):
-            self.log_run_params()
-            cases = self.filter_cases(self.load_data())
-            eval_dataset = self.build_eval_dataset(cases)
-            eval_results = mlflow.genai.evaluate(
-                data=eval_dataset,
-                predict_fn=self.predict_fn,
-                scorers=self.get_scorers(),
-                **self._get_evaluate_kwargs(),
-            )
-
-            results, metrics = self.get_report_results_and_metrics(eval_results)
-
-            self.log_run_metrics(metrics)
-            if results and metrics:
-                self.generate_report(results, metrics)
-            return results, metrics
+        self._loop_runner = AsyncLoopRunner()
+        self._loop_runner.start()
+        try:
+            with mlflow.start_run(run_name=run_name, nested=nested):
+                self.log_run_params()
+                cases = self.filter_cases(self.load_data())
+                eval_dataset = self.build_eval_dataset(cases)
+                eval_results = mlflow.genai.evaluate(
+                    data=eval_dataset,
+                    predict_fn=self.predict_fn,
+                    scorers=self.get_scorers(),
+                    **self._get_evaluate_kwargs(),
+                )
+
+                results, metrics = self.get_report_results_and_metrics(eval_results)
+
+                self.log_run_metrics(metrics)
+                if results and metrics:
+                    self.generate_report(results, metrics)
+                return results, metrics
+        finally:
+            self._loop_runner.close()
+            self._loop_runner = None

adt_eval/text_type.py

@@ -84,7 +84,6 @@ def text_type_per_page_scorer(inputs: Dict[str, Any], outputs: Dict[str, Any]) -
         rationale=json.dumps(matches),
     )
 
-
 class TextTypeEvaluator(MLflowEvaluatorBase):
     """Evaluator for text type accuracy."""
 
@@ -155,49 +154,6 @@ def build_eval_dataset(self, cases: List[Dict[str, Any]]) -> List[Dict[str, Any]
             )
         return records
 
-    def _run_coro(self, coro):
-        try:
-            loop = asyncio.get_running_loop()
-        except RuntimeError:
-            loop = None
-
-        if loop and loop.is_running():
-            # If an event loop is already running (e.g., in notebooks), offload to a thread
-            # and wait for the result to avoid nested loop issues.
-            result_container: Dict[str, Any] = {}
-            error_container: Dict[str, BaseException] = {}
-
-            def _runner():
-                try:
-                    result_container["value"] = asyncio.run(coro)
-                except BaseException as exc:  # pragma: no cover - re-raise below
-                    error_container["error"] = exc
-
-            thread = threading.Thread(target=_runner, daemon=True)
-            thread.start()
-            thread.join()
-
-            if "error" in error_container:
-                raise error_container["error"]
-            return result_container.get("value")
-
-        result_container: Dict[str, Any] = {}
-        error_container: Dict[str, BaseException] = {}
-
-        def _runner():
-            try:
-                result_container["value"] = asyncio.run(coro)
-            except BaseException as exc:  # pragma: no cover - re-raise below
-                error_container["error"] = exc
-
-        thread = threading.Thread(target=_runner, daemon=True)
-        thread.start()
-        thread.join()
-
-        if "error" in error_container:
-            raise error_container["error"]
-        return result_container.get("value")
-
     def predict_fn(self, **inputs: Any) -> Dict[str, Any]:
         page_text = inputs["page_text"]
         page_image_path = inputs["page_image_path"]