From 9c4eceee05804715038871247fb491a659517129 Mon Sep 17 00:00:00 2001 From: David Traina <44659830+DavidTraina@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:23:59 -0500 Subject: [PATCH] fix(media): don't hang flush() on media logged after shutdown() shutdown() stops the media upload consumers, but a client used afterwards (the same client, Langfuse(...) with the same key, or get_client()) still queues media found in span input, output and metadata. No consumer takes it, so flush() and every later shutdown() block forever in Queue.join(). Media logged while shutdown() stops the consumers can be stranded the same way. shutdown() now stops accepting media right after its own flush(), uploads what was queued since, stops the consumers, then marks done anything a thread that had already passed the check queued meanwhile. Later media is dropped with a warning, like scores created after shutdown(). --- langfuse/_client/resource_manager.py | 6 ++ langfuse/_task_manager/media_manager.py | 39 +++++++++ tests/unit/test_resource_manager.py | 104 +++++++++++++++++++++++- 3 files changed, 147 insertions(+), 2 deletions(-) diff --git a/langfuse/_client/resource_manager.py b/langfuse/_client/resource_manager.py index a274075e7..36c23c1e0 100644 --- a/langfuse/_client/resource_manager.py +++ b/langfuse/_client/resource_manager.py @@ -617,7 +617,13 @@ def shutdown(self) -> None: atexit.unregister(self.shutdown) self.flush() + # No consumer takes media queued after the consumers stop, and flush() + # would then wait for it forever. Upload what was queued since the + # flush() above, then stop the consumers and drop any late arrivals. + self._media_manager.stop_accepting_uploads() + self._media_upload_queue.join() self._stop_and_join_consumer_threads() + self._media_manager.drop_pending_uploads() def _init_tracer_provider( diff --git a/langfuse/_task_manager/media_manager.py b/langfuse/_task_manager/media_manager.py index 3ad9ac62f..732fac7b8 100644 --- a/langfuse/_task_manager/media_manager.py +++ b/langfuse/_task_manager/media_manager.py @@ -45,6 +45,7 @@ def __init__( self._enabled = os.environ.get( LANGFUSE_MEDIA_UPLOAD_ENABLED, "True" ).lower() not in ("false", "0") + self._accepting_uploads = True def reinitialize( self, @@ -57,6 +58,35 @@ def reinitialize( self._httpx_client = httpx_client self._queue = media_upload_queue + def stop_accepting_uploads(self) -> None: + """Skip media found from now on, as the consumers are about to stop.""" + self._accepting_uploads = False + + def drop_pending_uploads(self) -> None: + """Acknowledge uploads queued after the consumers stopped. + + A thread that checked the flag just before stop_accepting_uploads() + can still queue its upload, and flush() would wait for it forever. + """ + dropped = 0 + while True: + try: + upload_job = self._queue.get_nowait() + except Empty: + break + + # signal_shutdown() already marked its sentinels done. + if upload_job is not _SHUTDOWN_SENTINEL: + self._queue.task_done() + dropped += 1 + + if dropped: + logger.warning( + "Media upload dropped: Langfuse client has been shut down. " + "%s queued upload(s).", + dropped, + ) + def process_next_media_upload(self) -> None: try: upload_job = self._queue.get(block=True, timeout=1) @@ -343,6 +373,15 @@ def _process_media( field=field, ) + if not self._accepting_uploads: + logger.warning( + "Media upload dropped: Langfuse client has been shut down. " + "media_id=%s trace_id=%s", + media._media_id, + trace_id, + ) + return + self._queue.put( item=upload_media_job, block=False, diff --git a/tests/unit/test_resource_manager.py b/tests/unit/test_resource_manager.py index 27e600cdc..810ef90d6 100644 --- a/tests/unit/test_resource_manager.py +++ b/tests/unit/test_resource_manager.py @@ -2,13 +2,15 @@ import logging import threading +import time from queue import Queue from types import SimpleNamespace -from typing import Sequence +from typing import Callable, Sequence from unittest.mock import Mock +import httpx import pytest -from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult from langfuse import Langfuse @@ -17,6 +19,7 @@ from langfuse._task_manager.media_manager import MediaManager from langfuse._task_manager.media_upload_consumer import MediaUploadConsumer from langfuse._task_manager.score_consumer import ScoreConsumer +from langfuse.media import LangfuseMedia from langfuse.types import MaskOtelSpansResult @@ -506,3 +509,100 @@ def test_flush_returns_after_repeated_shutdown(concurrent): media_queue.task_done() assert not hung + + +def _media_client(**kwargs) -> Langfuse: + return Langfuse( + public_key="pk-media-shutdown", + secret_key="sk-media-shutdown", + span_exporter=NoOpSpanExporter(), + # The global provider may carry a sampler set by another test, and + # media is only processed on recorded spans. + tracer_provider=TracerProvider(), + **kwargs, + ) + + +def _log_media(client: Langfuse, content: bytes) -> None: + with client.start_as_current_observation( + name="media", + input=LangfuseMedia(content_bytes=content, content_type="image/png"), + ): + pass + + +def _hangs(target: Callable[[], None], media_queue: Queue, timeout: float) -> bool: + thread = threading.Thread(target=target, daemon=True) + thread.start() + thread.join(timeout=timeout) + hung = thread.is_alive() + # Acknowledge media that no consumer took, so a hung thread and the + # reset() in the autouse fixture can finish. + while media_queue.unfinished_tasks: + media_queue.get_nowait() + media_queue.task_done() + + return hung + + +@pytest.mark.filterwarnings("error::pytest.PytestUnhandledThreadExceptionWarning") +def test_flush_returns_when_media_is_logged_after_shutdown(caplog): + client = _media_client() + client.shutdown() + with caplog.at_level(logging.WARNING, logger="langfuse"): + _log_media(client, b"after-shutdown") + + assert not _hangs(client.flush, client._resources._media_upload_queue, timeout=2) + assert any("shut down" in r.getMessage() for r in caplog.records) + + +@pytest.mark.filterwarnings("error::pytest.PytestUnhandledThreadExceptionWarning") +def test_shutdown_uploads_media_logged_after_its_flush(monkeypatch): + media_requests = [] + + def handler(request: httpx.Request) -> httpx.Response: + media_requests.append(request) + # Slow enough that the second upload is still queued if shutdown() + # stops the consumer without waiting for it. + time.sleep(0.05) + # No uploadUrl: the server already has the media, so the job ends here. + return httpx.Response(200, json={"mediaId": "media-id", "uploadUrl": None}) + + client = _media_client( + httpx_client=httpx.Client(transport=httpx.MockTransport(handler)), + media_upload_thread_count=1, + ) + resources = client._resources + flush = resources.flush + + def flush_then_log_media() -> None: + flush() + # Stands in for other threads logging media while shutdown() runs. + _log_media(client, b"first") + _log_media(client, b"second") + + monkeypatch.setattr(resources, "flush", flush_then_log_media) + + assert not _hangs(client.shutdown, resources._media_upload_queue, timeout=5) + assert len(media_requests) == 2 + + +@pytest.mark.filterwarnings("error::pytest.PytestUnhandledThreadExceptionWarning") +def test_flush_returns_when_media_is_queued_while_consumers_stop(monkeypatch): + client = _media_client() + resources = client._resources + media_queue = resources._media_upload_queue + stop_consumers = resources._stop_and_join_consumer_threads + + def stop_consumers_then_queue_media() -> None: + stop_consumers() + # Stands in for a thread that passed the _accepting_uploads check + # before stop_accepting_uploads() and queues its upload only now. + media_queue.put({"media_id": "late"}) + + monkeypatch.setattr( + resources, "_stop_and_join_consumer_threads", stop_consumers_then_queue_media + ) + client.shutdown() + + assert not _hangs(client.flush, media_queue, timeout=2)