diff --git a/litellm/proxy/prisma_migration.py b/litellm/proxy/prisma_migration.py index 3561f190808..75b5116cac6 100644 --- a/litellm/proxy/prisma_migration.py +++ b/litellm/proxy/prisma_migration.py @@ -3,12 +3,16 @@ A failed migration fails the entrypoint, the same way it fails proxy startup. A failed 'prisma generate' is log-only: every shipped image bakes the client at build time, and refreshing it writes into site-packages, which an arbitrary non-root uid or a read-only -root filesystem cannot do. +root filesystem cannot do. When the prisma package is not writable it is skipped instead, +because its subprocesses roughly double the job's peak memory for a step that cannot succeed. """ +import importlib.util import os import subprocess import sys +from collections.abc import Callable +from importlib.machinery import ModuleSpec sys.path.insert(0, os.path.abspath("./")) @@ -20,9 +24,29 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy.proxy_cli import run_server -def main() -> int: +def _is_writable_dir(path: str) -> bool: + return os.access(path, os.W_OK) + + +def prisma_package_is_writable( + find_spec: Callable[[str], ModuleSpec | None] = importlib.util.find_spec, + is_writable: Callable[[str], bool] = _is_writable_dir, +) -> bool: + """'prisma generate' writes the client into the installed prisma package, so it can only succeed here.""" + spec: Final = find_spec("prisma") + locations: Final = tuple(spec.submodule_search_locations or ()) if spec is not None else () + return bool(locations) and all(is_writable(location) for location in locations) + + +def main(can_generate: Callable[[], bool] = prisma_package_is_writable) -> int: run_server(("--skip_server_startup",), standalone_mode=False) + if not can_generate(): + verbose_proxy_logger.info( + "Skipping 'prisma generate' because the prisma package is not writable, using the client baked at image build time" + ) + return 0 + verbose_proxy_logger.info("Running 'prisma generate'...") result: Final = subprocess.run(resolve_prisma_argv(("prisma", "generate")), capture_output=True, text=True) verbose_proxy_logger.info("'prisma generate' stdout: %s", result.stdout) diff --git a/tests/unit/proxy/test_prisma_migration.py b/tests/unit/proxy/test_prisma_migration.py index b7de849b3b4..5fbaed7eb1b 100644 --- a/tests/unit/proxy/test_prisma_migration.py +++ b/tests/unit/proxy/test_prisma_migration.py @@ -1,5 +1,6 @@ import os import sys +from importlib.machinery import ModuleSpec from pathlib import Path from unittest.mock import MagicMock, patch @@ -18,7 +19,7 @@ class TestPrismaMigration: mock_subprocess_run.return_value = MagicMock(returncode=0, stdout="", stderr="") with patch.dict(os.environ, env, clear=True): - assert prisma_migration.main() == 0 + assert prisma_migration.main(can_generate=lambda: True) == 0 mock_run_server.assert_called_once_with(("--skip_server_startup",), standalone_mode=False) @@ -34,7 +35,7 @@ class TestPrismaMigration: ) with patch.dict(os.environ, {}, clear=True): - assert prisma_migration.main() == 0 + assert prisma_migration.main(can_generate=lambda: True) == 0 @patch("litellm.proxy.prisma_migration.subprocess.run") @patch("litellm.proxy.prisma_migration.run_server") @@ -59,7 +60,7 @@ class TestPrismaMigration: empty_bin.mkdir() with patch.dict(os.environ, {"PATH": str(empty_bin)}, clear=True): - assert prisma_migration.main() == 0 + assert prisma_migration.main(can_generate=lambda: True) == 0 assert mock_subprocess_run.call_args.args[0] == (sys.executable, "-m", "prisma", "generate") @@ -76,6 +77,56 @@ class TestPrismaMigration: script.chmod(0o755) with patch.dict(os.environ, {"PATH": str(bin_dir)}, clear=True): - assert prisma_migration.main() == 0 + assert prisma_migration.main(can_generate=lambda: True) == 0 assert mock_subprocess_run.call_args.args[0] == ("prisma", "generate") + + @patch("litellm.proxy.prisma_migration.subprocess.run") # test-quality-ok: the spawned argv is the behavior under test + @patch("litellm.proxy.prisma_migration.run_server") # test-quality-ok: run_server boots the whole proxy + def test_prisma_generate_is_skipped_when_the_prisma_package_is_not_writable( + self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock + ) -> None: + assert prisma_migration.main(can_generate=lambda: False) == 0 + + mock_run_server.assert_called_once_with(("--skip_server_startup",), standalone_mode=False) + mock_subprocess_run.assert_not_called() + + @pytest.mark.parametrize( + ("locations", "writable", "expected"), + [ + (("/site-packages/prisma",), frozenset({"/site-packages/prisma"}), True), + (("/site-packages/prisma",), frozenset[str](), False), + (("/a/prisma", "/b/prisma"), frozenset({"/a/prisma"}), False), + ((), frozenset[str](), False), + ], + ids=("writable", "read-only", "one-of-two-read-only", "no-package-dir"), + ) + def test_prisma_package_is_writable_only_when_every_package_dir_is( + self, locations: tuple[str, ...], writable: frozenset[str], expected: bool + ) -> None: + spec = ModuleSpec("prisma", None, is_package=True) + spec.submodule_search_locations = list(locations) + + assert ( + prisma_migration.prisma_package_is_writable( + find_spec=lambda _name: spec, is_writable=lambda path: path in writable + ) + is expected + ) + + def test_prisma_package_is_not_writable_when_prisma_is_not_installed(self) -> None: + assert ( + prisma_migration.prisma_package_is_writable(find_spec=lambda _name: None, is_writable=lambda _path: True) + is False + ) + + def test_prisma_package_is_writable_checks_the_real_filesystem_by_default(self, tmp_path: Path) -> None: + missing_dir: Path = tmp_path / "missing" + + def spec_at(location: Path) -> ModuleSpec: + spec = ModuleSpec("prisma", None, is_package=True) + spec.submodule_search_locations = [str(location)] + return spec + + assert prisma_migration.prisma_package_is_writable(find_spec=lambda _name: spec_at(tmp_path)) is True + assert prisma_migration.prisma_package_is_writable(find_spec=lambda _name: spec_at(missing_dir)) is False