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)