speed up unit shutdown without weakening assertions

Author: hassiebp

langfuse/_task_manager/media_manager.py

@@ -18,6 +18,7 @@
 
 T = TypeVar("T")
 P = ParamSpec("P")
+_SHUTDOWN_SENTINEL = object()
 
 
 class MediaManager:
@@ -40,6 +41,11 @@ def __init__(
     def process_next_media_upload(self) -> None:
         try:
             upload_job = self._queue.get(block=True, timeout=1)
+
+            if upload_job is _SHUTDOWN_SENTINEL:
+                self._queue.task_done()
+                return
+
             logger.debug(
                 f"Media: Processing upload for media_id={upload_job['media_id']} in trace_id={upload_job['trace_id']}"
             )
@@ -54,6 +60,14 @@ def process_next_media_upload(self) -> None:
             )
             self._queue.task_done()
 
+    def signal_shutdown(self) -> None:
+        try:
+            self._queue.put(_SHUTDOWN_SENTINEL, block=False)
+        except Full:
+            # If the queue is full, the consumer will keep draining work and
+            # observe the paused flag on the next loop iteration.
+            pass
+
     def _find_and_process_media(
         self,
         *,

langfuse/_task_manager/media_upload_consumer.py

@@ -42,3 +42,4 @@ def pause(self) -> None:
             f"Thread: Pausing media upload consumer thread #{self._identifier}"
         )
         self.running = False
+        self._media_manager.signal_shutdown()

langfuse/_task_manager/score_ingestion_consumer.py

@@ -2,7 +2,7 @@
 import os
 import threading
 import time
-from queue import Empty, Queue
+from queue import Empty, Full, Queue
 from typing import Any, List, Optional
 
 import backoff
@@ -17,6 +17,7 @@
 
 MAX_EVENT_SIZE_BYTES = int(os.environ.get("LANGFUSE_MAX_EVENT_SIZE_BYTES", 1_000_000))
 MAX_BATCH_SIZE_BYTES = int(os.environ.get("LANGFUSE_MAX_BATCH_SIZE_BYTES", 2_500_000))
+_SHUTDOWN_SENTINEL = object()
 
 
 class ScoreIngestionMetadata(BaseModel):
@@ -71,6 +72,10 @@ def _next(self) -> list:
                     block=True, timeout=self._flush_interval - elapsed
                 )
 
+                if event is _SHUTDOWN_SENTINEL:
+                    self._ingestion_queue.task_done()
+                    break
+
                 # convert pydantic models to dicts
                 if "body" in event and isinstance(event["body"], BaseModel):
                     event["body"] = event["body"].model_dump(exclude_none=True)
@@ -139,6 +144,12 @@ def upload(self) -> None:
     def pause(self) -> None:
         """Pause the consumer."""
         self.running = False
+        try:
+            self._ingestion_queue.put(_SHUTDOWN_SENTINEL, block=False)
+        except Full:
+            # If the queue is full, the consumer will wake up naturally while
+            # draining items, so a dedicated shutdown signal is not required.
+            pass
 
     def _upload_batch(self, batch: List[Any]) -> None:
         logger.debug(

tests/unit/test_otel.py

@@ -3054,7 +3054,7 @@ def test_metrics_and_timing(self, langfuse_client, memory_exporter):
         span = langfuse_client.start_observation(name="timing-test-span")
 
         # Add a small delay
-        time.sleep(0.01)
+        time.sleep(0.1)
 
         # End the span
         span.end()
@@ -3096,10 +3096,10 @@ def test_metrics_and_timing(self, langfuse_client, memory_exporter):
         ) / 1_000_000_000
         assert span_duration_seconds > 0, "Span duration should be positive"
 
-        # Since we slept for 0.01 seconds, the span duration should be at least 0.005 seconds
+        # Since we slept for 0.1 seconds, the span duration should be at least 0.05 seconds
         # but we'll be generous with the upper bound due to potential system delays
-        assert span_duration_seconds >= 0.005, (
-            f"Span duration ({span_duration_seconds}s) should be at least 0.005s"
+        assert span_duration_seconds >= 0.05, (
+            f"Span duration ({span_duration_seconds}s) should be at least 0.05s"
         )
 
 

tests/unit/test_prompt_atexit.py

@@ -20,14 +20,14 @@ def test_prompts_atexit():
 print("Adding prompt cache", PromptCache)
 prompt_cache = PromptCache(max_prompt_refresh_workers=10)
 
-# example task that stays in flight briefly while the process exits
-def wait_briefly():
-    time.sleep(0.1)
+# example task that takes 2 seconds but we will force it to exit earlier
+def wait_2_sec():
+    time.sleep(2)
 
 # 8 times
 for i in range(8):
     prompt_cache.add_refresh_prompt_task(
-        f"key_wait_briefly_i_{i}", lambda: wait_briefly()
+        f"key_wait_2_sec_i_{i}", lambda: wait_2_sec()
     )
 """
 
@@ -76,13 +76,13 @@ async def main():
     print("Adding prompt cache", PromptCache)
     prompt_cache = PromptCache(max_prompt_refresh_workers=10)
 
-    # example task that stays in flight briefly while the process exits
-    def wait_briefly():
-        time.sleep(0.1)
+    # example task that takes 2 seconds but we will force it to exit earlier
+    def wait_2_sec():
+        time.sleep(2)
 
     async def add_new_prompt_refresh(i: int):
         prompt_cache.add_refresh_prompt_task(
-            f"key_wait_briefly_i_{i}", lambda: wait_briefly()
+            f"key_wait_2_sec_i_{i}", lambda: wait_2_sec()
         )
     
     # 8 times

tests/unit/test_resource_manager.py

@@ -1,8 +1,14 @@
 """Test the LangfuseResourceManager and get_client() function."""
 
+from queue import Queue
+from unittest.mock import Mock
+
 from langfuse import Langfuse
 from langfuse._client.get_client import get_client
 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.score_ingestion_consumer import ScoreIngestionConsumer
 
 
 def test_get_client_preserves_all_settings(monkeypatch):
@@ -101,3 +107,34 @@ def should_export_b(span):
 
     client_a.shutdown()
     client_b.shutdown()
+
+
+def test_score_ingestion_consumer_pause_wakes_blocked_thread():
+    consumer = ScoreIngestionConsumer(
+        ingestion_queue=Queue(),
+        identifier=0,
+        client=Mock(),
+        public_key="pk-test",
+        flush_interval=30,
+    )
+
+    consumer.start()
+    consumer.pause()
+    consumer.join(timeout=0.5)
+
+    assert not consumer.is_alive()
+
+
+def test_media_upload_consumer_pause_wakes_blocked_thread():
+    media_manager = MediaManager(
+        api_client=Mock(),
+        httpx_client=Mock(),
+        media_upload_queue=Queue(),
+    )
+    consumer = MediaUploadConsumer(identifier=0, media_manager=media_manager)
+
+    consumer.start()
+    consumer.pause()
+    consumer.join(timeout=0.5)
+
+    assert not consumer.is_alive()