fix(scores): make flush() and shutdown() work under gevent and eventlet

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

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:

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")