perf: memory usage on repeat instantiation

Author: hassiebp

langfuse/_task_manager/task_manager.py

@@ -191,6 +191,9 @@ def shutdown(self):
         """Flush all messages and cleanly shutdown the client."""
         self._log.debug("shutdown initiated")
 
+        # Unregister the atexit handler first
+        atexit.unregister(self.shutdown)
+
         self.flush()
         self.join()
 

langfuse/client.py

@@ -36,6 +36,7 @@
     UpdateGenerationBody,
 )
 from langfuse.api.resources.ingestion.types.update_span_body import UpdateSpanBody
+from langfuse.api.resources.media import GetMediaResponse
 from langfuse.api.resources.observations.types.observations_views import (
     ObservationsViews,
 )
@@ -49,7 +50,6 @@
 from langfuse.api.resources.utils.resources.pagination.types.meta_response import (
     MetaResponse,
 )
-from langfuse.api.resources.media import GetMediaResponse
 from langfuse.model import (
     ChatMessageDict,
     ChatPromptClient,
@@ -74,8 +74,8 @@
 from langfuse.api.client import FernLangfuse
 from langfuse.environment import get_common_release_envs
 from langfuse.logging import clean_logger
-from langfuse.model import Dataset, MapValue, Observation, TraceWithFullDetails
 from langfuse.media import LangfuseMedia
+from langfuse.model import Dataset, MapValue, Observation, TraceWithFullDetails
 from langfuse.request import LangfuseClient
 from langfuse.types import MaskFunction, ScoreDataType, SpanLevel
 from langfuse.utils import (
@@ -2018,6 +2018,14 @@ def shutdown(self):
         As the SDK calls join() already on shutdown, refer to flush() to ensure all events arive at the Langfuse API.
         """
         try:
+            self.prompt_cache._task_manager.shutdown()
+
+            # In logging.py, a handler is attached to the httpx logger.
+            # To avoid a memory leak on singleton reset, remove all handlers
+            httpx_logger = logging.getLogger("httpx")
+            for handler in httpx_logger.handlers:
+                httpx_logger.removeHandler(handler)
+
             return self.task_manager.shutdown()
         except Exception as e:
             self.log.exception(e)

langfuse/decorators/langfuse_decorator.py

@@ -1,51 +1,50 @@
 import asyncio
+import inspect
+import json
+import logging
 from collections import defaultdict
 from contextvars import ContextVar
 from datetime import datetime
 from functools import wraps
-import httpx
-import inspect
-import json
-import logging
 from typing import (
     Any,
+    AsyncGenerator,
     Callable,
     DefaultDict,
+    Dict,
+    Generator,
+    Iterable,
     List,
-    Optional,
-    Union,
     Literal,
-    Dict,
+    Optional,
     Tuple,
-    Iterable,
-    AsyncGenerator,
-    Generator,
     TypeVar,
+    Union,
     cast,
     overload,
 )
 
+import httpx
+from pydantic import BaseModel
 from typing_extensions import ParamSpec
 
 from langfuse.api import UsageDetails
 from langfuse.client import (
     Langfuse,
-    StatefulSpanClient,
-    StatefulTraceClient,
-    StatefulGenerationClient,
-    PromptClient,
-    ModelUsage,
     MapValue,
+    ModelUsage,
+    PromptClient,
     ScoreDataType,
+    StatefulGenerationClient,
+    StatefulSpanClient,
+    StatefulTraceClient,
     StateType,
 )
 from langfuse.serializer import EventSerializer
 from langfuse.types import ObservationParams, SpanLevel
 from langfuse.utils import _get_timestamp
-from langfuse.utils.langfuse_singleton import LangfuseSingleton
 from langfuse.utils.error_logging import catch_and_log_errors
-
-from pydantic import BaseModel
+from langfuse.utils.langfuse_singleton import LangfuseSingleton
 
 _observation_stack_context: ContextVar[
     List[Union[StatefulTraceClient, StatefulSpanClient, StatefulGenerationClient]]

langfuse/prompt_cache.py

@@ -1,15 +1,14 @@
 """@private"""
 
-from datetime import datetime
-from typing import List, Optional, Dict, Set
-from threading import Thread
 import atexit
 import logging
+from datetime import datetime
 from queue import Empty, Queue
+from threading import Thread
+from typing import Dict, List, Optional, Set
 
 from langfuse.model import PromptClient
 
-
 DEFAULT_PROMPT_CACHE_TTL_SECONDS = 60
 
 DEFAULT_PROMPT_CACHE_REFRESH_WORKERS = 1
@@ -114,6 +113,8 @@ def shutdown(self):
             f"Shutting down prompt refresh task manager, {len(self._consumers)} consumers,..."
         )
 
+        atexit.unregister(self.shutdown)
+
         for consumer in self._consumers:
             consumer.pause()
 

langfuse/utils/langfuse_singleton.py

@@ -1,7 +1,7 @@
-import httpx
 import threading
 from typing import Optional
 
+import httpx
 
 from langfuse import Langfuse
 from langfuse.types import MaskFunction