From d52769248e77296892174b675e6a2b52124a2103 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 15 Sep 2026 01:00:22 +0000 Subject: [PATCH] fix(langfuse): renew the client when eviction lands before the callback lease The cache can evict a logger between handing it to the callback and the callback taking its lease. Such a lease now hands back a fresh client acquired through the same parameters, so that callback exports through a live tracer provider instead of one teardown already shut down. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/langfuse/langfuse.py | 22 ++++++- litellm/integrations/langfuse/langfuse_sdk.py | 28 +++++--- .../langfuse/test_langfuse_sdk.py | 66 ++++++++++++++++--- .../integrations/test_langfuse.py | 40 +++++++++++ 4 files changed, 135 insertions(+), 21 deletions(-) diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 01838ed2a1b..28fb4f90f2b 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -355,7 +355,7 @@ class LangFuseLogger: self.langfuse_client = self._http_handler.client self.is_mock_mode = False - parameters: Final = { + self.langfuse_client_parameters: Final[dict[str, object]] = { "public_key": self.public_key, "secret_key": self.secret_key, "base_url": self.langfuse_host, @@ -365,7 +365,7 @@ class LangFuseLogger: "httpx_client": self.langfuse_client, "environment": self.langfuse_environment, } - self.Langfuse: Langfuse = self.safe_init_langfuse_client(parameters) + self.Langfuse: Langfuse = self.safe_init_langfuse_client(self.langfuse_client_parameters) # set the current langfuse project id in the environ # this is used by Alerting to link to the correct project @@ -416,6 +416,21 @@ class LangFuseLogger: verbose_logger.debug("Created langfuse client number %s", litellm.initialized_langfuse_clients) return langfuse_client + def _renew_langfuse_client(self) -> Langfuse: + """Replace a client the cache evicted after handing this logger to the callback. + + Bypasses the initialized-client ceiling: eviction already released this logger's slot, and the + replacement is never evicted itself, so its provider is retired with the logger instead. + """ + from litellm.integrations.langfuse.langfuse_sdk import acquire_langfuse_client + + return acquire_langfuse_client( + parameters=self.langfuse_client_parameters, + environment=self.langfuse_environment, + release=self.langfuse_release, + mock_mode=self.is_mock_mode, + ) + @staticmethod def add_metadata_from_header(litellm_params: dict, metadata: dict) -> dict[str, object]: """ @@ -512,7 +527,8 @@ class LangFuseLogger: verbose_logger.debug("OUTPUT IN LANGFUSE: %s; original: %s", output, response_obj) from litellm.integrations.langfuse.langfuse_sdk import lease_langfuse_client - with lease_langfuse_client(self.Langfuse): + with lease_langfuse_client(self.Langfuse, self._renew_langfuse_client) as leased: + self.Langfuse = leased trace_id, generation_id = self._log_langfuse_v2( user_id=user_id, metadata=metadata, diff --git a/litellm/integrations/langfuse/langfuse_sdk.py b/litellm/integrations/langfuse/langfuse_sdk.py index 4c70ad1263a..fe0f7332503 100644 --- a/litellm/integrations/langfuse/langfuse_sdk.py +++ b/litellm/integrations/langfuse/langfuse_sdk.py @@ -4,7 +4,7 @@ import os import re import threading from base64 import b64encode -from collections.abc import Generator, Mapping +from collections.abc import Callable, Generator, Mapping from contextlib import contextmanager from contextvars import ContextVar from datetime import datetime @@ -259,14 +259,20 @@ class _LangfuseLifecycleState: self.teardown_in_progress = False self.teardown_owner: int | None = None self.pending_clients: set[Langfuse] = set() # mutable-ok: eviction and callback threads queue into it + self.retired: WeakSet[Langfuse] = WeakSet() # mutable-ok: eviction marks its clients here from its own thread - def open_lease(self) -> None: + def open_lease(self, client: Langfuse) -> bool: + """Take a lease on ``client``; False when eviction already reached it, so a lease would guard a dead client.""" with self.lock: + if client in self.retired: + return False self.active_leases += 1 + return True def claim_for_teardown(self, client: Langfuse) -> bool: """Whether this thread owns ``client``'s teardown; a lease or another teardown in flight queues it instead.""" with self.lock: + self.retired.add(client) if self.active_leases > 0 or self.teardown_in_progress: self.pending_clients.add(client) return False @@ -326,21 +332,25 @@ def _lifecycle_state(client: Langfuse) -> _LangfuseLifecycleState: @contextmanager -def lease_langfuse_client(client: Langfuse) -> Generator[None]: +def lease_langfuse_client(client: Langfuse, renew: Callable[[], Langfuse]) -> Generator[Langfuse]: """Hold off cache eviction's teardown of ``client`` while the export inside is in flight. Eviction reaches a client the cache handed a callback moments earlier, so closing the SDK client and its tracer provider there drops the spans that callback is still writing. The lease protects exactly the window it wraps: an eviction arriving inside it is deferred to the last lease exit. - Taking a lease never blocks; a teardown already running keeps running, because the spans of a - lease taken that late were lost before the lease began, and stalling every other callback in the - process would not bring them back. A client the registry hands out during the deferral registers - as a holder, and the reference count keeps its bundle alive from there. + Taking a lease never blocks. When eviction already claimed ``client`` between the cache lookup + and this call, the lease is taken on ``renew()``'s fresh client instead and that client is what + the caller must export through: the registry hands it the live bundle when one remains, where + it registers as a holder and the reference count degrades the queued teardown to a flush, or a + fresh bundle once the old one is gone. """ state: Final = _lifecycle_state(client) - state.open_lease() + if not state.open_lease(client): + with lease_langfuse_client(renew(), renew) as leased: + yield leased + return try: - yield + yield client finally: _run_teardowns(state, state.release_lease()) diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py index 32cde376139..800502732f7 100644 --- a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py +++ b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py @@ -607,6 +607,10 @@ def _exports(client, exporter, name): return any(span.name == name for span in exporter.get_finished_spans()) +def _never_renew(): + raise AssertionError("the lease renewed a client eviction never reached") + + def test_garbage_collected_throwaway_clients_do_not_hold_shared_resources_open(): """A health probe or alerting lookup builds a client it never shuts down. @@ -658,7 +662,7 @@ def test_eviction_defers_teardown_until_active_callback_finishes(): def evict() -> None: shutdown_langfuse_client(client) - with lease_langfuse_client(client): + with lease_langfuse_client(client, _never_renew): evictor = threading.Thread(target=evict) evictor.start() evictor.join(timeout=5) @@ -713,7 +717,7 @@ def test_teardown_failure_does_not_strand_queued_clients(monkeypatch): raise RuntimeError("teardown failed") monkeypatch.setattr("litellm.integrations.langfuse.langfuse_sdk._teardown_langfuse_client", teardown) - with lease_langfuse_client(clients[0]): + with lease_langfuse_client(clients[0], _never_renew): for client in clients: shutdown_langfuse_client(client) @@ -738,19 +742,63 @@ def test_interrupt_during_deferred_teardown_propagates_and_requeues_the_client(m monkeypatch.setattr("litellm.integrations.langfuse.langfuse_sdk._teardown_langfuse_client", teardown) with pytest.raises(KeyboardInterrupt): - with lease_langfuse_client(client): + with lease_langfuse_client(client, _never_renew): shutdown_langfuse_client(client) assert state.pending_clients == {client} assert not state.teardown_in_progress - with lease_langfuse_client(client): - pass + replacement = Langfuse(public_key=PUBLIC_KEY, secret_key="sk-original", host="http://127.0.0.1:1") + with lease_langfuse_client(client, lambda: replacement) as leased: + assert leased is replacement assert calls == [client, client] assert not state.pending_clients +def test_lease_on_a_client_evicted_after_the_cache_lookup_exports_through_a_renewed_one(): + """Eviction can land between the cache handing out the logger and the callback taking its lease. + + That callback must not export into a shut-down provider; the lease has to hand it a live client. + """ + exporter = InMemorySpanExporter() + provider = build_isolated_tracer_provider(environment=None, release=None) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + evicted = Langfuse( + public_key=PUBLIC_KEY, + secret_key="sk-original", + host="http://127.0.0.1:1", + tracer_provider=provider, + span_exporter=exporter, + ) + register_langfuse_client(evicted) + shutdown_langfuse_client(evicted) + assert exporter._stopped + + renewed_exporter = InMemorySpanExporter() + renewed_provider = build_isolated_tracer_provider(environment=None, release=None) + renewed_provider.add_span_processor(SimpleSpanProcessor(renewed_exporter)) + + def renew(): + renewed = Langfuse( + public_key=PUBLIC_KEY, + secret_key="sk-original", + host="http://127.0.0.1:1", + tracer_provider=renewed_provider, + span_exporter=renewed_exporter, + ) + register_langfuse_client(renewed) + return renewed + + with lease_langfuse_client(evicted, renew) as leased: + assert leased is not evicted + assert _exports(leased, renewed_exporter, "after-lookup-eviction") + shutdown_langfuse_client(leased) + assert not renewed_exporter._stopped + + assert renewed_exporter._stopped + + def test_queued_eviction_waits_for_the_last_of_two_overlapping_leases(): exporter = InMemorySpanExporter() provider = build_isolated_tracer_provider(environment=None, release=None) @@ -764,8 +812,8 @@ def test_queued_eviction_waits_for_the_last_of_two_overlapping_leases(): ) register_langfuse_client(client) - with lease_langfuse_client(client): - with lease_langfuse_client(client): + with lease_langfuse_client(client, _never_renew): + with lease_langfuse_client(client, _never_renew): shutdown_langfuse_client(client) assert not exporter._stopped @@ -788,7 +836,7 @@ def test_a_client_adopted_during_deferred_teardown_keeps_exporting(): ) register_langfuse_client(evicted) - with lease_langfuse_client(evicted): + with lease_langfuse_client(evicted, _never_renew): shutdown_langfuse_client(evicted) adopter = Langfuse(public_key=PUBLIC_KEY, secret_key="sk-original", host="http://127.0.0.1:1") assert adopter._resources is evicted._resources @@ -804,7 +852,7 @@ def test_leases_on_one_client_do_not_serialise_callbacks(): both_inside = threading.Barrier(2, timeout=5) def hold_lease() -> None: - with lease_langfuse_client(client): + with lease_langfuse_client(client, _never_renew): both_inside.wait() holders = tuple(threading.Thread(target=hold_lease) for _ in range(2)) diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 64c211cad1b..d0acc010b15 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -1405,6 +1405,46 @@ def test_log_event_holds_a_client_lease_during_export(): assert state.active_leases == 0 +def test_log_event_renews_a_client_the_cache_evicted_before_the_lease(): + """The cache can evict this logger after handing it to the callback and before the lease opens. + + v2 lost that callback's events to a shut-down client; the callback must export through a fresh one. + """ + from litellm.integrations.langfuse.langfuse_sdk import register_langfuse_client, shutdown_langfuse_client + + logger, _ = _steering_logger() + evicted = logger.Langfuse + logger.langfuse_client_parameters = { + "public_key": "pk-steering-test", + "secret_key": "sk-steering-test", + "base_url": "http://127.0.0.1:1", + } + logger.langfuse_environment = None + logger.langfuse_release = None + logger.is_mock_mode = True + register_langfuse_client(evicted) + shutdown_langfuse_client(evicted) + + now = datetime.datetime.now() + returned = logger.log_event_on_langfuse( + kwargs={ + "call_type": "completion", + "litellm_params": {"metadata": {}}, + "messages": [{"role": "user", "content": "the-input"}], + "optional_params": {}, + }, + response_obj=litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "the-output"}}]), + start_time=now, + end_time=now, + ) + + assert returned["trace_id"] is not None + assert logger.Langfuse is not evicted + assert logger.Langfuse._resources is not evicted._resources + assert logger.Langfuse not in _lifecycle_state(logger.Langfuse).retired + shutdown_langfuse_client(logger.Langfuse) + + def _exported_span(logger, exporter): logger.Langfuse.flush() return exporter.get_finished_spans()[-1]