litellm/tests/rust-python-harness/shared/native_build.py
Yuneng Jiang dbf8fe0f4e
test: repair four chronically failing CI tests
test_no_linear_scans_in_router: #39468 added config_deployments() and
heuristic_v2_router_limit_violation(), which both scan the whole model_list
from admin-only paths (model add/upsert), so add them to the allowlist. The
allowlist becomes a mapping so each exemption carries its reason as data.

test_missing_model_parameter_curl: a request with no model is rejected by the
proxy when nothing can serve it and by the router when a wildcard or default
deployment exists, and by the upstream provider when a wildcard forwards it,
so the message text is not a stable contract. Assert the contract that holds
in every case: HTTP 400 with a non-empty error message.

test_model_group_info_e2e: /model_group/info resolves wildcards, so it can
never return "anthropic/*" verbatim. cc3f9cd65b rewrote the assertion to
expect the raw pattern after claude-3-5-haiku-20241022 left the price map,
which made it unsatisfiable. Assert the expansion instead.

test_should_derive_ocr_mapping_status_from_live_tests: the audit needs a
native bridge built with the trace-parity feature, which CI never builds, so
skip with the harness's own diagnostic instead of erroring. Extract that
check out of ensure_trace_bridge as trace_bridge_error so a pytest run
reports the state without kicking off a maturin rebuild.
2026-09-04 10:10:47 -07:00

97 lines
3.3 KiB
Python

from __future__ import annotations
import importlib.util
import os
import subprocess
import sys
from collections.abc import Iterator
from pathlib import Path
from typing import Final
from litellm.rust_bridge import get_native_bridge, reset_native_bridge_cache
MATURIN_SPEC: Final = "maturin==1.15.0"
BRIDGE_FEATURE: Final = "trace-parity"
_RUST_ROOT: Final = "litellm-rust"
_LOCKFILE: Final = "Cargo.lock"
_SOURCE_SUFFIXES: Final = frozenset({".rs", ".toml"})
_FAILURE_OUTPUT_LINES: Final = 15
def needs_rebuild(native_mtime: float | None, newest_source_mtime: float | None) -> bool:
if native_mtime is None:
return True
if newest_source_mtime is None:
return False
return newest_source_mtime > native_mtime
def _source_files(rust_root: Path) -> Iterator[Path]:
for path in rust_root.rglob("*"):
relative: Final = path.relative_to(rust_root)
if "target" in relative.parts or not path.is_file():
continue
if path.name == _LOCKFILE or path.suffix in _SOURCE_SUFFIXES:
yield path
def _newest_source_mtime(repo_root: Path) -> float | None:
rust_root: Final = repo_root / _RUST_ROOT
if not rust_root.is_dir():
return None
return max((path.stat().st_mtime for path in _source_files(rust_root)), default=None)
def _native_module_path() -> Path | None:
try:
spec: Final = importlib.util.find_spec("litellm.rust_bridge._native")
except (ImportError, ValueError):
return None
origin: Final = getattr(spec, "origin", None)
return Path(origin) if origin else None
def _drop_imported_bridge() -> None:
reset_native_bridge_cache()
for name in tuple(sys.modules):
if name.startswith("litellm.rust_bridge._native"):
del sys.modules[name]
def _rebuild(repo_root: Path) -> tuple[bool, str]:
command: Final = ("uvx", "--from", MATURIN_SPEC, "maturin", "develop", "--features", BRIDGE_FEATURE)
completed: Final = subprocess.run(
command,
cwd=repo_root,
env={**os.environ, "VIRTUAL_ENV": sys.prefix},
capture_output=True,
text=True,
check=False,
)
output: Final = f"{completed.stdout}\n{completed.stderr}".strip()
lines: Final = tuple(output.splitlines())
return completed.returncode == 0, "\n".join(lines[-_FAILURE_OUTPUT_LINES:])
def trace_bridge_error() -> str | None:
"""Why the installed bridge cannot serve trace parity, or None when it can. Never rebuilds."""
bridge: Final = get_native_bridge()
if bridge is None:
return "native Rust bridge is not importable"
if getattr(bridge, "_trace", None) is None:
return f"native Rust bridge does not expose _trace; it must be built with the {BRIDGE_FEATURE} feature"
return None
def ensure_trace_bridge(repo_root: Path) -> str | None:
native_path: Final = _native_module_path()
native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None
if needs_rebuild(native_mtime, _newest_source_mtime(repo_root)):
print(f"Rebuilding native Rust bridge ({BRIDGE_FEATURE} feature)...", flush=True)
succeeded: Final
output: Final
succeeded, output = _rebuild(repo_root)
if not succeeded:
return f"native Rust bridge rebuild failed:\n{output}"
_drop_imported_bridge()
return trace_bridge_error()