From 6c35c5f8bff3c475368ca02ce78040ad0ff908ab Mon Sep 17 00:00:00 2001 From: David Traina <44659830+DavidTraina@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:23:16 -0500 Subject: [PATCH] fix(scores): make flush() and shutdown() work under gevent and eventlet _join_score_queue() waited on queue.Queue's all_tasks_done Condition. gevent (25.4.1+) and eventlet replace queue.Queue with classes that have no all_tasks_done, so flush() and shutdown() raised AttributeError and shutdown() left the score consumer running. The score consumer now notifies a Condition owned by the resource manager after task_done(), and flush() waits on that instead, so it works with any queue type. flush() still wakes the consumer, returns as soon as the queue drains, and stops waiting once no consumer is alive. --- langfuse/_client/resource_manager.py | 32 ++++++++----- langfuse/_task_manager/score_consumer.py | 20 +++++--- tests/unit/test_score_consumer.py | 58 ++++++++++++++++++++++++ 3 files changed, 93 insertions(+), 17 deletions(-) diff --git a/langfuse/_client/resource_manager.py b/langfuse/_client/resource_manager.py index e633cdfaa..cf3064473 100644 --- a/langfuse/_client/resource_manager.py +++ b/langfuse/_client/resource_manager.py @@ -376,10 +376,12 @@ def _init_consumer_threads(self) -> None: self._media_upload_consumers.append(media_upload_consumer) self._score_queue: Queue[Any] = Queue(100_000) + self._score_tasks_done = threading.Condition() self._score_consumers = [] score_consumer = ScoreConsumer( score_queue=self._score_queue, + tasks_done=self._score_tasks_done, identifier=0, api_client=self.api, flush_at=self.flush_at, @@ -596,7 +598,7 @@ def flush(self) -> None: for score_consumer in self._score_consumers: if score_consumer.is_alive(): score_consumer.wake() - + self._join_score_queue() langfuse_logger.debug("Successfully flushed score queue") @@ -608,20 +610,28 @@ def _join_score_queue(self) -> None: Queue.join() would block forever when the consumers stop first, e.g. when shutdown() runs concurrently and a wake signal is never taken. + This waits on the consumers' own task_done() signal instead of the + queue's all_tasks_done, which gevent's and eventlet's patched Queues + do not have. """ - score_queue = self._score_queue + with self._score_tasks_done: + while self._should_wait_for_score_queue(): + self._score_tasks_done.wait(timeout=0.1) + + def _should_wait_for_score_queue(self) -> bool: + unfinished_tasks = self._score_queue.unfinished_tasks + if not unfinished_tasks: + return False - with score_queue.all_tasks_done: - while score_queue.unfinished_tasks: - if not any(c.is_alive() for c in self._score_consumers): - langfuse_logger.debug( - "Score consumers stopped with %s queue item(s) unprocessed", - score_queue.unfinished_tasks, - ) + if any(c.is_alive() for c in self._score_consumers): + return True - return + langfuse_logger.debug( + "Score consumers stopped with %s queue item(s) unprocessed", + unfinished_tasks, + ) - score_queue.all_tasks_done.wait(timeout=0.1) + return False def shutdown(self) -> None: self._shutdown = True diff --git a/langfuse/_task_manager/score_consumer.py b/langfuse/_task_manager/score_consumer.py index 21d9f5052..c46ee54e0 100644 --- a/langfuse/_task_manager/score_consumer.py +++ b/langfuse/_task_manager/score_consumer.py @@ -62,6 +62,7 @@ def __init__( flush_at: Optional[int] = None, flush_interval: Optional[float] = None, max_retries: Optional[int] = None, + tasks_done: Optional[threading.Condition] = None, ): """Create a consumer thread.""" super().__init__() @@ -73,6 +74,8 @@ def __init__( # Make consumer a daemon thread so that it doesn't block program exit self.daemon = True self._score_queue = score_queue + # Notified after every task_done(), so flush() can wait on any queue type. + self._tasks_done = threading.Condition() if tasks_done is None else tasks_done self._identifier = identifier self._api_client = api_client @@ -113,7 +116,7 @@ def _next(self) -> List[Dict[str, Any]]: ) if score is _WAKE_SIGNAL: - self._score_queue.task_done() + self._task_done() break try: @@ -124,7 +127,7 @@ def _next(self) -> List[Dict[str, Any]]: "dropped. Error: %s", e, ) - self._score_queue.task_done() + self._task_done() continue @@ -138,7 +141,7 @@ def _next(self) -> List[Dict[str, Any]]: item_size, MAX_BATCH_SIZE_BYTES, ) - self._score_queue.task_done() + self._task_done() continue @@ -167,7 +170,7 @@ def _next(self) -> List[Dict[str, Any]]: str(e), exc_info=True, ) - self._score_queue.task_done() + self._task_done() return scores @@ -209,8 +212,13 @@ def upload(self) -> None: ) finally: # mark items as acknowledged from queue - for _ in batch: - self._score_queue.task_done() + self._task_done(len(batch)) + + def _task_done(self, count: int = 1) -> None: + for _ in range(count): + self._score_queue.task_done() + with self._tasks_done: + self._tasks_done.notify_all() def wake(self) -> None: try: diff --git a/tests/unit/test_score_consumer.py b/tests/unit/test_score_consumer.py index cc6611d29..258d010de 100644 --- a/tests/unit/test_score_consumer.py +++ b/tests/unit/test_score_consumer.py @@ -15,6 +15,7 @@ from opentelemetry.sdk.trace.sampling import ALWAYS_OFF, ParentBased, TraceIdRatioBased from langfuse import Langfuse +from langfuse._client import resource_manager as resource_manager_module from langfuse._client.resource_manager import LangfuseResourceManager from langfuse._task_manager import score_consumer as score_consumer_module from langfuse._task_manager.score_consumer import ScoreConsumer @@ -417,6 +418,24 @@ def test_create_score_posts_to_scores_endpoint(client_factory): assert sent[1]["timestamp"].endswith("Z") +class GeventLikeQueue: + """queue.Queue as gevent's monkey patching replaces it: no all_tasks_done.""" + + def __init__(self, maxsize: int = 0) -> None: + self._queue: Queue = Queue(maxsize) + + def __getattr__(self, name: str) -> Any: + if name == "all_tasks_done": + raise AttributeError(name) + return getattr(self._queue, name) + + +@pytest.fixture(params=[Queue, GeventLikeQueue], ids=["stdlib", "gevent"]) +def queue_class(request, monkeypatch): + monkeypatch.setattr(resource_manager_module, "Queue", request.param) + + +@pytest.mark.usefixtures("queue_class") def test_flush_sends_pending_score_without_waiting_for_flush_interval( client_factory, ): @@ -432,6 +451,7 @@ def test_flush_sends_pending_score_without_waiting_for_flush_interval( assert elapsed < 1 +@pytest.mark.usefixtures("queue_class") def test_flush_after_shutdown_returns(client_factory): client = client_factory(Recorder(_responder(ACCEPTED))) client.shutdown() @@ -443,6 +463,7 @@ def test_flush_after_shutdown_returns(client_factory): assert not flush.is_alive() +@pytest.mark.usefixtures("queue_class") def test_flush_returns_when_consumer_stops_after_liveness_check(client_factory): client = client_factory(Recorder(_responder(ACCEPTED))) consumer = client._resources._score_consumers[0] @@ -472,6 +493,43 @@ def stop_then_wake() -> None: assert not flush_hung +@pytest.mark.usefixtures("queue_class") +def test_flush_returns_when_consumer_stops_while_flush_waits(client_factory): + request_started = threading.Event() + release_request = threading.Event() + accept = _responder(ACCEPTED) + + def respond(request: httpx.Request) -> httpx.Response: + request_started.set() + release_request.wait(timeout=5) + return accept(request) + + client = client_factory(Recorder(respond), flush_interval=5) + consumer = client._resources._score_consumers[0] + score_queue = client._resources._score_queue + client.create_score(name="a", value=1.0, trace_id="a" * 32) + + flush = threading.Thread(target=client.flush, daemon=True) + flush.start() + assert request_started.wait(timeout=5) + # flush() waits on the in-flight score. The consumer then stops without a + # final notify, as when its greenlet is killed, and never takes the item + # queued after it. + consumer.running = False + consumer._tasks_done = threading.Condition() + score_queue.put(_score("orphan")) + release_request.set() + flush.join(timeout=2) + flush_hung = flush.is_alive() + # Acknowledge the orphan so a hung flush() cannot block teardown. + while not score_queue.empty(): + score_queue.get_nowait() + score_queue.task_done() + + assert not consumer.is_alive() + assert not flush_hung + + def test_create_score_payload_matches_request_model_serialization(client_factory): recorder = Recorder(_responder(ACCEPTED)) client = client_factory(recorder, environment="staging")