diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5ce6cd93a22..639b02d1199 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -124,6 +124,7 @@ from .types_utils.utils import get_instance_fn, validate_custom_validate_return_ if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from litellm.repositories.database import Database from litellm.tracing import TraceReceiver Span = _Span | Any @@ -133,6 +134,7 @@ else: class ProxyLifespanState(TypedDict): tracing_receiver: ReadOnly["TraceReceiver | None"] + database: ReadOnly["Database | None"] class ReconcileOutcome(NamedTuple): diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 52a0be3d49e..6ac274583b0 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -233,8 +233,10 @@ class RoutingPrismaWrapper: await self._reader.start_token_refresh_task() async def stop_token_refresh_task(self) -> None: - await self._writer.stop_token_refresh_task() - await self._reader.stop_token_refresh_task() + try: + await self._writer.stop_token_refresh_task() + finally: + await self._reader.stop_token_refresh_task() async def recreate_prisma_client( self, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d6ec62b25c2..620f16099fb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -168,6 +168,7 @@ from litellm.proxy.common_utils.token_count_settings import ( ) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields +from litellm.repositories import database as database_runtime from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, STRATEGY_ROUTER_PARAM_FIELDS, @@ -1300,13 +1301,7 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N if worker_heartbeat is not None and prisma_client: await worker_heartbeat.deregister() if prisma_client: - # Drain the SGR fold first: it lives in memory, so an un-drained interval - # is lost, and a write attempted after disconnect raises - # ClientNotConnectedError rather than persisting anything. Ordering this - # inside the same guard is what keeps the two from drifting apart. await asyncio.gather(_drain_audit_tasks_on_shutdown(), _flush_request_accumulators(prisma_client)) - verbose_proxy_logger.debug("Disconnecting from Prisma") - await prisma_client.disconnect() else: await _drain_audit_tasks_on_shutdown() @@ -1429,16 +1424,37 @@ async def initialize_shared_aiohttp_session() -> "ClientSession | None": _initialize_shared_aiohttp_session: Final = initialize_shared_aiohttp_session -async def _connect_to_count_stored_values() -> SupportsRawQueries: +async def _connect_to_count_stored_values(resources: AsyncExitStack) -> SupportsRawQueries: client: Final = prisma_client or PrismaClient( database_url=str(get_secret("DATABASE_URL")), proxy_logging_obj=proxy_logging_obj ) + if prisma_client is None: + _own_database(resources, client) await client.connect() return client.writer_db +def _own_database(resources: AsyncExitStack, client: PrismaClient | None) -> database_runtime.Database | None: + if client is None: + return None + database: Final = database_runtime.Database(prisma=client) + resources.push_async_callback(database.aclose) + return database + + +def _forget_prisma_client() -> None: + global prisma_client # noqa: PLW0603 # routes still read this global until they move to DatabaseDependency + prisma_client = None # rebind-ok: a closed client must not survive into the next lifespan entry + + @asynccontextmanager async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]: + async with AsyncExitStack() as resources, _proxy_lifespan(app, resources) as state: + yield state + + +@asynccontextmanager +async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGenerator[ProxyLifespanState, None]: global \ prisma_client, \ master_key, \ @@ -1457,6 +1473,9 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState shared_aiohttp_session import json + resources.callback(_forget_prisma_client) + existing_database: Final = _own_database(resources, prisma_client) + init_verbose_loggers() running_loop: Final = asyncio.get_running_loop() @@ -1535,21 +1554,24 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState if decoded_worker_config is not None: await initialize_from_worker_config(decoded_worker_config) - enforce_master_key_boot_verdict( - await with_stored_secrets_counted( - master_key_boot_verdict( - master_key=master_key, - environment_master_key=os.getenv(MASTER_KEY_ENV_VAR), - general_settings=general_settings, - config_file_path=user_config_file_path, - override_env_is_on=get_secret_bool(WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR) is True, - salt_key_is_set=os.getenv(SALT_KEY_ENV_VAR) is not None, - database_is_configured=prisma_client is not None or get_secret("DATABASE_URL", None) is not None, + async with AsyncExitStack() as boot_database: + enforce_master_key_boot_verdict( + await with_stored_secrets_counted( + master_key_boot_verdict( + master_key=master_key, + environment_master_key=os.getenv(MASTER_KEY_ENV_VAR), + general_settings=general_settings, + config_file_path=user_config_file_path, + override_env_is_on=get_secret_bool(WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR) is True, + salt_key_is_set=os.getenv(SALT_KEY_ENV_VAR) is not None, + database_is_configured=prisma_client is not None or get_secret("DATABASE_URL", None) is not None, + ), + count_values_encrypted_with=partial( + count_values_encrypted_with_or_none, partial(_connect_to_count_stored_values, boot_database) + ), ), - count_values_encrypted_with=partial(count_values_encrypted_with_or_none, _connect_to_count_stored_values), - ), - announce=announce_on_stderr_at_exit, - ) + announce=announce_on_stderr_at_exit, + ) # check if DATABASE_URL in environment - load from there if prisma_client is None: @@ -1562,6 +1584,8 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState ) ) + database: Final = existing_database or _own_database(resources, prisma_client) + await migrate_if_requested( environ=os.environ, master_key=master_key, @@ -1571,15 +1595,16 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState ) if prisma_client is not None: + startup_prisma: Final = prisma_client - async def _run_pw_migration(): + async def _run_pw_migration() -> None: try: - result: Final = await migrate_passwords_to_scrypt_async(prisma_client) + result: Final = await migrate_passwords_to_scrypt_async(startup_prisma) verbose_proxy_logger.info("Password migration: %s", result) except Exception as e: verbose_proxy_logger.warning("Password migration skipped: %s", e) - asyncio.create_task(_run_pw_migration()) + await resources.enter_async_context(database_runtime.background_task(_run_pw_migration())) async def _run_agent_grant_id_migration() -> None: from litellm.proxy.agent_endpoints.agent_registry import ( @@ -1590,7 +1615,7 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState for attempt in range(3): try: result = await global_agent_registry.migrate_legacy_grant_ids( - table=object_permission_table(prisma_client) + table=object_permission_table(startup_prisma) ) if result.rewritten: verbose_proxy_logger.info( @@ -1610,7 +1635,7 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState if attempt < 2: await asyncio.sleep(5) - asyncio.create_task(_run_agent_grant_id_migration()) + await resources.enter_async_context(database_runtime.background_task(_run_agent_grant_id_migration())) ## A coordination_redis block saved from the admin UI lives in the database, ## which is only reachable once the prisma client exists. Apply it here, before @@ -1796,7 +1821,7 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState enabled=tracing_enabled, settings=tracing_settings, ) as receiver: - state: Final[ProxyLifespanState] = {"tracing_receiver": receiver} + state: Final[ProxyLifespanState] = {"tracing_receiver": receiver, "database": database} from litellm.proxy.admin_mcp import admin_mcp_lifespan telemetry_settings: Final = load_telemetry_settings() @@ -1850,30 +1875,6 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState except Exception as e: verbose_proxy_logger.error("Error closing shared aiohttp session: %s", e) - # Shutdown event - stop RDS IAM token refresh background task - if ( - prisma_client is not None - and hasattr(prisma_client, "db") - and hasattr(prisma_client.db, "stop_token_refresh_task") - ): - try: - await prisma_client.db.stop_token_refresh_task() - except Exception as e: - verbose_proxy_logger.error("Error stopping token refresh task: %s", e) - - # Shutdown event - stop Prisma DB health watchdog task - if prisma_client is not None and hasattr(prisma_client, "stop_db_health_watchdog_task"): - try: - await prisma_client.stop_db_health_watchdog_task() - except Exception as e: - verbose_proxy_logger.error("Error stopping DB health watchdog task: %s", e) - - if prisma_client is not None and hasattr(prisma_client, "stop_view_setup_task"): - try: - await prisma_client.stop_view_setup_task() - except Exception as e: - verbose_proxy_logger.error("Error stopping the spend view setup task: %s", e) - await _drain_spend_event_producer_on_shutdown() if telemetry_start is not None: telemetry_start.cancel() @@ -12050,17 +12051,20 @@ class ProxyStartupEvent: proxy_logging_obj: ProxyLogging, user_api_key_cache: UserApiKeyCache, ) -> PrismaClient | None: - """ - - Sets up prisma client - - Adds necessary views to proxy - """ + if database_url is None: + return None + client: Final = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) + async with AsyncExitStack() as setup_resources: + setup_resources.push_async_callback(database_runtime.Database(prisma=client).aclose) + connected: Final = await cls._initialize_prisma_client(client) + if connected is not None: + setup_resources.pop_all() + return connected + + @staticmethod + async def _initialize_prisma_client(prisma_client: PrismaClient) -> PrismaClient | None: connected_client: PrismaClient | None = None try: - if database_url is None: - return None - - prisma_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) - try: await prisma_client.connect() except Exception as e: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0ddfc5ce62f..ad501559786 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5365,6 +5365,8 @@ class PrismaClient: else: self.db = writer_wrapper # Client to connect to Prisma db self._db_reconnect_lock = asyncio.Lock() + self._closing: bool = False + self._engine_tasks: frozenset[asyncio.Task[object]] = frozenset() self._db_health_watchdog_task: asyncio.Task | None = None self._view_setup_task: asyncio.Task[_ViewSetupOutcome] | None = None self._db_last_reconnect_attempt_ts: float = 0.0 @@ -5406,6 +5408,7 @@ class PrismaClient: self._watching_engine: bool = False self._engine_confirmed_dead: bool = False self._engine_wait_thread: threading.Thread | None = None + self._engine_wait_threads: tuple[threading.Thread, ...] = () verbose_proxy_logger.debug("Success - Created Prisma Client") @property @@ -6614,6 +6617,9 @@ class PrismaClient: start_time: Final = time.time() try: await self.db.disconnect() + if self._closing: + await asyncio.gather(*(asyncio.to_thread(thread.join) for thread in self._engine_wait_threads)) + self._engine_wait_threads = () except Exception as e: import traceback @@ -6632,6 +6638,39 @@ class PrismaClient: ) raise e + def _spawn_engine_task(self, work: Coroutine[object, object, object]) -> None: + if self._closing: + work.close() + return + task: Final = asyncio.create_task(work) + self._engine_tasks = self._engine_tasks | {task} + task.add_done_callback(self._engine_task_finished) + + def _engine_task_finished(self, task: asyncio.Task[object]) -> None: + self._engine_tasks = self._engine_tasks - {task} + if not task.cancelled() and (error := task.exception()) is not None: + verbose_proxy_logger.error("Prisma engine background task failed: %s", error) + + async def _stop_engine_tasks(self) -> None: + tasks: Final = tuple(self._engine_tasks) + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + # Request-driven reconnects are not engine tasks. Wait for their bounded + # reconnect cycle to release the lock before disconnecting, otherwise a + # late reconnect can reopen the client after shutdown. + async with self._db_reconnect_lock: + pass + + async def stop_background_tasks(self) -> None: + self._closing = True + self._stop_engine_watcher() + async with contextlib.AsyncExitStack() as cleanup: + cleanup.push_async_callback(self._stop_engine_tasks) + cleanup.push_async_callback(self.db.stop_token_refresh_task) + cleanup.push_async_callback(self.stop_view_setup_task) + cleanup.push_async_callback(self.stop_db_health_watchdog_task) + def _get_engine_pid(self) -> int: """Get the PID of the writer's engine subprocess, or 0 if unavailable. @@ -6725,7 +6764,7 @@ class PrismaClient: self._engine_confirmed_dead = True self._reap_all_zombies() self._cleanup_engine_watcher() - asyncio.create_task( + self._spawn_engine_task( self.attempt_db_reconnect( reason="engine_process_death", force=True, @@ -6745,6 +6784,9 @@ class PrismaClient: name=f"prisma-engine-waitpid-{pid}", ) thread.start() + self._engine_wait_threads = tuple(thread for thread in self._engine_wait_threads if thread.is_alive()) + ( + thread, + ) self._engine_wait_thread = thread return True @@ -6806,7 +6848,7 @@ class PrismaClient: self._engine_confirmed_dead = True self._reap_all_zombies() self._cleanup_engine_watcher() - asyncio.create_task( + self._spawn_engine_task( self.attempt_db_reconnect( reason="engine_process_death", force=True, @@ -6868,7 +6910,7 @@ class PrismaClient: self._engine_confirmed_dead = True self._reap_all_zombies() self._cleanup_engine_watcher() - asyncio.create_task( + self._spawn_engine_task( self.attempt_db_reconnect( reason="engine_process_death", force=True, @@ -6941,7 +6983,7 @@ class PrismaClient: waitpid thread nor pidfd are available. """ - if self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None: + if self._closing or self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None: return pid: Final = self._get_engine_pid() if pid == 0: @@ -6968,7 +7010,7 @@ class PrismaClient: pid, ) self._watching_engine = True - asyncio.create_task(self._poll_engine_proc()) + self._spawn_engine_task(self._poll_engine_proc()) def _stop_engine_watcher(self) -> None: """Stop watching the engine process and clean up all resources.""" @@ -6990,7 +7032,7 @@ class PrismaClient: """ self._engine_confirmed_dead = False self._cleanup_engine_watcher() - asyncio.create_task(self._start_engine_watcher()) + self._spawn_engine_task(self._start_engine_watcher()) async def _run_reconnect_cycle( self, @@ -7197,6 +7239,8 @@ class PrismaClient: force_recreate: bool = False, stale_read_engine: "_StaleReadEngine | None" = None, ) -> bool: + if self._closing: + return False now: Final = time.time() if ( force is False @@ -7387,7 +7431,7 @@ class PrismaClient: - A periodic SELECT 1 probe that triggers reconnect on network/connection failure. - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling. """ - if self._db_health_watchdog_enabled is not True: + if self._closing or self._db_health_watchdog_enabled is not True: verbose_proxy_logger.debug("Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED") return if self._db_health_watchdog_task is not None: @@ -7421,7 +7465,7 @@ class PrismaClient: verbose_proxy_logger.info("Stopped Prisma DB health watchdog") def start_view_setup_task(self) -> None: - if self._view_setup_task is not None: + if self._closing or self._view_setup_task is not None: return self._view_setup_task = asyncio.create_task(self._run_view_setup()) diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py new file mode 100644 index 00000000000..c7291ad8adc --- /dev/null +++ b/litellm/repositories/database.py @@ -0,0 +1,32 @@ +import asyncio +from collections.abc import AsyncGenerator, Coroutine +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_proxy_logger + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +@dataclass(frozen=True, slots=True) +class Database: + prisma: "PrismaClient" + + async def aclose(self) -> None: + verbose_proxy_logger.debug("Disconnecting from Prisma") + try: + await self.prisma.stop_background_tasks() + finally: + await self.prisma.disconnect() + + +@asynccontextmanager +async def background_task(work: Coroutine[object, object, None]) -> AsyncGenerator[None, None]: + task: Final = asyncio.create_task(work) + try: + yield + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) diff --git a/tests/integration/database/test_proxy_database_lifespan.py b/tests/integration/database/test_proxy_database_lifespan.py new file mode 100644 index 00000000000..4a38cc4caf0 --- /dev/null +++ b/tests/integration/database/test_proxy_database_lifespan.py @@ -0,0 +1,153 @@ +import asyncio +import os +import subprocess +import sys +import threading +import uuid +from pathlib import Path +from typing import Final, Literal, TypeAlias + +import psutil +import pytest + +from tests.integration._support.database import read_rows, scratch_database + +pytestmark: Final = pytest.mark.timeout(180) +ROOT: Final = Path(__file__).resolve().parents[3] +PUBLICLY_KNOWN_KEY: Final = "sk-" + "1234" +Mode: TypeAlias = Literal["late", "unsafe"] +CONFIG: Final = ( + "model_list: []\n" + "general_settings:\n" + " disable_prisma_schema_update: true\n" + " disable_spend_logs: true\n" + " disable_model_info_refresh: true\n" +) + + +def _assert_database_idle() -> None: + assert ( + read_rows( + "SELECT pid FROM pg_stat_activity WHERE datname = current_database() AND pid <> pg_backend_pid()", + (), + ) + == [] + ) + assert psutil.Process().children(recursive=True) == [] + assert tuple(thread.name for thread in threading.enumerate() if thread.name.startswith("prisma-engine-")) == () + + +async def _assert_lifespan_closed() -> None: + from litellm.proxy import proxy_server + + await asyncio.sleep(0) + assert proxy_server.prisma_client is None + assert ( + tuple( + task.get_coro().__qualname__ + for task in asyncio.all_tasks() + if task.get_coro().__qualname__.startswith( + ("PrismaClient.", "PrismaWrapper.", "_proxy_lifespan.._run_") + ) + ) + == () + ) + await asyncio.to_thread(_assert_database_idle) + + +async def _exercise_lifespans(mode: Mode) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError + + expected: Final = UnsafeMasterKeyError if mode == "unsafe" else ValueError + match: Final = "master key" if mode == "unsafe" else "use_redis_transaction_buffer" + with pytest.raises(expected, match=match): + async with proxy_server.proxy_startup_event(proxy_server.app): + pytest.fail("Invalid configuration started the proxy") + await _assert_lifespan_closed() + + Path(os.environ["CONFIG_FILE_PATH"]).write_text(CONFIG) + with pytest.MonkeyPatch.context() as environment: + environment.setenv("LITELLM_MASTER_KEY", f"sk-{uuid.uuid4().hex}") + for cancel in (False, True, False): + try: + async with proxy_server.proxy_startup_event(proxy_server.app) as state: + database: Final = state["database"] + assert database is not None + assert await database.prisma.writer_db.query_raw("SELECT 42 AS answer") == [{"answer": 42}] + assert await database.prisma.read_db.query_raw("SELECT 43 AS answer") == [{"answer": 43}] + assert psutil.Process().children(recursive=True) + if cancel: + raise asyncio.CancelledError + except asyncio.CancelledError: + if not cancel: + raise + await _assert_lifespan_closed() + + +@pytest.mark.parametrize("mode", ("late", "unsafe")) +@pytest.mark.parametrize("replica", (False, True), ids=("writer", "writer-and-reader")) +def test_failed_startup_then_retry_releases_database_resources(tmp_path: Path, mode: Mode, replica: bool) -> None: + config: Final = tmp_path / "config.yaml" + config.write_text(CONFIG + (" use_redis_transaction_buffer: true\n" if mode == "late" else "")) + inherited: Final = { + key: value + for key, value in os.environ.items() + if not key.startswith("REDIS_") + and key + not in ( + "WORKER_CONFIG", + "DATABASE_URL_READ_REPLICA", + "LITELLM_SALT_KEY", + "LITELLM_MIGRATE_FROM_MASTER_KEY", + "LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", + "INTEGRATION_UPSTREAM_URL", + ) + } + with scratch_database("lifespan") as database_url: + environment: Final = { + **inherited, + "DATABASE_URL": database_url, + "DIRECT_URL": database_url, + "CONFIG_FILE_PATH": str(config), + "LITELLM_MASTER_KEY": PUBLICLY_KNOWN_KEY if mode == "unsafe" else f"sk-{uuid.uuid4().hex}", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "LITELLM_TELEMETRY_ENABLED": "false", + "PRISMA_HEALTH_WATCHDOG_ENABLED": "true", + "PYTHONPATH": os.pathsep.join((str(ROOT), str(ROOT / "enterprise"))), + **({"DATABASE_URL_READ_REPLICA": database_url} if replica else {}), + } + schema: Final = subprocess.run( + ( + sys.executable, + "-P", + "-m", + "prisma", + "db", + "push", + "--skip-generate", + "--schema", + str(ROOT / "litellm/proxy/schema.prisma"), + ), + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=60, + ) + assert schema.returncode == 0, schema.stderr + probe: Final = subprocess.run( + ( + sys.executable, + "-P", + "-c", + "import asyncio; from tests.integration.database.test_proxy_database_lifespan import _exercise_lifespans; " + f"asyncio.run(_exercise_lifespans({mode!r}))", + ), + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=90, + ) + assert probe.returncode == 0, probe.stdout + probe.stderr diff --git a/tests/unit/proxy/db/test_routing_prisma_wrapper.py b/tests/unit/proxy/db/test_routing_prisma_wrapper.py index 911a6d42ed3..73225acd2c8 100644 --- a/tests/unit/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/unit/proxy/db/test_routing_prisma_wrapper.py @@ -254,6 +254,55 @@ def test_token_refresh_delegates_to_both_writer_and_reader(): reader.stop_token_refresh_task.assert_awaited_once() +def test_reader_token_refresh_stops_even_when_the_writer_task_failed() -> None: + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + writer: Final = MagicMock() + reader: Final = MagicMock() + writer.stop_token_refresh_task = AsyncMock(side_effect=RuntimeError("refresh failed")) + reader.stop_token_refresh_task = AsyncMock() + routing: Final = RoutingPrismaWrapper(writer=writer, reader=reader) + + with pytest.raises(RuntimeError, match="refresh failed"): + asyncio.run(routing.stop_token_refresh_task()) + + reader.stop_token_refresh_task.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", (RuntimeError("refresh failed"), asyncio.CancelledError())) +async def test_reader_cleanup_finishes_before_writer_stop_failure_propagates(error: BaseException) -> None: + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + reader_stopping: Final = asyncio.Event() + release_reader: Final = asyncio.Event() + reader_stopped: Final = asyncio.Event() + + async def stop_reader() -> None: + reader_stopping.set() + await release_reader.wait() + reader_stopped.set() + + writer: Final = MagicMock() + reader: Final = MagicMock() + writer.stop_token_refresh_task = AsyncMock(side_effect=error) + reader.stop_token_refresh_task = AsyncMock(side_effect=stop_reader) + routing: Final = RoutingPrismaWrapper(writer=writer, reader=reader) + stopping: Final = asyncio.create_task(routing.stop_token_refresh_task()) + try: + await reader_stopping.wait() + await asyncio.sleep(0) + assert not stopping.done() + release_reader.set() + with pytest.raises(type(error)): + await stopping + assert reader_stopped.is_set() + finally: + release_reader.set() + stopping.cancel() + await asyncio.gather(stopping, return_exceptions=True) + + def test_routed_actions_falls_back_to_writer_for_unknown_methods(): from litellm.proxy.db.routing_prisma_wrapper import _RoutedActions diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 8b1d7c1222a..f25dfed196d 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -24,6 +24,7 @@ import logging import os import subprocess from collections.abc import Awaitable, Callable +from contextlib import aclosing from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -50,6 +51,7 @@ from litellm.proxy.proxy_server import ( run_ollama_serve, save_worker_config, ) +from litellm.repositories import database as database_runtime from .conftest import normalize @@ -100,9 +102,10 @@ def test_cleanup_router_config_variables_fails_on_unknown_attr_raises(): @pytest.mark.asyncio -async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch): +async def test_database_lifespan_disconnects_prisma_after_proxy_shutdown(monkeypatch): fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock() + fake_prisma.stop_background_tasks = AsyncMock() monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) monkeypatch.setattr(ps, "master_key", "sk-x", raising=False) @@ -116,7 +119,9 @@ async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch): monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - await proxy_shutdown_event() + async with aclosing(database_runtime.Database(prisma=fake_prisma)): + await proxy_shutdown_event() + fake_prisma.disconnect.assert_not_awaited() observed = { "disconnect_called": fake_prisma.disconnect.await_count == 1, @@ -203,6 +208,7 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect")) + fake_prisma.stop_background_tasks = AsyncMock() monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) async def _record_flush(client, accumulator): @@ -221,7 +227,8 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - await proxy_shutdown_event() + async with aclosing(database_runtime.Database(prisma=fake_prisma)): + await proxy_shutdown_event() assert calls == ["flush", "disconnect"] @@ -238,6 +245,7 @@ async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_befor fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect")) + fake_prisma.stop_background_tasks = AsyncMock() monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) async def _record_flush(client, accumulator): @@ -261,7 +269,8 @@ async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_befor monkeypatch.setattr(litellm, "success_callback", [], raising=False) track_audit_task(asyncio.create_task(_audit_write_that_finishes_after_the_flush_starts())) - await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) + async with aclosing(database_runtime.Database(prisma=fake_prisma)): + await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) assert calls == ["flush", "audit", "disconnect"] @@ -317,9 +326,10 @@ async def test_proxy_shutdown_without_a_database_drains_audit_callbacks_before_t @pytest.mark.asyncio -async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch): +async def test_database_lifespan_propagates_disconnect_error_after_shutdown(monkeypatch): fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock(side_effect=RuntimeError("db gone")) + fake_prisma.stop_background_tasks = AsyncMock() monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) fake_jwt = MagicMock() @@ -332,7 +342,8 @@ async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch): monkeypatch.setattr(litellm, "success_callback", [], raising=False) with pytest.raises(RuntimeError, match="db gone"): - await proxy_shutdown_event() + async with aclosing(database_runtime.Database(prisma=fake_prisma)): + await proxy_shutdown_event() # --------------------------------------------------------------------------- @@ -1007,7 +1018,9 @@ async def test_proxy_shutdown_commits_spend_before_waiting_on_a_slow_logging_tas task_may_finish: Final = asyncio.Event() task_finished: Final = asyncio.Event() - flushes: Final[list[tuple[str, bool]]] = [] # mutable-ok: records what each shutdown flush saw, which is the assertion + flushes: Final[ + list[tuple[str, bool]] + ] = [] # mutable-ok: records what each shutdown flush saw, which is the assertion async def slow_logging_task() -> None: await task_may_finish.wait() @@ -1059,7 +1072,7 @@ def test_otel_global_provider_published_after_callback_init(): span was orphaned. The publish (``publish_global_otel_v2_provider``) must therefore appear after ``_initialize_startup_logging`` in the lifespan source. """ - wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event) + wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan) source = inspect.getsource(wrapped) init_pos = source.find("ProxyStartupEvent.initialize_startup_logging(") publish_pos = source.find("publish_global_otel_v2_provider(") @@ -1145,7 +1158,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): The call must follow Prisma setup so DB-backed deployments do not false-positive. Direct ``_warn_budget_without_db`` tests cover the warning behavior itself. """ - wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event) + wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan) source = inspect.getsource(wrapped) budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:") warn_pos = source.find("warn_budget_without_db(") diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 976ed4f28a5..9b11b296aa9 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -115,7 +115,7 @@ async def test_admin_mcp_failure_still_closes_proxy_resources( message: Final = "LITELLM_LICENSE" if failure_phase == "license" else f"{failure_phase} failed" async def run_lifespan() -> None: async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state: - assert state == {"tracing_receiver": None} + assert state == {"tracing_receiver": None, "database": None} if failure_phase == "serving": raise RuntimeError("serving failed") if failure_phase == "cancelled": @@ -164,7 +164,7 @@ async def test_proxy_shutdown_drains_active_admin_tool_before_closing_connector( transport=httpx2.ASGITransport(app=admin_mcp_proxy), base_url="http://localhost:4000" ) as client: async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state: - assert state == {"tracing_receiver": None} + assert state == {"tracing_receiver": None, "database": None} request: Final = asyncio.create_task(client.post( "/admin/mcp", headers={"Authorization": "Bearer admin", "Accept": "application/json, text/event-stream"}, diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 55d1587669a..cd2578b4e3d 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -1651,6 +1651,9 @@ class MockPrisma: async def disconnect(self): pass + async def stop_background_tasks(self): + pass + mock_prisma = MockPrisma() @@ -1761,7 +1764,7 @@ async def test_proxy_startup_refuses_an_unsafe_master_key_even_when_the_database from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError from litellm.proxy.proxy_server import proxy_startup_event - async def unreachable(): + async def unreachable(resources): raise ConnectionError("database is down") _, announced = _boot_with_general_settings(monkeypatch, tmp_path, general_settings) @@ -1806,7 +1809,7 @@ async def test_proxy_startup_asks_to_migrate_only_when_the_database_holds_values database = _DatabaseWithOneStoredCredential(encrypt_value_helper("sk-provider", new_encryption_key=encrypted_with)) - async def connected(): + async def connected(resources): return database _, announced = _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": PUBLICLY_KNOWN_KEY}) @@ -1841,6 +1844,15 @@ async def test_proxy_startup_says_a_lingering_migrate_from_variable_can_be_delet class _PrismaClientWhoseDatabaseRejectsQueries: + def __init__(self) -> None: + self.disconnect_calls = 0 + + async def disconnect(self) -> None: + self.disconnect_calls += 1 + + async def stop_background_tasks(self) -> None: + pass + class _Database: async def query_raw(self, query, *args): raise RuntimeError("permission denied for table LiteLLM_CredentialsTable") @@ -1855,7 +1867,8 @@ async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatc from litellm.proxy.proxy_server import proxy_startup_event _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries()) + prisma: Final = _PrismaClientWhoseDatabaseRejectsQueries() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY) with ( @@ -1868,6 +1881,25 @@ async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatc notices = [record.getMessage() for record in caplog.records if "LITELLM_MIGRATE_FROM_MASTER_KEY" in record.message] assert len(notices) == 1 assert "Could not migrate stored values" in notices[0] + assert prisma.disconnect_calls == 1 + + +@pytest.mark.asyncio +async def test_proxy_startup_after_a_failed_migration_does_not_reuse_the_closed_client(monkeypatch, tmp_path): + from fastapi import FastAPI + + from litellm.proxy.proxy_server import proxy_startup_event + + _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries()) + monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY) + with pytest.raises(RuntimeError, match="permission denied"): + async with proxy_startup_event(FastAPI()): + pass + + monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", "") + async with proxy_startup_event(FastAPI()) as state: + assert state["database"] is None @pytest.mark.asyncio @@ -1914,11 +1946,11 @@ async def test_proxy_startup_boots_an_unsafe_master_key_under_the_override(monke class _ShutdownAwarePrisma(MockPrisma): def __init__(self): super().__init__() - self.stop_view_setup_task = AsyncMock() + self.stop_background_tasks = AsyncMock() @pytest.mark.asyncio -async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path): +async def test_proxy_shutdown_stops_database_background_tasks(monkeypatch, tmp_path): import yaml from fastapi import FastAPI @@ -1933,11 +1965,11 @@ async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path): monkeypatch.setattr(proxy_server_module, "store_model_in_db", False) async with proxy_startup_event(FastAPI()): - stopped_while_serving = fake_prisma.stop_view_setup_task.await_count + stopped_while_serving = fake_prisma.stop_background_tasks.await_count actual = { "stopped_while_serving": stopped_while_serving, - "stopped_after_shutdown": fake_prisma.stop_view_setup_task.await_count, + "stopped_after_shutdown": fake_prisma.stop_background_tasks.await_count, } assert actual == {"stopped_while_serving": 0, "stopped_after_shutdown": 1} @@ -13957,6 +13989,8 @@ async def test_team_window_spend_row_carries_the_request_start_time(): def _mock_startup_prisma_client(health_check_error=None, connect_error=None): client = MagicMock() client.connect = AsyncMock(side_effect=connect_error) + client.disconnect = AsyncMock() + client.stop_background_tasks = AsyncMock() client.db.start_token_refresh_task = AsyncMock() client.check_view_exists = AsyncMock() client._set_spend_logs_row_count_in_proxy_state = AsyncMock() @@ -14069,6 +14103,8 @@ async def test_setup_prisma_client_raises_when_db_unavailable_is_not_allowed(mon mock_client = _mock_startup_prisma_client(health_check_error=httpx.ReadTimeout("startup health check timed out")) with pytest.raises(httpx.ReadTimeout): await _run_setup_prisma_client(mock_client) + mock_client.stop_background_tasks.assert_awaited_once() + mock_client.disconnect.assert_awaited_once() @pytest.mark.asyncio @@ -14089,6 +14125,31 @@ async def test_setup_prisma_client_returns_none_when_connect_itself_fails(monkey assert result is None assert mock_client.start_db_health_watchdog_task.await_count == 0 assert mock_client.health_check.await_count == 0 + mock_client.stop_background_tasks.assert_awaited_once() + mock_client.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_cancelling_prisma_setup_closes_the_client_before_publishing_it(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "False") + checking: Final = asyncio.Event() + client: Final = _mock_startup_prisma_client() + + async def health_check() -> None: + checking.set() + await asyncio.Event().wait() + + client.health_check = health_check + setup: Final = asyncio.create_task(_run_setup_prisma_client(client)) + try: + await asyncio.wait_for(checking.wait(), timeout=2) + finally: + setup.cancel() + with pytest.raises(asyncio.CancelledError): + await setup + + client.stop_background_tasks.assert_awaited_once() + client.disconnect.assert_awaited_once() async def _run_scheduled_background_jobs(): diff --git a/tests/unit/proxy/utils/prisma_and_spend/conftest.py b/tests/unit/proxy/utils/prisma_and_spend/conftest.py index aefb95516af..afc282a14d4 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/unit/proxy/utils/prisma_and_spend/conftest.py @@ -139,6 +139,7 @@ def mock_prisma_client() -> MagicMock: client.db.is_connected = MagicMock(return_value=False) client.db.connect = AsyncMock() client.db.disconnect = AsyncMock() + client.db.stop_token_refresh_task = AsyncMock() client.db.query_raw = AsyncMock(return_value=[{"?column?": 1}]) client.db.execute_raw = AsyncMock() client.db.tx = MagicMock() diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py index 9aa57c7a19b..1c94075ca53 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py @@ -14,12 +14,106 @@ from __future__ import annotations import asyncio import logging -from typing import Any -from unittest.mock import AsyncMock, MagicMock +from typing import Any, Final +from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.proxy.utils import PrismaClient +from litellm.repositories.database import Database + + +@pytest.mark.asyncio +@pytest.mark.parametrize("engine_dead", (False, True)) +async def test_shutdown_waits_for_hung_request_reconnect_to_time_out( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch, engine_dead: bool +) -> None: + monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") + reconnecting: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + + async def hung_recreate(*args: object, **kwargs: object) -> None: + reconnecting.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + prisma_client.db.recreate_prisma_client = AsyncMock(side_effect=hung_recreate) + prisma_client._engine_confirmed_dead = engine_dead + with patch.object(asyncio.get_running_loop(), "time", return_value=1000.0) as clock: + reconnect: Final = asyncio.create_task( + prisma_client.attempt_db_reconnect(reason="request", force=True, force_recreate=True) + ) + await reconnecting.wait() + closing: Final = asyncio.create_task(Database(prisma=prisma_client).aclose()) + try: + await asyncio.sleep(0) + assert not closing.done() + prisma_client.db.disconnect.assert_not_awaited() + + clock.return_value = 1000.0 + prisma_client._db_watchdog_reconnect_timeout_seconds + 1 + assert await reconnect is False + await closing + finally: + reconnect.cancel() + closing.cancel() + await asyncio.gather(reconnect, closing, return_exceptions=True) + + assert cancelled.is_set() + assert not prisma_client._db_reconnect_lock.locked() + prisma_client.db.disconnect.assert_awaited_once() + assert await prisma_client.attempt_db_reconnect(reason="request", force=True) is False + prisma_client.db.recreate_prisma_client.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_shutdown_rejects_a_reconnect_already_waiting_for_the_lock(prisma_client: PrismaClient) -> None: + await prisma_client._db_reconnect_lock.acquire() + reconnect: Final = asyncio.create_task(prisma_client.attempt_db_reconnect(reason="engine_process_death", force=True)) + await asyncio.sleep(0) + closing: Final = asyncio.create_task(prisma_client.stop_background_tasks()) + await asyncio.sleep(0) + prisma_client._db_reconnect_lock.release() + + await asyncio.wait_for(closing, timeout=2) + assert await reconnect is False + assert await prisma_client.attempt_db_reconnect(reason="engine_process_death", force=True) is False + prisma_client.db.connect.assert_not_awaited() + prisma_client.db.disconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_spawned_engine_work_and_prevents_rearming(prisma_client: PrismaClient) -> None: + running: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + released: Final = asyncio.Event() + finished: Final = asyncio.Event() + + async def engine_work() -> None: + running.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + await released.wait() + finished.set() + + prisma_client._spawn_engine_task(engine_work()) + await asyncio.wait_for(running.wait(), timeout=2) + closing: Final = asyncio.create_task(prisma_client.stop_background_tasks()) + try: + await asyncio.wait_for(cancelled.wait(), timeout=2) + assert not finished.is_set() + assert not closing.done() + finally: + released.set() + await asyncio.wait_for(closing, timeout=2) + + assert finished.is_set() + prisma_client._handle_writer_engine_replaced() + await asyncio.sleep(0) + assert not prisma_client._engine_tasks _PROBE_SQL = "SELECT to_regclass($1) IS NOT NULL AS present" diff --git a/tests/unit/repositories/test_database.py b/tests/unit/repositories/test_database.py new file mode 100644 index 00000000000..008eb1124e8 --- /dev/null +++ b/tests/unit/repositories/test_database.py @@ -0,0 +1,18 @@ +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.repositories.database import Database + + +@pytest.mark.asyncio +async def test_database_disconnects_even_when_stopping_background_tasks_fails() -> None: + prisma: Final = MagicMock() + prisma.stop_background_tasks = AsyncMock(side_effect=RuntimeError("watchdog failed")) + prisma.disconnect = AsyncMock() + + with pytest.raises(RuntimeError, match="watchdog failed"): + await Database(prisma=prisma).aclose() + + prisma.disconnect.assert_awaited_once()