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

langfuse/_client/resource_manager.py

@@ -487,6 +487,13 @@ def add_score_task(self, score: dict, *, force_sample: bool = False) -> None:
                 (for example ``traceId``).
             force_sample: Enqueue the score even if its trace is sampled out.
         """
+        if self._shutdown:
+            langfuse_logger.warning(
+                "Score dropped: Langfuse client has been shut down."
+            )
+
+            return
+
         try:
             trace_id = score.get("traceId")
             tracer_provider = (
@@ -596,33 +603,13 @@ def flush(self) -> None:
         for score_consumer in self._score_consumers:
             if score_consumer.is_alive():
                 score_consumer.wake()
-                
-        self._join_score_queue()
+
+        self._score_queue.join()
         langfuse_logger.debug("Successfully flushed score queue")
 
         self._media_upload_queue.join()
         langfuse_logger.debug("Successfully flushed media upload queue")
 
-    def _join_score_queue(self) -> None:
-        """Wait until the score queue is drained or no consumer is left to drain it.
-
-        Queue.join() would block forever when the consumers stop first, e.g.
-        when shutdown() runs concurrently and a wake signal is never taken.
-        """
-        score_queue = self._score_queue
-
-        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,
-                    )
-
-                    return
-
-                score_queue.all_tasks_done.wait(timeout=0.1)
-
     def shutdown(self) -> None:
         self._shutdown = True
 

langfuse/_task_manager/score_consumer.py

@@ -36,6 +36,7 @@
         _SCORES_BODY_LIMIT_BYTES,
     )
 MAX_BATCH_SIZE_BYTES = min(_configured_batch_size_bytes, _SCORES_BODY_LIMIT_BYTES)
+# Ends the batch being collected. Never counted as an unfinished queue task.
 _WAKE_SIGNAL = object()
 
 
@@ -113,7 +114,6 @@ def _next(self) -> List[Dict[str, Any]]:
                 )
 
                 if score is _WAKE_SIGNAL:
-                    self._score_queue.task_done()
                     break
 
                 try:
@@ -216,7 +216,11 @@ def wake(self) -> None:
         try:
             self._score_queue.put(_WAKE_SIGNAL, block=False)
         except Full:
-            pass
+            return
+
+        # Mark it done now, since no consumer may take it (e.g. when shutdown()
+        # stops the consumer first) and queue.join() would then never return.
+        self._score_queue.task_done()
 
     def pause(self) -> None:
         """Pause the consumer."""

pyproject.toml

@@ -69,6 +69,7 @@ dev = [
     "autoevals>=0.0.130,<0.1",
     "opentelemetry-instrumentation-threading>=0.59b0,<1",
     "tenacity>=9.1.4",
+    "gevent>=24.2.1,<27",
 ]
 docs = [
     "pdoc>=15.0.4,<16",