From 275247716815c24c26d698d655b37863823587f1 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 3 Nov 2026 11:35:57 +0000 Subject: [PATCH] build(model-catalog): load the bundled cost map at runtime instead of embedding it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/smoke_test_native_wheel.py | 31 ++++++++++-------- .github/workflows/test-unit.yml | 2 +- backend/Dockerfile | 1 - gateway/Dockerfile | 1 - litellm-rust/Cargo.lock | 1 - .../crates/model-catalog/src/bundled.rs | 3 -- litellm-rust/crates/model-catalog/src/lib.rs | 2 -- .../crates/model-catalog/tests/catalog.rs | 19 +++++------ .../crates/model-catalog/tests/schema.rs | 6 ++-- litellm-rust/crates/python-bridge/Cargo.toml | 1 - .../python-bridge/src/model_registry.rs | 32 +++++++------------ migrations/Dockerfile | 1 - packaging/litellm-core/pyproject.toml | 1 - pyproject.toml | 2 -- scripts/build_core_distribution.py | 1 - .../base_sdk_tests/test_core_distribution.py | 2 +- tests/unit/rust_bridge/test_bindings.py | 10 ++++++ 17 files changed, 52 insertions(+), 64 deletions(-) delete mode 100644 litellm-rust/crates/model-catalog/src/bundled.rs diff --git a/.github/scripts/smoke_test_native_wheel.py b/.github/scripts/smoke_test_native_wheel.py index af4054ea92e..53af909cabf 100644 --- a/.github/scripts/smoke_test_native_wheel.py +++ b/.github/scripts/smoke_test_native_wheel.py @@ -7,19 +7,14 @@ 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 importlib.resources import files from pathlib import Path import sys -native_path = Path(sys.argv[1]) -spec = spec_from_file_location("litellm.rust_bridge._native", native_path) -if spec is None or spec.loader is None: - raise RuntimeError("cannot create native extension import specification") -module = module_from_spec(spec) -spec.loader.exec_module(module) +sys.path.insert(0, sys.argv[1]) + +import litellm.rust_bridge._native as module before = module.gil_stats() if not isinstance(before.get("releases"), int): @@ -37,8 +32,9 @@ 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") +wheel_cost_map = files("litellm").joinpath("model_prices_and_context_window_backup.json").read_bytes() +if module.bundled_model_cost_map() != wheel_cost_map: + raise AssertionError("the native cost map does not match the wheel package resource") """ @@ -58,10 +54,17 @@ def main() -> int: sys.stderr.write(f"expected one native extension, found {len(native_members)}\n") return 1 - native_path: Final = Path(temporary_directory) / Path(native_members[0].filename).name - native_path.write_bytes(archive.read(native_members[0])) + package_path: Final = Path(temporary_directory) / "litellm" + bridge_path: Final = package_path / "rust_bridge" + bridge_path.mkdir(parents=True) + (package_path / "__init__.py").write_bytes(b"") + (bridge_path / "__init__.py").write_bytes(b"") + (bridge_path / Path(native_members[0].filename).name).write_bytes(archive.read(native_members[0])) + (package_path / "model_prices_and_context_window_backup.json").write_bytes( + archive.read("litellm/model_prices_and_context_window_backup.json") + ) result: Final = subprocess.run( - (sys.executable, "-c", CHILD_SCRIPT, str(native_path), str(COST_MAP)), check=False + (sys.executable, "-c", CHILD_SCRIPT, temporary_directory), check=False ) if result.returncode != 0: diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 0a490bb8f45..f0beacc16ea 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -42,7 +42,7 @@ jobs: id: rust-bridge-key if: steps.changes.outputs.decision != 'skip' env: - RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'model_prices_and_context_window.json', 'pyproject.toml') }} + RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'pyproject.toml') }} run: echo "key=$RUST_BRIDGE_CACHE_KEY" >> "$GITHUB_OUTPUT" - name: Set up Python diff --git a/backend/Dockerfile b/backend/Dockerfile index aad59ba17c6..1c22f914100 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -24,7 +24,6 @@ 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 658bed831ba..002c129756f 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -40,7 +40,6 @@ 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 f91deaaa401..5424905ba26 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4356,7 +4356,6 @@ dependencies = [ "litellm-inference-transcription", "litellm-llms", "litellm-llms-types", - "litellm-model-catalog", "litellm-router-types", "litellm-secret-value", "litellm-secrets", diff --git a/litellm-rust/crates/model-catalog/src/bundled.rs b/litellm-rust/crates/model-catalog/src/bundled.rs deleted file mode 100644 index d71a2d57ebb..00000000000 --- a/litellm-rust/crates/model-catalog/src/bundled.rs +++ /dev/null @@ -1,3 +0,0 @@ -/// The registry document this build serves when no other one is loaded. -pub const BUNDLED_REGISTRY: &[u8] = - include_bytes!("../../../../model_prices_and_context_window.json"); diff --git a/litellm-rust/crates/model-catalog/src/lib.rs b/litellm-rust/crates/model-catalog/src/lib.rs index fe419dde054..f375a148e83 100644 --- a/litellm-rust/crates/model-catalog/src/lib.rs +++ b/litellm-rust/crates/model-catalog/src/lib.rs @@ -35,7 +35,6 @@ macro_rules_attribute::attribute_alias! { #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]; } -mod bundled; mod catalog; mod error; mod fallback; @@ -49,7 +48,6 @@ mod store; mod vocabulary; pub mod wire; -pub use bundled::*; pub use catalog::*; pub use error::*; pub use fallback::*; diff --git a/litellm-rust/crates/model-catalog/tests/catalog.rs b/litellm-rust/crates/model-catalog/tests/catalog.rs index 31f907689d3..1a37f171bdb 100644 --- a/litellm-rust/crates/model-catalog/tests/catalog.rs +++ b/litellm-rust/crates/model-catalog/tests/catalog.rs @@ -1,11 +1,11 @@ use std::collections::BTreeSet; -use litellm_model_catalog::{ - BUNDLED_REGISTRY, Catalog, Error, ErrorKind, IntegrityLimits, Mode, Rejects, Written, -}; +use litellm_model_catalog::{Catalog, Error, ErrorKind, IntegrityLimits, Mode, Rejects, Written}; use rstest::{fixture, rstest}; use serde_json::{Map, Value, json}; +const REGISTRY: &[u8] = include_bytes!("../../../../model_prices_and_context_window.json"); + const ALPHA_FIXTURE: &[u8] = br#"{ "sample_spec":{"explanation":"example","input_cost_per_token":"not checked"}, "fallback_generalizations":{"rules":[{"name":"family","pattern":"^new-","model_info":{"mode":"chat"}}]}, @@ -272,7 +272,7 @@ fn timezone_errors_reject_providerless_rows_without_rejecting_other_providerless #[rstest] fn bundled_per_row_errors_match_document_errors() { - let parsed = Catalog::parse(BUNDLED_REGISTRY).unwrap(); + let parsed = Catalog::parse(REGISTRY).unwrap(); let row_errors: Vec<_> = parsed .value .written() @@ -367,21 +367,18 @@ fn written_is_the_document_in_order_including_sections_and_bad_entries() { #[rstest] fn the_bundled_registry_has_no_problems() { - let problems: Vec<_> = Catalog::problems(BUNDLED_REGISTRY) - .iter() - .map(report) - .collect(); + let problems: Vec<_> = Catalog::problems(REGISTRY).iter().map(report).collect(); assert!(problems.is_empty(), "{problems:#?}"); const REFERENCE_MODEL_COUNT: usize = 4303; - let catalog = Catalog::parse(BUNDLED_REGISTRY).unwrap().value; + let catalog = Catalog::parse(REGISTRY).unwrap().value; assert!(catalog.model_count() as f64 >= REFERENCE_MODEL_COUNT as f64 * 0.9); } #[rstest] fn the_bundled_registry_is_written_back_as_it_was_read() { - let document: Value = serde_json::from_slice(BUNDLED_REGISTRY).unwrap(); + let document: Value = serde_json::from_slice(REGISTRY).unwrap(); assert_eq!( - written(&Catalog::parse(BUNDLED_REGISTRY).unwrap().value), + written(&Catalog::parse(REGISTRY).unwrap().value), serde_json::to_string(&document).unwrap() ); } diff --git a/litellm-rust/crates/model-catalog/tests/schema.rs b/litellm-rust/crates/model-catalog/tests/schema.rs index 464cf069cac..16b02d2f1fd 100644 --- a/litellm-rust/crates/model-catalog/tests/schema.rs +++ b/litellm-rust/crates/model-catalog/tests/schema.rs @@ -1,15 +1,17 @@ #![cfg(feature = "schema")] -use litellm_model_catalog::{BUNDLED_REGISTRY, model_entry_json_schema, registry_json_schema}; +use litellm_model_catalog::{model_entry_json_schema, registry_json_schema}; use rstest::rstest; use serde_json::{Map, Value, json}; +const REGISTRY: &[u8] = include_bytes!("../../../../model_prices_and_context_window.json"); + fn schema() -> Value { serde_json::to_value(model_entry_json_schema()).expect("generated schema serializes") } fn checked_in_catalog() -> Value { - serde_json::from_slice(BUNDLED_REGISTRY).unwrap() + serde_json::from_slice(REGISTRY).unwrap() } fn registry_validator() -> jsonschema::Validator { diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index e3ef45f7876..575867071fd 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -50,7 +50,6 @@ 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/model_registry.rs b/litellm-rust/crates/python-bridge/src/model_registry.rs index b7b7cd1fbb3..c1583ac7ef0 100644 --- a/litellm-rust/crates/python-bridge/src/model_registry.rs +++ b/litellm-rust/crates/python-bridge/src/model_registry.rs @@ -3,28 +3,18 @@ use pyo3::sync::PyOnceLock; use pyo3::types::PyBytes; #[pyfunction] -pub fn bundled_model_cost_map(py: Python<'_>) -> Bound<'_, PyBytes> { +pub fn bundled_model_cost_map(py: Python<'_>) -> PyResult> { static BYTES: PyOnceLock> = PyOnceLock::new(); BYTES - .get_or_init(py, || { - PyBytes::new(py, litellm_model_catalog::BUNDLED_REGISTRY).unbind() + .get_or_try_init(py, || -> PyResult> { + let bytes = py + .import("importlib.resources")? + .call_method1("files", ("litellm",))? + .call_method1("joinpath", ("model_prices_and_context_window_backup.json",))? + .call_method0("read_bytes")? + .cast_into::() + .map_err(PyErr::from)?; + Ok(bytes.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); - }); - } + .map(|bytes| bytes.bind(py).clone()) } diff --git a/migrations/Dockerfile b/migrations/Dockerfile index 3ffa9140543..2b1fd029143 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -43,7 +43,6 @@ 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 1141e69c0aa..893eddd7ed8 100644 --- a/packaging/litellm-core/pyproject.toml +++ b/packaging/litellm-core/pyproject.toml @@ -56,7 +56,6 @@ 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 7a653b28f81..8fecda4666e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -335,7 +335,6 @@ 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", @@ -358,7 +357,6 @@ 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 41d7208d2c3..98cf851f434 100644 --- a/scripts/build_core_distribution.py +++ b/scripts/build_core_distribution.py @@ -14,7 +14,6 @@ SOURCES: Final = ( "litellm-rust", ".cargo", "rust-toolchain.toml", - "model_prices_and_context_window.json", "README.md", "LICENSE", ) diff --git a/tests/base_sdk_tests/test_core_distribution.py b/tests/base_sdk_tests/test_core_distribution.py index c22fa4e54fc..51fcfa29d9a 100644 --- a/tests/base_sdk_tests/test_core_distribution.py +++ b/tests/base_sdk_tests/test_core_distribution.py @@ -133,7 +133,7 @@ def test_core_artifacts_exclude_local_configuration(distributions: tuple[Path, P @pytest.mark.parametrize( - "relative_path", ["rust-toolchain.toml", ".cargo/config.toml", "model_prices_and_context_window.json"] + "relative_path", ["rust-toolchain.toml", ".cargo/config.toml"] ) 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")) diff --git a/tests/unit/rust_bridge/test_bindings.py b/tests/unit/rust_bridge/test_bindings.py index 086be65ee26..5c02377d3d2 100644 --- a/tests/unit/rust_bridge/test_bindings.py +++ b/tests/unit/rust_bridge/test_bindings.py @@ -1,6 +1,7 @@ import asyncio from collections.abc import Coroutine from dataclasses import dataclass +from importlib.resources import files from types import SimpleNamespace from typing import Final, cast # noqa: TID251 # narrows the async binding's untyped coroutine @@ -17,6 +18,15 @@ from litellm.rust_bridge.responses import entrypoints as responses from litellm.rust_bridge.transcription import native as transcription +def test_bundled_model_cost_map_uses_package_resource_and_reuses_bytes() -> None: + first: Final = _native.bundled_model_cost_map() + second: Final = _native.bundled_model_cost_map() + expected: Final = files("litellm").joinpath("model_prices_and_context_window_backup.json").read_bytes() + + assert first is second + assert first == expected + + def test_binding_distinguishes_disable_from_reset() -> None: def route() -> str: return "native"