From 7dcdc2ffc930fac68f8a7bf0c02780afe390c035 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Thu, 29 Oct 2026 07:55:02 -0700 Subject: [PATCH] build(model-catalog): ship the cost map inside the native extension (#48907) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/smoke_test_native_wheel.py | 9 +++++- backend/Dockerfile | 1 + gateway/Dockerfile | 1 + litellm-rust/Cargo.lock | 1 + litellm-rust/crates/python-bridge/AGENTS.md | 12 ++++++++ litellm-rust/crates/python-bridge/Cargo.toml | 1 + litellm-rust/crates/python-bridge/src/lib.rs | 4 +++ .../python-bridge/src/model_registry.rs | 30 +++++++++++++++++++ litellm/rust_bridge/_native.pyi | 2 ++ migrations/Dockerfile | 1 + packaging/litellm-core/pyproject.toml | 1 + pyproject.toml | 2 ++ scripts/build_core_distribution.py | 10 ++++++- .../base_sdk_tests/test_core_distribution.py | 4 ++- 14 files changed, 76 insertions(+), 3 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/model_registry.rs diff --git a/.github/scripts/smoke_test_native_wheel.py b/.github/scripts/smoke_test_native_wheel.py index 577bb32fcf0..af4054ea92e 100644 --- a/.github/scripts/smoke_test_native_wheel.py +++ b/.github/scripts/smoke_test_native_wheel.py @@ -7,6 +7,8 @@ import zipfile from pathlib import Path from typing import Final +COST_MAP: Final = Path(__file__).resolve().parents[2] / "model_prices_and_context_window.json" + CHILD_SCRIPT: Final = """ from importlib.util import module_from_spec, spec_from_file_location from pathlib import Path @@ -34,6 +36,9 @@ else: after = module.gil_stats() if not isinstance(after.get("releases"), int): raise AssertionError(f"native module unusable after panic: {after!r}") + +if module.bundled_model_cost_map() != Path(sys.argv[2]).read_bytes(): + raise AssertionError("the wheel embeds a different model_prices_and_context_window.json than this checkout") """ @@ -55,7 +60,9 @@ def main() -> int: native_path: Final = Path(temporary_directory) / Path(native_members[0].filename).name native_path.write_bytes(archive.read(native_members[0])) - result: Final = subprocess.run((sys.executable, "-c", CHILD_SCRIPT, str(native_path)), check=False) + result: Final = subprocess.run( + (sys.executable, "-c", CHILD_SCRIPT, str(native_path), str(COST_MAP)), check=False + ) if result.returncode != 0: sys.stderr.write(f"native wheel smoke test exited with status {result.returncode}\n") diff --git a/backend/Dockerfile b/backend/Dockerfile index 1c22f914100..aad59ba17c6 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -24,6 +24,7 @@ ENV UV_PYTHON_DOWNLOADS=0 COPY pyproject.toml README.md rust-toolchain.toml ./ COPY .cargo ./.cargo COPY litellm-rust ./litellm-rust +COPY model_prices_and_context_window.json ./ RUN --mount=type=cache,target=/root/.cache/uv \ mkdir -p litellm/rust_bridge /wheels && \ diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 6dc8bab45ce..9d228627543 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -40,6 +40,7 @@ ENV UV_PYTHON_DOWNLOADS=0 COPY pyproject.toml README.md rust-toolchain.toml ./ COPY .cargo ./.cargo COPY litellm-rust ./litellm-rust +COPY model_prices_and_context_window.json ./ RUN --mount=type=cache,target=/root/.cache/uv \ mkdir -p litellm/rust_bridge /wheels && \ diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f7d7e5f4dab..e959f77917b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4345,6 +4345,7 @@ dependencies = [ "litellm-inference-transcription", "litellm-llms", "litellm-llms-types", + "litellm-model-catalog", "litellm-router-types", "litellm-secrets", "litellm-secrets-aws", diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index d90fe073b20..26858b2afe2 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -106,6 +106,18 @@ GIL handling to `litellm-host-python`. ## Data Handling +- Reuse large immutable bundled data and parsed catalogs across calls, including in + replacement APIs. Share Rust storage and cache repeated Python projections instead + of reparsing, serializing or copying the full payload on every access +- Cache owned Python objects with `PyOnceLock>`, then bind and clone references + while attached. Do not use blocking `OnceLock::get_or_init` while attached + ([PyO3 initialization](https://pyo3.rs/v0.29.2/faq.html), + [reference ownership](https://pyo3.rs/v0.29.2/types.html)) +- Never reuse Python objects across interpreters or interpreter reinitialization. + Once-only caches are for immutable data. Reloadable/custom data needs versioned + entries or invalidation, and caller-mutable Python containers must not be shared +- Test unchanged contents and allocation reuse, not just equal values. Keep this + coverage when replacing a bundled-data accessor with a catalog API - OCR payloads can contain personal data and large base64 images. Do not log payloads or provider responses. - Avoid copying large payloads more than needed. The current JSON round-trip is diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 9d69cd8c14b..a4e1c57ab1e 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -49,6 +49,7 @@ litellm-secrets-types.workspace = true litellm-llms-types.workspace = true litellm-router-types.workspace = true litellm-host-python.workspace = true +litellm-model-catalog.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true prost.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 12827b76ff3..683e90d02a2 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -8,6 +8,7 @@ mod execution; mod http; mod logger; mod marshal; +mod model_registry; mod preflight; mod python_settings; mod routes; @@ -30,6 +31,8 @@ mod _native { #[pymodule_export] use crate::logger::NativeDiagnosticProcessor; #[pymodule_export] + use crate::model_registry::bundled_model_cost_map; + #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] use crate::routes::chat_completions::{acompletion, completion}; @@ -110,6 +113,7 @@ mod tests { "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "trace_encode_error", + "bundled_model_cost_map", "TokenCounter", "Tokenizer", "gil_stats", diff --git a/litellm-rust/crates/python-bridge/src/model_registry.rs b/litellm-rust/crates/python-bridge/src/model_registry.rs new file mode 100644 index 00000000000..b7b7cd1fbb3 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_registry.rs @@ -0,0 +1,30 @@ +use pyo3::prelude::*; +use pyo3::sync::PyOnceLock; +use pyo3::types::PyBytes; + +#[pyfunction] +pub fn bundled_model_cost_map(py: Python<'_>) -> Bound<'_, PyBytes> { + static BYTES: PyOnceLock> = PyOnceLock::new(); + BYTES + .get_or_init(py, || { + PyBytes::new(py, litellm_model_catalog::BUNDLED_REGISTRY).unbind() + }) + .bind(py) + .clone() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[rstest::rstest] + fn repeated_calls_reuse_the_bundled_registry_bytes() { + Python::initialize(); + let first = Python::attach(|py| bundled_model_cost_map(py).unbind()); + Python::attach(|py| { + let second = bundled_model_cost_map(py); + assert!(first.bind(py).is(&second)); + assert_eq!(second.as_bytes(), litellm_model_catalog::BUNDLED_REGISTRY); + }); + } +} diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 7364a1d72f2..62b48621c05 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -14,6 +14,7 @@ class ForkedAfterNativeRuntimeStarted(RuntimeError): ... class ProcessReservedForForking(RuntimeError): ... def trace_encode_error(message: str) -> bytes: ... +def bundled_model_cost_map() -> bytes: ... @final class NativeClickHouseSpendConfig: @@ -220,6 +221,7 @@ __all__ = [ "aocr", "aresponses", "atranscription", + "bundled_model_cost_map", "completion", "embedding", "gil_stats", diff --git a/migrations/Dockerfile b/migrations/Dockerfile index 2b1fd029143..3ffa9140543 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -43,6 +43,7 @@ ENV UV_PYTHON_DOWNLOADS=0 COPY pyproject.toml README.md rust-toolchain.toml ./ COPY .cargo ./.cargo COPY litellm-rust ./litellm-rust +COPY model_prices_and_context_window.json ./ RUN --mount=type=cache,target=/root/.cache/uv \ mkdir -p litellm/rust_bridge /wheels && \ diff --git a/packaging/litellm-core/pyproject.toml b/packaging/litellm-core/pyproject.toml index 27681d9116c..db50589a675 100644 --- a/packaging/litellm-core/pyproject.toml +++ b/packaging/litellm-core/pyproject.toml @@ -56,6 +56,7 @@ editable-profile = "dev" include = [ { path = "rust-toolchain.toml", format = "sdist" }, { path = ".cargo/config.toml", format = "sdist" }, + { path = "model_prices_and_context_window.json", format = "sdist" }, "litellm/router_strategy/complexity_router/artifacts/*.json", "litellm/router_strategy/complexity_router/fuse_presets.json", "litellm/proxy/model_insights_tasks.json", diff --git a/pyproject.toml b/pyproject.toml index 200bb463313..e6d34a3d90c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -335,6 +335,7 @@ include = [ "litellm/proxy/public_endpoints/whats_new.json", "litellm/proxy/common_utils/codex_base_instructions.md", "litellm/proxy/common_utils/codex_bundled_models_0.159.3.json", + { path = "model_prices_and_context_window.json", format = "sdist" }, ] exclude = [ "litellm/proxy/enterprise", @@ -357,6 +358,7 @@ cache-keys = [ { file = "litellm-rust/Cargo.lock" }, { file = "litellm-rust/Cargo.toml" }, { file = "litellm-rust/crates/**/*" }, + { file = "model_prices_and_context_window.json" }, ] constraint-dependencies = [ "tornado>=6.5.8", diff --git a/scripts/build_core_distribution.py b/scripts/build_core_distribution.py index 36ec20a9031..41d7208d2c3 100644 --- a/scripts/build_core_distribution.py +++ b/scripts/build_core_distribution.py @@ -9,7 +9,15 @@ from pathlib import Path from typing import Final ROOT: Final = Path(__file__).resolve().parents[1] -SOURCES: Final = ("litellm", "litellm-rust", ".cargo", "rust-toolchain.toml", "README.md", "LICENSE") +SOURCES: Final = ( + "litellm", + "litellm-rust", + ".cargo", + "rust-toolchain.toml", + "model_prices_and_context_window.json", + "README.md", + "LICENSE", +) def stage_core_distribution(source: Path, destination: Path) -> None: diff --git a/tests/base_sdk_tests/test_core_distribution.py b/tests/base_sdk_tests/test_core_distribution.py index ba68d08e4d1..c22fa4e54fc 100644 --- a/tests/base_sdk_tests/test_core_distribution.py +++ b/tests/base_sdk_tests/test_core_distribution.py @@ -132,7 +132,9 @@ def test_core_artifacts_exclude_local_configuration(distributions: tuple[Path, P assert not any(member.name.endswith(f"/{name}") for member in archive.getmembers()) -@pytest.mark.parametrize("relative_path", ["rust-toolchain.toml", ".cargo/config.toml"]) +@pytest.mark.parametrize( + "relative_path", ["rust-toolchain.toml", ".cargo/config.toml", "model_prices_and_context_window.json"] +) def test_core_sdist_preserves_native_build_configuration(distribution_directory: Path, relative_path: str) -> None: sdists: Final = tuple(distribution_directory.glob("litellm_core-*.tar.gz")) assert len(sdists) == 1