scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

(perf) - fixes for aiohttp handler to hit 1K RPS (#7590)

Browse source 
* fix getting aiohttp sesson

* fix _get_async_client_session
This commit is contained in:
Ishaan Jaff 2025-01-06 15:41:39 -08:00 • committed by GitHub
parent 23685e93f3
commit 61d67cfa43
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 23 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

39
litellm/llms/custom_httpx/aiohttp_handler.py
View file

@ -1,8 +1,9 @@
import json
from typing import TYPE_CHECKING, Any, Optional, Tuple, Union
import aiohttp # Add this import
import aiohttp
import httpx # type: ignore
from aiohttp import ClientSession
import litellm
import litellm.litellm_core_utils
@ -13,7 +14,6 @@ from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
_get_httpx_client,
get_async_httpx_client,
)
from litellm.utils import CustomStreamWrapper, ModelResponse, ProviderConfigManager
@ -32,9 +32,21 @@ class BaseLLMAIOHTTPHandler:
def __init__(self):
self.client_session: Optional[aiohttp.ClientSession] = None
def _get_async_client_session(
self, dynamic_client_session: Optional[ClientSession] = None
) -> ClientSession:
if dynamic_client_session:
return dynamic_client_session
elif self.client_session:
return self.client_session
else:
# init client session, and then return new session
self.client_session = aiohttp.ClientSession()
return self.client_session
async def _make_common_async_call(
self,
async_httpx_client: AsyncHTTPHandler,
async_client_session: Optional[ClientSession],
provider_config: BaseConfig,
api_base: str,
headers: dict,
@ -49,12 +61,13 @@ class BaseLLMAIOHTTPHandler:
)
response: Optional[aiohttp.ClientResponse] = None
if self.client_session is None:
self.client_session = aiohttp.ClientSession()
async_client_session = self._get_async_client_session(
dynamic_client_session=async_client_session
)
for i in range(max(max_retry_on_unprocessable_entity_error, 1)):
try:
response = await self.client_session.post(
response = await async_client_session.post(
url=api_base,
headers=headers,
json=data,
@ -146,17 +159,11 @@ class BaseLLMAIOHTTPHandler:
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
client: Optional[AsyncHTTPHandler] = None,
client: Optional[ClientSession] = None,
):
if client is None:
async_httpx_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders(custom_llm_provider)
)
else:
async_httpx_client = client
_response = await self._make_common_async_call(
async_httpx_client=async_httpx_client,
async_client_session=client,
provider_config=provider_config,
api_base=api_base,
headers=headers,
@ -186,7 +193,7 @@ class BaseLLMAIOHTTPHandler:
fake_stream: bool = False,
api_key: Optional[str] = None,
headers: Optional[dict] = {},
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler, ClientSession]] = None,
):
provider_config = ProviderConfigManager.get_provider_chat_config(
model=model, provider=litellm.LlmProviders(custom_llm_provider)
@ -245,7 +252,7 @@ class BaseLLMAIOHTTPHandler:
encoding=encoding,
client=(
client
if client is not None and isinstance(client, AsyncHTTPHandler)
if client is not None and isinstance(client, ClientSession)
else None
),
)