ci: benchmark and gate an installed release wheel

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

.github/actions/cache-cargo-build/action.yml

@@ -15,6 +15,12 @@ description: >-
   cache the same directory for different workloads, and a shared key would let
   whichever ran first deny the others a save.
 
+inputs:
+  profile:
+    description: "Cargo profile the build uses (dev or release)"
+    required: false
+    default: "dev"
+
 runs:
   using: composite
   steps:
@@ -25,6 +31,6 @@ runs:
           ~/.cargo/registry
           ~/.cargo/git
           litellm-rust/target
-        key: ${{ runner.os }}-maturin-dev-${{ hashFiles('litellm-rust/Cargo.lock') }}
+        key: ${{ runner.os }}-maturin-${{ inputs.profile }}-${{ hashFiles('litellm-rust/Cargo.lock') }}
         restore-keys: |
-          ${{ runner.os }}-maturin-dev-
+          ${{ runner.os }}-maturin-${{ inputs.profile }}-

.github/scripts/verify_installed_wheel_imports.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import importlib.metadata
+import json
+import sys
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final, cast
+
+
+def _direct_url_failures(direct_url: str | None) -> tuple[str, ...]:
+    if direct_url is None:
+        return ()
+    try:
+        direct_url_data: Final = cast(object, json.loads(direct_url))
+    except json.JSONDecodeError:
+        return ("litellm distribution direct_url.json is invalid JSON",)
+    if not isinstance(direct_url_data, Mapping):
+        return ()
+    direct_url_mapping: Final = cast(Mapping[str, object], direct_url_data)
+    dir_info: Final = direct_url_mapping.get("dir_info")
+    if not isinstance(dir_info, Mapping):
+        return ()
+    dir_info_mapping: Final = cast(Mapping[str, object], dir_info)
+    return ("litellm distribution is editable",) if dir_info_mapping.get("editable") is True else ()
+
+
+def _native_bridge_available() -> bool:
+    try:
+        from litellm.rust_bridge import native_bridge_available
+    except Exception:
+        return False
+    return native_bridge_available()
+
+
+def provenance_failures(
+    *,
+    prefix: Path,
+    checkout: Path,
+    module_files: Mapping[str, Path],
+    direct_url: str | None,
+    native_available: bool,
+) -> tuple[str, ...]:
+    resolved_prefix: Final = prefix.resolve()
+    resolved_checkout: Final = checkout.resolve()
+    resolved_module_files: Final = tuple(
+        (module_name, module_path.resolve()) for module_name, module_path in module_files.items()
+    )
+    prefix_failures: Final = tuple(
+        f"{module_name} is not under sys.prefix: {module_path}"
+        for module_name, module_path in resolved_module_files
+        if not module_path.is_relative_to(resolved_prefix)
+    )
+    checkout_failures: Final = tuple(
+        f"{module_name} is under the checkout: {module_path}"
+        for module_name, module_path in resolved_module_files
+        if module_path.is_relative_to(resolved_checkout)
+    )
+    direct_url_failures: Final = _direct_url_failures(direct_url)
+    native_failures: Final = ("litellm.rust_bridge.native_bridge_available() is false",) if not native_available else ()
+    return prefix_failures + checkout_failures + direct_url_failures + native_failures
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    arguments: Final = tuple(sys.argv if argv is None else argv)
+    if len(arguments) != 2:
+        sys.stderr.write(f"usage: {Path(arguments[0]).name} CHECKOUT\n")
+        return 2
+
+    checkout: Final = Path(arguments[1])
+    try:
+        import litellm
+    except Exception as error:
+        sys.stderr.write(f"litellm import failed: {error}\n")
+        return 1
+
+    try:
+        import litellm.rust_bridge._native as native
+    except Exception as error:
+        sys.stderr.write(f"litellm.rust_bridge._native import failed: {error}\n")
+        return 1
+
+    imported_modules: Final = (
+        ("litellm", litellm),
+        ("litellm.rust_bridge._native", native),
+    )
+    missing_module_files: Final = tuple(
+        f"{module_name} has no __file__" for module_name, module in imported_modules if module.__file__ is None
+    )
+    module_files: Final = MappingProxyType(
+        {
+            module_name: Path(module.__file__).resolve()
+            for module_name, module in imported_modules
+            if module.__file__ is not None
+        }
+    )
+    try:
+        direct_url: Final = importlib.metadata.distribution("litellm").read_text("direct_url.json")
+    except Exception as error:
+        sys.stderr.write(f"litellm distribution lookup failed: {error}\n")
+        return 1
+
+    native_available: Final = _native_bridge_available()
+
+    failures: Final = missing_module_files + provenance_failures(
+        prefix=Path(sys.prefix),
+        checkout=checkout,
+        module_files=module_files,
+        direct_url=direct_url,
+        native_available=native_available,
+    )
+    if failures:
+        sys.stderr.write("".join(f"{failure}\n" for failure in failures))
+        return 1
+
+    sys.stdout.write("".join(f"{module_name}: {module_path}\n" for module_name, module_path in module_files.items()))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

