[Performance] RPS Improvement +500 RPS when sending the user field (#14616)

* perf tool

* fix: cache type issue

* fix: exception hanging & cache setting

1. Removed unhandled exceptions
2. Set cache value to dict
This commit is contained in:
Alexsander Hamir
2025-09-16 16:18:23 -07:00
committed by GitHub
parent 68105ce1a7
commit 02db2e8ae8
3 changed files with 134 additions and 12 deletions
+8 -11
View File
@@ -469,16 +469,13 @@ async def get_end_user_object(
# check if in cache
cached_user_obj = await user_api_key_cache.async_get_cache(key=_key)
if cached_user_obj is not None:
if isinstance(cached_user_obj, dict):
return_obj = LiteLLM_EndUserTable(**cached_user_obj)
check_in_budget(end_user_obj=return_obj)
return return_obj
elif isinstance(cached_user_obj, LiteLLM_EndUserTable):
return_obj = cached_user_obj
check_in_budget(end_user_obj=return_obj)
return return_obj
# Convert cached dict to LiteLLM_EndUserTable instance
return_obj = LiteLLM_EndUserTable(**cached_user_obj)
check_in_budget(end_user_obj=return_obj)
return return_obj
# else, check db
try:
try:
response = await prisma_client.db.litellm_endusertable.find_unique(
where={"user_id": end_user_id},
include={"litellm_budget_table": True},
@@ -487,9 +484,9 @@ async def get_end_user_object(
if response is None:
raise Exception
# save the end-user object to cache
# save the end-user object to cache (always store as dict for consistency)
await user_api_key_cache.async_set_cache(
key="end_user_id:{}".format(end_user_id), value=response
key="end_user_id:{}".format(end_user_id), value=response.dict()
)
_response = LiteLLM_EndUserTable(**response.dict())
@@ -0,0 +1,125 @@
"""
Performance utilities for LiteLLM proxy server.
This module provides performance monitoring and profiling functionality for endpoint
performance analysis using cProfile with configurable sampling rates.
"""
import asyncio
import cProfile
import functools
import threading
from pathlib import Path as PathLib
from litellm._logging import verbose_proxy_logger
# Global profiling state
_profile_lock = threading.Lock()
_profiler = None
_last_profile_file_path = None
_sample_counter = 0
_sample_counter_lock = threading.Lock()
def _should_sample(profile_sampling_rate: float) -> bool:
"""Determine if current request should be sampled based on sampling rate."""
if profile_sampling_rate >= 1.0:
return True # Always sample
elif profile_sampling_rate <= 0.0:
return False # Never sample
# Use deterministic sampling based on counter for consistent rate
global _sample_counter
with _sample_counter_lock:
_sample_counter += 1
# Sample based on rate (e.g., 0.1 means sample every 10th request)
should_sample = (_sample_counter % int(1.0 / profile_sampling_rate)) == 0
return should_sample
def _start_profiling(profile_sampling_rate: float) -> None:
"""Start cProfile profiling once globally."""
global _profiler
with _profile_lock:
if _profiler is None:
_profiler = cProfile.Profile()
_profiler.enable()
verbose_proxy_logger.info(f"Profiling started with sampling rate: {profile_sampling_rate}")
def _start_profiling_for_request(profile_sampling_rate: float) -> bool:
"""Start profiling for a specific request (if sampling allows)."""
if _should_sample(profile_sampling_rate):
_start_profiling(profile_sampling_rate)
return True
return False
def _save_stats(profile_file: PathLib) -> None:
"""Save current stats directly to file."""
with _profile_lock:
if _profiler is None:
return
try:
# Disable profiler temporarily to dump stats
_profiler.disable()
_profiler.dump_stats(str(profile_file))
# Re-enable profiler to continue profiling
_profiler.enable()
verbose_proxy_logger.debug(f"Profiling stats saved to {profile_file}")
except Exception as e:
verbose_proxy_logger.error(f"Error saving profiling stats: {e}")
# Make sure profiler is re-enabled even if there's an error
try:
_profiler.enable()
except Exception:
pass
def profile_endpoint(sampling_rate: float = 1.0):
"""Decorator to sample endpoint hits and save to a profile file.
Args:
sampling_rate: Rate of requests to profile (0.0 to 1.0)
- 1.0: Profile all requests (100%)
- 0.1: Profile 1 in 10 requests (10%)
- 0.0: Profile no requests (0%)
"""
def decorator(func):
def set_last_profile_path(path: PathLib) -> None:
global _last_profile_file_path
_last_profile_file_path = path
if asyncio.iscoroutinefunction(func):
@functools.wraps(func)
async def async_wrapper(*args, **kwargs):
is_sampling = _start_profiling_for_request(sampling_rate)
file_path_obj = PathLib("endpoint_profile.pstat")
set_last_profile_path(file_path_obj)
try:
result = await func(*args, **kwargs)
if is_sampling:
_save_stats(file_path_obj)
return result
except Exception:
if is_sampling:
_save_stats(file_path_obj)
raise
return async_wrapper
else:
@functools.wraps(func)
def sync_wrapper(*args, **kwargs):
is_sampling = _start_profiling_for_request(sampling_rate)
file_path_obj = PathLib("endpoint_profile.pstat")
set_last_profile_path(file_path_obj)
try:
result = func(*args, **kwargs)
if is_sampling:
_save_stats(file_path_obj)
return result
except Exception:
if is_sampling:
_save_stats(file_path_obj)
raise
return sync_wrapper
return decorator
+1 -1
View File
@@ -48,7 +48,7 @@ async def test_get_end_user_object(customer_spend, customer_budget):
)
_cache = DualCache()
_key = "end_user_id:{}".format(end_user_id)
_cache.set_cache(key=_key, value=end_user_obj)
_cache.set_cache(key=_key, value=end_user_obj.model_dump())
try:
await get_end_user_object(
end_user_id=end_user_id,