diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index e4f6eb5bb2b..c1ac7a278cf 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -6,7 +6,7 @@ import traceback from collections.abc import Callable, Iterable, Mapping from datetime import datetime from functools import lru_cache -from importlib.metadata import version +from importlib.metadata import PackageNotFoundError, version from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable @@ -278,13 +278,13 @@ class LangFuseLogger: allow_env_credentials: bool = True, ): try: - from litellm.integrations.langfuse.langfuse_sdk import acquire_langfuse_tracing - except Exception as e: + self.langfuse_sdk_version: str = installed_langfuse_version() + except PackageNotFoundError as e: raise Exception( - f"\033[91mLangfuse not installed, try running 'pip install langfuse' to fix this error: {e}\n{traceback.format_exc()}\033[0m" - ) - self.langfuse_sdk_version: str = installed_langfuse_version() + f"\033[91mLangfuse not installed, try running 'pip install langfuse' to fix this error: {e}\033[0m" + ) from e raise_if_unsupported_langfuse_version(self.langfuse_sdk_version) + from litellm.integrations.langfuse.langfuse_sdk import acquire_langfuse_tracing self.public_key, self.secret_key, self.langfuse_host = resolve_langfuse_credentials( langfuse_public_key=langfuse_public_key, @@ -362,6 +362,12 @@ class LangFuseLogger: """Push every queued observation to Langfuse before the process goes away.""" self.tracing.flush() + def stop(self) -> None: + """Give the export channel back; ``DynamicLoggingCache`` calls this when a per-key logger expires.""" + from litellm.integrations.langfuse.langfuse_sdk import release_langfuse_tracing + + release_langfuse_tracing(self.tracing) + @staticmethod def add_metadata_from_header(litellm_params: dict, metadata: dict) -> dict[str, object]: """ diff --git a/litellm/integrations/langfuse/langfuse_sdk.py b/litellm/integrations/langfuse/langfuse_sdk.py index 3bdb318e512..9e592f058e1 100644 --- a/litellm/integrations/langfuse/langfuse_sdk.py +++ b/litellm/integrations/langfuse/langfuse_sdk.py @@ -6,7 +6,7 @@ import threading from base64 import b64encode from collections.abc import Iterable, Mapping, Sequence from contextvars import ContextVar -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime from hashlib import sha256 from importlib.metadata import version @@ -49,6 +49,7 @@ __all__ = ( "configured_sample_rate", "flush_langfuse_tracing", "observation_attributes", + "release_langfuse_tracing", "resolve_observation_id", "resolve_trace_id", "start_child_span", @@ -62,6 +63,7 @@ _OBSERVATION_ID_PATTERN: Final = re.compile(r"^(?=.*[1-9a-f])[0-9a-f]{16}$") _TRACER_NAME: Final = "langfuse-sdk" _MAX_QUEUE_SIZE: Final = 100_000 _DEFAULT_FLUSH_AT: Final = 512 +_CHANNEL_RETIRE_GRACE_SECONDS: Final = 60.0 _SPAN_LIMITS: Final = SpanLimits( max_attributes=SpanLimits.UNSET, max_events=128, @@ -561,6 +563,9 @@ class LangfuseTracing: def flush(self, timeout_millis: int = 30_000) -> bool: return self.provider.force_flush(timeout_millis) + def shutdown(self) -> None: + self.provider.shutdown() + @dataclass(frozen=True, slots=True) class _TracingKey: @@ -575,10 +580,14 @@ class _TracingKey: mock_mode: bool +@dataclass(frozen=True, slots=True) +class _Lease: + tracing: LangfuseTracing + holders: int + + _TRACING_LOCK: Final = threading.Lock() -_TRACING: Final[ - dict[_TracingKey, LangfuseTracing] -] = {} # mutable-ok: process-wide channel cache, guarded by _TRACING_LOCK +_TRACING: Final[dict[_TracingKey, _Lease]] = {} # mutable-ok: process-wide channel cache, guarded by _TRACING_LOCK def acquire_langfuse_tracing( @@ -593,8 +602,8 @@ def acquire_langfuse_tracing( ) -> LangfuseTracing: """One export channel per credential set, shared by every logger built for it. - Channels live for the process: a provider owns a batch export thread, and tearing one down - while another logger for the same credentials still exports through it would drop its spans. + A provider owns a batch export thread, so a channel lives while any logger holds it and is + retired through ``release_langfuse_tracing`` once the last holder lets go. """ key: Final = _TracingKey( public_key=public_key, @@ -610,7 +619,8 @@ def acquire_langfuse_tracing( with _TRACING_LOCK: cached: Final = _TRACING.get(key) if cached is not None: - return cached + _TRACING[key] = replace(cached, holders=cached.holders + 1) + return cached.tracing created: Final = build_langfuse_tracing( exporter=DiscardingSpanExporter() if mock_mode @@ -621,10 +631,44 @@ def acquire_langfuse_tracing( flush_at=key.flush_at, flush_interval_millis=key.flush_interval_millis, ) - _TRACING[key] = created + _TRACING[key] = _Lease(tracing=created, holders=1) return created +def release_langfuse_tracing(tracing: LangfuseTracing, *, grace_seconds: float = _CHANNEL_RETIRE_GRACE_SECONDS) -> None: + """Let go of one logger's hold on its channel; a channel nobody holds is retired ``grace_seconds`` later. + + The grace covers a callback that fetched its logger from the cache just before the entry expired, + and a logger rebuilt for the same credentials in the meantime picks the channel back up instead. + """ + with _TRACING_LOCK: + held: Final = next(((key, lease) for key, lease in _TRACING.items() if lease.tracing is tracing), None) + if held is None: + return + key, lease = held + if lease.holders <= 0: + return + _TRACING[key] = replace(lease, holders=lease.holders - 1) + if lease.holders > 1: + return + if grace_seconds <= 0: + _retire_if_unheld(key, tracing) + return + retire: Final = threading.Timer(grace_seconds, _retire_if_unheld, args=(key, tracing)) + retire.name = "langfuse-retire" + retire.daemon = True + retire.start() + + +def _retire_if_unheld(key: _TracingKey, tracing: LangfuseTracing) -> None: + with _TRACING_LOCK: + lease: Final = _TRACING.get(key) + if lease is None or lease.tracing is not tracing or lease.holders > 0: + return + del _TRACING[key] + tracing.shutdown() + + class _FlushWorker(threading.Thread): """Daemon, so a channel still blocked at the deadline cannot hold up interpreter exit.""" @@ -645,7 +689,7 @@ def flush_langfuse_tracing(timeout_millis: int = 30_000) -> bool: finish in the background rather than pushing the deadline out for the channels after it. """ with _TRACING_LOCK: - channels: Final = tuple(_TRACING.values()) + channels: Final = tuple(lease.tracing for lease in _TRACING.values()) workers: Final = tuple(_FlushWorker(channel, timeout_millis) for channel in channels) deadline: Final = monotonic() + timeout_millis / 1000 for worker in workers: diff --git a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py b/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py index c29ea566a31..0dca6ab7550 100644 --- a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py +++ b/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py @@ -23,8 +23,8 @@ class LangfuseInMemoryCache(InMemoryCache): The counter is a soft budget: loggers built concurrently for one credential set before the first lands in the cache each take a slot, and only the cached one gives it back on expiry. - Export channels are shared per credential set and outlive the logger, so nothing else is - torn down here (https://github.com/BerriAI/litellm/issues/11169). + The logger's ``stop()`` below hands its shared export channel back + (https://github.com/BerriAI/litellm/issues/11169). """ def _remove_key(self, key: str) -> None: diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py index 7e05ff0053b..fc1343ec75c 100644 --- a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py +++ b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py @@ -39,6 +39,7 @@ from litellm.integrations.langfuse.langfuse_sdk import ( configured_sample_rate, flush_langfuse_tracing, observation_attributes, + release_langfuse_tracing, resolve_observation_id, resolve_trace_id, start_child_span, @@ -664,6 +665,78 @@ def test_changed_credentials_or_settings_get_their_own_channel(override): assert _acquire() is not _acquire(**override) +class _RecordsShutdown(InMemorySpanExporter): + def __init__(self) -> None: + super().__init__() + self.shutdowns = 0 + + def shutdown(self) -> None: + self.shutdowns += 1 + super().shutdown() + + +def _acquire_recorded(monkeypatch: pytest.MonkeyPatch, public_key: str) -> tuple[LangfuseTracing, _RecordsShutdown]: + exporter = _RecordsShutdown() + monkeypatch.setattr("litellm.integrations.langfuse.langfuse_sdk._build_span_exporter", lambda **_: exporter) + return _acquire(public_key=public_key, mock_mode=False, flush_interval=600.0), exporter + + +def test_channel_is_retired_only_after_its_last_holder_releases_it(monkeypatch: pytest.MonkeyPatch): + """Two loggers on one credential set share the channel: the first release must leave it + exporting for the second, and the last release must shut the batch thread down and drop + the registry entry so the next logger gets a fresh channel instead of a dead one.""" + first, exporter = _acquire_recorded(monkeypatch, "pk-lease-test") + second = _acquire(public_key="pk-lease-test", mock_mode=False, flush_interval=600.0) + assert second is first + + release_langfuse_tracing(first, grace_seconds=0.0) + second.tracer.start_span("generation").end() + assert exporter.shutdowns == 0 + assert flush_langfuse_tracing() is True + assert len(exporter.get_finished_spans()) == 1 + + release_langfuse_tracing(second, grace_seconds=0.0) + assert exporter.shutdowns == 1 + assert _acquire(public_key="pk-lease-test", mock_mode=False, flush_interval=600.0) is not first + + +def test_release_flushes_the_queued_spans_before_the_channel_goes_away(monkeypatch: pytest.MonkeyPatch): + tracing, exporter = _acquire_recorded(monkeypatch, "pk-lease-flush-test") + tracing.tracer.start_span("generation").end() + + release_langfuse_tracing(tracing, grace_seconds=0.0) + + assert len(exporter.get_finished_spans()) == 1 + + +def test_channel_reacquired_within_the_grace_is_kept(monkeypatch: pytest.MonkeyPatch): + """A logger rebuilt for the same credentials right after the old one expired, and a callback + that fetched the old logger just before expiry, both keep exporting through the same channel.""" + tracing, exporter = _acquire_recorded(monkeypatch, "pk-lease-grace-test") + + release_langfuse_tracing(tracing, grace_seconds=0.2) + assert _acquire(public_key="pk-lease-grace-test", mock_mode=False, flush_interval=600.0) is tracing + + threading.Event().wait(0.5) + tracing.tracer.start_span("generation").end() + assert exporter.shutdowns == 0 + assert flush_langfuse_tracing() is True + assert len(exporter.get_finished_spans()) == 1 + + +def test_release_of_a_channel_the_registry_never_handed_out_is_a_no_op(): + exporter = InMemorySpanExporter() + tracing = build_langfuse_tracing( + exporter=exporter, environment=None, release=None, sample_rate=1.0, flush_interval_millis=10 + ) + + release_langfuse_tracing(tracing, grace_seconds=0.0) + tracing.tracer.start_span("generation").end() + + assert tracing.flush() is True + assert len(exporter.get_finished_spans()) == 1 + + def test_flush_langfuse_tracing_exports_the_queued_spans_of_every_channel(monkeypatch: pytest.MonkeyPatch): """The proxy shutdown hook flushes through this, so a span finished just before a graceful restart must reach the exporter without waiting for the batch interval.""" diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 3f47a793056..50db3569c8e 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -2166,6 +2166,60 @@ def test_version_gate_rejects_v5_prereleases(): langfuse_module.raise_if_unsupported_langfuse_version("5.0.0") +def test_old_sdk_fails_with_the_upgrade_message_before_the_otel_module_is_imported(monkeypatch): + """On a v2 install `langfuse_sdk` itself fails to import, so the version gate must run first + or the caller is told the package is missing when it only needs upgrading.""" + import sys + + monkeypatch.setattr(langfuse_module, "installed_langfuse_version", lambda: "2.59.7") + monkeypatch.setitem(sys.modules, "litellm.integrations.langfuse.langfuse_sdk", None) + + with pytest.raises(ImportError) as raised: + _build_langfuse_logger(monkeypatch, langfuse_public_key="pk-old-sdk") + + assert "2.59.7" in str(raised.value) + assert "langfuse_otel" in str(raised.value) + assert "not installed" not in str(raised.value) + + +def test_missing_sdk_is_reported_as_not_installed(monkeypatch): + from importlib.metadata import PackageNotFoundError + + def not_installed() -> str: + raise PackageNotFoundError("langfuse") + + monkeypatch.setattr(langfuse_module, "installed_langfuse_version", not_installed) + + with pytest.raises(Exception, match="Langfuse not installed"): + _build_langfuse_logger(monkeypatch, langfuse_public_key="pk-no-sdk") + + +def test_stopped_logger_hands_its_export_channel_back(monkeypatch): + """`DynamicLoggingCache` calls `stop()` on expiry; the channel must be retired once every + logger that held it has stopped, or each credential rotation leaks a batch export thread.""" + from litellm.integrations.langfuse.langfuse_sdk import acquire_langfuse_tracing, release_langfuse_tracing + + logger = _build_langfuse_logger(monkeypatch, langfuse_public_key="pk-stop-releases") + + def acquire_same_credentials(): + return acquire_langfuse_tracing( + public_key="pk-stop-releases", + secret_key="sk-lit5228", + base_url=_UNREACHABLE_HOST, + environment=logger.langfuse_environment, + release=logger.langfuse_release, + flush_interval=logger.langfuse_flush_interval, + mock_mode=False, + ) + + logger.stop() + reacquired = acquire_same_credentials() + assert reacquired is logger.tracing, "the channel stays up while another logger still holds it" + + release_langfuse_tracing(reacquired, grace_seconds=0.0) + assert acquire_same_credentials() is not logger.tracing, "stop() did not give the logger's hold back" + + def test_int_steering_values_reach_langfuse_as_strings(): """Langfuse models user, session and version as strings; v2's pydantic coerced ints for the caller.""" rig = _steering_logger() diff --git a/tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py b/tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py index fd0192534cd..d3a30e998d1 100644 --- a/tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py +++ b/tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py @@ -47,23 +47,28 @@ class TestLangfuseInMemoryCache: assert litellm.initialized_langfuse_clients == initial_count - 1 @patch("litellm.initialized_langfuse_clients", 3) - def test_evicted_logger_keeps_its_export_channel_and_api_client_alive(self): - """Export channels are shared per credential set and outlive the logger, so eviction only - releases the initialized-client slot: prompts still resolve and the channel still flushes.""" + def test_evicted_logger_releases_its_hold_on_the_shared_export_channel(self): + """Export channels are shared per credential set: eviction gives this logger's hold back + while a sibling logger keeps exporting, and the channel is retired once the last hold goes.""" from litellm.integrations.langfuse.langfuse import LangFuseLogger - from litellm.integrations.langfuse.langfuse_sdk import DiscardingSpanExporter, build_langfuse_tracing + from litellm.integrations.langfuse.langfuse_sdk import acquire_langfuse_tracing, release_langfuse_tracing + + def acquire(): + return acquire_langfuse_tracing( + public_key="pk-eviction-test", + secret_key="sk", + base_url="http://127.0.0.1:1", + environment=None, + release=None, + flush_interval=1.0, + mock_mode=True, + ) logger = LangFuseLogger.__new__(LangFuseLogger) logger.api_client = MagicMock() logger.api_client.get_prompt.return_value = "prompt-after-eviction" - logger.tracing = build_langfuse_tracing( - exporter=DiscardingSpanExporter(), - environment=None, - release=None, - sample_rate=1.0, - flush_at=512, - flush_interval_millis=1000, - ) + logger.tracing = acquire() + sibling = acquire() self.cache.cache_dict["test_key"] = logger self.cache.ttl_dict["test_key"] = time.time() + 100 @@ -71,6 +76,9 @@ class TestLangfuseInMemoryCache: assert litellm.initialized_langfuse_clients == 2 assert logger.api_client.get_prompt("greeting") == "prompt-after-eviction" - with logger.tracing.tracer.start_as_current_span("still-open"): + with sibling.tracer.start_as_current_span("still-open"): pass - assert logger.tracing.flush(1000) is True + assert sibling.flush(1000) is True + + release_langfuse_tracing(sibling, grace_seconds=0.0) + assert acquire() is not logger.tracing, "eviction did not release the evicted logger's hold"