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