diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 43aead76630..ad501559786 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6656,6 +6656,9 @@ class PrismaClient: 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 diff --git a/tests/integration/database/test_proxy_database_lifespan.py b/tests/integration/database/test_proxy_database_lifespan.py index 014d6c59b94..4a38cc4caf0 100644 --- a/tests/integration/database/test_proxy_database_lifespan.py +++ b/tests/integration/database/test_proxy_database_lifespan.py @@ -14,6 +14,7 @@ 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" @@ -109,7 +110,7 @@ def test_failed_startup_then_retry_releases_database_resources(tmp_path: Path, m "DATABASE_URL": database_url, "DIRECT_URL": database_url, "CONFIG_FILE_PATH": str(config), - "LITELLM_MASTER_KEY": "sk-1234" if mode == "unsafe" else f"sk-{uuid.uuid4().hex}", + "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", diff --git a/tests/unit/proxy/db/test_routing_prisma_wrapper.py b/tests/unit/proxy/db/test_routing_prisma_wrapper.py index 4dd43d9643b..73225acd2c8 100644 --- a/tests/unit/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/unit/proxy/db/test_routing_prisma_wrapper.py @@ -269,6 +269,40 @@ def test_reader_token_refresh_stops_even_when_the_writer_task_failed() -> None: 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/utils/prisma_and_spend/test_prisma_client_lifecycle.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py index f6f9b2ab42a..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 @@ -15,11 +15,56 @@ from __future__ import annotations import asyncio import logging from typing import Any, Final -from unittest.mock import AsyncMock, MagicMock +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