fix(cost): keep off_peak_pricing scoped to its deployment

register_model inserted the first deployment's off_peak_pricing dict by
reference into the shared backend cost-map entry, and later deployments
sharing that backend merged their schedules into the same object,
corrupting the first deployment's schedule and polluting the built-in
entry. Nested dicts now merge copy-on-write, and off_peak_pricing stays
off the shared backend keys.
This commit is contained in:
mateo-berri 2026-09-01 11:17:55 -07:00
parent 7abed91523
commit 1ba13fcc25
3 changed files with 111 additions and 9 deletions

View file

@ -3506,17 +3506,22 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
return {k: v for k, v in model_info.items() if k not in cls.model_fields}
SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = frozenset(
ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__
) - frozenset(CustomPricingLiteLLMParams.model_fields)
DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"})
SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = (
frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__)
- frozenset(CustomPricingLiteLLMParams.model_fields)
- DEPLOYMENT_SCOPED_PRICING_FIELDS
)
def shared_backend_model_info(model_info: dict[str, Any]) -> dict[str, Any]:
"""Return only the fields safe to register under a shared ``{provider}/{model}``
key in ``litellm.model_cost``: cost-map schema fields (``ModelInfoBase``) minus
per-deployment pricing overrides. Per-deployment metadata (``id``,
``access_via_team_ids``, arbitrary custom keys) never belongs on the shared key;
it stays under the deployment's unique model id.
per-deployment pricing overrides and deployment-scoped pricing blocks such as
``off_peak_pricing``. Per-deployment metadata (``id``, ``access_via_team_ids``,
arbitrary custom keys) never belongs on the shared key; it stays under the
deployment's unique model id.
"""
return {k: v for k, v in model_info.items() if k in SHARED_BACKEND_MODEL_INFO_FIELDS}

View file

@ -2851,10 +2851,9 @@ def _update_dictionary(existing_dict: dict, new_dict: dict) -> dict:
elif isinstance(v, dict):
existing_nested_dict = existing_dict.get(k)
if isinstance(existing_nested_dict, dict):
existing_nested_dict.update(v)
existing_dict[k] = existing_nested_dict
existing_dict[k] = {**existing_nested_dict, **v} # mutable-ok: copy-on-write merge
else:
existing_dict[k] = v
existing_dict[k] = dict(v) # mutable-ok: detached copy, never the caller's dict by reference
else:
existing_dict[k] = v

View file

@ -793,3 +793,101 @@ def test_embedding_direct_sdk_custom_pricing_still_registers_shared_key():
finally:
litellm.model_cost.pop(model_key, None)
_invalidate_model_cost_lowercase_map()
def test_update_dictionary_merges_nested_dicts_without_aliasing():
"""A nested dict must be merged copy-on-write: the pre-existing nested dict
object stays untouched, and the caller's incoming nested dict is never
inserted by reference into the merged result.
"""
from litellm.utils import _update_dictionary
existing_nested = {"hours_utc": "01:00-02:00"}
existing = {"off_peak_pricing": existing_nested}
incoming_nested = {"windows": [{"hours_utc": "16:00-19:00", "weekdays": [2]}]}
incoming = {"off_peak_pricing": incoming_nested}
merged = _update_dictionary(existing, incoming)
assert merged["off_peak_pricing"] == {
"hours_utc": "01:00-02:00",
"windows": [{"hours_utc": "16:00-19:00", "weekdays": [2]}],
}
assert existing_nested == {"hours_utc": "01:00-02:00"}
assert merged["off_peak_pricing"] is not incoming_nested
fresh = _update_dictionary({}, incoming)
assert fresh["off_peak_pricing"] == incoming_nested
assert fresh["off_peak_pricing"] is not incoming_nested
def test_router_deployments_sharing_backend_keep_their_own_off_peak_pricing():
"""Two deployments of the same backend model with different
``off_peak_pricing`` blocks must each keep their own schedule under their
unique model id, and neither block may leak onto the shared backend keys.
Before the fix, ``register_model`` inserted the first deployment's block by
reference into the built-in ``gpt-4o-mini`` entry, and the second
deployment's registration merged its keys into that same object, corrupting
the first deployment's schedule and polluting the built-in entry.
"""
from litellm import Router
active_block = {
"windows": [{"hours_utc": "16:00-19:00", "weekdays": [2]}],
"input_cost_per_token": 5e-07,
"output_cost_per_token": 1e-06,
}
inactive_block = {
"hours_utc": "05:00-06:00",
"input_cost_per_token": 5e-07,
"output_cost_per_token": 1e-06,
}
shared_keys = ["gpt-4o-mini", "openai/gpt-4o-mini"]
deployment_ids = ["offpeak-alias-dep-1", "offpeak-alias-dep-2"]
original_entries = _snapshot_model_cost_entries(shared_keys)
router = Router(
model_list=[
{
"model_name": "offpeak-active-weekday",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "fake-key-for-registration",
},
"model_info": {
"id": deployment_ids[0],
"input_cost_per_token": 1e-06,
"output_cost_per_token": 2e-06,
"off_peak_pricing": dict(active_block),
},
},
{
"model_name": "offpeak-inactive-hours",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "fake-key-for-registration",
},
"model_info": {
"id": deployment_ids[1],
"input_cost_per_token": 1e-06,
"output_cost_per_token": 2e-06,
"off_peak_pricing": dict(inactive_block),
},
},
]
)
try:
registered_first = litellm.model_cost[deployment_ids[0]]["off_peak_pricing"]
registered_second = litellm.model_cost[deployment_ids[1]]["off_peak_pricing"]
assert registered_first == active_block
assert registered_second == inactive_block
for shared_key in shared_keys:
shared_entry = litellm.model_cost.get(shared_key) or {}
assert not shared_entry.get("off_peak_pricing")
finally:
for deployment_id in deployment_ids:
litellm.model_cost.pop(deployment_id, None)
_restore_model_cost_entries(original_entries)
del router