Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
177 changes: 177 additions & 0 deletions litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,177 @@
"""Prepare the Node toolchain the Prisma CLI needs, separately from migrations.

The Prisma CLI is a Node program. The first invocation inside a fresh
container installs a private Node runtime and npm-installs the CLI itself,
which can take minutes on a cold or slow machine. Sharing one timeout between
that one-time bootstrap and the migration commands makes a slow bootstrap
indistinguishable from a slow migration, so the bootstrap gets killed long
before it can finish.

A killed bootstrap does not correct itself. The installer leaves its cache
directory behind, and Prisma decides whether to install by testing that
directory for existence alone, so every later attempt skips the install and
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.
"""

import math
import os
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from litellm_proxy_extras._logging import logger

try:
from prisma import config as prisma_config
except ImportError:
prisma_config = None

PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT"
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT"
NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR"

DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0
DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0

BOOTSTRAP_ARG = "--version"


@dataclass(frozen=True)
class ToolchainBootstrap:
"""Outcome of preparing the Prisma toolchain."""

healed_incomplete_cache: bool
ready: bool


def _timeout_from_env(env_var: str, default: float) -> float:
raw = os.getenv(env_var)
if raw is None:
return default
try:
seconds = float(raw)
except ValueError:
logger.warning(
"%s=%r is not a number, falling back to %ss", env_var, raw, default
)
return default
if not math.isfinite(seconds) or seconds <= 0:
logger.warning(
"%s=%r is not a finite positive number, falling back to %ss",
env_var,
raw,
default,
)
return default
return seconds


def prisma_command_timeout() -> float:
"""Seconds any single Prisma command may run for."""
return _timeout_from_env(
PRISMA_COMMAND_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_COMMAND_TIMEOUT
)


def prisma_bootstrap_timeout() -> float:
"""Seconds the one-time Node toolchain install may run for."""
return _timeout_from_env(
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT
)


def nodeenv_cache_dir() -> Optional[Path]:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Optional annotation uses legacy syntax

The new nodeenv_cache_dir API uses Optional[Path] instead of the repository-standard Path | None, adding an avoidable style inconsistency and unnecessary Optional import.

Rule Used: In this repo, prefer str | None (PEP 604 union s... (source)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

"""Where Prisma installs its private Node runtime, or None if unknowable."""
override = os.getenv(NODEENV_CACHE_DIR_ENV_VAR)
if override:
return Path(override).absolute()
if prisma_config is not None:
try:
return Path(prisma_config.nodeenv_cache_dir).absolute()
except (OSError, ValueError) as e:
logger.warning("Could not read the Prisma nodeenv cache dir: %s", e)
try:
return Path.home() / ".cache" / "prisma-python" / "nodeenv"
except RuntimeError:
logger.warning(
"No resolvable home directory, cannot locate the Prisma nodeenv cache"
)
return None


def node_binary_path(cache_dir: Path) -> Path:
"""Path the Node binary occupies once the toolchain is fully installed."""
if os.name == "nt":
return cache_dir / "Scripts" / "node.exe"
return cache_dir / "bin" / "node"


def heal_incomplete_nodeenv_cache() -> bool:
"""Delete a nodeenv cache directory left without a Node binary.

Returns True when a half-installed toolchain was removed, so the next
Prisma invocation reinstalls it instead of failing on a missing binary.
"""
cache_dir = nodeenv_cache_dir()
if cache_dir is None or not cache_dir.is_dir():
return False
if node_binary_path(cache_dir).exists():
return False
logger.warning(
"Node toolchain at %s has no %s, so a previous install was interrupted. "
"Removing it so it can be reinstalled.",
cache_dir,
node_binary_path(cache_dir).name,
)
try:
shutil.rmtree(cache_dir)
except OSError as e:
logger.warning("Could not remove %s: %s", cache_dir, e)
return False
return True


def ensure_prisma_toolchain(
prisma_command: str, prisma_env: dict[str, str]
) -> ToolchainBootstrap:
"""Install whatever the Prisma CLI needs to run, under its own timeout.

Never raises. A toolchain that cannot be prepared is reported so the
caller can go on and let the real Prisma command produce the real error.
"""
healed = heal_incomplete_nodeenv_cache()
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,
)
except subprocess.TimeoutExpired:
logger.warning(
"Preparing the Prisma CLI toolchain timed out after %ss. Raise %s "
"if this machine needs longer to install it.",
timeout,
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR,
)
return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False)
except subprocess.CalledProcessError as e:
logger.warning("Preparing the Prisma CLI toolchain failed: %s", e.stderr)
return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False)
except OSError as e:
logger.warning("Could not run the Prisma CLI: %s", e)
return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False)
logger.info("Prisma CLI toolchain ready")
return ToolchainBootstrap(healed_incomplete_cache=healed, ready=True)
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
from pathlib import Path

