mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
fix(langfuse): retire the provider a credential rotation abandons
Rotating the secret or host for one public key evicted the registry entry but left the tracer provider litellm built for it running, so repeated rotations leaked an export thread each time. The evicted bundle is now retired after the registry lock is released, and only when no live client is still on it.
This commit is contained in:
parent
d4a1cd312f
commit
442d85a23c
2 changed files with 80 additions and 17 deletions
|
|
@ -213,25 +213,51 @@ class DiscardingSpanExporter(SpanExporter):
|
|||
return True
|
||||
|
||||
|
||||
_LIVE_CLIENTS_LOCK: Final = threading.Lock()
|
||||
# litellm clients still using each SDK resource bundle; the bundle is torn down with the last one.
|
||||
# Both sides are weak so a throwaway client (a health probe, an alerting lookup) that is simply
|
||||
# garbage-collected stops holding the bundle open rather than inflating a counter forever.
|
||||
_live_clients: Final[WeakKeyDictionary[LangfuseResourceManager, WeakSet]] = WeakKeyDictionary()
|
||||
|
||||
|
||||
def _evict_if_stale_locked(
|
||||
*, public_key: object, secret_key: object, base_url: object
|
||||
) -> LangfuseResourceManager | None:
|
||||
"""Assumes ``LangfuseResourceManager._lock`` is held; returns the still-valid cached bundle, if any."""
|
||||
) -> tuple[LangfuseResourceManager | None, LangfuseResourceManager | None]:
|
||||
"""Assumes ``LangfuseResourceManager._lock`` is held; returns the still-valid bundle and the evicted one."""
|
||||
if not public_key:
|
||||
return None
|
||||
return None, None
|
||||
cached: Final = LangfuseResourceManager._instances.get(public_key) # pyright: ignore[reportPrivateUsage] # registry has no public accessor
|
||||
if cached is None:
|
||||
return None
|
||||
return None, None
|
||||
if getattr(cached, "secret_key", None) == secret_key and getattr(cached, "base_url", None) == base_url:
|
||||
return cached
|
||||
LangfuseResourceManager._instances.pop(public_key, None) # pyright: ignore[reportPrivateUsage] # registry has no public accessor
|
||||
return None
|
||||
return cached, None
|
||||
return None, LangfuseResourceManager._instances.pop(public_key, None) # pyright: ignore[reportPrivateUsage] # registry has no public accessor
|
||||
|
||||
|
||||
def _shutdown_abandoned_provider(resources: LangfuseResourceManager | None) -> None:
|
||||
"""Retire the provider litellm built for an evicted bundle, once no live client is still on it.
|
||||
|
||||
Runs after the registry lock is released: provider shutdown flushes and joins the export
|
||||
thread, which must never happen under the lock every other client init contends on. A
|
||||
client still holding these resources tears them down itself in ``shutdown_langfuse_client``.
|
||||
"""
|
||||
if resources is None:
|
||||
return
|
||||
with _LIVE_CLIENTS_LOCK:
|
||||
holders: Final = _live_clients.get(resources)
|
||||
if holders is not None and len(holders) > 0:
|
||||
return
|
||||
_live_clients.pop(resources, None)
|
||||
provider: Final = getattr(resources, "tracer_provider", None)
|
||||
if provider is not None and provider in _litellm_built_providers:
|
||||
provider.shutdown()
|
||||
|
||||
|
||||
def evict_stale_langfuse_resources(*, public_key: str | None, secret_key: str | None, base_url: str | None) -> None:
|
||||
"""Drop a cached client whose credentials no longer match the ones being requested."""
|
||||
with LangfuseResourceManager._lock: # pyright: ignore[reportPrivateUsage] # registry has no public accessor
|
||||
_evict_if_stale_locked(public_key=public_key, secret_key=secret_key, base_url=base_url)
|
||||
_, abandoned = _evict_if_stale_locked(public_key=public_key, secret_key=secret_key, base_url=base_url)
|
||||
_shutdown_abandoned_provider(abandoned)
|
||||
|
||||
|
||||
def _build_verified_span_exporter(*, public_key: object, secret_key: object, base_url: object) -> SpanExporter | None:
|
||||
|
|
@ -300,7 +326,7 @@ def acquire_langfuse_client(
|
|||
)
|
||||
)
|
||||
with LangfuseResourceManager._lock: # pyright: ignore[reportPrivateUsage] # registry has no public accessor
|
||||
cached: Final = _evict_if_stale_locked(
|
||||
cached, abandoned = _evict_if_stale_locked(
|
||||
public_key=public_key,
|
||||
secret_key=parameters.get("secret_key"),
|
||||
base_url=parameters.get("base_url"),
|
||||
|
|
@ -313,16 +339,10 @@ def acquire_langfuse_client(
|
|||
span_exporter=span_exporter,
|
||||
)
|
||||
register_langfuse_client(client)
|
||||
_shutdown_abandoned_provider(abandoned)
|
||||
return client
|
||||
|
||||
|
||||
_LIVE_CLIENTS_LOCK: Final = threading.Lock()
|
||||
# litellm clients still using each SDK resource bundle; the bundle is torn down with the last one.
|
||||
# Both sides are weak so a throwaway client (a health probe, an alerting lookup) that is simply
|
||||
# garbage-collected stops holding the bundle open rather than inflating a counter forever.
|
||||
_live_clients: Final[WeakKeyDictionary[LangfuseResourceManager, WeakSet]] = WeakKeyDictionary()
|
||||
|
||||
|
||||
def register_langfuse_client(client: Langfuse) -> None:
|
||||
"""Track the client against the SDK resources it ended up with.
|
||||
|
||||
|
|
|
|||
|
|
@ -23,9 +23,9 @@ from litellm.integrations.langfuse.langfuse import (
|
|||
)
|
||||
from litellm.integrations.langfuse.langfuse_sdk import (
|
||||
AS_ROOT_ATTRIBUTE,
|
||||
_litellm_built_providers,
|
||||
PUBLIC_ATTRIBUTE,
|
||||
RELEASE_ATTRIBUTE,
|
||||
_litellm_built_providers,
|
||||
build_isolated_tracer_provider,
|
||||
evict_stale_langfuse_resources,
|
||||
open_trace_context,
|
||||
|
|
@ -572,6 +572,49 @@ def test_shutdown_of_a_stale_client_does_not_deregister_the_live_one():
|
|||
assert LangfuseResourceManager._instances.get(PUBLIC_KEY) is live._resources
|
||||
|
||||
|
||||
def _rotation_provider():
|
||||
"""A litellm-built provider on the lifecycle public key, exporting in memory."""
|
||||
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)
|
||||
assert _exports(client, exporter, "before-rotation")
|
||||
return client, exporter
|
||||
|
||||
|
||||
def test_rotation_retires_the_provider_no_client_is_left_on():
|
||||
"""Prompt management builds throwaway clients from request credentials.
|
||||
|
||||
Alternating the secret for one public key evicts a bundle nobody holds any
|
||||
more, and its export thread has to go with it or every rotation leaks one.
|
||||
"""
|
||||
import gc
|
||||
|
||||
client, exporter = _rotation_provider()
|
||||
del client
|
||||
gc.collect()
|
||||
|
||||
evict_stale_langfuse_resources(public_key=PUBLIC_KEY, secret_key="sk-rotated", base_url="http://127.0.0.1:2")
|
||||
|
||||
assert exporter._stopped
|
||||
|
||||
|
||||
def test_rotation_keeps_a_still_live_client_exporting():
|
||||
"""The evicted bundle is only retired when nothing is on it; a live logger must survive."""
|
||||
client, exporter = _rotation_provider()
|
||||
|
||||
evict_stale_langfuse_resources(public_key=PUBLIC_KEY, secret_key="sk-rotated", base_url="http://127.0.0.1:2")
|
||||
|
||||
assert _exports(client, exporter, "after-rotation")
|
||||
|
||||
|
||||
def test_ssl_exporter_is_only_built_with_custom_tls_material(monkeypatch, tmp_path):
|
||||
"""v4 exports over its own OTLP channel, so litellm's CA bundle must be rebuilt onto it."""
|
||||
import litellm
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue