mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
fix(proxy-extras): kill the whole Prisma process group when a command times out
Every Prisma CLI call now goes through one runner that starts the command in its own session and SIGKILLs the process group on timeout, so the Node process and the Rust schema engine die together with the Python wrapper instead of being reparented to pid 1, where they kept applying migrations after the proxy had given up and held the Prisma advisory lock against every retry and every later boot. Tests that faked subprocess.run now fake the runner, and the fake Prisma CLI in the migration tests forks a grandchild that must not outlive a timed-out migrate deploy.
This commit is contained in:
parent
ff1f21aea9
commit
105f99cea2
7 changed files with 139 additions and 78 deletions
|
|
@ -18,10 +18,15 @@ recoverable one.
|
|||
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.
|
||||
slow CI runner. Migrate deploy therefore runs under its own budget.
|
||||
|
||||
The Python ``prisma`` wrapper spawns Node, which spawns the Rust schema
|
||||
engine, so killing only the wrapper on timeout leaves the engine running with
|
||||
no parent: it keeps mutating the database after the proxy has given up, holds
|
||||
Prisma's advisory lock so every retry and every later boot queues behind it,
|
||||
and dies mid-migration once its pipes close, leaving a half-applied ledger row.
|
||||
Every Prisma command therefore runs in a process group of its own, and a
|
||||
timeout kills the whole group.
|
||||
|
||||
All three budgets are overridable so an operator can widen them without a
|
||||
release: ``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install,
|
||||
|
|
@ -35,10 +40,12 @@ the deploy override says otherwise.
|
|||
import math
|
||||
import os
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from typing import IO, Optional, Union
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
|
||||
|
|
@ -167,6 +174,49 @@ def heal_incomplete_nodeenv_cache() -> bool:
|
|||
return True
|
||||
|
||||
|
||||
def _kill_process_group(process: "subprocess.Popen[str]") -> None:
|
||||
if os.name == "nt":
|
||||
process.kill()
|
||||
return
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
return
|
||||
|
||||
|
||||
def run_prisma(
|
||||
argv: Sequence[str],
|
||||
*,
|
||||
timeout: float,
|
||||
env: Mapping[str, str],
|
||||
stdout: Union[IO[str], int, None] = subprocess.PIPE,
|
||||
stderr: Optional[int] = subprocess.PIPE,
|
||||
) -> "subprocess.CompletedProcess[str]":
|
||||
"""Run one Prisma CLI command in its own process group, bounded by ``timeout``.
|
||||
|
||||
Raises ``subprocess.TimeoutExpired`` once the budget is spent, after killing
|
||||
the command together with every process it spawned, and
|
||||
``subprocess.CalledProcessError`` on a non-zero exit. Output is captured as
|
||||
text unless ``stdout``/``stderr`` say otherwise.
|
||||
"""
|
||||
with subprocess.Popen(
|
||||
argv,
|
||||
env=env,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
text=True,
|
||||
start_new_session=True,
|
||||
) as process:
|
||||
try:
|
||||
out, err = process.communicate(timeout=timeout)
|
||||
except BaseException:
|
||||
_kill_process_group(process)
|
||||
raise
|
||||
if process.returncode:
|
||||
raise subprocess.CalledProcessError(process.returncode, process.args, out, err)
|
||||
return subprocess.CompletedProcess(process.args, process.returncode, out, err)
|
||||
|
||||
|
||||
def ensure_prisma_toolchain(
|
||||
prisma_command: str, prisma_env: dict[str, str]
|
||||
) -> ToolchainBootstrap:
|
||||
|
|
@ -179,14 +229,7 @@ def ensure_prisma_toolchain(
|
|||
timeout = prisma_bootstrap_timeout()
|
||||
logger.info("Preparing the Prisma CLI toolchain (timeout %ss)", timeout)
|
||||
try:
|
||||
subprocess.run(
|
||||
[prisma_command, BOOTSTRAP_ARG],
|
||||
timeout=timeout,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=prisma_env,
|
||||
)
|
||||
run_prisma([prisma_command, BOOTSTRAP_ARG], timeout=timeout, env=prisma_env)
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning(
|
||||
"Preparing the Prisma CLI toolchain timed out after %ss. Raise %s "
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ import tempfile
|
|||
from pathlib import Path
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.prisma_toolchain import prisma_command_timeout
|
||||
from litellm_proxy_extras.prisma_toolchain import prisma_command_timeout, run_prisma
|
||||
|
||||
REPLICA_IDENTITY_FULL_ENV_VAR = "LITELLM_SET_REPLICA_IDENTITY_FULL"
|
||||
|
||||
|
|
@ -66,7 +66,7 @@ def apply_replica_identity_full(
|
|||
with tempfile.TemporaryDirectory(prefix="litellm_replica_identity_") as tmp_dir:
|
||||
sql_path = Path(tmp_dir) / "replica_identity_full.sql"
|
||||
sql_path.write_text(REPLICA_IDENTITY_FULL_SQL)
|
||||
subprocess.run(
|
||||
run_prisma(
|
||||
[
|
||||
prisma_command,
|
||||
"db",
|
||||
|
|
@ -77,9 +77,6 @@ def apply_replica_identity_full(
|
|||
schema_path,
|
||||
],
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=prisma_env,
|
||||
)
|
||||
except subprocess.CalledProcessError as e:
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import time
|
|||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from litellm_proxy_extras import prisma_toolchain
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.replica_identity import (
|
||||
REPLICA_IDENTITY_FULL_ENV_VAR,
|
||||
|
|
@ -198,7 +199,7 @@ class ProxyExtrasDBManager:
|
|||
# 1. Generate migration SQL file by comparing empty state to current db state
|
||||
logger.info("Generating baseline migration...")
|
||||
migration_file = init_dir / "migration.sql"
|
||||
subprocess.run(
|
||||
prisma_toolchain.run_prisma(
|
||||
[
|
||||
_get_prisma_command(),
|
||||
"migrate",
|
||||
|
|
@ -209,14 +210,13 @@ class ProxyExtrasDBManager:
|
|||
"--script",
|
||||
],
|
||||
stdout=open(migration_file, "w"),
|
||||
check=True,
|
||||
timeout=prisma_command_timeout(),
|
||||
env=prisma_env,
|
||||
)
|
||||
|
||||
# 3. Mark the migration as applied since it represents current state
|
||||
logger.info("Marking baseline migration as applied...")
|
||||
subprocess.run(
|
||||
prisma_toolchain.run_prisma(
|
||||
[
|
||||
_get_prisma_command(),
|
||||
"migrate",
|
||||
|
|
@ -224,7 +224,6 @@ class ProxyExtrasDBManager:
|
|||
"--applied",
|
||||
"0_init",
|
||||
],
|
||||
check=True,
|
||||
timeout=prisma_command_timeout(),
|
||||
env=prisma_env,
|
||||
)
|
||||
|
|
@ -253,7 +252,7 @@ class ProxyExtrasDBManager:
|
|||
"""Mark a specific migration as rolled back"""
|
||||
# Set up environment for offline mode if configured
|
||||
prisma_env = _get_prisma_env()
|
||||
subprocess.run(
|
||||
prisma_toolchain.run_prisma(
|
||||
[
|
||||
_get_prisma_command(),
|
||||
"migrate",
|
||||
|
|
@ -262,8 +261,6 @@ class ProxyExtrasDBManager:
|
|||
migration_name,
|
||||
],
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
env=prisma_env,
|
||||
)
|
||||
|
||||
|
|
@ -315,11 +312,9 @@ class ProxyExtrasDBManager:
|
|||
def _resolve_specific_migration(migration_name: str):
|
||||
"""Mark a specific migration as applied"""
|
||||
prisma_env = _get_prisma_env()
|
||||
subprocess.run(
|
||||
prisma_toolchain.run_prisma(
|
||||
[_get_prisma_command(), "migrate", "resolve", "--applied", migration_name],
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
env=prisma_env,
|
||||
)
|
||||
|
||||
|
|
@ -403,7 +398,7 @@ class ProxyExtrasDBManager:
|
|||
try:
|
||||
logger.info("Generating migration diff between DB and schema.prisma...")
|
||||
with open(diff_sql_path, "w") as f:
|
||||
subprocess.run(
|
||||
prisma_toolchain.run_prisma(
|
||||
[
|
||||
_get_prisma_command(),
|
||||
"migrate",
|
||||
|
|
@ -414,7 +409,6 @@ class ProxyExtrasDBManager:
|
|||
schema_path,
|
||||
"--script",
|
||||
],
|
||||
check=True,
|
||||
timeout=prisma_command_timeout(),
|
||||
stdout=f,
|
||||
env=_get_prisma_env(),
|
||||
|
|
@ -437,7 +431,7 @@ class ProxyExtrasDBManager:
|
|||
migration_files = sorted(Path(migrations_dir).glob("*/migration.sql"))
|
||||
for mig_file in migration_files:
|
||||
try:
|
||||
subprocess.run(
|
||||
prisma_toolchain.run_prisma(
|
||||
[
|
||||
_get_prisma_command(),
|
||||
"db",
|
||||
|
|
@ -448,9 +442,6 @@ class ProxyExtrasDBManager:
|
|||
schema_path,
|
||||
],
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
logger.info(f"Applied migration: {mig_file.parent.name}")
|
||||
|
|
@ -483,7 +474,7 @@ class ProxyExtrasDBManager:
|
|||
applied_ok = False
|
||||
try:
|
||||
logger.info("Running prisma db execute to apply the migration diff...")
|
||||
result = subprocess.run(
|
||||
result = prisma_toolchain.run_prisma(
|
||||
[
|
||||
_get_prisma_command(),
|
||||
"db",
|
||||
|
|
@ -494,9 +485,6 @@ class ProxyExtrasDBManager:
|
|||
schema_path,
|
||||
],
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
logger.info(f"prisma db execute stdout: {result.stdout}")
|
||||
|
|
@ -525,7 +513,7 @@ class ProxyExtrasDBManager:
|
|||
for migration_name in migration_names:
|
||||
try:
|
||||
logger.info(f"Resolving migration: {migration_name}")
|
||||
subprocess.run(
|
||||
prisma_toolchain.run_prisma(
|
||||
[
|
||||
_get_prisma_command(),
|
||||
"migrate",
|
||||
|
|
@ -534,9 +522,6 @@ class ProxyExtrasDBManager:
|
|||
migration_name,
|
||||
],
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
logger.debug(f"Resolved migration: {migration_name}")
|
||||
|
|
@ -726,11 +711,12 @@ class ProxyExtrasDBManager:
|
|||
original_dir = os.getcwd()
|
||||
os.chdir(migrations_dir)
|
||||
try:
|
||||
subprocess.run(
|
||||
prisma_toolchain.run_prisma(
|
||||
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
env=_get_prisma_env(),
|
||||
stdout=None,
|
||||
stderr=None,
|
||||
)
|
||||
return True
|
||||
except (
|
||||
|
|
@ -752,12 +738,9 @@ class ProxyExtrasDBManager:
|
|||
try:
|
||||
for attempt in range(4):
|
||||
try:
|
||||
result = subprocess.run(
|
||||
result = prisma_toolchain.run_prisma(
|
||||
[_get_prisma_command(), "migrate", "deploy"],
|
||||
timeout=deploy_timeout,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
logger.info(f"prisma migrate deploy stdout: {result.stdout}")
|
||||
|
|
@ -1007,12 +990,9 @@ class ProxyExtrasDBManager:
|
|||
logger.info("Running prisma migrate deploy")
|
||||
try:
|
||||
# Set migrations directory for Prisma
|
||||
result = subprocess.run(
|
||||
result = prisma_toolchain.run_prisma(
|
||||
[_get_prisma_command(), "migrate", "deploy"],
|
||||
timeout=prisma_migrate_deploy_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
logger.info(f"prisma migrate deploy stdout: {result.stdout}")
|
||||
|
|
@ -1084,7 +1064,7 @@ class ProxyExtrasDBManager:
|
|||
f"Found failed migration: {failed_migration}, marking as rolled back"
|
||||
)
|
||||
# Mark the failed migration as rolled back
|
||||
subprocess.run(
|
||||
prisma_toolchain.run_prisma(
|
||||
[
|
||||
_get_prisma_command(),
|
||||
"migrate",
|
||||
|
|
@ -1093,9 +1073,6 @@ class ProxyExtrasDBManager:
|
|||
failed_migration,
|
||||
],
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
logger.info(
|
||||
|
|
@ -1220,10 +1197,12 @@ class ProxyExtrasDBManager:
|
|||
if ProxyExtrasDBManager.spend_logs_is_partitioned():
|
||||
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
|
||||
# Use prisma db push with increased timeout
|
||||
subprocess.run(
|
||||
prisma_toolchain.run_prisma(
|
||||
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
stdout=None,
|
||||
stderr=None,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
return True
|
||||
except subprocess.TimeoutExpired:
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
|
|||
"Error: P3018\nMigration name: 20250326162113_baseline\n"
|
||||
"Database error code: 42501\npermission denied for schema public"
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="permission"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
|
@ -60,7 +60,7 @@ def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
|
|||
"Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n"
|
||||
'Reason: syntax error at or near "BRKN" LINE 42'
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
|
@ -124,7 +124,7 @@ def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path):
|
|||
def fake_resolve(*args, **kwargs):
|
||||
resolve_called["n"] += 1
|
||||
|
||||
monkeypatch.setattr("subprocess.run", fake_run)
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", fake_run)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve)
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True) # v2 flag NOT set
|
||||
|
|
@ -139,7 +139,7 @@ def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_pat
|
|||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
stderr = "db push error"
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="prisma db push failed"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
|
||||
|
||||
|
|
@ -209,7 +209,7 @@ def test_v2_resolve_specific_migration_failure_raises_runtime_error(
|
|||
"Error: P3009\nMigration `20260101000000_some_migration` failed\n"
|
||||
"relation already exists"
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(
|
||||
RuntimeError, match="Failed to mark migration .* as applied"
|
||||
):
|
||||
|
|
@ -228,7 +228,7 @@ def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
|
|||
stdout = "Applied migration.\n"
|
||||
stderr = ""
|
||||
|
||||
monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult())
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", lambda *a, **kw: FakeResult())
|
||||
|
||||
resolve_called = {"n": 0}
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -296,7 +296,7 @@ def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path):
|
|||
"_resolve_specific_migration",
|
||||
lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
|
||||
)
|
||||
monkeypatch.setattr("subprocess.run", _succeed_after(1, _DEADLOCK_P3018_STDERR))
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, _DEADLOCK_P3018_STDERR))
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
|
|
@ -309,7 +309,7 @@ def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path):
|
|||
monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda name: None)
|
||||
|
||||
with patch(
|
||||
"subprocess.run",
|
||||
"litellm_proxy_extras.prisma_toolchain.run_prisma",
|
||||
side_effect=_fake_migrate_deploy_failure(1, _DEADLOCK_P3018_STDERR),
|
||||
):
|
||||
with pytest.raises(RuntimeError, match="after 4 attempts"):
|
||||
|
|
@ -343,7 +343,7 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_
|
|||
"_resolve_specific_migration",
|
||||
lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
|
||||
)
|
||||
monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr))
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
|
|
@ -372,7 +372,7 @@ def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path
|
|||
"_resolve_specific_migration",
|
||||
lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
|
||||
)
|
||||
monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr))
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
|
|
@ -395,7 +395,7 @@ def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path):
|
|||
"_roll_back_migration",
|
||||
lambda name: pytest.fail("an unreadable ledger must not trigger a retry"),
|
||||
)
|
||||
monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr))
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
|
||||
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
|
@ -417,7 +417,7 @@ def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path):
|
|||
lambda name: 'ERROR: syntax error at or near "BRKN"',
|
||||
)
|
||||
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
|
@ -427,7 +427,7 @@ def test_v2_bare_deadlock_stderr_retries(monkeypatch, tmp_path):
|
|||
waiter as victim) is retried, not fatal."""
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
monkeypatch.setattr(
|
||||
"subprocess.run", _succeed_after(1, "Database error: deadlock detected")
|
||||
"litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, "Database error: deadlock detected")
|
||||
)
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
|
@ -446,7 +446,7 @@ def test_v2_advisory_lock_timeout_retries(monkeypatch, tmp_path):
|
|||
"""v2: the advisory-lock waiter that times out while a peer's retry holds
|
||||
the lock retries instead of dying."""
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
monkeypatch.setattr("subprocess.run", _succeed_after(2, _P1002_ADVISORY_LOCK_STDERR))
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(2, _P1002_ADVISORY_LOCK_STDERR))
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
|
|
@ -456,7 +456,7 @@ def test_v2_p1002_without_advisory_lock_context_still_raises(monkeypatch, tmp_pa
|
|||
"""v2: a plain P1002 (database unreachable) stays fatal."""
|
||||
_stub_v2_env(monkeypatch, tmp_path)
|
||||
stderr = "Error: P1002\n\nThe database server at `db`:`5432` was reached but timed out."
|
||||
monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr))
|
||||
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
|
||||
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
|
|
|||
|
|
@ -567,7 +567,7 @@ class TestResolveAllMigrationsLedger:
|
|||
return _FakeCompleted()
|
||||
return _FakeCompleted()
|
||||
|
||||
monkeypatch.setattr(utils_module.subprocess, "run", fake_run)
|
||||
monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", fake_run)
|
||||
ProxyExtrasDBManager._resolve_all_migrations(str(tmp_path), "schema.prisma")
|
||||
return calls
|
||||
|
||||
|
|
@ -604,9 +604,9 @@ class TestPartitionedSpendLogsPushGuard:
|
|||
import litellm_proxy_extras.utils as utils_module
|
||||
|
||||
def fail_run(cmd, **kwargs):
|
||||
raise AssertionError(f"subprocess.run should not be called, got: {cmd}")
|
||||
raise AssertionError(f"run_prisma should not be called, got: {cmd}")
|
||||
|
||||
monkeypatch.setattr(utils_module.subprocess, "run", fail_run)
|
||||
monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", fail_run)
|
||||
|
||||
def test_v1_db_push_fails_fast_with_guidance(self, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ import ast
|
|||
import json
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
|
|
@ -47,6 +48,7 @@ FAKE_PRISMA = """#!{python}
|
|||
import json
|
||||
import os
|
||||
import pathlib
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
|
|
@ -66,6 +68,9 @@ with log_path.open("a") as log:
|
|||
time.sleep(float(os.environ.get("FAKE_PRISMA_SLEEP", "0")))
|
||||
if args[:2] == ["migrate", "deploy"]:
|
||||
if earlier_same_command == 0:
|
||||
if os.environ.get("FAKE_PRISMA_GRANDCHILD_PIDFILE"):
|
||||
grandchild = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(600)"])
|
||||
pathlib.Path(os.environ["FAKE_PRISMA_GRANDCHILD_PIDFILE"]).write_text(str(grandchild.pid))
|
||||
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)
|
||||
|
|
@ -272,6 +277,43 @@ def test_migrate_deploy_stops_at_its_own_timeout(
|
|||
assert elapsed < 30
|
||||
|
||||
|
||||
def _process_is_gone(pid: int, within_seconds: float) -> bool:
|
||||
deadline = time.monotonic() + within_seconds
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
except ProcessLookupError:
|
||||
return True
|
||||
time.sleep(0.05)
|
||||
return False
|
||||
|
||||
|
||||
def test_a_timed_out_migrate_deploy_takes_its_process_tree_with_it(
|
||||
toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
"""The real CLI forks Node and a schema engine; a timeout must not leave them running."""
|
||||
_, log_path = toolchain_env
|
||||
pidfile = tmp_path / "grandchild.pid"
|
||||
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")
|
||||
monkeypatch.setenv("FAKE_PRISMA_GRANDCHILD_PIDFILE", str(pidfile))
|
||||
|
||||
with pytest.raises(RuntimeError, match="insufficient permissions"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
grandchild_pid = int(pidfile.read_text())
|
||||
try:
|
||||
assert len(_deploy_calls(log_path)) == 2
|
||||
assert _process_is_gone(grandchild_pid, within_seconds=5)
|
||||
finally:
|
||||
try:
|
||||
os.kill(grandchild_pid, signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
|
||||
|
||||
def test_db_push_timeout_hint_names_the_per_command_budget(
|
||||
toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ def test_hands_the_alter_statement_to_the_prisma_cli():
|
|||
return subprocess.CompletedProcess(cmd, 0)
|
||||
|
||||
with patch(
|
||||
"litellm_proxy_extras.replica_identity.subprocess.run", side_effect=capture
|
||||
"litellm_proxy_extras.replica_identity.run_prisma", side_effect=capture
|
||||
):
|
||||
applied = apply_replica_identity_full(
|
||||
schema_path="/somewhere/schema.prisma",
|
||||
|
|
@ -60,7 +60,7 @@ def test_hands_the_alter_statement_to_the_prisma_cli():
|
|||
)
|
||||
def test_every_failure_is_reported_instead_of_raised(failure):
|
||||
with patch(
|
||||
"litellm_proxy_extras.replica_identity.subprocess.run", side_effect=failure
|
||||
"litellm_proxy_extras.replica_identity.run_prisma", side_effect=failure
|
||||
):
|
||||
assert (
|
||||
apply_replica_identity_full(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue