From 2495673d01ee26c73d54adb6b78e30ff87cf0bc2 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 28 Aug 2026 02:14:57 -0700 Subject: [PATCH] fix(proxy-extras): stop a db push timeout crashing the migration job subprocess.run leaves stderr as bytes on TimeoutExpired even under text=True, unlike CalledProcessError. Classifying both in one handler meant a real `prisma db push` timeout died on a TypeError, which proxy_cli.py's `except RuntimeError` does not catch, so the migrations Job container ended on an unhandled traceback instead of a clean exit. Give the timeout its own handler and retry it, matching what the migrate deploy loop beside it already does. That puts a fallthrough back into the loop, so the trailing raise removed in the previous commit is reachable again and comes back with it. Also drop a comment restating why the resolver cases exist and widen the db push test's docstring, which had stopped describing what it covers. --- .../litellm_proxy_extras/utils.py | 14 ++- .../test_setup_database_fail_fast.py | 86 +++++++++++++++++++ tests/test_litellm/proxy/test_proxy_cli.py | 7 +- 3 files changed, 98 insertions(+), 9 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index f751c4087eb..22e30dfa897 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -709,10 +709,13 @@ class ProxyExtrasDBManager: env=_get_prisma_env(), ) return True - except ( - subprocess.CalledProcessError, - subprocess.TimeoutExpired, - ) as e: + except subprocess.TimeoutExpired: + logger.info( + "prisma db push attempt %s timed out, retrying", + attempt + 1, + ) + time.sleep(random.randrange(5, 15)) + except subprocess.CalledProcessError as e: stderr = e.stderr or "" transient = ProxyExtrasDBManager._transient_prisma_failure( stderr @@ -732,6 +735,9 @@ class ProxyExtrasDBManager: stderr, ) time.sleep(random.randrange(5, 15)) + raise RuntimeError( + f"prisma db push failed after {_PRISMA_ATTEMPTS} attempts." + ) finally: os.chdir(original_dir) diff --git a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py index ed82037590a..3996b91c2a6 100644 --- a/tests/litellm-proxy-extras/test_setup_database_fail_fast.py +++ b/tests/litellm-proxy-extras/test_setup_database_fail_fast.py @@ -460,6 +460,92 @@ def test_v2_db_push_retries_are_bounded_and_report_the_prisma_error( assert "P1001" in str(exc.value) +def _db_push_only(push_side_effect): + """subprocess.run stand-in that only intercepts `prisma db push`.""" + pushes = {"n": 0} + + def _run(*args, **kwargs): + cmd = list(args[0] if args else kwargs.get("args", [])) + if cmd[-3:] != ["db", "push", "--accept-data-loss"]: + return _DeployApplied() + pushes["n"] += 1 + return push_side_effect(pushes["n"], cmd) + + return _run, pushes + + +def _timed_out_for_real(): + """Capture what subprocess.run really puts on a TimeoutExpired. + + Under text=True it still leaves stderr as bytes, unlike CalledProcessError, + so hardcoding a str here would test a shape production never sees. Derived + at import, before any test patches subprocess.run. + """ + try: + subprocess.run( + ["sh", "-c", "echo 'Error: P1001 unreachable' >&2; sleep 5"], + timeout=0.2, + check=True, + capture_output=True, + text=True, + ) + except subprocess.TimeoutExpired as e: + return e + raise AssertionError("the helper command was supposed to time out") + + +_TIMEOUT_TEMPLATE = _timed_out_for_real() + + +def _real_timeout_expired(cmd): + return subprocess.TimeoutExpired( + cmd=cmd, + timeout=_TIMEOUT_TEMPLATE.timeout, + output=_TIMEOUT_TEMPLATE.stdout, + stderr=_TIMEOUT_TEMPLATE.stderr, + ) + + +def test_v2_db_push_retries_a_timeout(monkeypatch, tmp_path): + """v2: a `prisma db push` that times out is retried, not turned into a + TypeError by classifying its bytes stderr as if it were text.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + if n == 1: + raise _real_timeout_expired(cmd) + return _DeployApplied() + + monkeypatch.setattr( + ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False + ) + run, pushes = _db_push_only(_side_effect) + with patch("subprocess.run", side_effect=run): + ok = ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True) + + assert ok is True + assert pushes["n"] == 2 + + +def test_v2_db_push_timeouts_are_bounded(monkeypatch, tmp_path): + """v2: a `prisma db push` that never stops timing out gives up as a + RuntimeError, which is the only exception proxy_cli.py exits cleanly on.""" + _prepare_v2_resolver(monkeypatch, tmp_path) + + def _side_effect(n, cmd): + raise _real_timeout_expired(cmd) + + monkeypatch.setattr( + ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False + ) + run, pushes = _db_push_only(_side_effect) + with patch("subprocess.run", side_effect=run): + with pytest.raises(RuntimeError, match=r"prisma db push failed after \d+"): + ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True) + + assert pushes["n"] == _PRISMA_ATTEMPTS + + def test_v2_unclassified_failure_is_not_treated_as_transient(monkeypatch, tmp_path): """v2: an unrecognised deploy failure still raises on the first attempt.""" _prepare_v2_resolver(monkeypatch, tmp_path) diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 4f205ada609..5e2dd358d75 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1737,7 +1737,8 @@ class TestRunServerDbSetup: mock_atexit_register, mock_subprocess_run, ): - """Test that use_prisma_db_push flag correctly controls PrismaManager.setup_database use_migrate parameter""" + """Which resolver and which migration mode run_server hands setup_database, + across the db push flag, the v2/legacy flag pair and USE_V2_MIGRATION_RESOLVER.""" from litellm.proxy.proxy_cli import run_server # Mock subprocess.run to simulate prisma being available @@ -1805,10 +1806,6 @@ class TestRunServerDbSetup: use_migrate=False, use_v2_resolver=True ) - # Test 3+: the resolver default and both routes back to v1. The flag - # covers a CLI boot; USE_V2_MIGRATION_RESOLVER covers deploys that - # cannot pass one, where prisma_migration.py fixes the argv. An - # explicit flag beats the env var. for argv, env_value, expected_v2 in ( ([], None, True), (["--use_v2_migration_resolver"], None, True),