mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(proxy): trace model list requests with the caller's traceparent
The proxy creates an OTEL SERVER span at the auth boundary whenever a request carries traceparent, but /v1/models never closed it. Unended spans are not exported, so callers saw nothing for the model list while inference requests showed up in the same trace Route the endpoint through the management span wrapper so the SERVER span is closed and parented to the caller's trace. The wrapper gains a variant that skips request and response payloads, since attaching the full model list would bloat every poll
This commit is contained in:
parent
5029d027a1
commit
ad5463d517
3 changed files with 101 additions and 10 deletions
|
|
@ -1,11 +1,11 @@
|
|||
# What is this?
|
||||
## Helper utils for the management endpoints (keys/users/teams)
|
||||
from collections.abc import Callable, Mapping, MutableMapping, Sequence
|
||||
from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from datetime import datetime
|
||||
from functools import wraps
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Protocol
|
||||
from typing import Any, Final, Protocol, TypeVar, cast
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
|
@ -767,19 +767,31 @@ async def _emit_management_endpoint_otel_span(
|
|||
)
|
||||
|
||||
|
||||
def management_endpoint_wrapper(func):
|
||||
_Endpoint: Final = TypeVar("_Endpoint", bound=Callable[..., Awaitable[Any]])
|
||||
|
||||
|
||||
def management_endpoint_wrapper(func: _Endpoint) -> _Endpoint:
|
||||
"""
|
||||
This wrapper does the following:
|
||||
|
||||
1. Log I/O, Exceptions to OTEL
|
||||
2. Create an Audit log for success calls
|
||||
"""
|
||||
return _wrap_management_endpoint(func, log_payload=True)
|
||||
|
||||
@wraps(func)
|
||||
|
||||
def management_endpoint_span_wrapper(func: _Endpoint) -> _Endpoint:
|
||||
"""``management_endpoint_wrapper`` without request/response payloads on the span, for
|
||||
endpoints that serialize large or high-frequency responses (e.g. the model list)."""
|
||||
return _wrap_management_endpoint(func, log_payload=False)
|
||||
|
||||
|
||||
def _wrap_management_endpoint(endpoint: _Endpoint, *, log_payload: bool) -> _Endpoint:
|
||||
@wraps(endpoint)
|
||||
async def wrapper(*args, **kwargs):
|
||||
start_time: Final = datetime.now()
|
||||
try:
|
||||
result: Final = await func(*args, **kwargs)
|
||||
result: Final = await endpoint(*args, **kwargs)
|
||||
end_time = datetime.now()
|
||||
try:
|
||||
user_api_key_dict: UserAPIKeyAuth = kwargs.get("user_api_key_dict") or UserAPIKeyAuth()
|
||||
|
|
@ -787,17 +799,17 @@ def management_endpoint_wrapper(func):
|
|||
await send_management_endpoint_alert(
|
||||
request_kwargs=kwargs,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
function_name=func.__name__,
|
||||
function_name=endpoint.__name__,
|
||||
)
|
||||
parent_otel_span: Span | None = getattr(user_api_key_dict, "parent_otel_span", None)
|
||||
if parent_otel_span is not None:
|
||||
await _emit_management_endpoint_otel_span(
|
||||
func=func,
|
||||
func=endpoint,
|
||||
kwargs=kwargs,
|
||||
parent_otel_span=parent_otel_span,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
result=result,
|
||||
result=result if log_payload else None,
|
||||
)
|
||||
|
||||
# Delete updated/deleted info from cache
|
||||
|
|
@ -818,7 +830,7 @@ def management_endpoint_wrapper(func):
|
|||
if parent_otel_span is not None:
|
||||
try:
|
||||
await _emit_management_endpoint_otel_span(
|
||||
func=func,
|
||||
func=endpoint,
|
||||
kwargs=kwargs,
|
||||
parent_otel_span=parent_otel_span,
|
||||
start_time=start_time,
|
||||
|
|
@ -835,4 +847,4 @@ def management_endpoint_wrapper(func):
|
|||
|
||||
raise e
|
||||
|
||||
return wrapper
|
||||
return cast(_Endpoint, wrapper) # cast-ok: wraps() preserves the endpoint's signature
|
||||
|
|
|
|||
|
|
@ -174,6 +174,7 @@ from litellm.proxy.common_utils.token_count_settings import (
|
|||
)
|
||||
from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog
|
||||
from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields
|
||||
from litellm.proxy.management_helpers.utils import management_endpoint_span_wrapper
|
||||
from litellm.repositories import database as database_runtime
|
||||
from litellm.router_strategy.auto_router.auto_router_model_naming import (
|
||||
GATED_AUTO_ROUTER_CAPABILITIES,
|
||||
|
|
@ -12343,6 +12344,7 @@ async def _deployment_hidden_by_listing_callbacks(deployment: Deployment, user_a
|
|||
@router.get(
|
||||
"/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"]
|
||||
) # if project requires model list
|
||||
@management_endpoint_span_wrapper
|
||||
async def model_list(
|
||||
request: Request = None, # pyright: ignore[reportArgumentType] # FastAPI always injects the Request; the None default only serves direct in-process callers
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
|
|
|
|||
|
|
@ -340,3 +340,80 @@ def test_management_wrapper_swallows_post_success_errors(server_span_factory, ot
|
|||
result = asyncio.run(fake_fn(data={}, user_api_key_dict=_real_user_api_key_dict(server_span)))
|
||||
|
||||
assert result == {"ok": True}
|
||||
|
||||
|
||||
def test_span_wrapper_omits_payload(server_span_factory, otel_with_exporter, monkeypatch):
|
||||
"""The span-only wrapper keeps the endpoint span but drops the response
|
||||
attributes — the models listing would otherwise serialize every model onto
|
||||
the span."""
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.management_helpers import utils as mgmt_utils
|
||||
|
||||
otel, exporter = otel_with_exporter
|
||||
monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False)
|
||||
monkeypatch.setattr(mgmt_utils, "send_management_endpoint_alert", _noop_alert)
|
||||
|
||||
server_span = server_span_factory(KEY_GENERATE_PATH)
|
||||
|
||||
@mgmt_utils.management_endpoint_span_wrapper
|
||||
async def fake_fn(data=None, user_api_key_dict=None):
|
||||
return {"data": [{"id": "gpt-4o"}]}
|
||||
|
||||
asyncio.run(fake_fn(data={}, user_api_key_dict=_real_user_api_key_dict(server_span)))
|
||||
|
||||
spans = [s for s in exporter.get_finished_spans() if s.name == "fake_fn"]
|
||||
assert len(spans) == 1
|
||||
assert not any(key.startswith("response.") for key in spans[0].attributes)
|
||||
|
||||
|
||||
MODEL_LIST_PATH = "/v1/models"
|
||||
INBOUND_TRACEPARENT = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
|
||||
|
||||
|
||||
def test_model_list_ends_server_span_in_inbound_trace(otel_with_exporter, monkeypatch):
|
||||
"""Regression: ``/v1/models`` created the OTEL SERVER span at the auth
|
||||
boundary but never ended it, so the span was never exported and a caller
|
||||
passing ``traceparent`` saw nothing for the models list.
|
||||
|
||||
Drives the real endpoint and asserts the SERVER span reaches the exporter
|
||||
under the caller's trace id."""
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
import litellm.proxy.utils as proxy_utils
|
||||
from litellm.proxy.management_helpers import utils as mgmt_utils
|
||||
|
||||
otel, exporter = otel_with_exporter
|
||||
monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False)
|
||||
monkeypatch.setattr(mgmt_utils, "send_management_endpoint_alert", _noop_alert)
|
||||
|
||||
async def _no_models_for_user(**kwargs):
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(proxy_utils, "get_available_models_for_user", _no_models_for_user)
|
||||
|
||||
server_span = otel.create_litellm_proxy_request_started_span(
|
||||
start_time=datetime.now(), headers={"traceparent": INBOUND_TRACEPARENT}
|
||||
)
|
||||
otel.set_proxy_request_route_attributes(server_span, url_path=MODEL_LIST_PATH, http_route=MODEL_LIST_PATH)
|
||||
|
||||
result = asyncio.run(
|
||||
proxy_server.model_list(
|
||||
user_api_key_dict=_real_user_api_key_dict(server_span),
|
||||
)
|
||||
)
|
||||
|
||||
assert result == {"data": [], "object": "list"}
|
||||
assert_server_span_attrs(
|
||||
exporter,
|
||||
expected_status=200,
|
||||
expected_url_path=MODEL_LIST_PATH,
|
||||
where="model list success",
|
||||
)
|
||||
exported = get_server_span(exporter)
|
||||
assert exported is not None
|
||||
assert exported.context.trace_id == int(INBOUND_TRACEPARENT.split("-")[1], 16)
|
||||
|
||||
endpoint_spans = [s for s in exporter.get_finished_spans() if s.name == "model_list"]
|
||||
assert len(endpoint_spans) == 1
|
||||
assert endpoint_spans[0].parent is not None
|
||||
assert endpoint_spans[0].parent.span_id == exported.context.span_id
|
||||
assert not any(key.startswith("response.") for key in endpoint_spans[0].attributes)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue