Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/middleware/in_flight_requests_middleware.py: 88%

1""" 

2Tracks the number of HTTP requests currently in-flight on this uvicorn worker. 

3 

4Used by /health/backlog to expose per-pod queue depth, and emitted as the 

5Prometheus gauge `litellm_in_flight_requests`. 

6""" 

7 

8import os 

9from typing import Any, Final 

10 

11from starlette.types import ASGIApp, Receive, Scope, Send 

12 

13 

14class InFlightRequestsMiddleware: 

15 """ 

16 ASGI middleware that increments a counter when a request arrives and 

17 decrements it when the response is sent (or an error occurs). 

18 

19 The counter is class-level and therefore scoped to a single uvicorn worker 

20 process — exactly the per-pod granularity we want. 

21 

22 Also updates the `litellm_in_flight_requests` Prometheus gauge if 

23 prometheus_client is installed. The gauge is lazily initialised on the 

24 first request so that PROMETHEUS_MULTIPROC_DIR is already set by the time 

25 we register the metric. Initialisation is attempted only once — if 

26 prometheus_client is absent the class remembers and never retries. 

27 """ 

28 

29 _in_flight: int = 0 

30 _gauge: Any | None = None 

31 _gauge_init_attempted: bool = False 

32 

33 def __init__(self, app: ASGIApp) -> None: 

34 self.app = app 

35 

36 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: 

37 if scope["type"] != "http": 

38 await self.app(scope, receive, send) 

39 return 

40 

41 InFlightRequestsMiddleware._in_flight += 1 

42 gauge: Final = InFlightRequestsMiddleware._get_gauge() 

43 if gauge is not None: 43 ↛ 45line 43 didn't jump to line 45 because the condition on line 43 was always true

44 gauge.inc() 

45 try: 

46 await self.app(scope, receive, send) 

47 finally: 

48 InFlightRequestsMiddleware._in_flight -= 1 

49 if gauge is not None: 49 ↛ exitline 49 didn't return from function '__call__' because the condition on line 49 was always true

50 gauge.dec() 

51 

52 @staticmethod 

53 def get_count() -> int: 

54 """Return the number of HTTP requests currently in-flight.""" 

55 return InFlightRequestsMiddleware._in_flight 

56 

57 @staticmethod 

58 def _get_gauge() -> Any | None: 

59 if InFlightRequestsMiddleware._gauge_init_attempted: 

60 return InFlightRequestsMiddleware._gauge 

61 InFlightRequestsMiddleware._gauge_init_attempted = True 

62 try: 

63 from prometheus_client import Gauge 

64 

65 if "PROMETHEUS_MULTIPROC_DIR" in os.environ: 65 ↛ 67line 65 didn't jump to line 67 because the condition on line 65 was never true

66 # livesum aggregates across all worker processes in the scrape response 

67 InFlightRequestsMiddleware._gauge = Gauge( 

68 "litellm_in_flight_requests", 

69 "Number of HTTP requests currently in-flight on this uvicorn worker", 

70 multiprocess_mode="livesum", 

71 ) 

72 else: 

73 InFlightRequestsMiddleware._gauge = Gauge( 

74 "litellm_in_flight_requests", 

75 "Number of HTTP requests currently in-flight on this uvicorn worker", 

76 ) 

77 except Exception: 

78 InFlightRequestsMiddleware._gauge = None 

79 return InFlightRequestsMiddleware._gauge 

80 

81 

82def get_in_flight_requests() -> int: 

83 """Module-level convenience wrapper used by the /health/backlog endpoint.""" 

84 return InFlightRequestsMiddleware.get_count()