From 07f62357306c44bc942fe8ded80b0552b35b6def Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 20 Aug 2025 21:46:33 -0700 Subject: [PATCH] [Performance] Improve LiteLLM Python SDK RPS by +200 RPS (#13839) * fix _proxy_from_env +100 RPS * fix: global_braintrust_http_handler * test_braintrust_logging --- litellm/integrations/braintrust_logging.py | 22 +++++++-------- .../llms/custom_httpx/aiohttp_transport.py | 28 +++++++++++++++++-- tests/local_testing/test_braintrust.py | 10 +++---- 3 files changed, 40 insertions(+), 20 deletions(-) diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 8149a6131e8..5cf403adee4 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -19,10 +19,6 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.utils import print_verbose -global_braintrust_http_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback -) -global_braintrust_sync_http_handler = HTTPHandler() API_BASE = "https://api.braintrustdata.com/v1" @@ -52,6 +48,10 @@ class BraintrustLogger(CustomLogger): self._project_id_cache: Dict[ str, str ] = {} # Cache mapping project names to IDs + self.global_braintrust_http_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) + self.global_braintrust_sync_http_handler = HTTPHandler() def validate_environment(self, api_key: Optional[str]): """ @@ -76,7 +76,7 @@ class BraintrustLogger(CustomLogger): return self._project_id_cache[project_name] try: - response = global_braintrust_sync_http_handler.post( + response = self.global_braintrust_sync_http_handler.post( f"{self.api_base}/project", headers=self.headers, json={"name": project_name}, @@ -96,7 +96,7 @@ class BraintrustLogger(CustomLogger): return self._project_id_cache[project_name] try: - response = await global_braintrust_http_handler.post( + response = await self.global_braintrust_http_handler.post( f"{self.api_base}/project/register", headers=self.headers, json={"name": project_name}, @@ -146,7 +146,7 @@ class BraintrustLogger(CustomLogger): return metadata async def create_default_project_and_experiment(self): - project = await global_braintrust_http_handler.post( + project = await self.global_braintrust_http_handler.post( f"{self.api_base}/project", headers=self.headers, json={"name": "litellm"} ) @@ -155,7 +155,7 @@ class BraintrustLogger(CustomLogger): self.default_project_id = project_dict["id"] def create_sync_default_project_and_experiment(self): - project = global_braintrust_sync_http_handler.post( + project = self.global_braintrust_sync_http_handler.post( f"{self.api_base}/project", headers=self.headers, json={"name": "litellm"} ) @@ -291,9 +291,9 @@ class BraintrustLogger(CustomLogger): try: print_verbose( - f"global_braintrust_sync_http_handler.post: {global_braintrust_sync_http_handler.post}" + f"self.global_braintrust_sync_http_handler.post: {self.global_braintrust_sync_http_handler.post}" ) - global_braintrust_sync_http_handler.post( + self.global_braintrust_sync_http_handler.post( url=f"{self.api_base}/project_logs/{project_id}/insert", json={"events": [request_data]}, headers=self.headers, @@ -446,7 +446,7 @@ class BraintrustLogger(CustomLogger): request_data["metrics"] = metrics try: - await global_braintrust_http_handler.post( + await self.global_braintrust_http_handler.post( url=f"{self.api_base}/project_logs/{project_id}/insert", json={"events": [request_data]}, headers=self.headers, diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 3ed7d04bde6..ab69ea1f8c3 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -3,7 +3,7 @@ import contextlib import os import typing import urllib.request -from typing import Callable, Dict, Union +from typing import Callable, Dict, Optional, Union import aiohttp import aiohttp.client_exceptions @@ -115,6 +115,12 @@ class AiohttpTransport(httpx.AsyncBaseTransport): ) -> None: self.client = client + ######################################################### + # Class variables for proxy settings + ######################################################### + self.proxy: Optional[str] = None + self.checked_proxy_env_settings: bool = False + async def aclose(self) -> None: if isinstance(self.client, ClientSession): await self.client.close() @@ -249,7 +255,22 @@ class LiteLLMAiohttpTransport(AiohttpTransport): def _proxy_from_env(self, url: httpx.URL) -> typing.Optional[str]: - """Return proxy URL from env for the given request URL.""" + """ + Return proxy URL from env for the given request URL + + Only check the proxy env settings once, this is a costly operation for CPU % usage + + .""" + ######################################################### + # Check if we've already checked the proxy env settings + ######################################################### + if self.checked_proxy_env_settings is True: + return self.proxy + + ######################################################### + # set self.checked_proxy_env_settings to True + ######################################################### + self.checked_proxy_env_settings = True proxies = urllib.request.getproxies() if urllib.request.proxy_bypass(url.host): return None @@ -257,4 +278,5 @@ class LiteLLMAiohttpTransport(AiohttpTransport): proxy = proxies.get(url.scheme) or proxies.get("all") if proxy and "://" not in proxy: proxy = f"http://{proxy}" - return proxy + self.proxy = proxy + return self.proxy diff --git a/tests/local_testing/test_braintrust.py b/tests/local_testing/test_braintrust.py index c898594683e..13b23a97585 100644 --- a/tests/local_testing/test_braintrust.py +++ b/tests/local_testing/test_braintrust.py @@ -35,9 +35,8 @@ def test_braintrust_logging(): http_client = HTTPHandler() - with patch.object( - litellm.integrations.braintrust_logging.global_braintrust_sync_http_handler, - "post", + with patch( + "litellm.integrations.braintrust_logging.HTTPHandler.post", new=MagicMock(), ) as mock_client: # set braintrust as a callback, litellm will send the data to braintrust @@ -57,9 +56,8 @@ def test_braintrust_logging_specific_project_id(): litellm.set_verbose = True - with patch.object( - litellm.integrations.braintrust_logging.global_braintrust_sync_http_handler, - "post", + with patch( + "litellm.integrations.braintrust_logging.HTTPHandler.post", new=MagicMock(), ) as mock_client: # set braintrust as a callback, litellm will send the data to braintrust