[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:
parent
b9621c760d
commit
07f6235730
@ -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,
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user