diff --git a/litellm/constants.py b/litellm/constants.py index b43f6d9f358..4ce9d97cda0 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -721,6 +721,7 @@ FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO: Final = 0.5 REDACTED_BY_LITELM_STRING: Final = "REDACTED_BY_LITELM" MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = get_env_int("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50) LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS: Final = 10_000 +OTEL_SHUTDOWN_FLUSH_TIMEOUT_MILLIS: Final = 10_000 # Backpressure + lifetime bounds for the /v1/messages streaming relay (see # BaseAnthropicMessagesStreamingIterator.async_sse_wrapper). The relay queue is # bounded so a slow client throttles the upstream pump instead of letting it diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py index 5a03b91d470..98a92080b97 100644 --- a/litellm/integrations/arize/arize.py +++ b/litellm/integrations/arize/arize.py @@ -104,12 +104,14 @@ class ArizeLogger(OpenTelemetry): if tracer_provider is not None: self.tracer = tracer_provider.get_tracer("litellm") + self._tracer_provider = tracer_provider self.span_kind = SpanKind return provider: Final = TracerProvider(resource=self._get_litellm_resource(self.config)) provider.add_span_processor(self._get_span_processor()) self.tracer = provider.get_tracer("litellm") + self._tracer_provider = provider self.span_kind = SpanKind def _handle_success( diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index 150d814f811..6db01176b36 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -48,6 +48,14 @@ ARIZE_HOSTED_PHOENIX_ENDPOINT: Final = "https://otlp.arize.com/v1/traces" _MAX_PROJECT_PROVIDERS: Final = 64 +def _force_flushed(flushable: "SpanProcessor | TracerProvider", timeout_millis: int) -> bool: + try: + return flushable.force_flush(timeout_millis) + except Exception as e: # noqa: BLE001 # one failed flush must not skip the rest at shutdown + verbose_logger.debug("ArizePhoenixLogger: force_flush failed: %s", e) + return False + + class ArizePhoenixLogger(OpenTelemetry): """ Arize Phoenix logger that sends traces to a Phoenix endpoint. @@ -91,7 +99,7 @@ class ArizePhoenixLogger(OpenTelemetry): self.config.exporter, ) - def flush_tracer_providers(self) -> None: + def flush_tracer_providers(self, timeout_millis: int) -> bool: """ Flush all cached per-project providers and the shared span processor. @@ -99,27 +107,17 @@ class ArizePhoenixLogger(OpenTelemetry): spans may still reference evicted providers. """ if getattr(self, "_use_injected_tracer_provider", False): - return + return True shared_processor: Final[SpanProcessor | None] = getattr(self, "_shared_span_processor", None) - if shared_processor is not None: - try: - shared_processor.force_flush() - except Exception as e: - verbose_logger.debug( - "ArizePhoenixLogger: shared span processor force_flush failed: %s", - e, - ) - with getattr(self, "_project_providers_lock", threading.Lock()): project_providers: Final[Mapping[str, TracerProvider]] = getattr(self, "_project_providers", {}) - providers: Final = list(project_providers.values()) + providers: Final = tuple(project_providers.values()) - for provider in providers: - try: - provider.force_flush() - except Exception as e: - verbose_logger.debug("ArizePhoenixLogger: TracerProvider force_flush failed: %s", e) + flushables: Final[tuple[SpanProcessor | TracerProvider, ...]] = ( + (shared_processor, *providers) if shared_processor is not None else providers + ) + return all(tuple(_force_flushed(flushable, timeout_millis) for flushable in flushables)) def _get_litellm_resource_for_project(self, project_name: str): """ diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 9bec220e073..5736621a39b 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -67,6 +67,7 @@ if TYPE_CHECKING: from opentelemetry.trace import Span as _Span from opentelemetry.trace import SpanKind as _SpanKind from opentelemetry.trace import Tracer as _Tracer + from opentelemetry.trace import TracerProvider as _APITracerProvider from litellm.proxy._types import ( ManagementEndpointLoggingPayload as _ManagementEndpointLoggingPayload, @@ -440,6 +441,8 @@ class OpenTelemetryConfig: class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): + _tracer_provider: "_APITracerProvider" + def __init__( self, config: OpenTelemetryConfig | None = None, @@ -1279,6 +1282,16 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): cache_key: Final = str(sorted(dynamic_headers.items())) return self._cached_dynamic_tracer(cache_key, _build, owns_exporter) + def flush_tracer_providers(self, timeout_millis: int) -> bool: + from opentelemetry.sdk.trace import TracerProvider + + with self._tracer_provider_cache_lock: + cached: Final = tuple(entry.provider for entry in self._tracer_provider_cache.values()) + providers: Final = tuple( + provider for provider in (self._tracer_provider, *cached) if isinstance(provider, TracerProvider) + ) + return all(tuple(provider.force_flush(timeout_millis) for provider in providers)) + def construct_dynamic_otel_headers( self, standard_callback_dynamic_params: StandardCallbackDynamicParams ) -> dict | None: diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index 41811f836c7..9e6c267438d 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -229,6 +229,10 @@ class OpenTelemetryV2(CustomLogger): """The provider this logger emits through, read-only to its callers.""" return self._tracer_provider + def flush_tracer_providers(self, timeout_millis: int) -> bool: + own: Final = self._tracer_provider.force_flush(timeout_millis) + return self._tenant_tracers.force_flush(timeout_millis) and own + def _init_metrics(self, meter_provider: "MeterProvider | None") -> "GenAIMetricRecorder | None": """Create the six GenAI histograms when metrics are enabled, else ``None``. diff --git a/litellm/integrations/otel/plumbing/routing.py b/litellm/integrations/otel/plumbing/routing.py index 9cf96a1e5e3..35a58eb6c5e 100644 --- a/litellm/integrations/otel/plumbing/routing.py +++ b/litellm/integrations/otel/plumbing/routing.py @@ -213,6 +213,11 @@ class TenantTracerCache: if drained: _shutdown_provider(provider) + def force_flush(self, timeout_millis: int) -> bool: + with self._lock: + providers: Final = (*self._providers.values(), *self._retired) + return all(tuple(provider.force_flush(timeout_millis) for provider in providers)) + def route_for( self, default: Tracer, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index fcc521cf6cf..c56dca84c11 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -349,6 +349,11 @@ else: _GENERIC_API_LOGGER_CLS: Final = GenericAPILogger _in_memory_loggers: Final[list[CustomLogger]] = [] + +def in_memory_loggers() -> tuple[CustomLogger, ...]: + return tuple(_in_memory_loggers) + + _STANDARD_LOGGING_METADATA_RESOLVED_KEYS: Final[frozenset[str]] = frozenset( ("used_client_oauth_token", "usage_object", KUBERNETES_POD_ROUTING_KEY, "is_silent_experiment") ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8c92a058f43..8d4da73277a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -55,6 +55,7 @@ from typing import ( get_args, get_origin, get_type_hints, + runtime_checkable, ) import anyio @@ -92,6 +93,7 @@ from litellm.constants import ( LITELLM_UI_ALLOW_HEADERS, LITELLM_UI_SESSION_DURATION, OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS, + OTEL_SHUTDOWN_FLUSH_TIMEOUT_MILLIS, RUNTIME_UPDATABLE_ROUTER_SETTINGS, VLLM_CHOICE_FIELDS, ) @@ -1418,6 +1420,48 @@ async def _flush_spend_around_logging_drain_on_shutdown( await flush_spend() +@runtime_checkable +class _TracerProviderOwner(Protocol): + def flush_tracer_providers(self, timeout_millis: int) -> bool: ... + + +def _tracer_provider_owners() -> tuple[_TracerProviderOwner, ...]: + from litellm.litellm_core_utils.litellm_logging import in_memory_loggers + + candidates: Final = (*litellm.logging_callback_manager.get_all_callbacks(), *in_memory_loggers()) + owners: Final = tuple(candidate for candidate in candidates if isinstance(candidate, _TracerProviderOwner)) + return tuple({id(owner): owner for owner in owners}.values()) + + +def _flushed_tracer_providers(owner: _TracerProviderOwner, timeout_millis: int) -> bool: + try: + return owner.flush_tracer_providers(timeout_millis) + except Exception as e: # noqa: BLE001 # one failed flush must not skip the other loggers at shutdown + verbose_proxy_logger.exception("Error flushing OpenTelemetry spans on shutdown: %s", e) + return False + + +def _flush_tracer_provider_owners(owners: tuple[_TracerProviderOwner, ...], timeout_millis: int) -> bool: + return all(tuple(_flushed_tracer_providers(owner, timeout_millis) for owner in owners)) + + +async def _flush_opentelemetry_on_shutdown(timeout_millis: int = OTEL_SHUTDOWN_FLUSH_TIMEOUT_MILLIS) -> None: + owners: Final = _tracer_provider_owners() + if not owners: + return + try: + flushed: Final = await asyncio.wait_for( + asyncio.to_thread(_flush_tracer_provider_owners, owners, timeout_millis), timeout=timeout_millis / 1000 + ) + except asyncio.TimeoutError: + verbose_proxy_logger.warning( + "OpenTelemetry shutdown flush did not finish within %dms; spans still queued may be lost", timeout_millis + ) + return + if not flushed: + verbose_proxy_logger.warning("OpenTelemetry shutdown flush incomplete; spans still queued may be lost") + + async def _drain_audit_tasks_on_shutdown() -> None: try: from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks @@ -1472,6 +1516,8 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N except Exception as e: # noqa: BLE001 # shutdown must continue even if the flush fails verbose_proxy_logger.exception("Error flushing Langfuse export channels on shutdown: %s", e) + await _flush_opentelemetry_on_shutdown() + ## RESET CUSTOM VARIABLES ## cleanup_router_config_variables() diff --git a/tests/unit/integrations/arize/test_arize.py b/tests/unit/integrations/arize/test_arize.py index 6accfdc9f93..b37b8f0a1ad 100644 --- a/tests/unit/integrations/arize/test_arize.py +++ b/tests/unit/integrations/arize/test_arize.py @@ -440,3 +440,13 @@ def test_get_arize_config_with_endpoints(mock_env_vars, monkeypatch): assert config.endpoint == "grpc://test.endpoint" assert config.protocol == "otlp_grpc" assert config.project_name == "custom-project" + + +def test_flush_tracer_providers_sends_spans_still_queued_in_arizes_own_provider(otlp_http_sink, monkeypatch): + monkeypatch.setenv("OTEL_BSP_SCHEDULE_DELAY", "3600000") + logger = ArizeLogger(config=OpenTelemetryConfig(exporter="otlp_http", endpoint=otlp_http_sink.url + "/v1/traces")) + logger.tracer.start_span("queued").end() + + assert otlp_http_sink.received.empty() + assert logger.flush_tracer_providers(timeout_millis=5_000) + assert otlp_http_sink.received.get_nowait() == "/v1/traces" diff --git a/tests/unit/integrations/arize/test_arize_phoenix.py b/tests/unit/integrations/arize/test_arize_phoenix.py index 5edf6e6d379..44519b26112 100644 --- a/tests/unit/integrations/arize/test_arize_phoenix.py +++ b/tests/unit/integrations/arize/test_arize_phoenix.py @@ -520,7 +520,7 @@ class TestPerProjectTracerProviderCache: mock_provider = MagicMock() logger._project_providers["proj"] = mock_provider - logger.flush_tracer_providers() + logger.flush_tracer_providers(timeout_millis=5_000) mock_processor.force_flush.assert_called_once() mock_provider.force_flush.assert_called_once() @@ -697,7 +697,7 @@ class TestTracerResolutionAndCache: callback_name="arize_phoenix", tracer_provider=provider, ) - logger.flush_tracer_providers() + logger.flush_tracer_providers(timeout_millis=5_000) exporter.shutdown() def test_standard_logging_metadata_wins_over_litellm_params(self): diff --git a/tests/unit/integrations/conftest.py b/tests/unit/integrations/conftest.py index 9278a280fe7..70bad89337c 100644 --- a/tests/unit/integrations/conftest.py +++ b/tests/unit/integrations/conftest.py @@ -25,6 +25,12 @@ class TlsSink: received: "queue.Queue[str]" +@dataclass(frozen=True, slots=True) +class OtlpHttpSink: + url: str + received: "queue.Queue[str]" + + class _RecordingOtelHandler(http.server.BaseHTTPRequestHandler): def __init__(self, *args: object, received: "queue.Queue[str]", **kwargs: object) -> None: self._received: Final = received @@ -117,6 +123,20 @@ def restore_process_environment() -> Iterator[None]: os.environ.update(original) +@pytest.fixture +def otlp_http_sink() -> Iterator[OtlpHttpSink]: + received: queue.Queue[str] = queue.Queue() + server: Final = http.server.ThreadingHTTPServer( + ("127.0.0.1", 0), functools.partial(_RecordingOtelHandler, received=received) + ) + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield OtlpHttpSink(url=f"http://127.0.0.1:{server.server_port}", received=received) + server.shutdown() + server.server_close() + thread.join(timeout=5) + + @pytest.fixture def tls_sink(tmp_path: Path) -> Iterator[TlsSink]: certificate_path, key_path = write_self_signed_cert(tmp_path, "sink") diff --git a/tests/unit/integrations/otel/test_otel_v2_dynamic.py b/tests/unit/integrations/otel/test_otel_v2_dynamic.py index 8163ba06317..fe25aa3e1d7 100644 --- a/tests/unit/integrations/otel/test_otel_v2_dynamic.py +++ b/tests/unit/integrations/otel/test_otel_v2_dynamic.py @@ -743,3 +743,31 @@ def test_signoz_keyless_team_endpoint_is_ignored_so_the_operator_key_never_reach owned = next(e for e in routed.exporters if e.owner == "signoz") assert owned.endpoint == "https://ingest.us.signoz.cloud:443" assert owned.headers == "signoz-ingestion-key=OPERATOR" + + +def test_force_flush_sends_spans_still_queued_in_a_live_tenant_provider(otlp_http_sink, monkeypatch): + monkeypatch.setenv("OTEL_BSP_SCHEDULE_DELAY", "3600000") + cache = _cache("arize", [ExporterSpec(kind="otlp_http", owner="arize", endpoint=otlp_http_sink.url)]) + route = cache.route_for(NoOpTracer(), {"arize_space_id": "S", "arize_api_key": "K"}) + route.tracer.start_span("queued").end() + cache.release(route.provider) + + assert otlp_http_sink.received.empty() + assert cache.force_flush(5_000) + assert otlp_http_sink.received.get_nowait() == "/v1/traces" + + +def test_force_flush_sends_spans_still_queued_in_a_retired_tenant_provider(otlp_http_sink, monkeypatch): + from litellm.integrations.otel.plumbing import routing as routing_mod + + monkeypatch.setenv("OTEL_BSP_SCHEDULE_DELAY", "3600000") + monkeypatch.setattr(routing_mod, "_MAX_CACHED_PROVIDERS", 1) + cache = _cache("arize", [ExporterSpec(kind="otlp_http", owner="arize", endpoint=otlp_http_sink.url)]) + held = cache.route_for(NoOpTracer(), {"arize_space_id": "held", "arize_api_key": "K"}) + held.tracer.start_span("queued").end() + cache.release(cache.route_for(NoOpTracer(), {"arize_space_id": "newer", "arize_api_key": "K"}).provider) + + assert otlp_http_sink.received.empty() + assert cache.force_flush(5_000) + assert otlp_http_sink.received.get_nowait() == "/v1/traces" + cache.release(held.provider) diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index 4ef95d3a79d..f00dc3d19a0 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -7822,3 +7822,16 @@ class TestOpentelemetryUnitTests: # Assert: error.message should be set from error_str using ErrorAttributes constant mock_span.set_attribute.assert_called_with(ErrorAttributes.ERROR_MESSAGE, "Fallback error message") + + +def test_flush_tracer_providers_sends_spans_still_queued_in_a_credential_scoped_provider(otlp_http_sink, monkeypatch): + monkeypatch.setenv("OTEL_BSP_SCHEDULE_DELAY", "3600000") + logger: Final = OpenTelemetry( + config=OpenTelemetryConfig(exporter="otlp_http", endpoint=otlp_http_sink.url + "/v1/traces"), + tracer_provider=TracerProvider(), + ) + logger._get_tracer_with_dynamic_headers({"api-key": "tenant-a"}).start_span("queued").end() + + assert otlp_http_sink.received.empty() + assert logger.flush_tracer_providers(timeout_millis=5_000) + assert otlp_http_sink.received.get_nowait() == "/v1/traces" diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index db5be07fb6f..0ce06d143f3 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -30,10 +30,16 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from apscheduler.schedulers.asyncio import AsyncIOScheduler +from opentelemetry.sdk.trace import SpanProcessor, TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from pydantic import BaseModel from typing_extensions import TypedDict import litellm.proxy.proxy_server as ps +from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig +from litellm.integrations.otel.logger import OpenTelemetryV2 +from litellm.integrations.otel.model.config import OpenTelemetryV2Config from litellm.proxy.management_helpers.audit_logs import track_audit_task from litellm.proxy.proxy_server import ( ProxyStartupEvent, @@ -202,6 +208,115 @@ async def test_proxy_shutdown_flushes_langfuse_off_the_event_loop_and_logs_a_tim assert any("Langfuse shutdown flush incomplete" in record.getMessage() for record in caplog.records) +_CALLBACK_LISTS: Final = ( + "callbacks", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", +) + + +def _isolate_shutdown(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + from litellm.litellm_core_utils import litellm_logging + + monkeypatch.setattr(ps, "prisma_client", None, raising=False) + monkeypatch.setattr(ps, "jwt_handler", MagicMock(close=AsyncMock()), raising=False) + monkeypatch.setattr(ps, "db_writer_client", None, raising=False) + monkeypatch.setattr(ps, "open_telemetry_logger", None, raising=False) + monkeypatch.setattr(litellm, "cache", None, raising=False) + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", []) + for name in _CALLBACK_LISTS: + monkeypatch.setattr(litellm, name, []) + + +def _batched_provider(exporter: InMemorySpanExporter) -> TracerProvider: + provider: Final = TracerProvider() + provider.add_span_processor(BatchSpanProcessor(exporter, schedule_delay_millis=3_600_000)) + return provider + + +def _legacy_otel_logger(provider: TracerProvider) -> OpenTelemetry: + return OpenTelemetry(config=OpenTelemetryConfig(exporter="console"), tracer_provider=provider) + + +def _v2_otel_logger(provider: TracerProvider) -> OpenTelemetryV2: + return OpenTelemetryV2(config=OpenTelemetryV2Config(exporter="in_memory"), tracer_provider=provider) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("make_logger", [_legacy_otel_logger, _v2_otel_logger], ids=["legacy", "v2"]) +@pytest.mark.parametrize("registered_in", ["callbacks", "in_memory_loggers"]) +async def test_proxy_shutdown_exports_spans_still_waiting_in_an_opentelemetry_batch( + monkeypatch, make_logger, registered_in +): + import litellm + from litellm.litellm_core_utils import litellm_logging + + _isolate_shutdown(monkeypatch) + exporter: Final = InMemorySpanExporter() + logger: Final = make_logger(_batched_provider(exporter)) + for name in _CALLBACK_LISTS: + monkeypatch.setattr(litellm, name, []) + if registered_in == "callbacks": + monkeypatch.setattr(litellm, "callbacks", [logger]) + else: + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [logger]) + logger.tracer.start_span("ended-before-shutdown").end() + + await proxy_shutdown_event() + + assert [span.name for span in exporter.get_finished_spans()] == ["ended-before-shutdown"] + + +@pytest.mark.asyncio +async def test_proxy_shutdown_keeps_flushing_other_loggers_when_one_raises(monkeypatch): + import litellm + + class _BrokenOwner: + def flush_tracer_providers(self, timeout_millis: int) -> bool: + raise RuntimeError("exporter gone") + + _isolate_shutdown(monkeypatch) + exporter: Final = InMemorySpanExporter() + logger: Final = _legacy_otel_logger(_batched_provider(exporter)) + monkeypatch.setattr(litellm, "callbacks", [_BrokenOwner(), logger]) + logger.tracer.start_span("ended-before-shutdown").end() + + await proxy_shutdown_event() + + assert [span.name for span in exporter.get_finished_spans()] == ["ended-before-shutdown"] + + +@pytest.mark.asyncio +async def test_opentelemetry_shutdown_flush_gives_up_at_its_deadline_and_warns(monkeypatch, caplog): + import threading + import time + + import litellm + + release: Final = threading.Event() + + class _StuckProcessor(SpanProcessor): + def force_flush(self, timeout_millis: int = 30_000) -> bool: + return release.wait(30) + + _isolate_shutdown(monkeypatch) + provider: Final = TracerProvider() + provider.add_span_processor(_StuckProcessor()) + monkeypatch.setattr(litellm, "callbacks", [_legacy_otel_logger(provider)]) + started: Final = time.monotonic() + + with caplog.at_level("WARNING", logger="LiteLLM Proxy"): + await ps._flush_opentelemetry_on_shutdown(timeout_millis=200) + elapsed: Final = time.monotonic() - started + release.set() + + assert elapsed < 5 + assert any("OpenTelemetry shutdown flush did not finish" in record.getMessage() for record in caplog.records) + + @pytest.mark.asyncio async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monkeypatch): """