mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(otel): fix thread leak in OpenTelemetry dynamic header path
Implemented TracerProvider caching in OpenTelemetry class to prevent linear thread growth when using dynamic headers (e.g., per-team Langfuse keys). Added formal unit test in tests/test_otel_thread_leak.py.
This commit is contained in:
parent
790a5ce0b3
commit
4b6b1da517
2 changed files with 99 additions and 0 deletions
|
|
@ -140,6 +140,7 @@ class OpenTelemetry(CustomLogger):
|
|||
self.OTEL_EXPORTER = self.config.exporter
|
||||
self.OTEL_ENDPOINT = self.config.endpoint
|
||||
self.OTEL_HEADERS = self.config.headers
|
||||
self._tracer_provider_cache: Dict[str, Any] = {}
|
||||
self._init_tracing(tracer_provider)
|
||||
|
||||
_debug_otel = str(os.getenv("DEBUG_OTEL", "False")).lower()
|
||||
|
|
@ -611,12 +612,20 @@ class OpenTelemetry(CustomLogger):
|
|||
"""Create a temporary tracer with dynamic headers for this request only."""
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
|
||||
# Prevents thread exhaustion by reusing providers for the same credential sets (e.g. per-team keys)
|
||||
cache_key = str(sorted(dynamic_headers.items()))
|
||||
if cache_key in self._tracer_provider_cache:
|
||||
return self._tracer_provider_cache[cache_key].get_tracer(LITELLM_TRACER_NAME)
|
||||
|
||||
# Create a temporary tracer provider with dynamic headers
|
||||
temp_provider = TracerProvider(resource=self._get_litellm_resource(self.config))
|
||||
temp_provider.add_span_processor(
|
||||
self._get_span_processor(dynamic_headers=dynamic_headers)
|
||||
)
|
||||
|
||||
# Store in cache for reuse
|
||||
self._tracer_provider_cache[cache_key] = temp_provider
|
||||
|
||||
return temp_provider.get_tracer(LITELLM_TRACER_NAME)
|
||||
|
||||
def construct_dynamic_otel_headers(
|
||||
|
|
|
|||
90
tests/test_otel_thread_leak.py
Normal file
90
tests/test_otel_thread_leak.py
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
import sys
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import pytest
|
||||
|
||||
# Add the project root to the path
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
def get_thread_count() -> int:
|
||||
"""Helper to get active thread count"""
|
||||
return threading.active_count()
|
||||
|
||||
@pytest.fixture
|
||||
def otel_logger():
|
||||
"""Fixture to provide a clean OTEL logger for each test"""
|
||||
config = OpenTelemetryConfig(
|
||||
exporter="console",
|
||||
enable_metrics=False,
|
||||
service_name="litellm-unit-test"
|
||||
)
|
||||
return OpenTelemetry(config=config)
|
||||
|
||||
def test_otel_thread_leak_dynamic_headers(otel_logger):
|
||||
"""
|
||||
Unit test to verify that calling get_tracer_to_use_for_request with
|
||||
dynamic headers doesn't cause a linear thread leak.
|
||||
|
||||
This test reproduces the issue where each unique team/key credential
|
||||
set causes a new TracerProvider (and its background threads) to be
|
||||
spawned but never closed.
|
||||
"""
|
||||
|
||||
# 1. Setup dynamic header simulation (monkey-patch)
|
||||
# This simulates what LangfuseOtelLogger does for per-team keys
|
||||
def mock_construct_dynamic_headers(standard_callback_dynamic_params):
|
||||
if standard_callback_dynamic_params:
|
||||
return {"Authorization": "Bearer fake_token"}
|
||||
return None
|
||||
|
||||
otel_logger.construct_dynamic_otel_headers = mock_construct_dynamic_headers
|
||||
|
||||
# 2. Establish Baseline
|
||||
initial_threads = get_thread_count()
|
||||
|
||||
# 3. Simulate requests
|
||||
num_requests = 10
|
||||
latencies = []
|
||||
|
||||
print("\n🚀 Simulating requests with dynamic headers:")
|
||||
for i in range(num_requests):
|
||||
kwargs = {
|
||||
"standard_callback_dynamic_params": StandardCallbackDynamicParams(
|
||||
langfuse_public_key=f"key_{i}",
|
||||
langfuse_secret_key=f"secret_{i}",
|
||||
)
|
||||
}
|
||||
|
||||
# Measure latency
|
||||
start_time = time.perf_counter()
|
||||
tracer = otel_logger.get_tracer_to_use_for_request(kwargs)
|
||||
end_time = time.perf_counter()
|
||||
|
||||
latency_ms = (end_time - start_time) * 1000
|
||||
latencies.append(latency_ms)
|
||||
print(f" Request {i+1:2d}: Latency = {latency_ms:6.2f} ms")
|
||||
|
||||
# Verify a tracer was actually returned
|
||||
assert tracer is not None
|
||||
|
||||
avg_latency = sum(latencies) / len(latencies)
|
||||
print(f"\n📊 Average Latency: {avg_latency:.2f} ms")
|
||||
|
||||
# 4. Check for leaks
|
||||
# Allow for a small constant increase (OTEL might start a few shared threads)
|
||||
# but a linear leak would result in +10 or more threads here.
|
||||
final_threads = get_thread_count()
|
||||
thread_delta = final_threads - initial_threads
|
||||
|
||||
print(f"\nThread growth: {thread_delta} threads across {num_requests} requests")
|
||||
|
||||
# ASSERTION: The growth should be significantly less than 1 thread per request.
|
||||
# If the bug exists, thread_delta will be >= num_requests.
|
||||
assert thread_delta < (num_requests / 2), (
|
||||
f"Thread leak detected! Threads grew by {thread_delta} over {num_requests} requests. "
|
||||
"Each request with dynamic headers appears to be leaking background threads."
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue