diff --git a/litellm/constants.py b/litellm/constants.py index 0dee5638ccb..bbe2ae3c8d5 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -910,9 +910,6 @@ MAX_SPENDLOG_ROWS_TO_QUERY = int( DEFAULT_SOFT_BUDGET = float( os.getenv("DEFAULT_SOFT_BUDGET", 50.0) ) # by default all litellm proxy keys have a soft budget of 50.0 -DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS = int( - os.getenv("DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS", 600) -) # 10 minutes timeout for client disconnect checking in proxy # makes it clear this is a rate limit error for a litellm virtual key RATE_LIMIT_ERROR_MESSAGE_FOR_VIRTUAL_KEY = "LiteLLM Virtual Key user_api_key_hash" diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a3a9c2cffc0..e900975f1cc 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1,7 +1,6 @@ import asyncio import json import logging -import time import traceback from datetime import datetime from typing import ( @@ -25,7 +24,6 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ( DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE, - DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS, STREAM_SSE_DATA_PREFIX, ) from litellm.litellm_core_utils.dd_tracing import tracer @@ -177,29 +175,6 @@ async def create_streaming_response( ) -async def _check_request_disconnection(request: Request, llm_api_call_task): - """ - Asynchronously checks if the request is disconnected at regular intervals. - If the request is disconnected - - cancel the litellm.router task - - Parameters: - - request: Request: The request object to check for disconnection. - Returns: - - None - """ - - # only run this function for configured timeout -> if these don't get cancelled -> we don't want the server to have many while loops - start_time = time.time() - while time.time() - start_time < DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS: - await asyncio.sleep(1) - message = await request.receive() - if message.get("type") == "http.disconnect": - # cancel the LLM API Call task if any passed - this is passed from individual providers - # Example OpenAI, Azure, VertexAI etc - llm_api_call_task.cancel() - return - class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict): self.data = data @@ -450,24 +425,12 @@ class ProxyBaseLLMRequestProcessing: ) tasks.append(llm_call) + # wait for call to end llm_responses = asyncio.gather( *tasks ) # run the moderation check in parallel to the actual llm api call - # Execute the task to detect disconnection - disconnect_task = asyncio.create_task(_check_request_disconnection(request, llm_responses)) - - try: - # wait for call to end - # Note: In the case of streaming, processing does not wait here, so disconnection detection is performed in StreamingResponse. - responses = await llm_responses - disconnect_task.cancel() - - except asyncio.CancelledError: - raise HTTPException( - status_code=499, - detail="Client disconnected the request", - ) + responses = await llm_responses response = responses[1] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e15d5401374..9f1566b2e00 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -997,6 +997,33 @@ db_writer_client: Optional[AsyncHTTPHandler] = None ### logger ### +async def check_request_disconnection(request: Request, llm_api_call_task): + """ + Asynchronously checks if the request is disconnected at regular intervals. + If the request is disconnected + - cancel the litellm.router task + - raises an HTTPException with status code 499 and detail "Client disconnected the request". + + Parameters: + - request: Request: The request object to check for disconnection. + Returns: + - None + """ + + # only run this function for 10 mins -> if these don't get cancelled -> we don't want the server to have many while loops + start_time = time.time() + while time.time() - start_time < 600: + await asyncio.sleep(1) + if await request.is_disconnected(): + # cancel the LLM API Call task if any passed - this is passed from individual providers + # Example OpenAI, Azure, VertexAI etc + llm_api_call_task.cancel() + + raise HTTPException( + status_code=499, + detail="Client disconnected the request", + ) + def _resolve_typed_dict_type(typ): """Resolve the actual TypedDict class from a potentially wrapped type.""" diff --git a/poetry.lock b/poetry.lock index 0ab437aec25..29d1a877087 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.1.4 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.1.2 and should not be changed by hand. [[package]] name = "aiohappyeyeballs" @@ -6122,15 +6122,15 @@ zstd = ["zstandard (>=0.18.0)"] [[package]] name = "uvicorn" -version = "0.32.1" +version = "0.29.0" description = "The lightning-fast ASGI server." optional = true python-versions = ">=3.8" groups = ["main"] markers = "python_version >= \"3.10\" and (extra == \"mlflow\" or extra == \"proxy\") or extra == \"proxy\"" files = [ - {file = "uvicorn-0.32.1-py3-none-any.whl", hash = "sha256:82ad92fd58da0d12af7482ecdb5f2470a04c9c9a53ced65b9bbb4a205377602e"}, - {file = "uvicorn-0.32.1.tar.gz", hash = "sha256:ee9519c246a72b1c084cea8d3b44ed6026e78a4a309cbedae9c37e4cb9fbb175"}, + {file = "uvicorn-0.29.0-py3-none-any.whl", hash = "sha256:2c2aac7ff4f4365c206fd773a39bf4ebd1047c238f8b8268ad996829323473de"}, + {file = "uvicorn-0.29.0.tar.gz", hash = "sha256:6a69214c0b6a087462412670b3ef21224fa48cae0e452b5883e8e8bdfdd11dd0"}, ] [package.dependencies] @@ -6139,7 +6139,7 @@ h11 = ">=0.8" typing-extensions = {version = ">=4.0", markers = "python_version < \"3.11\""} [package.extras] -standard = ["colorama (>=0.4) ; sys_platform == \"win32\"", "httptools (>=0.6.3)", "python-dotenv (>=0.13)", "pyyaml (>=5.1)", "uvloop (>=0.14.0,!=0.15.0,!=0.15.1) ; sys_platform != \"win32\" and sys_platform != \"cygwin\" and platform_python_implementation != \"PyPy\"", "watchfiles (>=0.13)", "websockets (>=10.4)"] +standard = ["colorama (>=0.4) ; sys_platform == \"win32\"", "httptools (>=0.5.0)", "python-dotenv (>=0.13)", "pyyaml (>=5.1)", "uvloop (>=0.14.0,!=0.15.0,!=0.15.1) ; sys_platform != \"win32\" and sys_platform != \"cygwin\" and platform_python_implementation != \"PyPy\"", "watchfiles (>=0.13)", "websockets (>=10.4)"] [[package]] name = "uvloop" @@ -6576,4 +6576,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.8.1,<4.0, !=3.9.7" -content-hash = "e48cc445bc012e020a9e311942e46833dda587b70a630a04bfc08b629746fe56" +content-hash = "f41e6359109c5c52dba2a28f301b04030d865265f408974082b390bf45568a01" diff --git a/pyproject.toml b/pyproject.toml index b1b11f5d21d..9f5d876cf2c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ pydantic = "^2.5.0" jsonschema = "^4.22.0" numpydoc = {version = "*", optional = true} # used in utils.py -uvicorn = {version = "^0.32.0", optional = true} +uvicorn = {version = "^0.29.0", optional = true} uvloop = {version = "^0.21.0", optional = true, markers="sys_platform != 'win32'"} gunicorn = {version = "^23.0.0", optional = true} fastapi = {version = "^0.115.5", optional = true} diff --git a/requirements.txt b/requirements.txt index 9b858e08a03..2d31819dc5b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,7 +5,7 @@ openai==1.99.5 # openai req. fastapi==0.115.5 # server dep backoff==2.2.1 # server dep pyyaml==6.0.2 # server dep -uvicorn==0.32.0 # server dep +uvicorn==0.29.0 # server dep gunicorn==23.0.0 # server dep fastuuid==0.12.0 # for uuid4 uvloop==0.21.0 # uvicorn dep, gives us much better performance under load diff --git a/tests/proxy_unit_tests/test_client_disconnection.py b/tests/proxy_unit_tests/test_client_disconnection.py deleted file mode 100644 index d894d7ad015..00000000000 --- a/tests/proxy_unit_tests/test_client_disconnection.py +++ /dev/null @@ -1,47 +0,0 @@ -""" -Test client disconnection detection functionality. -""" -import asyncio -import pytest -from unittest.mock import AsyncMock - -from litellm.proxy.common_request_processing import _check_request_disconnection - - -@pytest.mark.asyncio -async def test_check_request_disconnection_with_disconnect(): - """Test that _check_request_disconnection cancels task when client disconnects.""" - mock_request = AsyncMock() - mock_request.receive.side_effect = [ - {"type": "http.request"}, # First call - {"type": "http.disconnect"} # Second call - disconnect - ] - - mock_llm_task = AsyncMock() - - await _check_request_disconnection(mock_request, mock_llm_task) - - mock_llm_task.cancel.assert_called_once() - - -@pytest.mark.asyncio -async def test_check_request_disconnection_no_disconnect(): - """Test that _check_request_disconnection handles normal requests.""" - mock_request = AsyncMock() - mock_request.receive.return_value = {"type": "http.request"} - - mock_llm_task = AsyncMock() - - # This will timeout after 600 seconds, but we don't need to wait - # Just test that it doesn't crash immediately - task = asyncio.create_task(_check_request_disconnection(mock_request, mock_llm_task)) - await asyncio.sleep(0.1) # Let it run briefly - task.cancel() - - try: - await task - except asyncio.CancelledError: - pass - - # Task should not be cancelled during normal operation - mock_llm_task.cancel.assert_not_called() \ No newline at end of file