.github/workflows/codspeed.yml

@@ -13,6 +13,9 @@ on:
       - ".github/workflows/codspeed.yml"
       - ".github/actions/setup-uv-with-retries/**"
       - ".github/actions/cache-cargo-build/**"
+      - ".github/scripts/verify_installed_wheel_imports.py"
+      - ".github/scripts/verify_linux_native_wheel.py"
+      - ".github/scripts/uv_sync_with_retries.sh"
   pull_request:
     branches:
       - main
@@ -25,6 +28,9 @@ on:
       - ".github/workflows/codspeed.yml"
       - ".github/actions/setup-uv-with-retries/**"
       - ".github/actions/cache-cargo-build/**"
+      - ".github/scripts/verify_installed_wheel_imports.py"
+      - ".github/scripts/verify_linux_native_wheel.py"
+      - ".github/scripts/uv_sync_with_retries.sh"
   # Allow CodSpeed to trigger backtest performance analysis
   # in order to generate initial data
   workflow_dispatch:
@@ -59,19 +65,33 @@ jobs:
 
       - name: Cache the Rust build
         uses: ./.github/actions/cache-cargo-build
+        with:
+          profile: release
 
       # Build the wheel and resolve every dependency outside the CodSpeed
       # runner: the same maturin build took 42 minutes inside `codspeed run`
       # versus under 3 minutes as a plain step (LIT-6183)
