From 2f7b0db25eafc4db2ed5514b4f262c6257595a09 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 12 Aug 2026 18:56:34 -0700 Subject: [PATCH] fix(otel): take the provider hold inside route_for to close the eviction race pre_call can run on thread-pool workers, so between route_for returning a provider and the caller recording its open span, a concurrent request could overflow the LRU and shut that provider down with a zero span count, dropping the routed trace. route_for now increments the open-span count in the same locked critical section as the cache update and hands back an already-held provider; every caller releases it once its span has landed. The lock also makes the cache mutations safe under that same thread-pool concurrency. --- litellm/integrations/otel/logger.py | 105 +++++++++-------- litellm/integrations/otel/plumbing/routing.py | 106 ++++++++++++------ .../integrations/otel/test_otel_v2_dynamic.py | 19 ++-- 3 files changed, 139 insertions(+), 91 deletions(-) diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index dde9c33bfa6..d5014ec46ca 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -288,19 +288,23 @@ class OpenTelemetryV2(CustomLogger): # trace) — see ``TenantRoute.detached``. route: Final = self._tenant_tracers.route_for(self.tracer, call.dynamic_params, call.auth_metadata) parent_context: Final = resolve_request_span_context() - if is_recordable_span(get_current_span(parent_context)): - span = self._emitter.start_span( - SpanRole.LLM_CALL, - call.provisional_span_name, - parent_context=( - set_span_in_context(INVALID_SPAN, parent_context) if route.detached else parent_context - ), - start_time_ns=start_time_ns, - tracer=route.tracer, - links=_request_trace_links(parent_context) if route.detached else None, - ) - if span is not None and route.provider is not None: - self._tenant_tracers.hold(route.provider) + try: + if is_recordable_span(get_current_span(parent_context)): + span = self._emitter.start_span( + SpanRole.LLM_CALL, + call.provisional_span_name, + parent_context=( + set_span_in_context(INVALID_SPAN, parent_context) if route.detached else parent_context + ), + start_time_ns=start_time_ns, + tracer=route.tracer, + links=_request_trace_links(parent_context) if route.detached else None, + ) + finally: + # No span means route_for's hold has no owner (creation deferred to + # the close callback, which re-routes) — drop it here. + if span is None: + self._tenant_tracers.release(route.provider) self._open_llm_calls[call_id] = _LLMCallSpan( span=span, start_time_ns=start_time_ns, @@ -407,17 +411,20 @@ class OpenTelemetryV2(CustomLogger): if data.identity.call_id: self._release_carrier(self._open_llm_calls.pop(data.identity.call_id, None)) route: Final = self._tenant_tracers.route_for(self.tracer, None, auth_metadata(payload, kwargs)) - parent_context, links = resolve_mcp_span_context() - seeded: Final = self._seed_identity_baggage(data.identity, None, parent_context) - self._emitter.emit( - SpanRole.MCP_TOOL_CALL, - data, - parent_context=(set_span_in_context(INVALID_SPAN, seeded) if route.detached else seeded), - start_time_ns=to_ns(start_time), - end_time_ns=to_ns(end_time), - links=((*(links or ()), *(_request_trace_links(seeded) or ())) if route.detached else links), - tracer=route.tracer, - ) + try: + parent_context, links = resolve_mcp_span_context() + seeded: Final = self._seed_identity_baggage(data.identity, None, parent_context) + self._emitter.emit( + SpanRole.MCP_TOOL_CALL, + data, + parent_context=(set_span_in_context(INVALID_SPAN, seeded) if route.detached else seeded), + start_time_ns=to_ns(start_time), + end_time_ns=to_ns(end_time), + links=((*(links or ()), *(_request_trace_links(seeded) or ())) if route.detached else links), + tracer=route.tracer, + ) + finally: + self._tenant_tracers.release(route.provider) return True def _emit_mcp_list_tools( @@ -445,17 +452,20 @@ class OpenTelemetryV2(CustomLogger): if data.identity.call_id: self._release_carrier(self._open_llm_calls.pop(data.identity.call_id, None)) route: Final = self._tenant_tracers.route_for(self.tracer, None, auth_metadata(payload, kwargs)) - parent_context, links = resolve_mcp_span_context() - seeded: Final = self._seed_identity_baggage(data.identity, None, parent_context) - self._emitter.emit( - SpanRole.MCP_LIST_TOOLS, - data, - parent_context=(set_span_in_context(INVALID_SPAN, seeded) if route.detached else seeded), - start_time_ns=to_ns(start_time), - end_time_ns=to_ns(end_time), - links=((*(links or ()), *(_request_trace_links(seeded) or ())) if route.detached else links), - tracer=route.tracer, - ) + try: + parent_context, links = resolve_mcp_span_context() + seeded: Final = self._seed_identity_baggage(data.identity, None, parent_context) + self._emitter.emit( + SpanRole.MCP_LIST_TOOLS, + data, + parent_context=(set_span_in_context(INVALID_SPAN, seeded) if route.detached else seeded), + start_time_ns=to_ns(start_time), + end_time_ns=to_ns(end_time), + links=((*(links or ()), *(_request_trace_links(seeded) or ())) if route.detached else links), + tracer=route.tracer, + ) + finally: + self._tenant_tracers.release(route.provider) return True def _close_llm_call( @@ -486,7 +496,7 @@ class OpenTelemetryV2(CustomLogger): def _release_carrier(self, carrier: "_LLMCallSpan | None") -> None: """Release the routed provider a removed carrier was holding open.""" - if carrier is not None and carrier.provider is not None: + if carrier is not None: self._tenant_tracers.release(carrier.provider) def _finish_carrier( @@ -520,16 +530,19 @@ class OpenTelemetryV2(CustomLogger): # Baggage so the span — and the SDK path, which has none — is labeled # consistently. A detached route roots its own trace instead, linked back. route: Final = self._tenant_tracers.route_for(self.tracer, call.dynamic_params, call.auth_metadata) - parent_ctx = self._seed_identity_baggage(data.identity, data.request_model, resolve_request_span_context()) - return self._emitter.emit( - SpanRole.LLM_CALL, - data, - parent_context=(set_span_in_context(INVALID_SPAN, parent_ctx) if route.detached else parent_ctx), - start_time_ns=carrier.start_time_ns, - end_time_ns=end_time_ns, - tracer=route.tracer, - links=_request_trace_links(parent_ctx) if route.detached else None, - ) + try: + parent_ctx = self._seed_identity_baggage(data.identity, data.request_model, resolve_request_span_context()) + return self._emitter.emit( + SpanRole.LLM_CALL, + data, + parent_context=(set_span_in_context(INVALID_SPAN, parent_ctx) if route.detached else parent_ctx), + start_time_ns=carrier.start_time_ns, + end_time_ns=end_time_ns, + tracer=route.tracer, + links=_request_trace_links(parent_ctx) if route.detached else None, + ) + finally: + self._tenant_tracers.release(route.provider) # ====================================================================== # # Service hooks diff --git a/litellm/integrations/otel/plumbing/routing.py b/litellm/integrations/otel/plumbing/routing.py index 3a440bdb90c..0d6b007b188 100644 --- a/litellm/integrations/otel/plumbing/routing.py +++ b/litellm/integrations/otel/plumbing/routing.py @@ -10,6 +10,7 @@ tracer. This lets a single logger fan requests out to many tenants without needing a logger per tenant. """ +import threading from collections import OrderedDict from collections.abc import Mapping from dataclasses import dataclass @@ -92,9 +93,10 @@ class TenantRoute: tracer: Tracer detached: bool - #: The provider ``tracer`` came from, or ``None`` on the default route. A - #: caller that keeps a span open across callbacks must ``hold``/``release`` - #: it so LRU eviction can't shut the provider down under the open span. + #: The provider ``tracer`` came from, or ``None`` on the default route. It + #: is returned already held (counted as an open span, atomically with the + #: cache update), so LRU eviction can't shut it down before the caller's + #: span lands; the caller must ``release`` it exactly once when done. provider: TracerProvider | None = None @@ -110,6 +112,10 @@ class TenantTracerCache: self._config = config self._callback_name = callback_name self._tracer_name = tracer_name + # Guards the three mutable structures below: ``pre_call`` can run on + # thread-pool workers concurrently with the event loop, so cache + # updates, span counts, and retirement must be atomic. + self._lock: Final = threading.Lock() self._providers: OrderedDict[tuple[_HeaderItems, _HeaderItems], TracerProvider] = OrderedDict() self._open_span_counts: dict[TracerProvider, int] = {} # mutable-ok: live refcount state self._retired: set[TracerProvider] = set() # mutable-ok: evicted providers draining open spans @@ -119,34 +125,26 @@ class TenantTracerCache: ) self._warned_project_unroutable = False - def hold(self, provider: TracerProvider) -> None: - """Count an open span on ``provider`` so eviction defers its shutdown.""" - self._open_span_counts[provider] = self._open_span_counts.get(provider, 0) + 1 + def release(self, provider: TracerProvider | None) -> None: + """Drop one open-span count; shut a retired provider down once drained. - def release(self, provider: TracerProvider) -> None: - """Drop one open-span count; shut a retired provider down once drained.""" - remaining: Final = self._open_span_counts.get(provider, 0) - 1 - if remaining > 0: - self._open_span_counts[provider] = remaining - return - self._open_span_counts.pop(provider, None) - if provider in self._retired: - self._retired.discard(provider) - _shutdown_provider(provider) - - def _retire(self, provider: TracerProvider) -> None: - """Shut an evicted provider down, or defer while spans are still open. - - Shutting down immediately would stop the provider's processors while a - span opened at ``pre_call`` is still live, silently dropping it at end - instead of exporting it. A retired provider is pinned only by its open - spans, and the logger's open-call map is itself bounded, so a request - that never reaches a close callback still releases it eventually. + ``None`` (the default route) is a no-op so callers can release a + ``TenantRoute.provider`` unconditionally. The shutdown itself runs + outside the lock: it force-flushes over the network and must not stall + every concurrently routing request. """ - if self._open_span_counts.get(provider, 0) > 0: - self._retired.add(provider) + if provider is None: return - _shutdown_provider(provider) + with self._lock: + remaining: Final = self._open_span_counts.get(provider, 0) - 1 + if remaining > 0: + self._open_span_counts[provider] = remaining + return + self._open_span_counts.pop(provider, None) + drained: Final = provider in self._retired + self._retired.discard(provider) + if drained: + _shutdown_provider(provider) def route_for( self, @@ -160,6 +158,11 @@ class TenantTracerCache: project require a scoped tracer, in which case build (or reuse) one. The cache is a bounded LRU: the least-recently-used provider is flushed and shut down on overflow so its exporter threads don't accumulate. + + A routed provider is returned already held — its open-span count is + incremented in the same critical section as the cache update — so a + concurrent overflow eviction can't shut it down between selection and + the caller's span start. The caller must ``release`` it exactly once. """ credential_headers: Final = dynamic_otlp_headers(self._callback_name, dynamic_params) or _NO_HEADERS project_headers: Final = self._project_headers(auth_metadata) @@ -169,21 +172,50 @@ class TenantTracerCache: tuple(sorted(credential_headers.items())), tuple(sorted(project_headers.items())), ) - provider = self._providers.get(cache_key) - if provider is not None: - self._providers.move_to_end(cache_key) - else: - provider = build_tracer_provider(self._routed_config(credential_headers, project_headers)) - self._providers[cache_key] = provider - if len(self._providers) > _MAX_CACHED_PROVIDERS: - _, evicted = self._providers.popitem(last=False) - self._retire(evicted) + with self._lock: + provider: Final = self._cached_provider_locked(cache_key, credential_headers, project_headers) + self._open_span_counts[provider] = self._open_span_counts.get(provider, 0) + 1 + evicted: Final = self._evicted_on_overflow_locked() + if evicted is not None: + _shutdown_provider(evicted) return TenantRoute( tracer=get_tracer(provider, self._tracer_name), detached=bool(project_headers), provider=provider, ) + def _cached_provider_locked( + self, + cache_key: tuple[_HeaderItems, _HeaderItems], + credential_headers: Mapping[str, str], + project_headers: Mapping[str, str], + ) -> TracerProvider: + cached: Final = self._providers.get(cache_key) + if cached is not None: + self._providers.move_to_end(cache_key) + return cached + built: Final = build_tracer_provider(self._routed_config(credential_headers, project_headers)) + self._providers[cache_key] = built + return built + + def _evicted_on_overflow_locked(self) -> TracerProvider | None: + """Pop the LRU provider past the cap; return it if the caller must shut it down. + + A provider with open spans is retired to drain instead: stopping its + processors while a span opened at ``pre_call`` is still live would + silently drop that span at end instead of exporting it. A retired + provider is pinned only by its open spans, and the logger's open-call + map is itself bounded, so a request that never reaches a close callback + still releases it eventually. + """ + if len(self._providers) <= _MAX_CACHED_PROVIDERS: + return None + _, evicted = self._providers.popitem(last=False) + if self._open_span_counts.get(evicted, 0) > 0: + self._retired.add(evicted) + return None + return evicted + def _project_headers(self, auth_metadata: Mapping[str, str] | None) -> Mapping[str, str]: """The per-request project-routing headers, if this cache can apply them. diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py b/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py index 8d14a58e352..408f010dbd7 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py @@ -101,10 +101,11 @@ def test_provider_cache_is_bounded_and_evicts_lru(monkeypatch): def creds(space): return {"arize_space_id": space, "arize_api_key": "K"} - cache.route_for(default, creds("1")) - cache.route_for(default, creds("2")) - cache.route_for(default, creds("1")) # touch "1" → "2" is now LRU - cache.route_for(default, creds("3")) # overflow → evict "2" + # route_for returns a held provider; release models the span closing. + cache.release(cache.route_for(default, creds("1")).provider) + cache.release(cache.route_for(default, creds("2")).provider) + cache.release(cache.route_for(default, creds("1")).provider) # touch "1" → "2" is now LRU + cache.release(cache.route_for(default, creds("3")).provider) # overflow → evict "2" assert len(cache._providers) == 2 assert len(shut_down) == 1 # exactly the evicted provider was shut down @@ -303,7 +304,10 @@ def test_grpc_exporter_gets_no_project_routing(): def test_eviction_defers_shutdown_while_a_span_is_open(monkeypatch): # An LLM span opened at pre_call stays open until the close callback; LRU # eviction in that window must not stop the provider's processors, or the - # span is silently dropped at end instead of exported. + # span is silently dropped at end instead of exported. route_for itself + # takes the hold, atomically with the cache update, so a concurrent + # eviction can never shut a just-selected provider down before the caller + # records its span. from litellm.integrations.otel.plumbing import routing as routing_mod monkeypatch.setattr(routing_mod, "_MAX_CACHED_PROVIDERS", 1) @@ -316,9 +320,8 @@ def test_eviction_defers_shutdown_while_a_span_is_open(monkeypatch): route_a = cache.route_for(default, {"arize_space_id": "A", "arize_api_key": "K"}) assert route_a.provider is not None - cache.hold(route_a.provider) cache.route_for(default, {"arize_space_id": "B", "arize_api_key": "K"}) # evicts A - assert shut_down == [] # deferred: A still has an open span + assert shut_down == [] # deferred: A is still held by route_a cache.release(route_a.provider) assert shut_down == [route_a.provider] @@ -332,6 +335,6 @@ def test_release_without_eviction_keeps_provider_alive(monkeypatch): ) cache = _cache("arize") route = cache.route_for(NoOpTracer(), {"arize_space_id": "A", "arize_api_key": "K"}) - cache.hold(route.provider) cache.release(route.provider) assert shut_down == [] # still cached, never retired + cache.release(None) # default-route release is a no-op