[Performance] Improve LiteLLM Python SDK RPS by +200 RPS (#13839)

* fix _proxy_from_env +100 RPS

* fix: global_braintrust_http_handler

* test_braintrust_logging
This commit is contained in:
Ishaan Jaff 2025-08-20 21:46:33 -07:00 committed by GitHub
parent b9621c760d
commit 07f6235730
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 40 additions and 20 deletions

View File

@ -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,

View File

@ -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

View File

@ -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