From 1ae3216120a49fd2ac1ea6a2f43e42e8a33c276b Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 7 Sep 2026 12:05:43 -0700 Subject: [PATCH] fix(router): preserve default heuristic updates (#40007) --- .../auto_router_tuning_baseline.py | 14 ++++-- .../test_model_management_endpoints.py | 10 ++-- .../proxy/proxy_server/test_lifecycle.py | 21 ++++++++ .../test_auto_router_tuning_baseline.py | 48 +++++++++++++++++-- 4 files changed, 80 insertions(+), 13 deletions(-) diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_utils/auto_router_tuning_baseline.py index 1e6412856a6..b82269d5824 100644 --- a/litellm/router_utils/auto_router_tuning_baseline.py +++ b/litellm/router_utils/auto_router_tuning_baseline.py @@ -10,7 +10,7 @@ from pydantic import ValidationError from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig -TUNING_BASELINE_PARAM_NAME: Final = "auto_router_tuning_baseline" +TUNING_BASELINE_PARAM_NAME: Final = "auto_router_tuning_baseline_v2" HEURISTIC_V1_TUNING_FIELDS: Final = ( "tiers", @@ -29,6 +29,8 @@ HEURISTIC_V1_TUNING_FIELDS: Final = ( "keyword_tier_rules", ) +_TUNING_FIELD_SET: Final = frozenset(HEURISTIC_V1_TUNING_FIELDS) + _V1_SCORING_CLASSIFIER_TYPES: Final = frozenset({"heuristic", "heuristic_first", "hybrid"}) _AUTO_ROUTER_COMPLEXITY_PREFIX: Final = "auto_router/complexity_router" _EMPTY: Final[Mapping[str, object]] = MappingProxyType({}) @@ -40,12 +42,16 @@ def _mapping(value: object) -> Mapping[str, object]: def tuning_fingerprint(complexity_router_config: object) -> str | None: - """Digest of normalized heuristic-v1 tuning fields, or None when the config is invalid.""" + """Digest of explicitly supplied heuristic-v1 tuning fields, normalized, or None when invalid.""" + raw: Final = _mapping(complexity_router_config) try: - validated: Final = ComplexityRouterConfig.model_validate(_mapping(complexity_router_config)) + validated: Final = ComplexityRouterConfig.model_validate(raw) except ValidationError: return None - payload: Final = validated.model_dump(mode="json", include=frozenset(HEURISTIC_V1_TUNING_FIELDS)) + supplied: Final = ((_TUNING_FIELD_SET - frozenset(("tier_model_configs",))) & frozenset(raw)) | ( + frozenset(("tier_model_configs",)) if validated.tier_model_configs else frozenset() + ) + payload: Final = validated.model_dump(mode="json", include=supplied) return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()).hexdigest() diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 63c08529b4d..d90338f8480 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -4447,13 +4447,13 @@ class TestStrategyRouterWriteValidation: self.litellm_proxymodeltable = MagicMock( create=AsyncMock(), update=AsyncMock(), - find_many=AsyncMock( - return_value=tuple( - LiteLLM_ProxyModelTable.model_validate(row) for row in self.tuning_rows - ) - ), + find_many=AsyncMock(side_effect=self._find_many), ) + async def _find_many(self, where: object = None) -> tuple[LiteLLM_ProxyModelTable, ...]: + json.dumps(where) # prisma serializes the filter with json.dumps and rejects a mappingproxy + return tuple(LiteLLM_ProxyModelTable.model_validate(row) for row in self.tuning_rows) + @property def db(self) -> "TestStrategyRouterWriteValidation._FakeTx": return self diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index a518c9289ae..7db85dc6943 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -821,6 +821,27 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): ) +@pytest.mark.asyncio +async def test_tuning_baseline_v2_is_created_alongside_the_legacy_row(): + from litellm.router_utils.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT + + prisma_client = MagicMock() + prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_config.create = AsyncMock() + deployment = { + "model_name": "a", + "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": {}}, + } + + result = await ProxyStartupEvent._load_heuristic_v1_tuning_baselines(prisma_client, [deployment]) + + assert result == {'yaml:["a",[]]': DEFAULT_TUNING_FINGERPRINT} + assert prisma_client.db.litellm_config.create.await_args.kwargs["data"] == { + "param_name": "auto_router_tuning_baseline_v2", + "param_value": json.dumps(dict(result)), + } + + @pytest.mark.asyncio async def test_tuning_baseline_waits_for_a_complete_db_model_census(monkeypatch): prisma_client = MagicMock() diff --git a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py index c3395ab641f..fa7a96adb20 100644 --- a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py +++ b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py @@ -72,6 +72,9 @@ class TestTuningFingerprint: config["classifier_llm_config"] = {"model": "judge"} assert tuning_fingerprint(config) != DEFAULT_TUNING_FINGERPRINT + def test_explicit_empty_tier_model_configs_follow_omission(self) -> None: + assert tuning_fingerprint({"tier_model_configs": {}}) == DEFAULT_TUNING_FINGERPRINT + def test_tier_model_overrides_change_the_fingerprint(self) -> None: plain = tuning_fingerprint({"tiers": {"SIMPLE": "x"}}) with_override = tuning_fingerprint( @@ -80,11 +83,39 @@ class TestTuningFingerprint: assert plain != with_override def test_non_tuning_fields_do_not_change_the_fingerprint(self) -> None: - assert tuning_fingerprint({"return_raw_model_name": True, "session_affinity": True}) == DEFAULT_TUNING_FINGERPRINT + assert ( + tuning_fingerprint({"return_raw_model_name": True, "session_affinity": True}) == DEFAULT_TUNING_FINGERPRINT + ) def test_invalid_config_has_no_fingerprint(self) -> None: assert tuning_fingerprint({"tier_boundaries": "not-a-mapping"}) is None + def test_a_release_changing_a_shipped_default_is_not_an_operator_edit(self, monkeypatch) -> None: + """An omitted setting follows the shipped default and stays off the quota when that default moves; + only what the operator wrote is fingerprinted, so an explicit value equal to the old default still counts.""" + import litellm.router_strategy.complexity_router.config as config_module + + untouched = _router("a", {}) + tiers_only = _router("b", {"tiers": _TIERS}) + pinned = _router("c", {"dimension_weights": dict(config_module.DEFAULT_DIMENSION_WEIGHTS)}) + baselines = snapshot_tuning_baselines([untouched, tiers_only, pinned]) + assert tuning_fingerprint(pinned["litellm_params"]["complexity_router_config"]) != DEFAULT_TUNING_FINGERPRINT + + monkeypatch.setattr( + config_module, + "DEFAULT_DIMENSION_WEIGHTS", + {**config_module.DEFAULT_DIMENSION_WEIGHTS, "codePresence": 0.99}, + ) + monkeypatch.setattr( + config_module, + "DEFAULT_TIER_BOUNDARIES", + {**config_module.DEFAULT_TIER_BOUNDARIES, "simple_medium": 0.42}, + ) + + assert tuning_fingerprint({}) == DEFAULT_TUNING_FINGERPRINT + assert mutable_tuned_identities([untouched, tiers_only, pinned], baselines) == frozenset() + assert mutable_tuned_identities([untouched], snapshot_tuning_baselines([])) == frozenset() + class TestRouterIdentity: def test_db_rows_key_on_model_id_and_yaml_rows_on_name_and_tags(self) -> None: @@ -161,12 +192,18 @@ class TestQuota: new_c = _router("c", {"tiers": _TIERS}) assert tuning_quota_violation(candidate=edited_a, others=[legacy_b], baselines=baselines, limit=1) is None - assert tuning_quota_violation(candidate=edited_a, others=[edited_a, legacy_b], baselines=baselines, limit=1) is None + assert ( + tuning_quota_violation(candidate=edited_a, others=[edited_a, legacy_b], baselines=baselines, limit=1) + is None + ) assert tuning_quota_violation(candidate=legacy_a, others=[edited_b], baselines=baselines, limit=1) is None assert tuning_quota_violation(candidate=edited_b, others=[edited_a], baselines=baselines, limit=1) is not None assert tuning_quota_violation(candidate=new_c, others=[edited_a], baselines=baselines, limit=1) is not None assert tuning_quota_violation(candidate=new_c, others=[edited_a], baselines=baselines, limit=None) is None - assert tuning_quota_violation(candidate=legacy_a, others=[edited_a, edited_b], baselines=baselines, limit=1) is None + assert ( + tuning_quota_violation(candidate=legacy_a, others=[edited_a, edited_b], baselines=baselines, limit=1) + is None + ) def test_reverting_to_baseline_frees_the_quota(self) -> None: legacy_a = _router("a", {"tiers": _TIERS}) @@ -174,7 +211,10 @@ class TestQuota: baselines = snapshot_tuning_baselines([legacy_a, legacy_b]) edited_b = _router("b", {"tiers": _TIERS}) assert tuning_quota_violation(candidate=edited_b, others=[legacy_a], baselines=baselines, limit=1) is None - assert tuning_quota_violation(candidate=edited_b, others=[legacy_a, edited_b], baselines=baselines, limit=1) is None + assert ( + tuning_quota_violation(candidate=edited_b, others=[legacy_a, edited_b], baselines=baselines, limit=1) + is None + ) def test_violation_message_names_the_limit_and_remedy(self) -> None: message = tuning_limit_violation(held=2, limit=1)