diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 08bf56e3773..801f39de4a0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13,6 +13,14 @@ import time import traceback import warnings from datetime import datetime, timedelta, timezone + +# Patch Uvicorn so scope gets _uvicorn_received_at (request ready → app entry timing) +try: + from litellm.proxy.uvicorn_handoff_patch import apply_uvicorn_handoff_patch + + apply_uvicorn_handoff_patch() +except Exception: + pass from typing import ( TYPE_CHECKING, Any, @@ -1222,8 +1230,8 @@ current_dir = os.path.dirname(os.path.abspath(__file__)) # app.add_middleware(PrometheusAuthMiddleware) -class _UvicornHandoffTimingMiddleware: - """Raw ASGI middleware to timestamp when Uvicorn hands request to our app.""" +class _UvicornHandoffLoggingMiddleware: + """ASGI middleware: log time from Uvicorn request ready → app entry (when patch applied).""" def __init__(self, app: Any) -> None: self.app = app @@ -1232,7 +1240,13 @@ class _UvicornHandoffTimingMiddleware: self, scope: dict, receive: Any, send: Any ) -> None: if scope.get("type") == "http": - scope["_uvicorn_handoff_at"] = time.perf_counter() + t_received = scope.get("_uvicorn_received_at") + if t_received is not None: + delta_ms = (time.perf_counter() - t_received) * 1000 + print( + f"[proxy] uvicorn receive → app entry: {delta_ms:.2f}ms", + flush=True, + ) await self.app(scope, receive, send) @@ -1242,14 +1256,6 @@ class _ChatCompletionsEarlyReturnMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next): path = request.scope.get("path", "") if "chat/completions" in path and request.method == "POST": - t_handoff = request.scope.get("_uvicorn_handoff_at") - if t_handoff is not None: - delta_ms = (time.perf_counter() - t_handoff) * 1000 - print( - f"[proxy] Uvicorn handoff → middleware: {delta_ms:.2f}ms", - flush=True, - ) - print("[proxy] _ChatCompletionsEarlyReturnMiddleware: early return", flush=True) return JSONResponse( content={"message": "Hello, world!", "status": "success"}, status_code=200, @@ -1258,6 +1264,7 @@ class _ChatCompletionsEarlyReturnMiddleware(BaseHTTPMiddleware): app.add_middleware(_ChatCompletionsEarlyReturnMiddleware) +app.add_middleware(_UvicornHandoffLoggingMiddleware) def mount_swagger_ui(): @@ -11889,5 +11896,3 @@ app.mount(path=BASE_MCP_ROUTE, app=mcp_app) app.include_router(mcp_rest_endpoints_router) app.include_router(mcp_discoverable_endpoints_router) -# Wrap app so Uvicorn's handoff to our ASGI stack is timestamped for perf tracing -app = _UvicornHandoffTimingMiddleware(app) diff --git a/litellm/proxy/uvicorn_handoff_patch.py b/litellm/proxy/uvicorn_handoff_patch.py new file mode 100644 index 00000000000..c9929c72f0c --- /dev/null +++ b/litellm/proxy/uvicorn_handoff_patch.py @@ -0,0 +1,48 @@ +""" +Monkey-patch Uvicorn to record when a request is ready (headers complete) +so LiteLLM middleware can measure "Uvicorn receive → app entry" latency. + +Applied when litellm.proxy.proxy_server is imported (so it runs in the process +that serves requests, including gunicorn workers). Sets scope["_uvicorn_received_at"] +(perf_counter) in both httptools and h11 protocol paths when the request +scope is built and before the ASGI app is invoked. +""" +from __future__ import annotations + +import time + + +def _patch_httptools_impl() -> None: + import uvicorn.protocols.http.httptools_impl as m # noqa: PLC0415 + + _orig_on_headers_complete = m.HttpToolsProtocol.on_headers_complete + + def on_headers_complete(self: m.HttpToolsProtocol) -> None: + _orig_on_headers_complete(self) + # Scope is the same dict passed to RequestResponseCycle; stamp it so + # middleware can measure receive → app entry (timestamp is just after + # request was ready and task was scheduled). + self.scope["_uvicorn_received_at"] = time.perf_counter() # type: ignore[typeddict-unknown-key] + + m.HttpToolsProtocol.on_headers_complete = on_headers_complete # type: ignore[method-assign] + + +def _patch_h11_impl() -> None: + import uvicorn.protocols.http.h11_impl as m # noqa: PLC0415 + + # In h11, scope is built inside data_received; we stamp at start of run_asgi + # so "received" = when ASGI task began (slightly later than headers-complete). + _orig_run_asgi = m.RequestResponseCycle.run_asgi + + async def run_asgi(self: m.RequestResponseCycle, app: m.ASGI3Application) -> None: + if "_uvicorn_received_at" not in self.scope: + self.scope["_uvicorn_received_at"] = time.perf_counter() # type: ignore[typeddict-unknown-key] + await _orig_run_asgi(self, app) + + m.RequestResponseCycle.run_asgi = run_asgi # type: ignore[method-assign] + + +def apply_uvicorn_handoff_patch() -> None: + """Apply patches so scope gets _uvicorn_received_at for handoff timing.""" + _patch_httptools_impl() + _patch_h11_impl()