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, then stops the consumers. The flag check and put()
share a lock, so nothing can be queued after the consumers stop. Later
media is skipped with a warning, like scores created after shutdown().

Author: DavidTraina

langfuse/_client/resource_manager.py

@@ -410,6 +410,8 @@ def _at_fork_reinit(self) -> None:
         # the lock is class-level state needed by the child (e.g. to create a new client)
         # even if this particular instance was already shut down.
         LangfuseResourceManager._lock = threading.RLock()
+        # Media logged after shutdown() still takes this lock.
+        self._media_manager._enqueue_lock = threading.RLock()
 
         if self._shutdown:
             return
@@ -617,6 +619,11 @@ def shutdown(self) -> None:
         atexit.unregister(self.shutdown)
 
         self.flush()
+        # No consumer would take media queued after the consumers stop, and
+        # flush() would then wait for it forever. Upload what was queued since
+        # the flush() above before stopping them.
+        self._media_manager.stop_accepting_uploads()
+        self._media_upload_queue.join()
         self._stop_and_join_consumer_threads()
 
 

langfuse/_task_manager/media_manager.py

@@ -1,4 +1,5 @@
 import os
+import threading
 import time
 from queue import Empty, Full, Queue
 from typing import Any, Callable, Optional, TypeVar, cast
@@ -45,6 +46,10 @@ def __init__(
         self._enabled = os.environ.get(
             LANGFUSE_MEDIA_UPLOAD_ENABLED, "True"
         ).lower() not in ("false", "0")
+        # Reentrant, so shutdown() from a signal handler does not deadlock when
+        # the interrupted thread holds it.
+        self._enqueue_lock = threading.RLock()
+        self._accepting_uploads = True
 
     def reinitialize(
         self,
@@ -57,6 +62,11 @@ 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."""
+        with self._enqueue_lock:
+            self._accepting_uploads = False
+
     def process_next_media_upload(self) -> None:
         try:
             upload_job = self._queue.get(block=True, timeout=1)
@@ -343,10 +353,23 @@ def _process_media(
                 field=field,
             )
 
-            self._queue.put(
-                item=upload_media_job,
-                block=False,
-            )
+            with self._enqueue_lock:
+                accepting_uploads = self._accepting_uploads
+                if accepting_uploads:
+                    self._queue.put(
+                        item=upload_media_job,
+                        block=False,
+                    )
+
+            if not accepting_uploads:
+                logger.warning(
+                    "Media upload skipped: Langfuse client has been shut down. "
+                    "media_id=%s trace_id=%s",
+                    media._media_id,
+                    trace_id,
+                )
+                return
+
             logger.debug(
                 "Queue: Enqueued media ID %s for upload processing | trace_id=%s | "
                 "field=%s",

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 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
 
 
@@ -226,6 +229,7 @@ def test_at_fork_reinit_skips_when_shutdown(monkeypatch):
 
     old_score_queue = rm._score_queue
     old_lock = LangfuseResourceManager._lock
+    old_media_lock = rm._media_manager._enqueue_lock
 
     rm._shutdown = True
     rm._at_fork_reinit()
@@ -234,6 +238,7 @@ def test_at_fork_reinit_skips_when_shutdown(monkeypatch):
     # The class-level lock is shared state the child still needs (e.g. to create
     # a new client), so it must be replaced even when this instance is shut down.
     assert LangfuseResourceManager._lock is not old_lock
+    assert rm._media_manager._enqueue_lock is not old_media_lock
 
     client.shutdown()
 
@@ -506,3 +511,81 @@ def test_flush_returns_after_repeated_shutdown(concurrent):
         media_queue.task_done()
 
     assert not hung
+
+
+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
+
+
+@pytest.mark.filterwarnings("error::pytest.PytestUnhandledThreadExceptionWarning")
+def test_flush_returns_when_media_is_logged_after_shutdown(caplog):
+    client = Langfuse(
+        public_key="pk-media-after-shutdown",
+        secret_key="sk-media-after-shutdown",
+        span_exporter=NoOpSpanExporter(),
+        tracer_provider=TracerProvider(),
+    )
+    media_queue = client._resources._media_upload_queue
+    client.shutdown()
+    with caplog.at_level(logging.WARNING, logger="langfuse"):
+        _log_media(client, b"after-shutdown")
+
+    flush = threading.Thread(target=client.flush, daemon=True)
+    flush.start()
+    flush.join(timeout=2)
+    hung = flush.is_alive()
+    # Acknowledge media that no consumer took, so a hung flush() and the
+    # reset() in the autouse fixture can finish.
+    while media_queue.unfinished_tasks:
+        media_queue.get_nowait()
+        media_queue.task_done()
+
+    assert not hung
+    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)
+        # Keeps the consumer busy while shutdown() stops 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 = Langfuse(
+        public_key="pk-media-during-shutdown",
+        secret_key="sk-media-during-shutdown",
+        span_exporter=NoOpSpanExporter(),
+        tracer_provider=TracerProvider(),
+        httpx_client=httpx.Client(transport=httpx.MockTransport(handler)),
+    )
+    resources = client._resources
+    media_queue = resources._media_upload_queue
+    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)
+    shutdown = threading.Thread(target=client.shutdown, daemon=True)
+    shutdown.start()
+    shutdown.join(timeout=5)
+    hung = shutdown.is_alive()
+    # Acknowledge media that no consumer took, so a hung shutdown() and the
+    # reset() in the autouse fixture can finish.
+    while media_queue.unfinished_tasks:
+        media_queue.get_nowait()
+        media_queue.task_done()
+
+    assert not hung
+    assert len(media_requests) == 2