diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index c860f8e540d0..b17e055c7ea1 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -85,6 +85,22 @@ class CachingHandlerResponse(BaseModel): in_memory_cache_obj = InMemoryCache() +def _drop_logging_obj_from_kwargs(request_kwargs: dict[str, object]) -> dict[str, object]: + """ + The caching handler is stored on the Logging object + (``logging_obj._llm_caching_handler``), so keeping ``litellm_logging_obj`` + inside ``request_kwargs`` closes a reference cycle + (Logging -> LLMCachingHandler -> kwargs -> Logging) that keeps the full + request payload (messages included) alive until a generational GC pass + instead of being freed by refcount when the request ends. Nothing in the + caching layer reads the logging object from these kwargs; cache-key + generation ignores litellm-internal params. + """ + if "litellm_logging_obj" not in request_kwargs: + return request_kwargs + return {k: v for k, v in request_kwargs.items() if k != "litellm_logging_obj"} + + def _is_chat_completion_cached_dict(cached_result: dict) -> bool: cached_id = cached_result.get("id") if isinstance(cached_id, str) and cached_id.startswith("chatcmpl"): @@ -118,7 +134,7 @@ def __init__( self.async_streaming_chunks: List[ModelResponse] = [] self.sync_streaming_chunks: List[ModelResponse] = [] - self.request_kwargs = request_kwargs + self.request_kwargs = _drop_logging_obj_from_kwargs(request_kwargs) self.preset_cache_key: Optional[str] = None self.original_function = original_function self.start_time = start_time @@ -297,7 +313,7 @@ def _sync_get_cache( new_kwargs.pop("metadata", None) if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs: new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs) - self.request_kwargs = new_kwargs + self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs) print_verbose("Checking Sync Cache") cached_result = litellm.cache.get_cache(**new_kwargs) if cached_result is not None: @@ -693,7 +709,7 @@ async def _retrieve_from_cache( new_kwargs.pop("metadata", None) if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs: new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs) - self.request_kwargs = new_kwargs + self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs) cached_result: Optional[Any] = None if call_type == CallTypes.aembedding.value: if isinstance(new_kwargs["input"], str): diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 461ab62b8158..8c85659075d0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -925,7 +925,6 @@ def _pre_call(self, input, api_key, model=None, additional_args={}): def pre_call(self, input, api_key, model=None, additional_args={}): # Log the exact input to the LLM API - litellm.error_logs["PRE_CALL"] = locals() try: self._pre_call( input=input, @@ -1135,7 +1134,6 @@ def _get_masked_headers(self, headers: dict, ignore_sensitive_headers: bool = Fa def post_call(self, original_response, input=None, api_key=None, additional_args={}): # Log the exact result from the LLM API, for streaming - log the type of response received - litellm.error_logs["POST_CALL"] = locals() if isinstance(original_response, dict): original_response = json.dumps(original_response, default=str) try: diff --git a/litellm/router.py b/litellm/router.py index 46794ac7a247..f408d030b8b0 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2986,7 +2986,7 @@ def _update_kwargs_with_deployment( # here before it's wiped below, instead of relying on that attempt's # (possibly still-pending) failure event to do it. refund_stale_reservation_before_retry(self.cache, kwargs) - set_io_token_rate_limit_request_kwargs(kwargs) + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=deployment_has_io_token_limits(deployment)) ## DEPLOYMENT-LEVEL TAGS deployment_tags = deployment.get("litellm_params", {}).get("tags") diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 62c99a1c9f63..803fdc4b3530 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -43,14 +43,21 @@ OTPM_CACHE_KEY = "_litellm_otpm_cache_key" -def set_io_token_rate_limit_request_kwargs(kwargs: Optional[dict[str, Any]]) -> None: +def set_io_token_rate_limit_request_kwargs(kwargs: Optional[dict[str, Any]], store_in_context: bool = True) -> None: # The reservation sentinels are server-only, but `metadata` is caller # controlled on proxy requests. Strip any client-supplied copies here (this # runs before the router stashes its own reservation) so a forged # reservation can't drive the post-call reconcile/refund against an # arbitrary counter and bypass the configured limits. _clear_reservation_from_kwargs(kwargs) - _io_token_rate_limit_request_kwargs.set(kwargs) + # The context slot pins the entire request kwargs (messages included) for + # the lifetime of the surrounding context, which outlives the request when + # the context is captured by pooled resources (e.g. a redis connection + # created mid-request). Only ITPM/OTPM-limited deployments read it, so for + # every other deployment overwrite the slot with None instead of the + # kwargs; overwriting (rather than skipping) also releases a previous + # request's kwargs when a context is reused. + _io_token_rate_limit_request_kwargs.set(kwargs if store_in_context else None) def get_io_token_rate_limit_request_kwargs() -> Optional[dict[str, Any]]: diff --git a/tests/test_litellm/caching/test_caching_handler.py b/tests/test_litellm/caching/test_caching_handler.py index 01327529410d..1136a0b7e7b7 100644 --- a/tests/test_litellm/caching/test_caching_handler.py +++ b/tests/test_litellm/caching/test_caching_handler.py @@ -556,3 +556,31 @@ async def test_embedding_cache_falls_back_to_token_counter_for_legacy_entries(): assert cache_hit # token_counter over "hello world" yields a nonzero count — fallback path still runs assert response.usage.prompt_tokens > 0 + + +def test_request_kwargs_does_not_retain_logging_obj(): + """ + The caching handler lives on logging_obj._llm_caching_handler, so keeping + litellm_logging_obj inside request_kwargs closes a reference cycle + (Logging -> LLMCachingHandler -> kwargs -> Logging). That cycle keeps the + full request payload alive until a generational GC pass instead of being + freed by refcount when the request finishes; under bursts of large-token + requests this presents as stepwise RSS growth that never returns to + baseline. Other kwargs (messages included) must be preserved. + """ + logging_obj = MagicMock() + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hello"}], + "litellm_logging_obj": logging_obj, + } + + handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs=kwargs, + start_time=datetime.now(), + ) + + assert "litellm_logging_obj" not in handler.request_kwargs + assert handler.request_kwargs["messages"] == kwargs["messages"] + assert handler.request_kwargs["model"] == "gpt-4o" diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index ade2c6777458..19e6dcc33ed8 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3773,3 +3773,19 @@ def test_zero_token_video_usage_preserves_duration_seconds(logging_obj): assert payload["metadata"]["usage_object"]["duration_seconds"] == 4.0 assert payload["total_tokens"] == 0 assert payload["completion_tokens"] == 0 + + +def test_pre_call_does_not_pin_request_in_module_state(logging_obj): + """ + pre_call/post_call must not stash their locals (full messages, the Logging + object, complete_input_dict) into module-level state. That pinned the most + recent request's entire payload in memory for the life of the worker, + which with multi-hundred-KB requests is a permanent per-worker leak. + """ + litellm.error_logs.clear() + big_input = [{"role": "user", "content": "x" * 10_000}] + + logging_obj.pre_call(input=big_input, api_key="sk-test") + logging_obj.post_call(original_response='{"ok": true}', input=big_input, api_key="sk-test") + + assert litellm.error_logs == {} diff --git a/tests/test_litellm/test_router/test_io_token_rate_limits.py b/tests/test_litellm/test_router/test_io_token_rate_limits.py index 939e31895968..a5a682711114 100644 --- a/tests/test_litellm/test_router/test_io_token_rate_limits.py +++ b/tests/test_litellm/test_router/test_io_token_rate_limits.py @@ -977,3 +977,65 @@ async def test_model_group_info_aggregates_io_limits(self): assert info is not None assert info.itpm == 100 assert info.otpm == 20 + + +class TestContextSlotRetention: + def test_setter_stores_kwargs_only_for_io_limited_deployments(self): + """ + The context slot pins the entire request kwargs (messages included) + for the lifetime of the surrounding asyncio context, and pooled + resources created mid-request (e.g. redis connections) capture that + context, extending the pin far past the request. Only ITPM/OTPM + pre-call checks read the slot, so the setter must store None for + deployments without io token limits and still clear reservation + sentinels from kwargs either way. + """ + kwargs = { + "messages": [{"role": "user", "content": "x" * 1000}], + "metadata": {ITPM_RESERVED_KEY: 999, ITPM_CACHE_KEY: "forged"}, + } + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + assert get_io_token_rate_limit_request_kwargs() is None + assert ITPM_RESERVED_KEY not in kwargs["metadata"] + assert ITPM_CACHE_KEY not in kwargs["metadata"] + + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True) + assert get_io_token_rate_limit_request_kwargs() is kwargs + + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + assert get_io_token_rate_limit_request_kwargs() is None + + @pytest.mark.asyncio + async def test_router_does_not_pin_kwargs_without_io_limits(self): + router = Router( + model_list=[ + { + "model_name": "plain", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + } + ] + ) + set_io_token_rate_limit_request_kwargs(None) + kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + deployment = router.get_deployment_by_model_group_name("plain") + assert deployment is not None + router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) + assert get_io_token_rate_limit_request_kwargs() is None + + @pytest.mark.asyncio + async def test_router_pins_kwargs_for_io_limited_deployment(self): + router = Router( + model_list=[ + { + "model_name": "limited", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "itpm": 100}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + set_io_token_rate_limit_request_kwargs(None) + kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + deployment = router.get_deployment_by_model_group_name("limited") + assert deployment is not None + router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) + assert get_io_token_rate_limit_request_kwargs() is kwargs