diff --git a/langfuse/_client/resource_manager.py b/langfuse/_client/resource_manager.py index e633cdfaa..e920353a1 100644 --- a/langfuse/_client/resource_manager.py +++ b/langfuse/_client/resource_manager.py @@ -596,11 +596,11 @@ 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") - self._media_upload_queue.join() + self._join_media_upload_queue() langfuse_logger.debug("Successfully flushed media upload queue") def _join_score_queue(self) -> None: @@ -623,6 +623,27 @@ def _join_score_queue(self) -> None: score_queue.all_tasks_done.wait(timeout=0.1) + def _join_media_upload_queue(self) -> None: + """Wait until the media upload queue is drained or no consumer is left. + + Queue.join() would block forever when the consumers stop first, e.g. + when shutdown() runs concurrently or media is enqueued after shutdown. + """ + media_upload_queue = self._media_upload_queue + + with media_upload_queue.all_tasks_done: + while media_upload_queue.unfinished_tasks: + if not any(c.is_alive() for c in self._media_upload_consumers): + langfuse_logger.debug( + "Media upload consumers stopped with %s queue item(s) " + "unprocessed", + media_upload_queue.unfinished_tasks, + ) + + return + + media_upload_queue.all_tasks_done.wait(timeout=0.1) + def shutdown(self) -> None: self._shutdown = True diff --git a/tests/unit/test_resource_manager.py b/tests/unit/test_resource_manager.py index 232ec4f2f..944d0acaa 100644 --- a/tests/unit/test_resource_manager.py +++ b/tests/unit/test_resource_manager.py @@ -1,7 +1,8 @@ """Test the LangfuseResourceManager and get_client() function.""" import logging -from queue import Queue +import threading +from queue import Empty, Queue from types import SimpleNamespace from typing import Sequence from unittest.mock import Mock @@ -14,6 +15,7 @@ from langfuse._client.resource_manager import LangfuseResourceManager from langfuse._task_manager.media_manager import MediaManager from langfuse._task_manager.media_upload_consumer import MediaUploadConsumer +from langfuse._task_manager.media_upload_queue import UploadMediaJob from langfuse._task_manager.score_consumer import ScoreConsumer from langfuse.types import MaskOtelSpansResult @@ -471,3 +473,80 @@ def signal_shutdown(self, *, count): ("join", 0), ("join", 1), ] + + +def test_flush_returns_when_media_consumers_stopped_with_pending_items(): + """flush() must not block on the media queue when no consumer is left. + + Mirrors the score-queue guard: Queue.join() blocks forever when media + upload consumers have stopped with items still pending, e.g. when media is + enqueued after shutdown() or shutdown() runs concurrently. + """ + with LangfuseResourceManager._lock: + LangfuseResourceManager._instances.clear() + + client = Langfuse( + public_key="pk-flush-media-join", + secret_key="sk-flush-media-join", + span_exporter=NoOpSpanExporter(), + ) + rm = client._resources + assert rm is not None + assert len(rm._media_upload_consumers) > 0 + + # Simulate a shutdown race: consumers stopped with an item still pending. + # Consumers poll with a 1s timeout, so an unbounded join is safe here. + for consumer in rm._media_upload_consumers: + consumer.pause() + for consumer in rm._media_upload_consumers: + consumer.join() + assert not consumer.is_alive() + + rm._media_upload_queue.put( + UploadMediaJob( + media_id="test-media-id", + content_type="image/png", + content_length=3, + content_bytes=b"abc", + content_sha256_hash="", + trace_id="test-trace-id", + observation_id=None, + dataset_id=None, + dataset_item_id=None, + field="input", + ) + ) + + flushed = threading.Event() + + def _flush(): + rm.flush() + flushed.set() + + flush_thread = threading.Thread(target=_flush, daemon=True) + flush_thread.start() + flush_thread.join(timeout=5.0) + hung = not flushed.is_set() + + # Unblock flush() on unfixed code so the thread does not linger. Never + # blocks: on fixed code the guard already returned and this just drains. + try: + rm._media_upload_queue.get_nowait() + except Empty: + pass + else: + rm._media_upload_queue.task_done() + flush_thread.join(timeout=5.0) + + client.shutdown() + + # shutdown() enqueues one sentinel per consumer. Drain them so a later + # flush() (e.g. from the autouse fixture) cannot hang on unfixed code. + while True: + try: + rm._media_upload_queue.get_nowait() + except Empty: + break + rm._media_upload_queue.task_done() + + assert not hung, "flush() blocked on the media upload queue with no live consumer"