fix(proxy): skip prisma generate in the migration job when the prisma package is not writable

prisma_migration.py runs `prisma generate` after the migrations are done. The
generator writes the client into the installed prisma package, so with a
read-only root filesystem, or a uid that can't write site-packages, it always
fails and the failure is only logged. Its subprocesses still run alongside the
parent that holds the imported proxy, which roughly doubles the job's peak
memory (1156 MiB vs ~630 MiB on v1.105.1) and OOMKills the Helm migration Job
at a 1Gi limit after the migrations already succeeded

main() now checks that every prisma package directory is writable before
running the generator and logs the skip otherwise. Nothing changes where the
package is writable

Fixes #47076

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,45 @@ 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
+        )