Merge e57d0e692b into dddbaa803d

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)

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