Merge pull request #39365 from BerriAI/litellm_fix_v2_resolver_migrate_deploy_timeout

fix(proxy-extras): give prisma migrate deploy its own timeout budget

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)

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

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"])