litellm/litellm/proxy/common_utils/performance_utils.py

"""
Performance utilities for LiteLLM proxy server.

This module provides performance monitoring and profiling functionality for endpoint
performance analysis using cProfile with configurable sampling rates, and line_profiler
for line-by-line profiling.

See performance_utils.md for detailed usage examples and documentation.
"""

import atexit
import cProfile
import functools
import inspect
import threading
from collections.abc import Callable
from pathlib import Path as PathLib
from typing import Any, Final

from litellm._logging import verbose_proxy_logger

# Global profiling state
_profile_lock: Final = threading.Lock()
_profiler = None
_last_profile_file_path = None
_sample_counter = 0
_sample_counter_lock: Final = threading.Lock()

# Global line_profiler state
_line_profiler: Any | None = None
_line_profiler_lock: Final = threading.Lock()
_wrapped_functions: Final[dict[str, Callable]] = {}  # Store original functions


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: Final = (_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("Profiling started with sampling rate: %s", 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("Profiling stats saved to %s", profile_file)
        except Exception as e:
            verbose_proxy_logger.error("Error saving profiling stats: %s", 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 inspect.iscoroutinefunction(func):

            @functools.wraps(func)
            async def async_wrapper(*args, **kwargs):
                is_sampling: Final = _start_profiling_for_request(sampling_rate)
                file_path_obj: Final = PathLib("endpoint_profile.pstat")
                set_last_profile_path(file_path_obj)
                try:
                    result: Final = 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: Final = _start_profiling_for_request(sampling_rate)
                file_path_obj: Final = PathLib("endpoint_profile.pstat")
                set_last_profile_path(file_path_obj)
                try:
                    result: Final = 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


def enable_line_profiler() -> None:
    """Enable line_profiler for dynamic function wrapping.

    Raises:
        ImportError: If line_profiler is not available
    """
    global _line_profiler
    from line_profiler import LineProfiler  # Will raise ImportError if not available

    with _line_profiler_lock:
        if _line_profiler is None:
            _line_profiler = LineProfiler()
            verbose_proxy_logger.info("Line profiler enabled")


def wrap_function_with_line_profiler(module: Any, function_name: str) -> bool:
    """Dynamically wrap a function with line_profiler.

    Args:
        module: The module containing the function
        function_name: Name of the function to wrap

    Returns:
        True if wrapping was successful, False otherwise
    """
    try:
        enable_line_profiler()  # May raise ImportError if not available
    except ImportError:
        return False

    if _line_profiler is None:
        return False

    try:
        original_function: Final = getattr(module, function_name, None)
        if original_function is None:
            verbose_proxy_logger.warning("Function %s not found in module %s", function_name, module.__name__)
            return False

        # Store original function if not already wrapped
        if function_name not in _wrapped_functions:
            _wrapped_functions[function_name] = original_function

        # Wrap with line_profiler
        profiled_function: Final = _line_profiler(original_function)
        setattr(module, function_name, profiled_function)

        verbose_proxy_logger.info("Wrapped %s.%s with line_profiler", module.__name__, function_name)
        return True
    except Exception as e:
        verbose_proxy_logger.error("Error wrapping %s with line_profiler: %s", function_name, e)
        return False


def wrap_function_directly(func: Callable) -> Callable:
    """Wrap a function directly with line_profiler.

    This is the recommended way to profile functions, especially closures or
    functions created dynamically (like wrapper_async in litellm/utils.py).

    Args:
        func: The function to wrap

    Returns:
        The wrapped function that will be profiled when called

    Raises:
        ImportError: If line_profiler is not available
        RuntimeError: If line_profiler cannot be enabled or function cannot be wrapped
    """
    import warnings

    enable_line_profiler()  # Will raise ImportError if not available

    if _line_profiler is None:
        raise RuntimeError("Line profiler was not initialized")

    # Suppress warnings about __wrapped__ - we intentionally want to profile the wrapper
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", message=".*__wrapped__.*", category=UserWarning)
        # Add function to line_profiler and wrap it
        _line_profiler.add_function(func)
        profiled_function: Final = _line_profiler(func)

    verbose_proxy_logger.info("Wrapped function %s with line_profiler", func.__name__)
    return profiled_function


def collect_line_profiler_stats(output_file: str | None = None) -> None:
    """Collect and save line_profiler statistics.

    This can be called manually to collect stats at any time, or it's
    automatically called on shutdown if register_shutdown_handler() was used.

    Args:
        output_file: Optional path to save stats. If None, prints to stdout.
    """
    global _line_profiler

    with _line_profiler_lock:
        if _line_profiler is None:
            verbose_proxy_logger.debug("Line profiler not enabled, nothing to collect")
            return

        try:
            if output_file:
                # Save to file
                output_path: Final = PathLib(output_file)
                _line_profiler.dump_stats(str(output_path))
                verbose_proxy_logger.info("Line profiler stats saved to %s", output_path)
            else:
                # Print to stdout
                from io import StringIO

                stream: Final = StringIO()
                _line_profiler.print_stats(stream=stream)
                stats_output: Final = stream.getvalue()
                verbose_proxy_logger.info("Line profiler stats:\n" + stats_output)
        except Exception as e:
            verbose_proxy_logger.error("Error collecting line profiler stats: %s", e)


def register_shutdown_handler(output_file: str | None = None) -> None:
    """Register a shutdown handler to collect line_profiler stats.

    This registers an atexit handler that will automatically save profiling
    statistics when the Python process exits. Safe to call multiple times
    (only registers once).

    Args:
        output_file: Optional path to save stats on shutdown.
                     Defaults to 'line_profile_stats.lprof'
    """
    if output_file is None:
        output_file = "line_profile_stats.lprof"

    def shutdown_handler():
        collect_line_profiler_stats(output_file=output_file)

    atexit.register(shutdown_handler)
    verbose_proxy_logger.debug("Registered line_profiler shutdown handler for %s", output_file)