diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index f52323a89f5..dadc01e16f2 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -459,20 +459,23 @@ class LangFuseLogger: status_message=status_message, ) verbose_logger.debug("OUTPUT IN LANGFUSE: %s; original: %s", output, response_obj) - trace_id, generation_id = self._log_langfuse_v2( - user_id=user_id, - metadata=metadata, - litellm_params=litellm_params, - output=output, - start_time=start_time, - end_time=end_time, - kwargs=kwargs, - optional_params=optional_params, - input=input, - response_obj=response_obj, - level=level, - litellm_call_id=litellm_call_id, - ) + from litellm.integrations.langfuse.langfuse_sdk import lease_langfuse_client + + with lease_langfuse_client(self.Langfuse): + trace_id, generation_id = self._log_langfuse_v2( + user_id=user_id, + metadata=metadata, + litellm_params=litellm_params, + output=output, + start_time=start_time, + end_time=end_time, + kwargs=kwargs, + optional_params=optional_params, + input=input, + response_obj=response_obj, + level=level, + litellm_call_id=litellm_call_id, + ) verbose_logger.debug("Langfuse Layer Logging - final response object: %s", response_obj) verbose_logger.info("Langfuse Layer Logging - logging success") diff --git a/litellm/integrations/langfuse/langfuse_sdk.py b/litellm/integrations/langfuse/langfuse_sdk.py index e443b508dd8..1d4979e0924 100644 --- a/litellm/integrations/langfuse/langfuse_sdk.py +++ b/litellm/integrations/langfuse/langfuse_sdk.py @@ -4,7 +4,8 @@ import os import re import threading from base64 import b64encode -from collections.abc import Mapping +from collections.abc import Generator, Mapping +from contextlib import contextmanager from datetime import datetime from hashlib import sha256 from types import MappingProxyType @@ -28,6 +29,7 @@ __all__ = ( "acquire_langfuse_client", "build_isolated_tracer_provider", "evict_stale_langfuse_resources", + "lease_langfuse_client", "open_trace_context", "propagate_attributes", "register_langfuse_client", @@ -221,6 +223,137 @@ _LIVE_CLIENTS_LOCK: Final = threading.Lock() _live_clients: Final[WeakKeyDictionary[LangfuseResourceManager, WeakSet]] = WeakKeyDictionary() +class _LangfuseLifecycleState: + """How many callbacks are leasing one SDK resource bundle, and what eviction has queued behind them. + + ``lock`` is never held across a teardown, which takes the SDK's own registry lock. + """ + + def __init__(self) -> None: + self.lock = threading.Lock() + self.active_leases = 0 + 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 + + def open_lease(self) -> None: + with self.lock: + self.active_leases += 1 + + 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: + if self.active_leases > 0 or self.teardown_in_progress: + self.pending_clients.add(client) + return False + self.teardown_in_progress = True + self.teardown_owner = threading.get_ident() + return True + + def release_lease(self) -> tuple[Langfuse, ...]: + """Drop this lease and take ownership of the teardowns it was holding up, if it was the last one.""" + with self.lock: + self.active_leases -= 1 + if self.active_leases > 0 or self.teardown_in_progress or not self.pending_clients: + return () + claimed: Final = tuple(self.pending_clients) + self.pending_clients.clear() + self.teardown_in_progress = True + self.teardown_owner = threading.get_ident() + return claimed + + def next_teardown_batch(self) -> tuple[Langfuse, ...]: + """Whatever eviction queued while the last batch was draining, handing the ownership flag back when empty.""" + with self.lock: + if self.active_leases == 0 and self.pending_clients: + claimed: Final = tuple(self.pending_clients) + self.pending_clients.clear() + return claimed + self.teardown_in_progress = False + self.teardown_owner = None + return () + + def requeue(self, clients: tuple[Langfuse, ...]) -> None: + with self.lock: + self.pending_clients.update(clients) + + def end_teardown(self) -> None: + with self.lock: + if self.teardown_owner == threading.get_ident(): + self.teardown_in_progress = False + self.teardown_owner = None + + +_LIFECYCLE_STATES_LOCK: Final = threading.Lock() +_LIFECYCLE_STATES: Final[WeakKeyDictionary[object, _LangfuseLifecycleState]] = WeakKeyDictionary() + + +def _lifecycle_state(client: Langfuse) -> _LangfuseLifecycleState: + """One state per resource bundle, since teardown closes the provider every client on that bundle exports through.""" + resources: Final = getattr(client, "_resources", None) + key: Final = client if resources is None else resources + with _LIFECYCLE_STATES_LOCK: + existing: Final = _LIFECYCLE_STATES.get(key) + if existing is not None: + return existing + created: Final = _LangfuseLifecycleState() + _LIFECYCLE_STATES[key] = created + return created + + +@contextmanager +def lease_langfuse_client(client: Langfuse) -> Generator[None]: + """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. + """ + state: Final = _lifecycle_state(client) + state.open_lease() + try: + yield + finally: + _run_teardowns(state, state.release_lease(), propagate_base_exception=False) + + +def _run_teardowns( + state: _LangfuseLifecycleState, + clients: tuple[Langfuse, ...], + *, + propagate_base_exception: bool = True, +) -> None: + """Tear down ``clients``, then whatever eviction queued meanwhile, and hand the flag back. + + A failing ordinary teardown is logged and skipped rather than raised: the thread here is usually a + request callback that merely held the last lease, and its request must not fail on eviction's behalf. + Interrupts requeue the unfinished batch and normally propagate, while a callback exception already + in flight takes precedence over an eviction interrupt. + """ + batch = clients # rebind-ok: drains each batch queued while the previous one was being torn down + try: + from litellm._logging import verbose_logger + + while batch: + for index, client in enumerate(batch): + try: + _teardown_langfuse_client(client) + except Exception: + verbose_logger.exception("Langfuse client teardown failed during cache eviction") + except BaseException: + state.requeue(batch[index:]) + if propagate_base_exception: + raise + return + batch = state.next_teardown_batch() + finally: + state.end_teardown() + + def _evict_if_stale_locked( *, public_key: object, secret_key: object, base_url: object ) -> LangfuseResourceManager | None: @@ -398,6 +531,19 @@ def shutdown_langfuse_client(client: Langfuse) -> None: tracer provider's export thread running and leaves the client in the registry, so a later request for the same key gets a dead client back. + A callback holding a lease on the client's bundle postpones all of this to + the moment that lease ends, so eviction cannot close the provider out from + under an export the lease is wrapping. See ``lease_langfuse_client``. + """ + state: Final = _lifecycle_state(client) + if not state.claim_for_teardown(client): + return + _run_teardowns(state, (client,)) + + +def _teardown_langfuse_client(client: Langfuse) -> None: + """The blocking teardown behind ``shutdown_langfuse_client``. + A client that shares its resources with another live client only flushes: shutting the shared provider down here would silence the other client for the rest of its life, as it did before the reference count existed. diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py index 0693fb6ad9f..b980768657c 100644 --- a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py +++ b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py @@ -6,6 +6,7 @@ call would otherwise record its own duration instead of the call's. """ import json +import threading from datetime import datetime, timedelta, timezone import opentelemetry.trace as otel_trace @@ -25,9 +26,12 @@ from litellm.integrations.langfuse.langfuse_sdk import ( AS_ROOT_ATTRIBUTE, PUBLIC_ATTRIBUTE, RELEASE_ATTRIBUTE, + _lifecycle_state, _litellm_built_providers, + _teardown_langfuse_client, build_isolated_tracer_provider, evict_stale_langfuse_resources, + lease_langfuse_client, open_trace_context, register_langfuse_client, resolve_observation_id, @@ -102,7 +106,12 @@ def test_guardrail_span_with_float_timestamps_does_not_break_the_generation(clie context, claim_root = open_trace_context(client=lf, trace_id="9" * 32, parent_observation_id=None) guardrail_start = 1709294400.0 start_child_span( - client=lf, context=context, name="guardrail", start_time=guardrail_start, claim_trace_root=claim_root, attributes={} + client=lf, + context=context, + name="guardrail", + start_time=guardrail_start, + claim_trace_root=claim_root, + attributes={}, ).end(end_time=to_unix_nanos(guardrail_start + 2)) start_generation( client=lf, context=context, name="gen", start_time=CALL_START, claim_trace_root=claim_root, attributes={} @@ -143,7 +152,12 @@ def test_child_span_keeps_its_own_window_and_stays_a_sibling(client): context, claim_root = open_trace_context(client=lf, trace_id="d" * 32, parent_observation_id=None) guardrail_start = CALL_START + timedelta(seconds=1) start_child_span( - client=lf, context=context, name="guardrail", start_time=guardrail_start, claim_trace_root=claim_root, attributes={} + client=lf, + context=context, + name="guardrail", + start_time=guardrail_start, + claim_trace_root=claim_root, + attributes={}, ).end(end_time=to_unix_nanos(guardrail_start + timedelta(seconds=2))) start_generation( client=lf, context=context, name="gen", start_time=CALL_START, claim_trace_root=claim_root, attributes={} @@ -550,6 +564,153 @@ def test_evicting_a_client_that_shares_resources_keeps_the_other_exporting(): assert LangfuseResourceManager._instances.get(PUBLIC_KEY) is first._resources +def test_eviction_defers_teardown_until_active_callback_finishes(): + """A cached client must keep exporting while its callback lease is active.""" + exporter = InMemorySpanExporter() + provider = build_isolated_tracer_provider(environment=None, release=None) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + client = 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(client) + + def evict() -> None: + shutdown_langfuse_client(client) + + with lease_langfuse_client(client): + evictor = threading.Thread(target=evict) + evictor.start() + evictor.join(timeout=5) + assert not evictor.is_alive() + assert not exporter._stopped + + context, claim_root = open_trace_context(client=client, trace_id="a" * 32, parent_observation_id=None) + start_generation( + client=client, + context=context, + name="active-callback", + start_time=None, + claim_trace_root=claim_root, + attributes={}, + ).end() + client.flush() + assert any(span.name == "active-callback" for span in exporter.get_finished_spans()) + + evictor.join(timeout=5) + assert not evictor.is_alive() + assert exporter._stopped + assert LangfuseResourceManager._instances.get(PUBLIC_KEY) is not client._resources + + +def test_teardown_failure_does_not_strand_queued_clients(monkeypatch): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + first = Langfuse( + public_key=PUBLIC_KEY, + secret_key="sk-original", + host="http://127.0.0.1:1", + tracer_provider=provider, + span_exporter=exporter, + ) + clients = ( + first, + Langfuse(public_key=PUBLIC_KEY, secret_key="sk-original", host="http://127.0.0.1:1"), + Langfuse(public_key=PUBLIC_KEY, secret_key="sk-original", host="http://127.0.0.1:1"), + ) + assert len({client._resources for client in clients}) == 1 + for client in clients: + register_langfuse_client(client) + state = _lifecycle_state(clients[0]) + original_teardown = _teardown_langfuse_client + calls = [] + + def teardown(client): + calls.append(client) + original_teardown(client) + if len(calls) == 1: + raise RuntimeError("teardown failed") + + monkeypatch.setattr("litellm.integrations.langfuse.langfuse_sdk._teardown_langfuse_client", teardown) + with lease_langfuse_client(clients[0]): + for client in clients: + shutdown_langfuse_client(client) + + assert len(calls) == 3 + assert not state.pending_clients + assert not state.teardown_in_progress + + +def test_queued_eviction_waits_for_the_last_of_two_overlapping_leases(): + exporter = InMemorySpanExporter() + provider = build_isolated_tracer_provider(environment=None, release=None) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + client = 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(client) + + with lease_langfuse_client(client): + with lease_langfuse_client(client): + shutdown_langfuse_client(client) + assert not exporter._stopped + + assert exporter._stopped + + +def test_a_client_adopted_during_deferred_teardown_keeps_exporting(): + """The registry hands the same bundle back out while its teardown is queued behind a lease; + the holder count must degrade that teardown to a flush.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + _litellm_built_providers.add(provider) + 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) + + with lease_langfuse_client(evicted): + 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 + register_langfuse_client(adopter) + + assert _exports(adopter, exporter, "after-deferred-teardown") + assert LangfuseResourceManager._instances.get(PUBLIC_KEY) is adopter._resources + + +def test_leases_on_one_client_do_not_serialise_callbacks(): + """Every langfuse callback in the process shares one client, so leases must overlap.""" + client = _lifecycle_client() + both_inside = threading.Barrier(2, timeout=5) + + def hold_lease() -> None: + with lease_langfuse_client(client): + both_inside.wait() + + holders = tuple(threading.Thread(target=hold_lease) for _ in range(2)) + for holder in holders: + holder.start() + for holder in holders: + holder.join(timeout=5) + + assert not any(holder.is_alive() for holder in holders) + assert not both_inside.broken + + def test_last_client_on_shared_resources_tears_them_down(): first, second, exporter = _shared_resources_pair() shutdown_langfuse_client(second) diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 46485d7145e..10d611e1f9d 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -10,7 +10,7 @@ import pytest import litellm from litellm.integrations.langfuse import langfuse as langfuse_module from litellm.integrations.langfuse.langfuse import LangFuseLogger -from litellm.integrations.langfuse.langfuse_sdk import resolve_trace_id +from litellm.integrations.langfuse.langfuse_sdk import _lifecycle_state, resolve_trace_id # Import LangfuseUsageDetails directly from the module where it's defined @@ -1386,6 +1386,34 @@ def _steering_logger(): return logger, exporter +def test_log_event_holds_a_client_lease_during_export(): + logger, _ = _steering_logger() + state = _lifecycle_state(logger.Langfuse) + + def assert_lease_is_active(**_: object) -> tuple[str, str]: + assert state.active_leases == 1 + return "trace-id", "generation-id" + + now = datetime.datetime.now() + with patch.object(logger, "_log_langfuse_v2", side_effect=assert_lease_is_active): + 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": "trace-id", "generation_id": "generation-id"} + assert state.active_leases == 0 + + def _emit(rig, *, metadata=None, headers=None): """``log_event_on_langfuse`` is the entry point that folds ``langfuse_*`` headers into metadata.