diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py index f3b55fd4d96..2283814ab35 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py +++ b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py @@ -14,9 +14,22 @@ then fails on a Node binary that was never written. Deleting a cache directory that exists without a Node binary is what turns a killed bootstrap back into a recoverable one. -Both budgets are overridable so an operator can widen them without a release: -``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install and -``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every individual Prisma command. +``prisma migrate deploy`` is the other command whose runtime is not a +constant: it grows with the number of pending migrations, so a fresh database +that has to replay every migration this package ships overruns a per-command +budget sized for the short bookkeeping commands, on a laptop as much as on a +slow CI runner. The Python ``prisma`` wrapper spawns Node and the schema engine +as separate children, so killing the wrapper on timeout leaves them running: +the retry then contends with that orphan for Prisma's advisory lock and cannot +finish any sooner. Migrate deploy therefore runs under its own budget. + +All three budgets are overridable so an operator can widen them without a +release: ``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install, +``LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT`` for ``prisma migrate deploy`` and +``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every other Prisma command. The +per-command budget used to bound migrate deploy as well, so a deployment that +raised it above the deploy default keeps that larger budget for deploy unless +the deploy override says otherwise. """ import math @@ -36,10 +49,12 @@ except ImportError: PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT" PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT" +PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT" NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR" DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0 DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0 +DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0 BOOTSTRAP_ARG = "--version" @@ -88,6 +103,15 @@ def prisma_bootstrap_timeout() -> float: ) +def prisma_migrate_deploy_timeout() -> float: + """Seconds one ``prisma migrate deploy`` may run for, however many migrations are pending.""" + if os.getenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR) is not None: + return _timeout_from_env( + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT + ) + return max(DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT, prisma_command_timeout()) + + def nodeenv_cache_dir() -> Optional[Path]: """Where Prisma installs its private Node runtime, or None if unknowable.""" override = os.getenv(NODEENV_CACHE_DIR_ENV_VAR) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b8032dd0d28..f6647268624 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -15,8 +15,11 @@ from litellm_proxy_extras.replica_identity import ( apply_replica_identity_full, ) from litellm_proxy_extras.prisma_toolchain import ( + PRISMA_COMMAND_TIMEOUT_ENV_VAR, + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, ensure_prisma_toolchain, prisma_command_timeout, + prisma_migrate_deploy_timeout, ) @@ -698,12 +701,13 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) + deploy_timeout = prisma_migrate_deploy_timeout() try: for attempt in range(4): try: result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], - timeout=prisma_command_timeout(), + timeout=deploy_timeout, check=True, capture_output=True, text=True, @@ -713,8 +717,12 @@ class ProxyExtrasDBManager: return True except subprocess.TimeoutExpired: - logger.info( - f"prisma migrate deploy attempt {attempt + 1} timed out, retrying" + logger.warning( + "prisma migrate deploy attempt %s timed out after %ss, retrying. " + "Raise %s if this database needs longer to apply its pending migrations.", + attempt + 1, + deploy_timeout, + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, ) time.sleep(random.randrange(5, 15)) continue @@ -823,7 +831,8 @@ class ProxyExtrasDBManager: "Database migration failed after 4 attempts (retry loop " "exhausted by timeouts or repeated idempotent-recovery " "continues). Check database connectivity, load, and " - "_prisma_migrations ledger state." + "_prisma_migrations ledger state, and raise " + f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out." ) finally: os.chdir(original_dir) @@ -908,7 +917,7 @@ class ProxyExtrasDBManager: # Set migrations directory for Prisma result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], - timeout=prisma_command_timeout(), + timeout=prisma_migrate_deploy_timeout(), check=True, capture_output=True, text=True, @@ -1126,7 +1135,11 @@ class ProxyExtrasDBManager: ) return True except subprocess.TimeoutExpired: - logger.info(f"Attempt {attempt + 1} timed out") + logger.warning( + "Attempt %s timed out. Raise %s if this database needs longer to apply its schema.", + attempt + 1, + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR if use_migrate else PRISMA_COMMAND_TIMEOUT_ENV_VAR, + ) time.sleep(random.randrange(5, 15)) except subprocess.CalledProcessError as e: attempts_left = 3 - attempt diff --git a/tests/proxy_migration_tests/test_prisma_toolchain.py b/tests/proxy_migration_tests/test_prisma_toolchain.py index d12a2c4dd4e..733870f3239 100644 --- a/tests/proxy_migration_tests/test_prisma_toolchain.py +++ b/tests/proxy_migration_tests/test_prisma_toolchain.py @@ -7,26 +7,36 @@ attempt fails identically. These tests pin the two behaviours that keep a container recoverable: an incomplete cache is deleted before Prisma is invoked, and the install gets a budget of its own rather than sharing the one that bounds each migration command. + +``prisma migrate deploy`` gets a budget of its own for the same reason: its +runtime grows with the number of pending migrations, so a fresh database that +replays every migration overran the per-command budget on slow machines and +the proxy gave up after four identical timeouts. """ import ast import json +import logging import os import sys import time +from collections.abc import Callable from pathlib import Path import pytest from litellm_proxy_extras.prisma_toolchain import ( DEFAULT_PRISMA_COMMAND_TIMEOUT, + DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT, PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, PRISMA_COMMAND_TIMEOUT_ENV_VAR, + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, ensure_prisma_toolchain, heal_incomplete_nodeenv_cache, node_binary_path, prisma_bootstrap_timeout, prisma_command_timeout, + prisma_migrate_deploy_timeout, ) from litellm_proxy_extras.utils import ProxyExtrasDBManager @@ -42,14 +52,27 @@ import time args = sys.argv[1:] cache_dir = os.environ["PRISMA_NODEENV_CACHE_DIR"] -with pathlib.Path(os.environ["FAKE_PRISMA_LOG"]).open("a") as log: +log_path = pathlib.Path(os.environ["FAKE_PRISMA_LOG"]) +earlier_same_command = sum( + 1 + for line in (log_path.read_text().splitlines() if log_path.exists() else []) + if json.loads(line)["args"][:2] == args[:2] +) +with log_path.open("a") as log: log.write( json.dumps({{"args": args, "cache_dir_present": os.path.isdir(cache_dir)}}) + "\\n" ) time.sleep(float(os.environ.get("FAKE_PRISMA_SLEEP", "0"))) if args[:2] == ["migrate", "deploy"]: + if earlier_same_command == 0: + time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "0"))) + elif os.environ.get("FAKE_PRISMA_LATER_DEPLOY_STDERR"): + print(os.environ["FAKE_PRISMA_LATER_DEPLOY_STDERR"], file=sys.stderr) + sys.exit(1) print("No pending migrations to apply") +if args[:2] == ["db", "push"] and earlier_same_command == 0: + time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_PUSH_SLEEP", "0"))) sys.exit(0) """ @@ -80,9 +103,14 @@ def toolchain_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ['PATH']}") monkeypatch.delenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, raising=False) monkeypatch.delenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, raising=False) + monkeypatch.delenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, raising=False) return cache_dir, log_path +def _deploy_calls(log_path: Path) -> list[list[str]]: + return [call["args"] for call in _fake_prisma_calls(log_path) if call["args"][:2] == ["migrate", "deploy"]] + + def _make_incomplete_cache(cache_dir: Path) -> None: (cache_dir / "lib").mkdir(parents=True) (cache_dir / "bin").mkdir() @@ -209,25 +237,111 @@ def test_setup_database_prepares_the_toolchain_before_migrating( assert calls[0]["cache_dir_present"] is False +@pytest.mark.parametrize("use_v2_resolver", [False, True], ids=["v1", "v2"]) +def test_migrate_deploy_is_not_bounded_by_the_per_command_timeout( + toolchain_env: tuple[Path, Path], + monkeypatch: pytest.MonkeyPatch, + use_v2_resolver: bool, +) -> None: + """A fresh database replays every migration, which takes longer than any bookkeeping command.""" + _, log_path = toolchain_env + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "3") + + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True + assert _deploy_calls(log_path) == [["migrate", "deploy"]] + + +def test_migrate_deploy_stops_at_its_own_timeout( + toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + """The deploy budget still bounds a deploy that hangs, so boot cannot wait forever.""" + _, log_path = toolchain_env + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "60") + monkeypatch.setenv("FAKE_PRISMA_LATER_DEPLOY_STDERR", "Error: P3018 permission denied for schema public") + + started = time.monotonic() + with pytest.raises(RuntimeError, match="insufficient permissions"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + elapsed = time.monotonic() - started + + assert len(_deploy_calls(log_path)) == 2 + assert elapsed < 30 + + +def test_db_push_timeout_hint_names_the_per_command_budget( + toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """``db push`` keeps the per-command budget, so its timeout hint has to name that variable.""" + _, log_path = toolchain_env + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv("FAKE_PRISMA_FIRST_PUSH_SLEEP", "3") + + with caplog.at_level(logging.WARNING, logger="litellm_proxy_extras"): + assert ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=False) is True + + assert [call["args"][:2] for call in _fake_prisma_calls(log_path)].count(["db", "push"]) == 2 + assert [record.getMessage() for record in caplog.records if "timed out" in record.getMessage()] == [ + f"Attempt 1 timed out. Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer to apply its schema." + ] + + +@pytest.mark.parametrize( + ("command_timeout", "deploy_timeout", "expected"), + [ + ("900", None, 900.0), + ("12", None, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT), + ("900", "1200", 1200.0), + ("900", "300", 300.0), + ], + ids=["raised_command_budget_carries_over", "lowered_command_budget_does_not", "override_wins_upward", "override_wins_downward"], +) +def test_migrate_deploy_budget_keeps_a_raised_command_budget( + command_timeout: str, deploy_timeout: str | None, expected: float, monkeypatch: pytest.MonkeyPatch +) -> None: + """Deployments that raised the per-command budget to survive a long deploy keep that budget for deploy.""" + monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, command_timeout) + if deploy_timeout is None: + monkeypatch.delenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, raising=False) + else: + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, deploy_timeout) + + assert prisma_migrate_deploy_timeout() == expected + + @pytest.mark.parametrize( "raw", ["", "0", "-5", "not-a-number", "nan", "inf", "-inf", "1e400"], ) +@pytest.mark.parametrize( + ("env_var", "read_timeout", "default"), + [ + (PRISMA_COMMAND_TIMEOUT_ENV_VAR, prisma_command_timeout, DEFAULT_PRISMA_COMMAND_TIMEOUT), + (PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, prisma_migrate_deploy_timeout, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT), + ], + ids=["command", "migrate_deploy"], +) def test_unusable_timeout_override_falls_back_to_the_default( - raw: str, monkeypatch: pytest.MonkeyPatch + raw: str, env_var: str, read_timeout: Callable[[], float], default: float, monkeypatch: pytest.MonkeyPatch ) -> None: """A non-finite override would silently disable the timeout it configures.""" - monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, raw) + monkeypatch.setenv(env_var, raw) - assert prisma_command_timeout() == DEFAULT_PRISMA_COMMAND_TIMEOUT + assert read_timeout() == default def test_timeout_overrides_are_independent(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "12") monkeypatch.setenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, "900") + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1200") assert prisma_command_timeout() == 12 assert prisma_bootstrap_timeout() == 900 + assert prisma_migrate_deploy_timeout() == 1200 @pytest.mark.parametrize("module", ["utils.py", "replica_identity.py"])