from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.prisma_toolchain import prisma_command_timeout

REPLICA_IDENTITY_FULL_ENV_VAR = "LITELLM_SET_REPLICA_IDENTITY_FULL"

Expand Down Expand Up @@ -75,7 +76,7 @@ def apply_replica_identity_full(
"--schema",
schema_path,
],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
Expand Down
33 changes: 20 additions & 13 deletions litellm-proxy-extras/litellm_proxy_extras/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,10 @@
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
from litellm_proxy_extras.prisma_toolchain import (
ensure_prisma_toolchain,
prisma_command_timeout,
)


def str_to_bool(value: Optional[str]) -> bool:
Expand Down Expand Up @@ -142,7 +146,7 @@ def _create_baseline_migration(schema_path: str) -> bool:
],
stdout=open(migration_file, "w"),
check=True,
timeout=30,
timeout=prisma_command_timeout(),
env=prisma_env,
)

Expand All @@ -157,7 +161,7 @@ def _create_baseline_migration(schema_path: str) -> bool:
"0_init",
],
check=True,
timeout=30,
timeout=prisma_command_timeout(),
env=prisma_env,
)

Expand Down Expand Up @@ -193,7 +197,7 @@ def _roll_back_migration(migration_name: str):
"--rolled-back",
migration_name,
],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
env=prisma_env,
Expand All @@ -205,7 +209,7 @@ def _resolve_specific_migration(migration_name: str):
prisma_env = _get_prisma_env()
subprocess.run(
[_get_prisma_command(), "migrate", "resolve", "--applied", migration_name],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
env=prisma_env,
Expand Down Expand Up @@ -303,7 +307,7 @@ def _resolve_all_migrations(
"--script",
],
check=True,
timeout=60,
timeout=prisma_command_timeout(),
stdout=f,
env=_get_prisma_env(),
)
Expand Down Expand Up @@ -335,7 +339,7 @@ def _resolve_all_migrations(
"--schema",
schema_path,
],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
Expand Down Expand Up @@ -364,7 +368,7 @@ def _resolve_all_migrations(
"--schema",
schema_path,
],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
Expand Down Expand Up @@ -393,7 +397,7 @@ def _resolve_all_migrations(
"--applied",
migration_name,
],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
Expand Down Expand Up @@ -530,7 +534,7 @@ def _setup_database_v2(use_migrate: bool) -> bool:
try:
subprocess.run(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
env=_get_prisma_env(),
)
Expand All @@ -555,7 +559,7 @@ def _setup_database_v2(use_migrate: bool) -> bool:
try:
result = subprocess.run(
[_get_prisma_command(), "migrate", "deploy"],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
Expand Down Expand Up @@ -731,6 +735,9 @@ def setup_database(
Returns:
bool: True if setup was successful, False otherwise
"""
ensure_prisma_toolchain(
prisma_command=_get_prisma_command(), prisma_env=_get_prisma_env()
)
migrated = ProxyExtrasDBManager._run_migrations(
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
)
Expand All @@ -757,7 +764,7 @@ def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool:
# Set migrations directory for Prisma
result = subprocess.run(
[_get_prisma_command(), "migrate", "deploy"],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
Expand Down Expand Up @@ -840,7 +847,7 @@ def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool:
"--rolled-back",
failed_migration,
],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
Expand Down Expand Up @@ -968,7 +975,7 @@ def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool:
# Use prisma db push with increased timeout
subprocess.run(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
)
return True
Expand Down
Loading
Loading