-      - name: Build environment
+      - name: Build the release wheel
+        run: uv build --wheel --out-dir dist
+
+      - name: Verify the release wheel
+        run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
+        env:
+          RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
+
+      - name: Install the wheel into the benchmark environment
+        run: |
+          UV_PROJECT_ENVIRONMENT="${RUNNER_TEMP}/benchmark-venv" .github/scripts/uv_sync_with_retries.sh --frozen --no-default-groups --group benchmarks --no-install-project --python 3.12
+          uv pip install --python "${RUNNER_TEMP}/benchmark-venv/bin/python" --no-deps dist/*.whl
+          (cd "${RUNNER_TEMP}" && "${RUNNER_TEMP}/benchmark-venv/bin/python" -I "${GITHUB_WORKSPACE}/.github/scripts/verify_installed_wheel_imports.py" "${GITHUB_WORKSPACE}")
+
+      - name: Collect benchmarks
+        env:
+          PYTEST_DISABLE_PLUGIN_AUTOLOAD: "1"
+          LITELLM_REQUIRE_INSTALLED_WHEEL: "1"
         run: >
-          env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1
-          uv run --frozen --no-default-groups
-          --with pytest==8.3.5
-          --with pytest-codspeed==4.3.0
-          --with "mcp>=2.2.0,<3.0"
-          --with "a2a-sdk>=1.1.0,<2.0"
-          pytest
+          "${RUNNER_TEMP}/benchmark-venv/bin/python" -I -m pytest
+          --import-mode=importlib
           -p pytest_codspeed.plugin
           tests/benchmarks/
           --codspeed
@@ -82,13 +102,9 @@ jobs:
         with:
           mode: simulation
           run: >
-            env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1
-            uv run --frozen --no-default-groups
-            --with pytest==8.3.5
-            --with pytest-codspeed==4.3.0
-            --with "mcp>=2.2.0,<3.0"
-            --with "a2a-sdk>=1.1.0,<2.0"
-            pytest
+            env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 LITELLM_REQUIRE_INSTALLED_WHEEL=1
+            "${RUNNER_TEMP}/benchmark-venv/bin/python" -I -m pytest
+            --import-mode=importlib
             -p pytest_codspeed.plugin
             tests/benchmarks/
             --codspeed

.github/workflows/test-rust.yml

@@ -24,7 +24,9 @@ on:
       - ".github/actions/setup-uv-with-retries/**"
       - ".github/scripts/smoke_test_native_wheel.py"
       - ".github/scripts/verify_linux_native_wheel.py"
+      - ".github/scripts/verify_installed_wheel_imports.py"
       - "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
+      - "tests/test_litellm/rust_bridge/test_verify_installed_wheel_imports.py"
       - ".github/workflows/test-rust.yml"
   pull_request:
     branches:
@@ -54,7 +56,9 @@ on:
       - ".github/actions/setup-uv-with-retries/**"
       - ".github/scripts/smoke_test_native_wheel.py"
       - ".github/scripts/verify_linux_native_wheel.py"
+      - ".github/scripts/verify_installed_wheel_imports.py"
       - "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
+      - "tests/test_litellm/rust_bridge/test_verify_installed_wheel_imports.py"
       - ".github/workflows/test-rust.yml"
 
 permissions:
@@ -159,7 +163,7 @@ jobs:
       - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
 
       - name: Run pytest tests/test_litellm_rust with the compiled extension
-        run: make test-rust-extension
+        run: make test-rust-extension WHEEL="$(echo dist/*.whl)"
 
       - run: >-
           uv build --wheel --out-dir panic-dist

Makefile

@@ -55,7 +55,7 @@ help:
 	@echo "  make test-proxy-unit-b  - Run proxy_unit_tests (p-z, ~28 files)"
 	@echo "  make test-integration   - Run integration tests"
 	@echo "  make test-unit-helm     - Run helm unit tests"
-	@echo "  make test-rust-extension - Build the Rust extension and run its public Python tests"
+	@echo "  make test-rust-extension - Build the Rust extension and run its public Python tests (WHEEL=path/to.whl reuses a prebuilt wheel)"
 	@echo ""
 	@echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide"
 	@echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine."
@@ -294,11 +294,13 @@ pre-commit:
 test-rust-extension:
 	@temporary=$$(mktemp -d) && \
 	trap 'rm -rf "$$temporary"' EXIT HUP INT TERM && \
-	$(UV) build --python 3.12 --wheel --out-dir "$$temporary/wheels" && \
+	mkdir -p "$$temporary/wheels" && \
+	if [ -n "$(WHEEL)" ]; then cp $(WHEEL) "$$temporary/wheels/"; else $(UV) build --python 3.12 --wheel --out-dir "$$temporary/wheels"; fi && \
 	set -- "$$temporary"/wheels/*.whl && \
 	[ "$$#" -eq 1 ] && \
 	UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \
 	$(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \
+	(cd "$$temporary" && "$$temporary/venv/bin/python" -I "$(CURDIR)/.github/scripts/verify_installed_wheel_imports.py" "$(CURDIR)") && \
 	"$$temporary/venv/bin/python" -I -m mypy.stubtest \
 		--mypy-config-file tests/test_litellm/rust_bridge/stubtest.ini \
 		litellm.rust_bridge._native && \

pyproject.toml

@@ -287,6 +287,12 @@ healthcheck = [
     "httpx==0.28.1",
     "pyyaml==6.0.3",
 ]
+benchmarks = [
+    "pytest==9.0.3",
+    "pytest-codspeed==4.3.0",
+    "mcp>=2.2.0,<3",
+    "a2a-sdk==1.1.0",
+]
 
 [build-system]
 requires = ["maturin==1.15.0"]

tests/benchmarks/conftest.py

@@ -8,9 +8,14 @@ flipping results between runs. Running the executor inline keeps each
 benchmark's cost self-contained and deterministic.
 """
 
-from collections.abc import Callable, Iterator
+import importlib.metadata
+import json
+import os
+import sys
+from collections.abc import Callable, Iterator, Mapping
 from concurrent.futures import Future
-from typing import ParamSpec, TypeVar
+from pathlib import Path
+from typing import Final, ParamSpec, TypeVar, cast
 
 import pytest
 
@@ -20,6 +25,45 @@ P = ParamSpec("P")
 R = TypeVar("R")
 
 
+def pytest_configure(config: pytest.Config) -> None:
+    if os.environ.get("LITELLM_REQUIRE_INSTALLED_WHEEL") != "1":
+        return
+
+    import litellm
+    import litellm.rust_bridge._native as native
+
+    prefix: Final = Path(sys.prefix).resolve()
+    module_paths: Final = (
+        ("litellm", Path(litellm.__file__).resolve()),
+        ("litellm.rust_bridge._native", Path(native.__file__).resolve()),
+    )
+    for module_name, module_path in module_paths:
+        sys.stdout.write(f"{module_name}: {module_path}\n")
+
+    path_failures: Final = tuple(
+        f"{module_name} is not under sys.prefix: {module_path}"
+        for module_name, module_path in module_paths
+        if not module_path.is_relative_to(prefix)
+    )
+    if path_failures:
+        raise pytest.UsageError("; ".join(path_failures))
+
+    direct_url: Final = importlib.metadata.distribution("litellm").read_text("direct_url.json")
+    if direct_url is None:
+        return
+
+    try:
+        direct_url_data: Final = cast(object, json.loads(direct_url))
+    except json.JSONDecodeError as error:
+        raise pytest.UsageError("litellm distribution direct_url.json is invalid JSON") from error
+
+    if isinstance(direct_url_data, Mapping):
+        direct_url_mapping: Final = cast(Mapping[str, object], direct_url_data)
+        dir_info: Final = direct_url_mapping.get("dir_info")
+        if isinstance(dir_info, Mapping) and cast(Mapping[str, object], dir_info).get("editable") is True:
+            raise pytest.UsageError("litellm distribution is editable")
+
+
 def _submit_inline(fn: Callable[P, R], /, *args: P.args, **kwargs: P.kwargs) -> Future[R]:
     future: Future[R] = Future()
     try:

tests/test_litellm/rust_bridge/test_verify_installed_wheel_imports.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import importlib.util
+import sys
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Final, Protocol, cast
+
+
+class _VerifierModule(Protocol):
+    def provenance_failures(
+        self,
+        *,
+        prefix: Path,
+        checkout: Path,
+        module_files: Mapping[str, Path],
+        direct_url: str | None,
+        native_available: bool,
+    ) -> tuple[str, ...]: ...
+
+
+_REPO_ROOT: Final = Path(__file__).resolve().parents[3]
+_MODULE_PATH: Final = _REPO_ROOT / ".github" / "scripts" / "verify_installed_wheel_imports.py"
+_SPEC: Final = importlib.util.spec_from_file_location("verify_installed_wheel_imports", _MODULE_PATH)
+assert _SPEC is not None and _SPEC.loader is not None
+_LOADED_VERIFIER: Final = importlib.util.module_from_spec(_SPEC)
+sys.modules[_SPEC.name] = _LOADED_VERIFIER
+_SPEC.loader.exec_module(_LOADED_VERIFIER)
+verifier: Final = cast(_VerifierModule, _LOADED_VERIFIER)
+
+
+def _failures(
+    *,
+    prefix: Path,
+    checkout: Path,
+    module_files: Mapping[str, Path],
+    direct_url: str | None = '{"dir_info": {"editable": false}}',
+    native_available: bool = True,
+) -> tuple[str, ...]:
+    return verifier.provenance_failures(
+        prefix=prefix,
+        checkout=checkout,
+        module_files=module_files,
+        direct_url=direct_url,
+        native_available=native_available,
+    )
+
+
+def test_accepts_installed_non_editable_modules(tmp_path: Path) -> None:
+    prefix: Final = tmp_path / "venv" / "lib" / "python3.12" / "site-packages"
+    checkout: Final = tmp_path / "checkout"
+    module_files: Final = {
+        "litellm": prefix / "litellm" / "__init__.py",
+        "litellm.rust_bridge._native": prefix / "litellm" / "rust_bridge" / "_native.abi3.so",
+    }
+
+    assert _failures(prefix=prefix, checkout=checkout, module_files=module_files) == ()
+
+
+def test_rejects_module_under_checkout(tmp_path: Path) -> None:
+    prefix: Final = tmp_path
+    checkout: Final = tmp_path / "checkout"
+    module_files: Final = {"litellm": checkout / "litellm" / "__init__.py"}
+
+    assert _failures(prefix=prefix, checkout=checkout, module_files=module_files) == (
+        f"litellm is under the checkout: {(checkout / 'litellm' / '__init__.py').resolve()}",
+    )
+
+
+def test_rejects_editable_direct_url(tmp_path: Path) -> None:
+    prefix: Final = tmp_path / "site-packages"
+    checkout: Final = tmp_path / "checkout"
+    module_files: Final = {"litellm": prefix / "litellm" / "__init__.py"}
+
+    assert _failures(
+        prefix=prefix,
+        checkout=checkout,
+        module_files=module_files,
+        direct_url='{"dir_info": {"editable": true}}',
+    ) == ("litellm distribution is editable",)
+
+
+def test_rejects_unavailable_native_bridge(tmp_path: Path) -> None:
+    prefix: Final = tmp_path / "site-packages"
+    checkout: Final = tmp_path / "checkout"
+    module_files: Final = {"litellm": prefix / "litellm" / "__init__.py"}
+
+    assert _failures(
+        prefix=prefix,
+        checkout=checkout,
+        module_files=module_files,
+        native_available=False,
+    ) == ("litellm.rust_bridge.native_bridge_available() is false",)

uv.lock

@@ -4648,6 +4648,12 @@ utils = [
 ]
 
 [package.dev-dependencies]
+benchmarks = [
+    { name = "a2a-sdk" },
+    { name = "mcp" },
+    { name = "pytest" },
+    { name = "pytest-codspeed" },
+]
 ci = [
     { name = "aiodynamo" },
     { name = "anthropic" },
@@ -4849,6 +4855,12 @@ requires-dist = [
 provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "mcp", "saml", "semantic-router", "mlflow", "grpc", "stt-vertex-chirp", "stt-nvidia-riva", "google", "bedrock-realtime", "proxy-runtime"]
 
 [package.metadata.requires-dev]
+benchmarks = [
+    { name = "a2a-sdk", specifier = "==1.1.0" },
+    { name = "mcp", specifier = ">=2.2.0,<3" },
+    { name = "pytest", specifier = "==9.0.3" },
+    { name = "pytest-codspeed", specifier = "==4.3.0" },
+]
 ci = [
     { name = "aiodynamo", specifier = "==24.7" },
     { name = "anthropic", specifier = "==0.84.0" },