mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
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.
This commit is contained in:
mubashir1osmani
parent
cece0789ea
commit
2f7b0db25e
3 changed files with 139 additions and 91 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue