From 842d09973abf21a727975f008ed33a86875e14a1 Mon Sep 17 00:00:00 2001 From: Yahya Kayaal <117476621+kayaal34@users.noreply.github.com> Date: Sat, 19 Sep 2026 15:44:12 +0300 Subject: [PATCH] fix: make LLMDescriptionScorer work inside a running event loop The scorer ran its async batch with `run_until_complete` on a loop taken from `asyncio.get_event_loop()`. Called from async code (notebook, async web handler) that returned the running loop and raised "This event loop is already running". On Python 3.14 `get_event_loop()` raises when no loop is set, so every instance created a new loop that `clear_cache` never closed, leaking one loop per HPO trial. Each scorer now owns a private loop, which the generator's async client stays bound to. Batches run on it directly from sync code, or on a worker thread when the caller is already inside a running loop. The loop is closed on refit and in `clear_cache`. Closes #353 --- .../scoring/_description/llm_encoder.py | 41 ++++++++++++------ tests/modules/scoring/test_description_llm.py | 42 +++++++++++++++++++ 2 files changed, 71 insertions(+), 12 deletions(-) diff --git a/src/autointent/modules/scoring/_description/llm_encoder.py b/src/autointent/modules/scoring/_description/llm_encoder.py index 41adaddba..47116ef79 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 +from concurrent.futures import ThreadPoolExecutor 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,30 @@ 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. + 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() + + def _close_event_loop(self) -> None: + if hasattr(self, "_event_loop"): + self._event_loop.close() + delattr(self, "_event_loop") + + def _run_async(self, coro: Coroutine[Any, Any, _T]) -> _T: + """Run ``coro`` to completion on the scorer's loop, whether or not the caller is already in one.""" + try: + asyncio.get_running_loop() + except RuntimeError: + return self._event_loop.run_until_complete(coro) + # Called from async code (notebook, async web handler): this thread's loop is busy, + # so drive the scorer's loop from a worker thread and block until it finishes. + with ThreadPoolExecutor(max_workers=1) as executor: + return executor.submit(self._event_loop.run_until_complete, coro).result() def dump(self, path: str) -> None: dump_path = Path(path) diff --git a/tests/modules/scoring/test_description_llm.py b/tests/modules/scoring/test_description_llm.py index 3f678a4f4..fbc04e458 100644 --- a/tests/modules/scoring/test_description_llm.py +++ b/tests/modules/scoring/test_description_llm.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import tempfile from typing import TYPE_CHECKING, Any, cast @@ -112,3 +113,44 @@ 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 + + scorer.fit([], [], scorer._description_texts) + assert first_loop.is_closed() + assert scorer._event_loop is not first_loop + + second_loop = scorer._event_loop + scorer.clear_cache() + assert second_loop.is_closed() + assert not hasattr(scorer, "_event_loop")