diff --git a/.github/actions/cache-cargo-build/action.yml b/.github/actions/cache-cargo-build/action.yml index c3b8ce22c68..222fad637fb 100644 --- a/.github/actions/cache-cargo-build/action.yml +++ b/.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 }}- diff --git a/.github/scripts/verify_installed_wheel_imports.py b/.github/scripts/verify_installed_wheel_imports.py new file mode 100644 index 00000000000..7775dd333f0 --- /dev/null +++ b/.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()) diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index fd7513a3937..692445be7f6 100644 --- a/.github/workflows/codspeed.yml +++ b/.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 diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 70d828b5c8b..4d9f7140f70 100644 --- a/.github/workflows/test-rust.yml +++ b/.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 diff --git a/Makefile b/Makefile index 0e9d2bbf82c..27bb5abe540 100644 --- a/Makefile +++ b/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 && \ diff --git a/pyproject.toml b/pyproject.toml index 3b17a397f3c..32bb55d174c 100644 --- a/pyproject.toml +++ b/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"] diff --git a/tests/benchmarks/conftest.py b/tests/benchmarks/conftest.py index c9b31cfb7d7..659a27ce392 100644 --- a/tests/benchmarks/conftest.py +++ b/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: diff --git a/tests/test_litellm/rust_bridge/test_verify_installed_wheel_imports.py b/tests/test_litellm/rust_bridge/test_verify_installed_wheel_imports.py new file mode 100644 index 00000000000..729f56ce993 --- /dev/null +++ b/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",) diff --git a/uv.lock b/uv.lock index dfba6bfc258..8384fa21307 100644 --- a/uv.lock +++ b/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" },