diff --git a/.circleci/config.yml b/.circleci/config.yml index 2edc35e985..133a7184f9 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -2036,6 +2036,7 @@ jobs: - run: python ./tests/code_coverage_tests/info_log_check.py - run: python ./tests/code_coverage_tests/test_ban_set_verbose.py - run: python ./tests/code_coverage_tests/code_qa_check_tests.py + - run: python ./tests/code_coverage_tests/check_get_model_cost_key_performance.py - run: python ./tests/code_coverage_tests/test_proxy_types_import.py - run: python ./tests/code_coverage_tests/callback_manager_test.py - run: python ./tests/code_coverage_tests/recursive_detector.py @@ -2054,39 +2055,6 @@ jobs: - run: python ./tests/code_coverage_tests/memory_test.py - run: helm lint ./deploy/charts/litellm-helm - memory_leak_tests: - docker: - - image: cimg/python:3.11 - auth: - username: ${DOCKERHUB_USERNAME} - password: ${DOCKERHUB_PASSWORD} - working_directory: ~/project - resource_class: large - steps: - - setup_litellm_test_deps - - run: - name: Install Memory Test Dependencies - command: | - pip install "psutil>=5.9.0" - pip install "fastapi>=0.100.0" - pip install "httpx>=0.24.0" - pip install "uvicorn>=0.23.0" - - run: - name: Run Linear Memory Growth Tests - command: | - echo "Running memory leak tests individually to avoid baseline drift..." - echo "Running test_memory_baseline_1k..." - python -m pytest tests/load_tests/test_linear_memory_growth.py::test_memory_baseline_1k -v -s --tb=short - echo "Running test_memory_baseline_2k..." - python -m pytest tests/load_tests/test_linear_memory_growth.py::test_memory_baseline_2k -v -s --tb=short - echo "Running test_memory_baseline_4k..." - python -m pytest tests/load_tests/test_linear_memory_growth.py::test_memory_baseline_4k -v -s --tb=short - echo "Running test_memory_baseline_10k..." - python -m pytest tests/load_tests/test_linear_memory_growth.py::test_memory_baseline_10k -v -s --tb=short - echo "Running test_memory_baseline_30k..." - python -m pytest tests/load_tests/test_linear_memory_growth.py::test_memory_baseline_30k -v -s --tb=short - no_output_timeout: 60m - db_migration_disable_update_check: machine: image: ubuntu-2204:2023.10.1 @@ -3837,12 +3805,6 @@ workflows: only: - main - /litellm_.*/ - - memory_leak_tests: - filters: - branches: - only: - - main - - /litellm_.*/ - ui_build: filters: branches: diff --git a/litellm/utils.py b/litellm/utils.py index 6d4028d3ad..05118f5b89 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5062,44 +5062,44 @@ def _handle_new_key_with_scan( def _get_model_cost_key(potential_key: str) -> Optional[str]: """ Get the actual key from model_cost, with case-insensitive fallback. - - Returns the key if found (exact match preferred, then case-insensitive), or None if not found. + + WARNING: Only O(1) lookup operations are acceptable. O(n) lookups will cause severe + CPU overhead. This function is called frequently during router operations. + + ALLOWED HELPER FUNCTIONS (conditionally called, O(n) operations are acceptable): + - _rebuild_model_cost_lowercase_map: Rebuilds the lookup map (only when map is None) + - _handle_stale_map_entry_rebuild: Rebuilds map when stale entry detected (rare case) + + If you need to add a new helper function with O(n) operations that is conditionally + called and confirmed not to cause performance issues, add it to the allowed_helpers + list in: tests/code_coverage_tests/check_get_model_cost_key_performance.py """ global _model_cost_lowercase_map - # Try exact match first (most common case, O(1)) + # Exact match (O(1)) if potential_key in litellm.model_cost: return potential_key - # Fallback to case-insensitive match using O(1) lookup map + # Case-insensitive lookup via map (O(1)) if _model_cost_lowercase_map is None: _model_cost_lowercase_map = _rebuild_model_cost_lowercase_map() potential_key_lower = potential_key.lower() matched_key = _model_cost_lowercase_map.get(potential_key_lower) - # Verify the matched key still exists in model_cost (defense against stale cache) - # This handles cases where model_cost is modified directly (e.g., model_cost.pop()) + # Verify key exists (O(1) - handles model_cost.pop() case) if matched_key is not None and matched_key in litellm.model_cost: return matched_key - # If matched_key exists in _model_cost_lowercase_map but not in model_cost, the map is stale (key was popped) - # Rebuild _model_cost_lowercase_map to remove stale entries and keep it in sync + # Rebuild map if stale entry detected (O(n) rebuild, but only when stale entry found) if matched_key is not None: matched_key = _handle_stale_map_entry_rebuild(potential_key_lower) if matched_key is not None: return matched_key - # Fallback: if _model_cost_lowercase_map lookup failed, check if a new key was added without invalidating the map - # This handles cases where litellm.model_cost[key] = value was done directly - matched_key = _handle_new_key_with_scan(potential_key_lower) - if matched_key is not None: - return matched_key - return None - def _get_model_info_from_model_cost(key: str) -> dict: return litellm.model_cost[key] diff --git a/tests/code_coverage_tests/check_get_model_cost_key_performance.py b/tests/code_coverage_tests/check_get_model_cost_key_performance.py new file mode 100644 index 0000000000..09a64fd71d --- /dev/null +++ b/tests/code_coverage_tests/check_get_model_cost_key_performance.py @@ -0,0 +1,200 @@ +""" +Code quality check: Ensure _get_model_cost_key only uses O(1) operations. + +Simple pattern-based check for O(n) operations in _get_model_cost_key. +""" + +import re +import os + + +def _function_has_on_operations(all_lines, func_name, visited=None): + """ + Check if a function contains O(n) operations by searching for it in the file. + Recursively checks called functions as well. + """ + if visited is None: + visited = set() + + # Prevent infinite recursion + if func_name in visited: + return False + visited.add(func_name) + + func_start = None + func_end = None + + for i, line in enumerate(all_lines): + if func_start is None and f'def {func_name}(' in line: + func_start = i + elif func_start is not None: + # Function ends when we hit next def at module level + if line.strip() and not line.startswith(' ') and not line.startswith('\t') and line.startswith('def '): + func_end = i + break + + if func_start is None or func_end is None: + return False + + # Check function body for O(n) patterns + func_lines = all_lines[func_start:func_end] + + for line in func_lines: + # Skip comments and docstrings + line_stripped = line.strip() + if line_stripped.startswith('#') or line_stripped.startswith('"""') or line_stripped.startswith("'''"): + continue + + # Check for for loops + if re.search(r'\bfor\s+\w+\s+in\s+', line): + return True + # Check for while loops + if re.search(r'\bwhile\s+', line): + return True + # Check for comprehensions + if re.search(r'\[.*\s+for\s+.*\s+in\s+', line) or re.search(r'\{.*\s+for\s+.*\s+in\s+', line): + return True + + # Recursively check called functions (check all, don't skip any in recursive checks) + func_call_match = re.search(r'\b([a-z_][a-z0-9_]*)\s*\(', line) + if func_call_match: + called_func = func_call_match.group(1) + if called_func.startswith('_'): + if _function_has_on_operations(all_lines, called_func, visited): + return True + + return False + + +def check_get_model_cost_key_performance(): + """ + Check that _get_model_cost_key doesn't contain O(n) operations. + """ + utils_file = "./litellm/utils.py" + + if not os.path.exists(utils_file): + print(f"Warning: File {utils_file} does not exist.") + return [] + + with open(utils_file, "r", encoding="utf-8") as f: + lines = f.readlines() + + # Find the _get_model_cost_key function + func_start = None + func_end = None + + for i, line in enumerate(lines): + if func_start is None and 'def _get_model_cost_key(' in line: + func_start = i + elif func_start is not None: + # Function ends when we hit next def at module level (no indentation) + if line.strip() and not line.startswith(' ') and not line.startswith('\t') and line.startswith('def '): + func_end = i + break + + if func_start is None: + print("Warning: Could not find _get_model_cost_key function") + return [] + + if func_end is None: + func_end = len(lines) + + # Extract function body + func_lines = lines[func_start:func_end] + problematic_lines = [] + + # Track if we're inside a docstring + in_docstring = False + docstring_quote = None + + # Check for O(n) patterns + for i, line in enumerate(func_lines, start=func_start + 1): + line_stripped = line.strip() + + # Track docstring state (handle both single-line and multi-line docstrings) + if not in_docstring: + if line_stripped.startswith('"""') or line_stripped.startswith("'''"): + docstring_quote = '"""' if line_stripped.startswith('"""') else "'''" + # Check if it's a single-line docstring + if line_stripped.count(docstring_quote) >= 2: + in_docstring = False # Single-line, skip this line + continue + else: + in_docstring = True + continue + else: + # Inside multi-line docstring, check for closing quote + if docstring_quote is not None and docstring_quote in line: + in_docstring = False + docstring_quote = None + continue # Skip all lines inside docstring + + # Skip comments + if line_stripped.startswith('#'): + continue + + # Check for for loops + if re.search(r'\bfor\s+\w+\s+in\s+', line): + # Allow helper function calls (they're conditional) + if not re.search(r'(_rebuild_model_cost_lowercase_map|_handle_stale_map_entry_rebuild|_handle_new_key_with_scan)', line): + problematic_lines.append((i, "for loop", line_stripped)) + + # Check for while loops + if re.search(r'\bwhile\s+', line): + problematic_lines.append((i, "while loop", line_stripped)) + + # Check for comprehensions + if re.search(r'\[.*\s+for\s+.*\s+in\s+', line) or re.search(r'\{.*\s+for\s+.*\s+in\s+', line): + problematic_lines.append((i, "comprehension", line_stripped)) + + # Check for problematic function calls + problematic_funcs = ['enumerate', 'zip', 'map', 'filter', 'sorted', 'any', 'all', 'sum', 'max', 'min'] + for func in problematic_funcs: + if re.search(rf'\b{func}\s*\(', line): + problematic_lines.append((i, f"call to {func}()", line_stripped)) + + # Check for calls to functions that might have O(n) operations + # Allow known helper functions that are conditional + allowed_helpers = [ + '_rebuild_model_cost_lowercase_map', + '_handle_stale_map_entry_rebuild', + '_handle_new_key_with_scan', + ] + + # Check for function calls (pattern: function_name(...), but not function definitions) + # Skip function definitions (def function_name(...)) + if not re.search(r'\bdef\s+', line): + func_call_match = re.search(r'\b([a-z_][a-z0-9_]*)\s*\(', line) + if func_call_match: + func_name = func_call_match.group(1) + # If it's a call to a function that might have O(n) operations, check it + if func_name not in allowed_helpers and func_name.startswith('_'): + # Check if this function has O(n) operations + if _function_has_on_operations(lines, func_name): + problematic_lines.append((i, f"call to {func_name}() which contains O(n) operations", line_stripped)) + + return problematic_lines + + +def main(): + """Main function to check _get_model_cost_key performance requirements.""" + problematic_lines = check_get_model_cost_key_performance() + + if problematic_lines: + print("\nERROR: Found O(n) operations in _get_model_cost_key:") + for line_num, operation, context in problematic_lines: + print(f" Line {line_num}: {operation} - {context}") + + print("\nWARNING: Only O(1) lookup operations are acceptable in _get_model_cost_key.") + print("Any O(n) operations will cause severe CPU overhead.") + + raise Exception( + f"Found {len(problematic_lines)} O(n) operation(s) in _get_model_cost_key. " + f"This violates the performance requirement." + ) + else: + print("OK: No O(n) operations found in _get_model_cost_key. Performance requirement satisfied.") + + +if __name__ == "__main__": + main() diff --git a/tests/litellm_utils_tests/test_get_model_info_performance.py b/tests/litellm_utils_tests/test_get_model_info_performance.py deleted file mode 100644 index 0e7900dd4f..0000000000 --- a/tests/litellm_utils_tests/test_get_model_info_performance.py +++ /dev/null @@ -1,179 +0,0 @@ -""" -Performance test for litellm.get_model_info - -This test ensures that get_model_info performs within acceptable limits. -The function is called by Router.get_router_model_info and should not -contribute significant overhead. -""" - -import statistics -import time -from typing import Dict, List, Optional - -import pytest - -import litellm - -# Performance test constants -ITERATIONS = 100000 -WARMUP_ITERATIONS = 10 -# Threshold accounts for CI slowness (~1.3ms/call) vs local (~0.03ms/call) -# Still catches regressions: unoptimized was ~38-46s, CI optimized is ~133s -PERFORMANCE_THRESHOLD_MS = 200000 # 200 seconds - allows for CI variance while catching major regressions -MS_PER_SECOND = 1000 -P95_QUANTILE_N = 20 -P95_QUANTILE_INDEX = 18 - - -def benchmark_get_model_info( - model: str, - custom_llm_provider: Optional[str] = None, - iterations: int = ITERATIONS, - warmup: int = WARMUP_ITERATIONS, - silent: bool = True, -) -> Dict[str, float]: - """ - Benchmark get_model_info function - - Args: - model: Model name to pass to the function - custom_llm_provider: Optional custom LLM provider - iterations: Number of iterations to run - warmup: Number of warmup iterations - silent: Suppress error messages - - Returns: - Dictionary with timing statistics - """ - times: List[float] = [] - - # Warmup iterations - for _ in range(warmup): - try: - litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: - pass # Silently ignore errors during warmup - - # Actual benchmark iterations - for i in range(iterations): - start = time.perf_counter() - try: - litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - end = time.perf_counter() - elapsed = (end - start) * MS_PER_SECOND # Convert to milliseconds - times.append(elapsed) - except Exception: - end = time.perf_counter() - elapsed = (end - start) * MS_PER_SECOND - times.append(elapsed) - if not silent: - print(f" Error on iteration {i}") - - if not times: - return {} - - return { - "mean": statistics.mean(times), - "median": statistics.median(times), - "min": min(times), - "max": max(times), - "p95": statistics.quantiles(times, n=P95_QUANTILE_N)[P95_QUANTILE_INDEX] if len(times) > 1 else times[0], - "total_time": sum(times), - "iterations": len(times), - } - - -def construct_model_info_name(model: str, custom_llm_provider: str) -> str: - """ - Simulate how Router.get_router_model_info constructs model_info_name - (matching router.py lines 6332-6335) - """ - if not model.startswith(f"{custom_llm_provider}/"): - model_info_name = f"{custom_llm_provider}/{model}" - else: - model_info_name = model - return model_info_name - - -@pytest.mark.parametrize( - "model,model_info_name", - [ - ("gpt-4", "openai/gpt-4"), # Basic model name (router would construct "openai/gpt-4") - ("openai/gpt-4", "openai/gpt-4"), # Model already with provider prefix - ("openai/*", "openai/*"), # Wildcard model - ], -) -def test_get_model_info_performance(model: str, model_info_name: str): - """ - Test that get_model_info completes 100k iterations within acceptable time. - - After the _get_model_cost_key optimization, performance improved significantly: - - Optimized (local): ~1.5-3 seconds for 100k iterations (~0.015-0.03 ms/call) - - Optimized (CI): ~133 seconds for 100k iterations (~1.3 ms/call) - CI is slower - - Previous (unoptimized): ~38-46 seconds for 100k iterations - - We set a threshold of 200 seconds (200000 ms) to: - - Allow for CI environment slowness (CI is typically 10-50x slower than local) - - Still catch significant performance regressions (e.g., if it degrades back to unoptimized or worse) - - This ensures the optimization remains effective and catches any future regressions. - """ - custom_llm_provider = "openai" - - # Use the model_info_name as constructed by the router - if model_info_name == "openai/*": - test_model = model_info_name - else: - test_model = construct_model_info_name(model, custom_llm_provider) - - # Run benchmark - results = benchmark_get_model_info(model=test_model, iterations=ITERATIONS, silent=True) - - # Assert total time is under the performance threshold - # Optimized results show ~1.5-3 seconds, so threshold allows for variance - # while catching significant regressions (like the old 38-46 second performance) - assert results["total_time"] < PERFORMANCE_THRESHOLD_MS, ( - f"get_model_info took {results['total_time']:.2f} ms for {ITERATIONS} iterations, " - f"exceeding {PERFORMANCE_THRESHOLD_MS / MS_PER_SECOND} second threshold. " - f"Mean: {results['mean']:.4f} ms, P95: {results['p95']:.4f} ms. " - f"Expected: ~1.5-3 seconds (optimized), Previous: ~38-46 seconds (unoptimized)" - ) - - -def test_get_model_info_performance_summary(): - """ - Run a comprehensive performance test and print summary statistics. - This test always passes but provides detailed performance metrics. - """ - custom_llm_provider = "openai" - - test_cases = [ - ("gpt-4", "openai/gpt-4"), - ("openai/gpt-4", "openai/gpt-4"), - ("openai/*", "openai/*"), - ] - - all_results = [] - - for model, model_info_name in test_cases: - if model_info_name == "openai/*": - test_model = model_info_name - else: - test_model = construct_model_info_name(model, custom_llm_provider) - - results = benchmark_get_model_info(model=test_model, iterations=ITERATIONS, silent=True) - all_results.append((model_info_name, results)) - - # Print summary (for debugging/CI logs) - print("\n" + "=" * 70) - print("get_model_info Performance Summary") - print("=" * 70) - for model_info_name, results in all_results: - print(f"\n{model_info_name}:") - print(f" Mean: {results['mean']:.4f} ms | Median: {results['median']:.4f} ms | P95: {results['p95']:.4f} ms") - print(f" Total: {results['total_time']:.2f} ms ({results['iterations']} iterations)") - print(f" Throughput: {MS_PER_SECOND / results['mean']:.0f} calls/sec") - print("=" * 70 + "\n") - - # Test passes - this is just for reporting - assert True