diff --git a/litellm/constants.py b/litellm/constants.py index 32b47bc6b46..aa132a1c51e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1749,6 +1749,7 @@ INVALID_VIRTUAL_KEY_ERROR_MARKER: Final = "_litellm_invalid_virtual_key_error" # Python garbage collection threshold configuration # Format: "gen0,gen1,gen2" e.g., "1000,50,50" PYTHON_GC_THRESHOLD: Final = os.getenv("PYTHON_GC_THRESHOLD") +LITELLM_GC_FREEZE_AFTER_STARTUP: Final = os.getenv("LITELLM_GC_FREEZE_AFTER_STARTUP", "true").lower() == "true" # pass through route constansts BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES: Final = [ diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 884aacb9fe1..8b1c1a53af1 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -16,7 +16,7 @@ from typing_extensions import ReadOnly from litellm import get_secret_str from litellm._logging import verbose_proxy_logger -from litellm.constants import PYTHON_GC_THRESHOLD +from litellm.constants import LITELLM_GC_FREEZE_AFTER_STARTUP, PYTHON_GC_THRESHOLD from litellm.litellm_core_utils.bug_report import EnvironmentReport from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import UserAPIKeyAuth @@ -63,6 +63,21 @@ def configure_gc_thresholds(): configure_gc_thresholds() +def freeze_gc_after_startup() -> None: + """Freeze the objects that exist once startup is done (LITELLM_GC_FREEZE_AFTER_STARTUP). + + Full collections are stop-the-world and walk every tracked object; the startup heap + (modules, typing/Pydantic metadata, cost map, config) never becomes garbage, so moving + it to the permanent generation shortens those pauses. Startup garbage is collected first + so it is not frozen with the rest. + """ + if not LITELLM_GC_FREEZE_AFTER_STARTUP: + return + gc.collect() + gc.freeze() + verbose_proxy_logger.info("GC: froze %s objects after startup", gc.get_freeze_count()) + + @router.get( "/debug/asyncio-tasks", dependencies=[Depends(user_api_key_auth)], diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 698c1e6e752..eb294d47afe 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1,7 +1,6 @@ import asyncio import copy import enum -import gc import importlib import inspect import io @@ -518,7 +517,7 @@ from litellm.proxy.common_utils.config_includes import resolve_include_file_path from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.credential_hydration import decrypted_or_stored from litellm.proxy.common_utils.custom_openapi_spec import JsonObject -from litellm.proxy.common_utils.debug_utils import init_verbose_loggers +from litellm.proxy.common_utils.debug_utils import freeze_gc_after_startup, init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router from litellm.proxy.common_utils.discoverable_model_filter import ( discoverable_rows, @@ -1971,8 +1970,7 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener async with AsyncExitStack() as admin_mcp_stack: try: await admin_mcp_stack.enter_async_context(admin_mcp_lifespan(app)) - gc.freeze() - verbose_proxy_logger.warning("gc.freeze pinned %d startup objects", gc.get_freeze_count()) + freeze_gc_after_startup() yield state finally: if model_info_scheduler is not None and model_info_scheduler.running: diff --git a/tests/unit/proxy/common_utils/test_debug_utils.py b/tests/unit/proxy/common_utils/test_debug_utils.py index 456be4e8c00..d635091a6ff 100644 --- a/tests/unit/proxy/common_utils/test_debug_utils.py +++ b/tests/unit/proxy/common_utils/test_debug_utils.py @@ -1,3 +1,4 @@ +import gc import json import os import socket @@ -13,6 +14,7 @@ from litellm.proxy import proxy_server from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.bug_report_config import build_proxy_bug_report +from litellm.proxy.common_utils import debug_utils from litellm.proxy.common_utils.debug_utils import ( PSUTIL_MISSING_ERROR, _ProcFilesystemProcess, @@ -198,3 +200,28 @@ def test_memory_routes_count_prompt_caching_pins( assert response.status_code == 200 assert pin_count(response.json()) == expected + + +@pytest.fixture +def unfrozen_gc() -> Iterator[None]: + gc.unfreeze() + yield + gc.unfreeze() + + +def test_freeze_gc_after_startup_freezes_existing_objects(unfrozen_gc: None, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(debug_utils, "LITELLM_GC_FREEZE_AFTER_STARTUP", True) + startup_object = {"loaded": "at startup"} + + debug_utils.freeze_gc_after_startup() + + assert gc.get_freeze_count() > 0 + assert not any(obj is startup_object for obj in gc.get_objects()) + + +def test_freeze_gc_after_startup_does_nothing_when_disabled(unfrozen_gc: None, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(debug_utils, "LITELLM_GC_FREEZE_AFTER_STARTUP", False) + + debug_utils.freeze_gc_after_startup() + + assert gc.get_freeze_count() == 0 diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index db5be07fb6f..fd5b72cbe64 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -18,6 +18,7 @@ Pins covered: from __future__ import annotations import asyncio +import gc import inspect import json import logging @@ -34,6 +35,7 @@ from pydantic import BaseModel from typing_extensions import TypedDict import litellm.proxy.proxy_server as ps +from litellm.proxy.common_utils import debug_utils from litellm.proxy.management_helpers.audit_logs import track_audit_task from litellm.proxy.proxy_server import ( ProxyStartupEvent, @@ -1022,6 +1024,32 @@ async def test_proxy_startup_event_honors_disable_model_info_refresh( ) +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", [True, False]) +async def test_proxy_startup_event_freezes_gc_before_serving(enabled: bool) -> None: + """With LITELLM_GC_FREEZE_AFTER_STARTUP (#46517) the startup heap is already frozen + when the proxy starts serving.""" + scheduler = AsyncIOScheduler() + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} | { + "LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY": "true" + } + gc.unfreeze() + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.object(ps, "scheduler", scheduler), + patch.object(debug_utils, "LITELLM_GC_FREEZE_AFTER_STARTUP", enabled), + ): + try: + async with proxy_startup_event(app=None): + frozen_while_serving = gc.get_freeze_count() + finally: + gc.unfreeze() + if scheduler.running: + scheduler.shutdown(wait=False) + + assert (frozen_while_serving > 0) is enabled + + @pytest.mark.asyncio async def test_proxy_shutdown_commits_spend_before_waiting_on_a_slow_logging_task_and_again_after_it() -> None: """A worker killed while shutdown waits on a slow logging task keeps the spend it had already counted,