diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 94cfa6fc67..df2b7fc803 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -470,6 +470,7 @@ router_settings: | DEFAULT_MAX_RETRIES | Default maximum retry attempts. Default is 2 | DEFAULT_MAX_TOKENS | Default maximum tokens for LLM calls. Default is 4096 | DEFAULT_MAX_TOKENS_FOR_TRITON | Default maximum tokens for Triton models. Default is 2000 +| DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE | Default maximum size for redis batch cache. Default is 1000 | DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT | Default token count for mock response completions. Default is 20 | DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT | Default token count for mock response prompts. Default is 10 | DEFAULT_MODEL_CREATED_AT_TIME | Default creation timestamp for models. Default is 1677610602 @@ -717,6 +718,7 @@ router_settings: | PROXY_BATCH_POLLING_INTERVAL | Time in seconds to wait before polling a batch, to check if it's completed. Default is 6000s (1 hour) | PROXY_BUDGET_RESCHEDULER_MAX_TIME | Maximum time in seconds to wait before checking database for budget resets. Default is 605 | PROXY_BUDGET_RESCHEDULER_MIN_TIME | Minimum time in seconds to wait before checking database for budget resets. Default is 597 +| PYTHON_GC_THRESHOLD | GC thresholds ('gen0,gen1,gen2', e.g. '1000,50,50'); defaults to Python’s values. | PROXY_LOGOUT_URL | URL for logging out of the proxy service | QDRANT_API_BASE | Base URL for Qdrant API | QDRANT_API_KEY | API key for Qdrant service diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index ce07f7ce70..3edc3f4282 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -19,6 +19,7 @@ if TYPE_CHECKING: import litellm from litellm._logging import print_verbose, verbose_logger +from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE from .base_cache import BaseCache from .in_memory_cache import InMemoryCache @@ -60,7 +61,7 @@ class DualCache(BaseCache): default_in_memory_ttl: Optional[float] = None, default_redis_ttl: Optional[float] = None, default_redis_batch_cache_expiry: Optional[float] = None, - default_max_redis_batch_cache_size: int = 100, + default_max_redis_batch_cache_size: int = DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE, ) -> None: super().__init__() # If in_memory_cache is not provided, use the default InMemoryCache @@ -260,7 +261,7 @@ class DualCache(BaseCache): **kwargs, ): try: - result = [None for _ in range(len(keys))] + result = [None] * len(keys) if self.in_memory_cache is not None: in_memory_result = await self.in_memory_cache.async_batch_get_cache( keys, **kwargs @@ -283,20 +284,27 @@ class DualCache(BaseCache): redis_result = await self.redis_cache.async_batch_get_cache( sublist_keys, parent_otel_span=parent_otel_span ) + + # Update the last access time for ALL queried keys + # This includes keys with None values to throttle repeated Redis queries + for key in sublist_keys: + self.last_redis_batch_access_time[key] = current_time + + # Short-circuit if redis_result is None or contains only None values + if redis_result is None or all(v is None for v in redis_result.values()): + return result - if redis_result is not None: - # Update in-memory cache with the value from Redis - for key, value in redis_result.items(): - if value is not None: - await self.in_memory_cache.async_set_cache( - key, redis_result[key], **kwargs - ) - # Update the last access time for each key fetched from Redis - self.last_redis_batch_access_time[key] = current_time - + # Pre-compute key-to-index mapping for O(1) lookup + key_to_index = {key: i for i, key in enumerate(keys)} + + # Update both result and in-memory cache in a single loop for key, value in redis_result.items(): - index = keys.index(key) - result[index] = value + result[key_to_index[key]] = value + + if value is not None and self.in_memory_cache is not None: + await self.in_memory_cache.async_set_cache( + key, value, **kwargs + ) return result except Exception: diff --git a/litellm/constants.py b/litellm/constants.py index f4a5f92810..cf6b53fb3e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -199,6 +199,9 @@ JITTER = float(os.getenv("JITTER", 0.75)) DEFAULT_IN_MEMORY_TTL = int( os.getenv("DEFAULT_IN_MEMORY_TTL", 5) ) # default time to live for the in-memory cache +DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE = int( + os.getenv("DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE", 1000) +) # default max size for redis batch cache DEFAULT_POLLING_INTERVAL = float( os.getenv("DEFAULT_POLLING_INTERVAL", 0.03) ) # default polling interval for the scheduler @@ -970,6 +973,10 @@ DEFAULT_SOFT_BUDGET = float( # makes it clear this is a rate limit error for a litellm virtual key RATE_LIMIT_ERROR_MESSAGE_FOR_VIRTUAL_KEY = "LiteLLM Virtual Key user_api_key_hash" +# Python garbage collection threshold configuration +# Format: "gen0,gen1,gen2" e.g., "1000,50,50" +PYTHON_GC_THRESHOLD = os.getenv("PYTHON_GC_THRESHOLD") + # pass through route constansts BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES = [ "agents/", diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 16ab2cc805..0cb7f0058f 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -1,19 +1,46 @@ # Start tracing memory allocations import asyncio +import gc import json import os +import sys import tracemalloc from collections import Counter +from typing import Any, Dict, List, Optional, Tuple -from fastapi import APIRouter, Depends +from fastapi import APIRouter, Depends, HTTPException, Query from litellm import get_secret_str from litellm._logging import verbose_proxy_logger +from litellm.constants import PYTHON_GC_THRESHOLD from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router = APIRouter() +# Configure garbage collection thresholds from environment variables +def configure_gc_thresholds(): + """Configure Python garbage collection thresholds from environment variables.""" + gc_threshold_env = PYTHON_GC_THRESHOLD + if gc_threshold_env: + try: + # Parse threshold string like "1000,50,50" + thresholds = [int(x.strip()) for x in gc_threshold_env.split(",")] + if len(thresholds) == 3: + gc.set_threshold(*thresholds) + verbose_proxy_logger.info(f"GC thresholds set to: {thresholds}") + else: + verbose_proxy_logger.warning(f"GC threshold not set: {gc_threshold_env}. Expected format: 'gen0,gen1,gen2'") + except ValueError as e: + verbose_proxy_logger.warning(f"Failed to parse GC threshold: {gc_threshold_env}. Error: {e}") + + # Log current thresholds + current_thresholds = gc.get_threshold() + verbose_proxy_logger.info(f"Current GC thresholds: gen0={current_thresholds[0]}, gen1={current_thresholds[1]}, gen2={current_thresholds[2]}") + +# Initialize GC configuration +configure_gc_thresholds() + @router.get("/debug/asyncio-tasks") async def get_active_tasks_stats(): @@ -158,6 +185,490 @@ async def memory_usage_in_mem_cache_items( } +@router.get("/debug/memory/summary", include_in_schema=False) +async def get_memory_summary( + _: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> Dict[str, Any]: + """ + Get simplified memory usage summary for the proxy. + + Returns: + - worker_pid: Process ID + - status: Overall health based on memory usage + - memory: Process memory usage and RAM info + - caches: Cache item counts and descriptions + - garbage_collector: GC status and pending object counts + + Example usage: + curl http://localhost:4000/debug/memory/summary -H "Authorization: Bearer sk-1234" + + For detailed analysis, call GET /debug/memory/details + For cache management, use the cache management endpoints + """ + from litellm.proxy.proxy_server import ( + llm_router, + proxy_logging_obj, + user_api_key_cache, + ) + + # Get process memory info + process_memory = {} + health_status = "healthy" + + try: + import psutil + + process = psutil.Process() + memory_info = process.memory_info() + memory_mb = memory_info.rss / (1024 * 1024) + memory_percent = process.memory_percent() + + process_memory = { + "summary": f"{memory_mb:.1f} MB ({memory_percent:.1f}% of system memory)", + "ram_usage_mb": round(memory_mb, 2), + "system_memory_percent": round(memory_percent, 2), + } + + # Check memory health status + if memory_percent > 80: + health_status = "critical" + elif memory_percent > 60: + health_status = "warning" + else: + health_status = "healthy" + + except ImportError: + process_memory["error"] = "Install psutil for memory monitoring: pip install psutil" + except Exception as e: + process_memory["error"] = str(e) + + # Get cache information + caches: Dict[str, Any] = {} + total_cache_items = 0 + + try: + # User API key cache + user_cache_items = len(user_api_key_cache.in_memory_cache.cache_dict) + total_cache_items += user_cache_items + caches["user_api_keys"] = { + "count": user_cache_items, + "count_readable": f"{user_cache_items:,}", + "what_it_stores": "Validated API keys for faster authentication" + } + + # Router cache + if llm_router is not None: + router_cache_items = len(llm_router.cache.in_memory_cache.cache_dict) + total_cache_items += router_cache_items + caches["llm_responses"] = { + "count": router_cache_items, + "count_readable": f"{router_cache_items:,}", + "what_it_stores": "LLM responses for identical requests" + } + + # Proxy logging cache + logging_cache_items = len( + proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict + ) + total_cache_items += logging_cache_items + caches["usage_tracking"] = { + "count": logging_cache_items, + "count_readable": f"{logging_cache_items:,}", + "what_it_stores": "Usage metrics before database write" + } + + except Exception as e: + caches["error"] = str(e) + + # Get garbage collector stats + gc_enabled = gc.isenabled() + objects_pending = gc.get_count()[0] + uncollectable = len(gc.garbage) + + gc_info = { + "status": "enabled" if gc_enabled else "disabled", + "objects_awaiting_collection": objects_pending, + } + + # Add warning if garbage collection issues detected + if uncollectable > 0: + gc_info["warning"] = f"{uncollectable} uncollectable objects (possible memory leak)" + + return { + "worker_pid": os.getpid(), + "status": health_status, + "memory": process_memory, + "caches": { + "total_items": total_cache_items, + "breakdown": caches, + }, + "garbage_collector": gc_info, + } + + +def _get_gc_statistics() -> Dict[str, Any]: + """Get garbage collector statistics.""" + return { + "enabled": gc.isenabled(), + "thresholds": { + "generation_0": gc.get_threshold()[0], + "generation_1": gc.get_threshold()[1], + "generation_2": gc.get_threshold()[2], + "explanation": "Number of allocations before automatic collection for each generation" + }, + "current_counts": { + "generation_0": gc.get_count()[0], + "generation_1": gc.get_count()[1], + "generation_2": gc.get_count()[2], + "explanation": "Current number of allocated objects in each generation" + }, + "collection_history": [ + { + "generation": i, + "total_collections": stat["collections"], + "total_collected": stat["collected"], + "uncollectable": stat["uncollectable"], + } + for i, stat in enumerate(gc.get_stats()) + ], + } + + +def _get_object_type_counts(top_n: int) -> Tuple[int, List[Dict[str, Any]]]: + """Count objects by type and return total count and top N types.""" + type_counts: Counter = Counter() + total_objects = 0 + + for obj in gc.get_objects(): + total_objects += 1 + obj_type = type(obj).__name__ + type_counts[obj_type] += 1 + + top_object_types = [ + { + "type": obj_type, + "count": count, + "count_readable": f"{count:,}" + } + for obj_type, count in type_counts.most_common(top_n) + ] + + return total_objects, top_object_types + + +def _get_uncollectable_objects_info() -> Dict[str, Any]: + """Get information about uncollectable objects (potential memory leaks).""" + uncollectable = gc.garbage + return { + "count": len(uncollectable), + "sample_types": [type(obj).__name__ for obj in uncollectable[:10]], + "warning": "If count > 0, you may have reference cycles preventing garbage collection" if len(uncollectable) > 0 else None, + } + + +def _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache) -> Dict[str, Any]: + """Calculate memory usage for all caches.""" + cache_stats: Dict[str, Any] = {} + try: + # User API key cache + user_cache_size = sys.getsizeof(user_api_key_cache.in_memory_cache.cache_dict) + user_ttl_size = sys.getsizeof(user_api_key_cache.in_memory_cache.ttl_dict) + cache_stats["user_api_key_cache"] = { + "num_items": len(user_api_key_cache.in_memory_cache.cache_dict), + "cache_dict_size_bytes": user_cache_size, + "ttl_dict_size_bytes": user_ttl_size, + "total_size_mb": round((user_cache_size + user_ttl_size) / (1024 * 1024), 2), + } + + # Router cache + if llm_router is not None: + router_cache_size = sys.getsizeof(llm_router.cache.in_memory_cache.cache_dict) + router_ttl_size = sys.getsizeof(llm_router.cache.in_memory_cache.ttl_dict) + cache_stats["llm_router_cache"] = { + "num_items": len(llm_router.cache.in_memory_cache.cache_dict), + "cache_dict_size_bytes": router_cache_size, + "ttl_dict_size_bytes": router_ttl_size, + "total_size_mb": round((router_cache_size + router_ttl_size) / (1024 * 1024), 2), + } + + # Proxy logging cache + logging_cache_size = sys.getsizeof( + proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict + ) + logging_ttl_size = sys.getsizeof( + proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.ttl_dict + ) + cache_stats["proxy_logging_cache"] = { + "num_items": len( + proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict + ), + "cache_dict_size_bytes": logging_cache_size, + "ttl_dict_size_bytes": logging_ttl_size, + "total_size_mb": round((logging_cache_size + logging_ttl_size) / (1024 * 1024), 2), + } + + # Redis cache info + if redis_usage_cache is not None: + cache_stats["redis_usage_cache"] = { + "enabled": True, + "cache_type": type(redis_usage_cache).__name__, + } + # Try to get Redis connection pool info if available + try: + if hasattr(redis_usage_cache, 'redis_client') and redis_usage_cache.redis_client: + if hasattr(redis_usage_cache.redis_client, 'connection_pool'): + pool_info = redis_usage_cache.redis_client.connection_pool # type: ignore + cache_stats["redis_usage_cache"]["connection_pool"] = { + "max_connections": pool_info.max_connections if hasattr(pool_info, 'max_connections') else None, + "connection_class": pool_info.connection_class.__name__ if hasattr(pool_info, 'connection_class') else None, + } + except Exception as e: + verbose_proxy_logger.debug(f"Error getting Redis pool info: {e}") + else: + cache_stats["redis_usage_cache"] = {"enabled": False} + + except Exception as e: + verbose_proxy_logger.debug(f"Error calculating cache stats: {e}") + cache_stats["error"] = str(e) + + return cache_stats + + +def _get_router_memory_stats(llm_router) -> Dict[str, Any]: + """Get memory usage statistics for LiteLLM router.""" + litellm_router_memory: Dict[str, Any] = {} + try: + if llm_router is not None: + # Model list memory size + if hasattr(llm_router, 'model_list') and llm_router.model_list: + model_list_size = sys.getsizeof(llm_router.model_list) + litellm_router_memory["model_list"] = { + "num_models": len(llm_router.model_list), + "size_bytes": model_list_size, + "size_mb": round(model_list_size / (1024 * 1024), 4), + } + + # Model names set + if hasattr(llm_router, 'model_names') and llm_router.model_names: + model_names_size = sys.getsizeof(llm_router.model_names) + litellm_router_memory["model_names_set"] = { + "num_model_groups": len(llm_router.model_names), + "size_bytes": model_names_size, + "size_mb": round(model_names_size / (1024 * 1024), 4), + } + + # Deployment names list + if hasattr(llm_router, 'deployment_names') and llm_router.deployment_names: + deployment_names_size = sys.getsizeof(llm_router.deployment_names) + litellm_router_memory["deployment_names"] = { + "num_deployments": len(llm_router.deployment_names), + "size_bytes": deployment_names_size, + "size_mb": round(deployment_names_size / (1024 * 1024), 4), + } + + # Deployment latency map + if hasattr(llm_router, 'deployment_latency_map') and llm_router.deployment_latency_map: + latency_map_size = sys.getsizeof(llm_router.deployment_latency_map) + litellm_router_memory["deployment_latency_map"] = { + "num_tracked_deployments": len(llm_router.deployment_latency_map), + "size_bytes": latency_map_size, + "size_mb": round(latency_map_size / (1024 * 1024), 4), + } + + # Fallback configuration + if hasattr(llm_router, 'fallbacks') and llm_router.fallbacks: + fallbacks_size = sys.getsizeof(llm_router.fallbacks) + litellm_router_memory["fallbacks"] = { + "num_fallback_configs": len(llm_router.fallbacks), + "size_bytes": fallbacks_size, + "size_mb": round(fallbacks_size / (1024 * 1024), 4), + } + + # Total router object size + router_obj_size = sys.getsizeof(llm_router) + litellm_router_memory["router_object"] = { + "size_bytes": router_obj_size, + "size_mb": round(router_obj_size / (1024 * 1024), 4), + } + + else: + litellm_router_memory = {"note": "Router not initialized"} + except Exception as e: + verbose_proxy_logger.debug(f"Error getting router memory info: {e}") + litellm_router_memory = {"error": str(e)} + + return litellm_router_memory + + +def _get_process_memory_info(worker_pid: int, include_process_info: bool) -> Optional[Dict[str, Any]]: + """Get process-level memory information using psutil.""" + if not include_process_info: + return None + + try: + import psutil + + process = psutil.Process() + memory_info = process.memory_info() + ram_usage_mb = round(memory_info.rss / (1024 * 1024), 2) + virtual_memory_mb = round(memory_info.vms / (1024 * 1024), 2) + memory_percent = round(process.memory_percent(), 2) + + return { + "pid": worker_pid, + "summary": f"Worker PID {worker_pid} using {ram_usage_mb:.1f} MB of RAM ({memory_percent:.1f}% of system memory)", + "ram_usage": { + "megabytes": ram_usage_mb, + "description": "Actual physical RAM used by this process" + }, + "virtual_memory": { + "megabytes": virtual_memory_mb, + "description": "Total virtual memory allocated (includes swapped memory)" + }, + "system_memory_percent": { + "percent": memory_percent, + "description": "Percentage of total system RAM being used" + }, + "open_file_handles": { + "count": process.num_fds() if hasattr(process, "num_fds") else "N/A (Windows)", + "description": "Number of open file descriptors/handles" + }, + "threads": { + "count": process.num_threads(), + "description": "Number of active threads in this process" + } + } + except ImportError: + return { + "pid": worker_pid, + "error": "psutil not installed. Install with: pip install psutil" + } + except Exception as e: + verbose_proxy_logger.debug(f"Error getting process info: {e}") + return {"pid": worker_pid, "error": str(e)} + + +@router.get("/debug/memory/details", include_in_schema=False) +async def get_memory_details( + _: UserAPIKeyAuth = Depends(user_api_key_auth), + top_n: int = Query(20, description="Number of top object types to return"), + include_process_info: bool = Query(True, description="Include process memory info"), +) -> Dict[str, Any]: + """ + Get detailed memory diagnostics for deep debugging. + + Returns: + - worker_pid: Process ID + - process_memory: RAM usage, virtual memory, file handles, threads + - garbage_collector: GC thresholds, counts, collection history + - objects: Total tracked objects and top object types + - uncollectable: Objects that can't be garbage collected (potential leaks) + - cache_memory: Memory usage of user_api_key, router, and logging caches + - router_memory: Memory usage of router components (model_list, deployment_names, etc.) + + Query Parameters: + - top_n: Number of top object types to return (default: 20) + - include_process_info: Include process-level memory info using psutil (default: true) + + Example usage: + curl "http://localhost:4000/debug/memory/details?top_n=30" -H "Authorization: Bearer sk-1234" + + All memory sizes are reported in both bytes and MB. + """ + from litellm.proxy.proxy_server import ( + llm_router, + proxy_logging_obj, + user_api_key_cache, + redis_usage_cache, + ) + + worker_pid = os.getpid() + + # Collect all diagnostics using helper functions + gc_stats = _get_gc_statistics() + total_objects, top_object_types = _get_object_type_counts(top_n) + uncollectable_info = _get_uncollectable_objects_info() + cache_stats = _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache) + litellm_router_memory = _get_router_memory_stats(llm_router) + process_info = _get_process_memory_info(worker_pid, include_process_info) + + return { + "worker_pid": worker_pid, + "process_memory": process_info, + "garbage_collector": gc_stats, + "objects": { + "total_tracked": total_objects, + "total_tracked_readable": f"{total_objects:,}", + "top_types": top_object_types, + }, + "uncollectable": uncollectable_info, + "cache_memory": cache_stats, + "router_memory": litellm_router_memory, + } + + +@router.post("/debug/memory/gc/configure", include_in_schema=False) +async def configure_gc_thresholds_endpoint( + _: UserAPIKeyAuth = Depends(user_api_key_auth), + generation_0: int = Query(700, description="Generation 0 threshold (default: 700)"), + generation_1: int = Query(10, description="Generation 1 threshold (default: 10)"), + generation_2: int = Query(10, description="Generation 2 threshold (default: 10)"), +) -> Dict[str, Any]: + """ + Configure Python garbage collection thresholds. + + Lower thresholds mean more frequent GC cycles (less memory, more CPU overhead). + Higher thresholds mean less frequent GC cycles (more memory, less CPU overhead). + + Returns: + - message: Confirmation message + - previous_thresholds: Old threshold values + - new_thresholds: New threshold values + - objects_awaiting_collection: Current object count in gen-0 + - tip: Hint about when next collection will occur + + Query Parameters: + - generation_0: Number of allocations before gen-0 collection (default: 700) + - generation_1: Number of gen-0 collections before gen-1 collection (default: 10) + - generation_2: Number of gen-1 collections before gen-2 collection (default: 10) + + Example for more aggressive collection: + curl -X POST "http://localhost:4000/debug/memory/gc/configure?generation_0=500" -H "Authorization: Bearer sk-1234" + + Example for less aggressive collection: + curl -X POST "http://localhost:4000/debug/memory/gc/configure?generation_0=1000" -H "Authorization: Bearer sk-1234" + + Monitor memory usage with GET /debug/memory/summary after changes. + """ + # Get current thresholds for logging + old_thresholds = gc.get_threshold() + + # Set new thresholds with error handling + try: + gc.set_threshold(generation_0, generation_1, generation_2) + verbose_proxy_logger.info( + f"GC thresholds updated from {old_thresholds} to " + f"({generation_0}, {generation_1}, {generation_2})" + ) + except Exception as e: + verbose_proxy_logger.error(f"Failed to set GC thresholds: {e}") + raise HTTPException( + status_code=500, + detail=f"Failed to set GC thresholds: {str(e)}" + ) + + # Get current object count to show immediate impact + current_count = gc.get_count()[0] + + return { + "message": "GC thresholds updated", + "previous_thresholds": f"{old_thresholds[0]}, {old_thresholds[1]}, {old_thresholds[2]}", + "new_thresholds": f"{generation_0}, {generation_1}, {generation_2}", + "objects_awaiting_collection": current_count, + "tip": f"Next collection will run after {generation_0 - current_count} more allocations" + } + + @router.get("/otel-spans", include_in_schema=False) async def get_otel_spans(): from litellm.proxy.proxy_server import open_telemetry_logger diff --git a/litellm/router.py b/litellm/router.py index 1b691678bd..1d768415ff 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -973,7 +973,8 @@ class Router: ) self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data = deployment["litellm_params"].copy() + # No copy needed - data is only read and spread into new dict below + data = deployment["litellm_params"] model_name = data["model"] potential_model_client = self._get_client( deployment=deployment, kwargs=kwargs @@ -1280,7 +1281,8 @@ class Router: deployment=deployment, parent_otel_span=parent_otel_span ) self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data = deployment["litellm_params"].copy() + # No copy needed - data is only read and spread into new dict below + data = deployment["litellm_params"] model_name = data["model"] @@ -1944,21 +1946,15 @@ class Router: def _is_prompt_management_model(self, model: str) -> bool: model_list = self.get_model_list(model_name=model) - if model_list is None: - return False - if len(model_list) != 1: + if model_list is None or len(model_list) != 1: return False litellm_model = model_list[0]["litellm_params"].get("model", None) - - if litellm_model is None: + if litellm_model is None or "/" not in litellm_model: return False - if "/" in litellm_model: - split_litellm_model = litellm_model.split("/")[0] - if split_litellm_model in litellm._known_custom_logger_compatible_callbacks: - return True - return False + split_litellm_model = litellm_model.split("/")[0] + return split_litellm_model in litellm._known_custom_logger_compatible_callbacks async def _prompt_management_factory( self, @@ -6726,9 +6722,11 @@ class Router: f"Starting Pre-call checks for deployments in model={model}" ) - _returned_deployments = copy.deepcopy(healthy_deployments) + # Optimized: Use list() shallow copy instead of deepcopy + # We only pop from the list, not modify deployment dicts - 100x+ faster on hot path (every request) + _returned_deployments = list(healthy_deployments) - invalid_model_indices = [] + invalid_model_indices = set() # Use set for O(1) membership checks try: input_tokens = litellm.token_counter(messages=messages) @@ -6778,7 +6776,7 @@ class Router: isinstance(model_info["max_input_tokens"], int) and input_tokens > model_info["max_input_tokens"] ): - invalid_model_indices.append(idx) + invalid_model_indices.add(idx) _context_window_error = True _potential_error_str += ( "Model={}, Max Input Tokens={}, Got={}".format( @@ -6817,7 +6815,7 @@ class Router: isinstance(_litellm_params["rpm"], int) and _litellm_params["rpm"] <= current_request ): - invalid_model_indices.append(idx) + invalid_model_indices.add(idx) _rate_limit_error = True continue @@ -6833,7 +6831,7 @@ class Router: litellm_params=LiteLLM_Params(**_litellm_params), allowed_model_region=allowed_model_region, ): - invalid_model_indices.append(idx) + invalid_model_indices.add(idx) continue ## INVALID PARAMS ## -> catch 'gpt-3.5-turbo-16k' not supporting 'response_format' param @@ -6862,7 +6860,7 @@ class Router: verbose_router_logger.debug( f"INVALID MODEL INDEX @ REQUEST KWARG FILTERING, k={k}" ) - invalid_model_indices.append(idx) + invalid_model_indices.add(idx) if len(invalid_model_indices) == len(_returned_deployments): """ @@ -6885,8 +6883,10 @@ class Router: llm_provider="", ) if len(invalid_model_indices) > 0: - for idx in reversed(invalid_model_indices): - _returned_deployments.pop(idx) + # Single-pass filter using set for O(1) lookups (avoids O(n^2) from repeated pops) + _returned_deployments = [ + d for i, d in enumerate(_returned_deployments) if i not in invalid_model_indices + ] ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) if len(_returned_deployments) > 0: @@ -6986,9 +6986,11 @@ class Router: # check if default deployment is set if self.default_deployment is not None: - updated_deployment = copy.deepcopy( - self.default_deployment - ) # self.default_deployment + # Shallow copy with nested litellm_params copy (100x+ faster than deepcopy) + updated_deployment = self.default_deployment.copy() + updated_deployment["litellm_params"] = self.default_deployment[ + "litellm_params" + ].copy() updated_deployment["litellm_params"]["model"] = model return model, updated_deployment diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index aba8bfd72f..e92f114dd5 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -125,9 +125,9 @@ class CooldownCache: ) active_cooldowns: List[Tuple[str, CooldownCacheValue]] = [] - if results is None: + if results is None or all(v is None for v in results): return active_cooldowns - + # Process the results for model_id, result in zip(model_ids, results): if result and isinstance(result, dict): diff --git a/tests/local_testing/test_redis_batch_optimizations.py b/tests/local_testing/test_redis_batch_optimizations.py new file mode 100644 index 0000000000..4d8f4e6a04 --- /dev/null +++ b/tests/local_testing/test_redis_batch_optimizations.py @@ -0,0 +1,127 @@ +""" +Tests for Redis batch caching optimizations (commit 3f52e8c) + +Verifies: + +1. Batch cache size increased from 100 → 1000 (minimum 1k) +2. Repeated Redis queries for cache misses are throttled +""" + +import os +import sys +import time +from unittest.mock import AsyncMock, patch + +import pytest +from dotenv import load_dotenv + +load_dotenv() +sys.path.insert(0, os.path.abspath("../..")) + +import uuid +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache +from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE + + +@pytest.fixture +def cache_setup(): + """Create cache instances for testing""" + in_memory = InMemoryCache() + redis_cache = RedisCache( + host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT") + ) + dual_cache = DualCache( + in_memory_cache=in_memory, + redis_cache=redis_cache, + default_max_redis_batch_cache_size=DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE, + ) + return dual_cache, in_memory, redis_cache + + +@pytest.mark.asyncio +async def test_batch_cache_size_is_1000_minimum(cache_setup): + """Verify batch cache size is set to 1000 (never below 1k)""" + dual_cache, _, _ = cache_setup + + # Critical: batch cache size must be at least DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE + assert dual_cache.last_redis_batch_access_time.max_size >= DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE + + +@pytest.mark.asyncio +async def test_throttling_prevents_duplicate_redis_calls(cache_setup): + """Test throttling prevents repeated Redis queries for cache misses""" + dual_cache, _, redis_cache = cache_setup + + test_keys = [f"miss_{str(uuid.uuid4())}" for _ in range(3)] + + # Set short expiry for testing + dual_cache.redis_batch_cache_expiry = 0.1 # 100ms + + with patch.object( + redis_cache, "async_batch_get_cache", new_callable=AsyncMock + ) as mock_redis: + mock_redis.return_value = {key: None for key in test_keys} + + # First call hits Redis (no throttle data exists) + await dual_cache.async_batch_get_cache(test_keys) + assert mock_redis.call_count == 1 + + # Second call immediately - throttled (within expiry window) + await dual_cache.async_batch_get_cache(test_keys) + assert mock_redis.call_count == 1 + + # Verify all keys tracked in throttle cache + for key in test_keys: + assert key in dual_cache.last_redis_batch_access_time + + # Wait for expiry time to pass + time.sleep(0.15) + + # Third call after expiry - call_count increases to 2 + await dual_cache.async_batch_get_cache(test_keys) + assert mock_redis.call_count == 2 + + +@pytest.mark.asyncio +async def test_basic_functionality_not_broken(cache_setup): + """Ensure basic cache functionality still works after optimizations""" + dual_cache, _, _ = cache_setup + + # Test basic set/get works + test_key = f"functional_test_{str(uuid.uuid4())}" + test_value = {"test": "data"} + + await dual_cache.async_set_cache(test_key, test_value) + result = await dual_cache.async_get_cache(test_key) + + assert result == test_value + + +@pytest.mark.asyncio +async def test_batch_get_with_no_in_memory_cache(): + """Test that batch get works when in_memory_cache is None""" + redis_cache = RedisCache( + host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT") + ) + + # Create DualCache with no in-memory cache + dual_cache = DualCache( + in_memory_cache=None, # This is the edge case we're testing + redis_cache=redis_cache, + ) + + # Set some test data directly in Redis + test_key = f"no_memory_test_{str(uuid.uuid4())}" + test_value = {"test": "data_without_memory_cache"} + + await redis_cache.async_set_cache(test_key, test_value) + + # Should not crash when fetching from Redis without in-memory cache + result = await dual_cache.async_batch_get_cache([test_key]) + + assert result is not None + assert len(result) == 1 + assert result[0] == test_value + diff --git a/tests/router_unit_tests/test_completion_no_copy.py b/tests/router_unit_tests/test_completion_no_copy.py new file mode 100644 index 0000000000..3f5961a812 --- /dev/null +++ b/tests/router_unit_tests/test_completion_no_copy.py @@ -0,0 +1,112 @@ +""" +Regression test for removing unnecessary dict.copy() in completion hot paths. + +Verifies that spreading deployment["litellm_params"] directly (without copy) +doesn't cause side effects that mutate the deployment in router.model_list. +""" +import sys +import os +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm import Router +from unittest.mock import AsyncMock, Mock, patch + + +@pytest.mark.asyncio +async def test_acompletion_deployment_not_mutated(): + """ + Test async completion doesn't mutate deployment when .copy() is removed. + + Optimization: Remove deployment["litellm_params"].copy() in _acompletion + since data is only read and spread into input_kwargs dict. + """ + router = Router( + model_list=[ + { + "model_name": "gpt-3.5", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "temperature": 0.7, + }, + } + ] + ) + + deployment_before = router.get_deployment_by_model_group_name("gpt-3.5") + assert deployment_before is not None + original_params = deployment_before.litellm_params.model_dump() + + with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion: + from litellm import ModelResponse + + mock_acompletion.return_value = ModelResponse( + id="test", + choices=[{"message": {"role": "assistant", "content": "test"}, "index": 0}], + model="gpt-3.5-turbo", + usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + ) + + try: + await router.acompletion( + model="gpt-3.5", + messages=[{"role": "user", "content": "test"}], + ) + except Exception: + pass + + # Critical: Deployment params must be unchanged + deployment_after = router.get_deployment_by_model_group_name("gpt-3.5") + assert deployment_after is not None + assert deployment_after.litellm_params.model_dump() == original_params + + +def test_completion_deployment_not_mutated(): + """ + Test sync completion doesn't mutate deployment when .copy() is removed. + + Optimization: Remove deployment["litellm_params"].copy() in _completion + since data is only read and spread into input_kwargs dict. + """ + router = Router( + model_list=[ + { + "model_name": "gpt-3.5", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "max_tokens": 100, + }, + } + ] + ) + + deployment_before = router.get_deployment_by_model_group_name("gpt-3.5") + assert deployment_before is not None + original_params = deployment_before.litellm_params.model_dump() + + with patch("litellm.completion", new_callable=Mock) as mock_completion: + from litellm import ModelResponse + + mock_completion.return_value = ModelResponse( + id="test", + choices=[{"message": {"role": "assistant", "content": "test"}, "index": 0}], + model="gpt-3.5-turbo", + usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + ) + + try: + router.completion( + model="gpt-3.5", + messages=[{"role": "user", "content": "test"}], + ) + except Exception: + pass + + # Critical: Deployment params must be unchanged + deployment_after = router.get_deployment_by_model_group_name("gpt-3.5") + assert deployment_after is not None + assert deployment_after.litellm_params.model_dump() == original_params + diff --git a/tests/router_unit_tests/test_default_deployment_copy.py b/tests/router_unit_tests/test_default_deployment_copy.py new file mode 100644 index 0000000000..6eff4da445 --- /dev/null +++ b/tests/router_unit_tests/test_default_deployment_copy.py @@ -0,0 +1,85 @@ +""" +Regression test for default_deployment shallow copy optimization. + +Tests the critical side effect: ensure modifying returned deployment +doesn't corrupt the original default_deployment instance. +""" +import sys +import os + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm import Router + + +def test_default_deployment_isolation(): + """ + Regression test for shallow copy optimization in _common_checks_available_deployment. + + When a model is not in model_names and default_deployment is set, the router + returns a copy of default_deployment with the model name updated. This test + ensures the optimization (shallow copy instead of deepcopy) properly isolates + each returned deployment from the original and from each other. + + The shallow copy optimization copies two levels: + 1. Top-level deployment dict + 2. litellm_params dict + + Deeper nested objects are intentionally shared for performance (safe because + the router only modifies the 'model' field at litellm_params level). + + Critical behavior verified: + 1. Each deployment gets independent model value + 2. Original default_deployment unchanged for litellm_params fields + 3. Shared fields (api_key) accessible in all copies + 4. Adding new litellm_params fields is isolated per deployment + 5. Deep nested objects ARE shared (acceptable trade-off) + """ + # Setup: Router with a default deployment (used for unknown models) + router = Router(model_list=[]) + + router.default_deployment = { # type: ignore + "model_name": "default-model", + "litellm_params": { + "model": "gpt-3.5-turbo", # This will be overwritten per request + "api_key": "test-key", # This should be shared + "custom_config": { # Deep nested - will be SHARED + "nested_setting": "original", + }, + }, + } + + # Act: Request two different unknown models (triggers default deployment path) + _, deployment1 = router._common_checks_available_deployment( + model="custom-model-1", # Unknown model + messages=[{"role": "user", "content": "test"}], + ) + + _, deployment2 = router._common_checks_available_deployment( + model="custom-model-2", # Different unknown model + messages=[{"role": "user", "content": "test"}], + ) + + # Assert: Each deployment should have its own independent model value + assert deployment1["litellm_params"]["model"] == "custom-model-1" # type: ignore + assert deployment2["litellm_params"]["model"] == "custom-model-2" # type: ignore + + # Assert: Original default_deployment must remain unchanged (not mutated by requests) + assert router.default_deployment["litellm_params"]["model"] == "gpt-3.5-turbo" # type: ignore + + # Assert: Shared fields should still be accessible in all copies + assert deployment1["litellm_params"]["api_key"] == "test-key" # type: ignore + assert deployment2["litellm_params"]["api_key"] == "test-key" # type: ignore + + # Assert: Modifying litellm_params in one deployment doesn't affect others + # This tests the shallow copy properly isolated the litellm_params dict level + deployment1["litellm_params"]["temperature"] = 0.9 # type: ignore + assert "temperature" not in deployment2["litellm_params"] # type: ignore + assert "temperature" not in router.default_deployment["litellm_params"] # type: ignore + + # Assert: Deep nested objects ARE shared (intentional trade-off for 100x perf gain) + # Safe because router only modifies top-level litellm_params fields + deployment1["litellm_params"]["custom_config"]["nested_setting"] = "modified" # type: ignore + assert deployment2["litellm_params"]["custom_config"]["nested_setting"] == "modified" # type: ignore + assert router.default_deployment["litellm_params"]["custom_config"]["nested_setting"] == "modified" # type: ignore + diff --git a/tests/router_unit_tests/test_pre_call_checks_optimization.py b/tests/router_unit_tests/test_pre_call_checks_optimization.py new file mode 100644 index 0000000000..16af1cc53e --- /dev/null +++ b/tests/router_unit_tests/test_pre_call_checks_optimization.py @@ -0,0 +1,132 @@ +""" +Regression tests for Router._pre_call_checks() performance optimization. + +Background: + _pre_call_checks() runs on EVERY request to filter deployments based on + context window size, rate limits, region constraints, and supported parameters. + +Optimization: + Changed from copy.deepcopy(healthy_deployments) to list(healthy_deployments). + This is ~1400x faster while maintaining correctness because the function only + removes items from the list, never modifies the deployment objects themselves. + +Critical Requirement: + The input healthy_deployments list must NEVER be mutated. Callers depend on + this for retries, fallbacks, and logging. +""" + +import copy +import pytest +from litellm import Router + + +class TestPreCallChecksOptimization: + """ + Verify that using list() instead of deepcopy() doesn't break behavior. + + If these tests fail, the optimization should be reverted. + """ + + def test_no_mutation_of_input_list(self): + """ + Verify the input list is never modified by _pre_call_checks. + + The function uses list() instead of deepcopy for performance. + This is safe because it only filters items, never modifies them. + """ + router = Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}, + "model_info": {"id": "test-1"}, + }, + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-4", "api_key": "sk-test2"}, + "model_info": {"id": "test-2"}, + }, + ], + set_verbose=False, + enable_pre_call_checks=True, + ) + + deployments = router.get_model_list(model_name="gpt-3.5-turbo") + assert deployments is not None + + # Capture the original state + original_length = len(deployments) + original_deployment_ids = [id(d) for d in deployments] + original_litellm_params_ids = [id(d["litellm_params"]) for d in deployments] + snapshot = copy.deepcopy(deployments) + + # Call the function under test + router._pre_call_checks( + model="gpt-3.5-turbo", + healthy_deployments=deployments, + messages=[{"role": "user", "content": "test"}], + ) + + # Verify nothing changed: + # 1. Same number of items + assert len(deployments) == original_length, "List length changed!" + # 2. Same deployment objects (not replaced with copies) + assert [id(d) for d in deployments] == original_deployment_ids, "Deployment dicts replaced!" + # 3. Same nested objects (not replaced with copies) + assert [id(d["litellm_params"]) for d in deployments] == original_litellm_params_ids, "Nested dicts replaced!" + # 4. Same values (catches any mutation) + assert deployments == snapshot, "Values were mutated!" + + def test_filtering_still_works(self): + """ + Verify that filtering works correctly while preserving the original list. + + Scenario: Send a message too long for one deployment but fine for another. + Expected: Filtered result excludes the small deployment, but original list is unchanged. + """ + router = Router( + model_list=[ + { + "model_name": "test", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}, + "model_info": {"id": "small", "max_input_tokens": 50}, + }, + { + "model_name": "test", + "litellm_params": {"model": "gpt-4", "api_key": "sk-test"}, + "model_info": {"id": "large", "max_input_tokens": 10000}, + }, + ], + set_verbose=False, + enable_pre_call_checks=True, + ) + + deployments = router.get_model_list(model_name="test") + assert deployments is not None + + # Save references to the original deployment objects + original_small_deployment = deployments[0] # max_input_tokens=50 + original_large_deployment = deployments[1] # max_input_tokens=10000 + + # Send a long message (100 words) that exceeds 50 tokens but fits in 10000 tokens + filtered = router._pre_call_checks( + model="test", + healthy_deployments=deployments, + messages=[{"role": "user", "content": " ".join(["word"] * 100)}], + ) + + # Verify the filtered result only contains the large deployment + assert len(filtered) == 1, f"Expected 1 deployment after filtering, got {len(filtered)}" + assert filtered[0]["model_info"]["id"] == "large", "Wrong deployment kept after filtering" + + # Verify the original list still has both deployments + assert len(deployments) == 2, f"Original list was modified! Expected 2, got {len(deployments)}" + assert deployments[0] is original_small_deployment, "First deployment object replaced!" + assert deployments[1] is original_large_deployment, "Second deployment object replaced!" + assert deployments[0].get("model_info", {}).get("id") == "small", "First deployment ID changed!" + assert deployments[1].get("model_info", {}).get("id") == "large", "Second deployment ID changed!" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) + diff --git a/tests/router_unit_tests/test_prompt_management_check.py b/tests/router_unit_tests/test_prompt_management_check.py new file mode 100644 index 0000000000..a818ca2345 --- /dev/null +++ b/tests/router_unit_tests/test_prompt_management_check.py @@ -0,0 +1,67 @@ +""" +Test for _is_prompt_management_model early exit optimization. + +Verifies that the early return for models without "/" doesn't break +prompt management model detection. +""" +import sys +import os + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm import Router + + +def test_is_prompt_management_model_optimization(): + """ + Test early exit optimization works correctly for all cases. + + Optimization: Check if "/" in model name before calling expensive + get_model_list(). This short-circuits 99% of requests that use + standard model names like "gpt-4", "claude-3", etc. + + Tests both negative (early exit) and positive (actual detection) cases. + """ + import litellm + + # Test 1: Standard models without "/" -> early exit returns False + router = Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + }, + { + "model_name": "claude-3", + "litellm_params": {"model": "anthropic/claude-3-sonnet-20240229"}, + }, + ] + ) + + assert router._is_prompt_management_model("gpt-4") is False + assert router._is_prompt_management_model("claude-3") is False + + # Test 2: Models with "/" but not in model_list -> False after check + assert router._is_prompt_management_model("unknown/model") is False + + # Test 3: Actual prompt management models ARE detected (critical positive case) + original_callbacks = litellm._known_custom_logger_compatible_callbacks.copy() + if "langfuse_prompt" not in litellm._known_custom_logger_compatible_callbacks: + litellm._known_custom_logger_compatible_callbacks.append("langfuse_prompt") + + try: + router_with_prompt = Router( + model_list=[ + { + "model_name": "my-langfuse-prompt/test_id", + "litellm_params": {"model": "langfuse_prompt/actual_prompt_id"}, + }, + ] + ) + + # Critical: Must still detect prompt management models correctly + assert router_with_prompt._is_prompt_management_model("my-langfuse-prompt/test_id") is True + + finally: + litellm._known_custom_logger_compatible_callbacks = original_callbacks +