diff --git a/src/autointent/modules/scoring/_description/llm_encoder.py b/src/autointent/modules/scoring/_description/llm_encoder.py index 41adaddba..53fca3e2f 100644 --- a/src/autointent/modules/scoring/_description/llm_encoder.py +++ b/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 diff --git a/tests/modules/scoring/test_description_llm.py b/tests/modules/scoring/test_description_llm.py index 3f678a4f4..ab7332f77 100644 --- a/tests/modules/scoring/test_description_llm.py +++ b/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)