fix: make LLMDescriptionScorer work inside a running event loop

src/autointent/modules/scoring/_description/llm_encoder.py

@@ -5,10 +5,11 @@
 import asyncio
 import json
 import logging
+import threading
 from functools import partial
 from pathlib import Path
 from textwrap import dedent
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TypeVar
 
 import aiometer
 import numpy as np
@@ -23,12 +24,17 @@
 from .base import BaseDescriptionScorer
 
 if TYPE_CHECKING:
+    from collections.abc import Coroutine
+
     from numpy.typing import NDArray
 
     from autointent.configs import CrossEncoderConfig, EmbedderConfig
 
 logger = logging.getLogger(__name__)
 
+
+_T = TypeVar("_T")
+
 GENERATOR_CONFIG_FILENAME = "generator_config.json"
 
 
@@ -249,7 +255,7 @@ def _compute_similarities(self, utterances: list[str]) -> NDArray[np.float64]:
                 max_at_once=self.max_concurrent,
                 max_per_second=self.max_per_second,
             )
-            categorizations = self._event_loop.run_until_complete(task)  # type: ignore[arg-type]
+            categorizations = self._run_async(task)  # type: ignore[arg-type]
 
         for i, categorization in enumerate(categorizations):
             if isinstance(categorization, IntentCategorization):
@@ -276,19 +282,38 @@ def clear_cache(self) -> None:
         # Generator doesn't have a clear_ram method, so we just set it to None
         if hasattr(self, "_generator"):
             delattr(self, "_generator")
-        if hasattr(self, "_event_loop"):
-            delattr(self, "_event_loop")
+        self._close_event_loop()
 
     def _init_event_loop(self) -> None:
+        # A private loop, never the caller's: the generator's async client keeps its
+        # connections bound to the loop it first ran on, so all batches reuse this one.
+        # The loop runs forever on one dedicated thread, so it is only ever driven from
+        # that thread, whichever thread (sync or async, one or several) submits work.
+        self._close_event_loop()
         if self.max_concurrent is not None:
-            try:
-                loop = asyncio.get_event_loop()
-            except RuntimeError:
-                loop = asyncio.new_event_loop()
-            else:
-                if loop.is_closed():
-                    loop = asyncio.new_event_loop()
-            self._event_loop = loop
+            self._event_loop = asyncio.new_event_loop()
+            self._event_loop_thread = threading.Thread(
+                target=self._event_loop.run_forever, name="LLMDescriptionScorer-loop", daemon=True
+            )
+            self._event_loop_thread.start()
+
+    def _close_event_loop(self) -> None:
+        if hasattr(self, "_event_loop"):
+            if hasattr(self, "_event_loop_thread"):
+                self._event_loop.call_soon_threadsafe(self._event_loop.stop)
+                self._event_loop_thread.join()
+                delattr(self, "_event_loop_thread")
+            self._event_loop.close()
+            delattr(self, "_event_loop")
+
+    def _run_async(self, coro: Coroutine[Any, Any, _T]) -> _T:
+        """Run ``coro`` on the scorer's loop thread and block until it finishes.
+
+        Works the same whether or not the caller is already inside a running event loop
+        (a notebook, an async web handler), and concurrent callers are safe: the loop
+        schedules their batches itself instead of being re-entered from several threads.
+        """
+        return asyncio.run_coroutine_threadsafe(coro, self._event_loop).result()
 
     def dump(self, path: str) -> None:
         dump_path = Path(path)
@@ -298,7 +323,7 @@ def dump(self, path: str) -> None:
         # it via GeneratorDumper — load() recreates it from generator_config.
         generator = self.__dict__.pop("_generator", None)
         try:
-            Dumper.dump(self, dump_path, exclude=[asyncio.BaseEventLoop, dict])
+            Dumper.dump(self, dump_path, exclude=[asyncio.BaseEventLoop, threading.Thread, dict])
         finally:
             if generator is not None:
                 self._generator = generator

tests/modules/scoring/test_description_llm.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import asyncio
 import tempfile
+from concurrent.futures import ThreadPoolExecutor
 from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
@@ -112,3 +114,63 @@ def test_llm_description_in_pipeline(dataset: Dataset, patch_llm_scorer_generato
     pipeline.fit(dataset)
     predictions = pipeline.predict(["test utterance"])
     assert len(predictions) == 1
+
+
+def _fit_llm_scorer(dataset: Dataset) -> LLMDescriptionScorer:
+    data_handler = DataHandler(dataset)
+    scorer = LLMDescriptionScorer(generator_config={"temperature": 0})
+    labels = data_handler.train_labels(0)
+    assert is_strict_labels(labels)
+    descriptions = data_handler.intent_descriptions
+    assert all(d is not None for d in descriptions)
+    scorer.fit(data_handler.train_utterances(0), labels, cast("list[str]", descriptions))
+    return scorer
+
+
+def test_description_scorer_llm_predict_inside_running_loop(
+    dataset: Dataset, patch_llm_scorer_generator: Generator
+) -> None:
+    """predict() is sync, but must also work when called from async code (notebooks, async servers)."""
+    scorer = _fit_llm_scorer(dataset)
+    utterances = ["What is the balance on my account?", "How do I reset my online banking password?"]
+    expected = scorer.predict(utterances)
+
+    async def predict_from_coroutine() -> npt.NDArray[Any]:
+        return scorer.predict(utterances)
+
+    np.testing.assert_array_equal(asyncio.run(predict_from_coroutine()), expected)
+    # the scorer's loop is still usable from sync code afterwards
+    np.testing.assert_array_equal(scorer.predict(utterances), expected)
+
+
+def test_description_scorer_llm_closes_its_event_loop(dataset: Dataset, patch_llm_scorer_generator: Generator) -> None:
+    scorer = _fit_llm_scorer(dataset)
+    first_loop = scorer._event_loop
+
+    first_thread = scorer._event_loop_thread
+
+    scorer.fit([], [], scorer._description_texts)
+    assert first_loop.is_closed()
+    assert not first_thread.is_alive()
+    assert scorer._event_loop is not first_loop
+
+    second_loop, second_thread = scorer._event_loop, scorer._event_loop_thread
+    scorer.clear_cache()
+    assert second_loop.is_closed()
+    assert not second_thread.is_alive()
+    assert not hasattr(scorer, "_event_loop")
+
+
+def test_description_scorer_llm_predict_from_several_threads(
+    dataset: Dataset, patch_llm_scorer_generator: Generator
+) -> None:
+    """Concurrent predict() calls share the scorer's loop without re-entering it from several threads."""
+    scorer = _fit_llm_scorer(dataset)
+    utterances = ["What is the balance on my account?", "How do I reset my online banking password?"]
+    expected = scorer.predict(utterances)
+
+    with ThreadPoolExecutor(max_workers=4) as executor:
+        results = list(executor.map(lambda _: scorer.predict(utterances), range(8)))
+
+    for result in results:
+        np.testing.assert_array_equal(result, expected)