diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 0f19e9ce149..742b9d9817f 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -119,6 +119,7 @@ from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, GenericLiteLLMParams, + LiteLLM_Params, ModelInfo, updateDeployment, ) @@ -728,6 +729,44 @@ def _ptu_priced_deployment(model_params: Deployment) -> Deployment: ) +_OWNERSHIP_FIELDS: Final = frozenset( + { + "db_model", + "team_id", + "team_public_model_name", + "access_groups", + "created_at", + "created_by", + "updated_at", + "updated_by", + "blocked", + } +) + +_STORED_REQUIRED_FIELDS: Final = frozenset( + name for model in (LiteLLM_Params, ModelInfo) for name, field in model.model_fields.items() if field.is_required() +) + +_NULL_CLEAR_IGNORED_FIELDS: Final = _OWNERSHIP_FIELDS | _STORED_REQUIRED_FIELDS | frozenset(PTU_MODEL_INFO_FIELDS) + + +def _explicitly_cleared_fields(patch: BaseModel | None) -> frozenset[str]: + """The keys a patch sends as an explicit null, which update_db_model removes from the + stored blob (JSON Merge Patch). Ownership keys are left alone, as are the keys the stored + models require, since clearing one writes a row no reload can rebuild through + LiteLLM_Params / ModelInfo. The PTU keys are handled by _explicitly_cleared_ptu_fields, + whose clear is gated on the feature flag. Applied after both blobs merge, so a model_info blob + the UI echoes back cannot resurrect a pricing key the litellm_params patch clears. + """ + if patch is None: + return frozenset() + return frozenset( + field + for field in patch.model_fields_set + if field not in _NULL_CLEAR_IGNORED_FIELDS and getattr(patch, field) is None + ) + + def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: if updated_patch.model_info is not None: _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) @@ -748,25 +787,15 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if updated_patch.model_info: merged_model_info.update(updated_patch.model_info.model_dump(exclude_none=True)) - # Honor explicit-null clears LAST, after both merges, so a model_info blob the UI - # passes through (which today re-sends the OLD pricing on every save) cannot - # silently undo a litellm_params clear via .update(). - # - # Restricted to SPECIAL_MODEL_INFO_PARAMS (input/output cost per token/character - # and cache read/write costs) so this path cannot be used to null out privileged - # model_info fields like team_id or access groups. SPECIAL_MODEL_INFO_PARAMS are - # mirrored between litellm_params and model_info by Deployment.__init__, so the - # clear propagates to both blobs. - if updated_patch.litellm_params: - for field in updated_patch.litellm_params.model_fields_set: - if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None: - merged_litellm_params.pop(field, None) - merged_model_info.pop(field, None) + for field in _explicitly_cleared_fields(updated_patch.litellm_params): + merged_litellm_params.pop(field, None) + if field in SPECIAL_MODEL_INFO_PARAMS: + merged_model_info.pop(field, None) + for field in _explicitly_cleared_fields(updated_patch.model_info): + merged_model_info.pop(field, None) + if field in SPECIAL_MODEL_INFO_PARAMS: + merged_litellm_params.pop(field, None) if updated_patch.model_info: - for field in updated_patch.model_info.model_fields_set: - if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: - merged_model_info.pop(field, None) - merged_litellm_params.pop(field, None) for field in _explicitly_cleared_ptu_fields(updated_patch.model_info): merged_model_info.pop(field, None) @@ -816,8 +845,9 @@ async def patch_model( """ PATCH Endpoint for partial model updates. - Only updates the fields specified in the request while preserving other existing values. - Follows proper PATCH semantics by only modifying provided fields. + JSON Merge Patch semantics over `litellm_params` and `model_info`: a key absent from the + body is unchanged, a key sent as null is removed from the stored row, and a key sent with a + value is set (identity and ownership keys such as `id` and `team_id` ignore a null). Args: model_id: The ID of the model to update diff --git a/litellm/router.py b/litellm/router.py index 934a4ac86a9..1252d7e7487 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -628,6 +628,15 @@ RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset( RETRY_BREADCRUMB_LIMIT: Final = 4 +# Cost-map keys created by _register_deployment_in_model_cost, which shares one flat +# namespace with the built-in model catalog. Only a key it created may be evicted, or a +# deployment whose id names a real model would strip that model's pricing and +# capabilities for every other deployment of it. delete_deployment gives a key back once no +# live router still serves that id, so a later catalog refresh that starts serving the name +# is not treated as a deployment's own. +_DEPLOYMENT_COST_MAP_KEYS: Final[set[str]] = set() # mutable-ok: ownership of shared cost-map keys + + class FallbackAwareStreamWrapper(CustomStreamWrapper): """Base for the Router's chat-completion stream wrappers, which are built around the attempt the Router picked first and have to repoint themselves when a fallback takes over.""" @@ -9761,6 +9770,7 @@ class Router: """ idx: Final = len(self.model_list) self.model_list.append(model) + _live_routers.add(self) # mutable-ok: track dynamic routers without extending their lifetimes self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() @@ -9929,7 +9939,12 @@ class Router: """Write a deployment's metadata into ``litellm.model_cost``. Runs when a deployment is added and again after a price data reload, so - the entries a refresh rebuilds are the ones a fresh boot would produce. + the entries a refresh rebuilds are the ones a fresh boot would produce. An + entry this function created is replaced rather than merged, so a price cleared + from the deployment does not linger from an earlier registration and keep + billing at the old rate. An entry it did not create is left to merge, because + a deployment id that collides with a catalog model name shares that model's + entry with every other deployment of it. Nothing is recorded for replay: a refresh walks the live routers instead, so a deleted, repointed or never-added deployment, and a discarded router, drop out of the rebuild on their own. @@ -9946,6 +9961,10 @@ class Router: } if model_id is not None: + if model_id in _DEPLOYMENT_COST_MAP_KEYS: + litellm.model_cost.pop(model_id, None) # mutable-ok: remove cleared prices from the shared entry + elif model_id not in litellm.model_cost: + _DEPLOYMENT_COST_MAP_KEYS.add(model_id) # mutable-ok: retain shared ownership across reloads litellm.register_model( model_cost={model_id: model_info}, persist_across_reloads=False, @@ -10042,6 +10061,11 @@ class Router: _budget_limiter: Final = self._get_router_deployment_budget_limiter() if _budget_limiter is not None: _budget_limiter.unregister_deployment_budget(model_id=id) + if not any( + router is not self and id in router.model_id_to_deployment_index_map + for router in tuple(_live_routers) + ): + _DEPLOYMENT_COST_MAP_KEYS.discard(id) # mutable-ok: the last owning router released this key try: self._unregister_pre_routing_strategy_for_deployment( deployment=item if isinstance(item, Deployment) else Deployment(**item) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index c8d7037d2fd..83f1711a245 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -76,6 +76,8 @@ - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} - {id: mgmt.model.test_connection.happy_path, module: mgmt, tier: P0, surface: api, assertions: [happy_path], source: "_health_endpoints.py:1785", rationale: "Test Connection for a responses-mode Bedrock Mantle deployment reaches the live provider and reports success; this exact shape 500ed on an acompletion partial before v1.91.0", fail_before_fix: proven} +- {id: mgmt.model.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "model_management_endpoints.py:731", rationale: "A partial PATCH changes only the keys it names; every other stored key reads back byte-for-byte, and the new rate reaches billing"} +- {id: mgmt.model.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "model_management_endpoints.py:731", fail_before_fix: proven, rationale: "An explicit null on PATCH removes the key from the stored row and billing falls back to the cost map; before the fix only mirrored pricing keys could be cleared"} - {id: mgmt.mcp_server.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1577", rationale: "Every field of an admin-created MCP server reads back verbatim, by id and in the list, on every replica"} - {id: mgmt.mcp_server.list.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1112", rationale: "The MCP page grid lists a created server with the same field values its detail view reports"} - {id: mgmt.mcp_server.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:2665", rationale: "A dashboard edit of one field leaves the others intact and is visible on every replica after one save; edits that took several saves to stick were a customer defect"} diff --git a/tests/e2e/management/test_model_lifecycle_e2e.py b/tests/e2e/management/test_model_lifecycle_e2e.py new file mode 100644 index 00000000000..99044efeb46 --- /dev/null +++ b/tests/e2e/management/test_model_lifecycle_e2e.py @@ -0,0 +1,365 @@ +"""Live e2e: the lifecycle of a DB-stored deployment through the model management +routes, read back on every gateway replica. + +Each test registers its own gpt-4o-mini mock deployment through /model/new (deleted +on teardown) with non-default pricing, context window, mode, and api_base pinned, then +walks the lifecycle up to the step it proves: the create reads back field for field, +a partial PATCH changes only the key it names, an explicit null on PATCH removes the +key from the stored row (JSON Merge Patch), a call after the price clear is billed at +the cost map's rate rather than the cleared override, and a delete removes the +deployment from /model/info and makes the model name unknown to /chat/completions. + +The stored row is read back from /model/info, a control-plane route with one answer +behind it. What every gateway must agree on is which models it serves, so the create +and delete steps poll /v1/models on every URL in PROXY_REPLICA_URLS through +ProxyClient.read_model_back_everywhere, failing by name on the gateway that never converged. +""" + +from __future__ import annotations + +import math +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import Final + +import pytest +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from management_client import ManagementClient +from models import ( + ChatBody, + ChatMessage, + Clear, + LiteLLMParamsBody, + LiteLLMParamsPatch, + ModelInfoBody, + ModelInfoEntry, + ModelInfoResponse, + ModelNewBody, + ModelPatchBody, + ModelsListResponse, + SpendLogRow, +) + +pytestmark = pytest.mark.e2e + +BACKEND_MODEL: Final = "gpt-4o-mini" +PINNED_API_BASE: Final = "https://pinned.example.invalid/v1" +PINNED_MAX_INPUT_TOKENS: Final = 4096 +PINNED_INPUT_RATE: Final = 1e-05 +UPDATED_INPUT_RATE: Final = 2e-05 +PINNED_OUTPUT_RATE: Final = 3e-05 + +# A PATCH lands on the control plane, and each gateway picks it up on its own config +# reload, so the first call after the write can still be billed at the old rate. There +# is no price on the gateway's data-plane surface to poll, so the billing steps drive +# calls until the new rate shows up in the spend row and let the deadline be what fails. +BILLING_CONVERGENCE_TIMEOUT: Final = 90.0 +BILLING_CONVERGENCE_INTERVAL: Final = 5.0 + + +@dataclass(frozen=True, slots=True) +class Registered: + model_name: str + model_id: str + + +class _ErrorDetail(BaseModel): + message: str + + +class _ErrorEnvelope(BaseModel): + error: _ErrorDetail + + +def _register(client: ManagementClient, resources: ResourceManager) -> Registered: + """Register a mock gpt-4o-mini deployment with every field under test pinned to a + non-default value, deleted on teardown. max_input_tokens is pinned in + litellm_params only: a value in model_info is copied into the shared cost-map + entry for the backend model, which would leak into every other gpt-4o-mini + deployment on the proxy.""" + model_name: Final = f"e2e-lifecycle-{unique_marker()}" + model_id: Final = client.proxy.register_model( + ModelNewBody( + model_name=model_name, + litellm_params=LiteLLMParamsBody( + model=BACKEND_MODEL, + mock_response="ok", + api_base=PINNED_API_BASE, + input_cost_per_token=PINNED_INPUT_RATE, + output_cost_per_token=PINNED_OUTPUT_RATE, + max_input_tokens=PINNED_MAX_INPUT_TOKENS, + ), + model_info=ModelInfoBody(mode="chat"), + ) + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return Registered(model_name=model_name, model_id=model_id) + + +def _entry(body: ModelInfoResponse, model_name: str) -> ModelInfoEntry | None: + return next((entry for entry in body.data if entry.model_name == model_name), None) + + +def _stored_entry( + client: ManagementClient, + model_name: str, + *, + converged: Callable[[ModelInfoEntry], bool], +) -> ModelInfoEntry: + """The stored /model/info row for `model_name`, once it satisfies `converged`. + + /model/info is a control-plane route: the gateways named in PROXY_REPLICA_URLS + serve the LLM surface only, so the stored row has one answer, not one per + gateway. What every gateway must agree on is which models it serves, and + `_assert_served_everywhere` / `_assert_absent_everywhere` poll /v1/models for + that.""" + + def has_converged(body: ModelInfoResponse) -> bool: + entry: Final = _entry(body, model_name) + return entry is not None and converged(entry) + + body: Final = client.proxy.read_model_back("/model/info", ModelInfoResponse, predicate=has_converged) + entry: Final = _entry(body, model_name) + assert entry is not None, f"/model/info stopped listing {model_name!r} between the poll and the read" + return entry + + +def _serves(body: ModelsListResponse, model_name: str) -> bool: + return any(entry.id == model_name for entry in body.data) + + +def _assert_served_everywhere(client: ManagementClient, model_name: str) -> None: + _ = client.proxy.read_model_back_everywhere( + "/v1/models", ModelsListResponse, predicate=lambda body: _serves(body, model_name) + ) + + +def _assert_absent_everywhere(client: ManagementClient, model_name: str) -> None: + _ = client.proxy.read_model_back_everywhere( + "/v1/models", ModelsListResponse, predicate=lambda body: not _serves(body, model_name) + ) + + +def _assert_untouched_keys_as_created(entry: ModelInfoEntry, replica: str) -> None: + """The keys no later step names read back byte-for-byte as /model/new wrote them.""" + params: Final = entry.litellm_params + assert params.model == BACKEND_MODEL, f"{replica}: litellm_params.model {params.model!r} != {BACKEND_MODEL!r}" + assert params.api_base == PINNED_API_BASE, f"{replica}: api_base {params.api_base!r} != {PINNED_API_BASE!r}" + assert params.output_cost_per_token == PINNED_OUTPUT_RATE, ( + f"{replica}: output_cost_per_token {params.output_cost_per_token} != {PINNED_OUTPUT_RATE}" + ) + assert entry.model_info.mode == "chat", f"{replica}: model_info.mode {entry.model_info.mode!r} != 'chat'" + + +def _approx_equal(actual: float, expected: float) -> bool: + return math.isclose(actual, expected, rel_tol=1e-2, abs_tol=1e-9) + + +def _priced(rows: list[SpendLogRow]) -> bool: + return any(row.metadata and row.metadata.cost_breakdown and row.metadata.cost_breakdown.input_cost for row in rows) + + +def _billed_input_cost(client: ManagementClient, model_name: str, key: str) -> tuple[int, float]: + """Drive one chat completion through `model_name` and return the prompt tokens and + input cost its spend row recorded, so a test can assert the rate the gateway actually + billed rather than only the rate it stored.""" + chat: Final = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model_name, + messages=[ChatMessage(role="user", content=f"reply with one word {unique_marker()}")], + max_tokens=16, + ), + ) + ) + assert chat.id is not None, f"chat completion carried no id to find its spend row by: {chat}" + + rows: Final = client.proxy.poll_logs_for_request_id(chat.id, predicate=_priced) + row: Final = next((row for row in rows if row.request_id == chat.id), None) + assert row is not None and row.metadata and row.metadata.cost_breakdown, ( + f"no priced spend row for request {chat.id} before the deadline: {rows}" + ) + prompt_tokens: Final = row.prompt_tokens or 0 + input_cost: Final = row.metadata.cost_breakdown.input_cost + assert prompt_tokens > 0 and input_cost is not None, f"spend row logged no prompt tokens or input cost: {row}" + return prompt_tokens, input_cost + + +def _await_billed_input_cost( + client: ManagementClient, model_name: str, key: str, *, expected_rate: float +) -> tuple[int, float]: + """Drive calls through `model_name` until one is billed at `expected_rate`, and + return the prompt tokens and input cost of the last spend row either way. + + Only the deadline ends the wait unsatisfied: a rate that never reaches the gateway + comes back as the stale cost for the caller to assert on, so the rate the caller + expects is still what decides the test.""" + deadline: Final = time.monotonic() + BILLING_CONVERGENCE_TIMEOUT + while True: + prompt_tokens, input_cost = _billed_input_cost(client, model_name, key) + if _approx_equal(input_cost, prompt_tokens * expected_rate) or time.monotonic() >= deadline: + return prompt_tokens, input_cost + time.sleep(BILLING_CONVERGENCE_INTERVAL) + + +class TestModelLifecycle: + @pytest.mark.covers("mgmt.model.add.persists") + def test_create_reads_back_every_field_and_serves_on_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + registered = _register(client, resources) + + entry = _stored_entry(client, registered.model_name, converged=lambda _entry: True) + stored = "/model/info" + + _assert_untouched_keys_as_created(entry, stored) + assert entry.litellm_params.input_cost_per_token == PINNED_INPUT_RATE, ( + f"{stored}: input_cost_per_token {entry.litellm_params.input_cost_per_token} != {PINNED_INPUT_RATE}" + ) + assert entry.litellm_params.max_input_tokens == PINNED_MAX_INPUT_TOKENS, ( + f"{stored}: max_input_tokens {entry.litellm_params.max_input_tokens} != {PINNED_MAX_INPUT_TOKENS}" + ) + assert entry.model_info.id == registered.model_id, ( + f"{stored}: model_info.id {entry.model_info.id!r} != {registered.model_id!r}" + ) + + _assert_served_everywhere(client, registered.model_name) + + @pytest.mark.covers("mgmt.model.update.preserves_unrelated_fields") + def test_partial_update_changes_only_the_named_key( + self, client: ManagementClient, resources: ResourceManager, scoped_key: str + ) -> None: + registered = _register(client, resources) + + stored = client.proxy.patch_model( + registered.model_id, + ModelPatchBody(litellm_params=LiteLLMParamsPatch(input_cost_per_token=UPDATED_INPUT_RATE)), + ) + assert stored.litellm_params.input_cost_per_token == UPDATED_INPUT_RATE, ( + f"PATCH response stores input_cost_per_token {stored.litellm_params.input_cost_per_token}, " + f"sent {UPDATED_INPUT_RATE}" + ) + + entry = _stored_entry( + client, + registered.model_name, + converged=lambda entry: entry.litellm_params.input_cost_per_token == UPDATED_INPUT_RATE, + ) + stored = "/model/info" + + _assert_untouched_keys_as_created(entry, stored) + assert entry.litellm_params.max_input_tokens == PINNED_MAX_INPUT_TOKENS, ( + f"{stored}: max_input_tokens {entry.litellm_params.max_input_tokens} != {PINNED_MAX_INPUT_TOKENS}" + ) + assert entry.model_info.input_cost_per_token == UPDATED_INPUT_RATE, ( + f"{stored}: model_info.input_cost_per_token {entry.model_info.input_cost_per_token} " + f"did not mirror the updated {UPDATED_INPUT_RATE}" + ) + + prompt_tokens, input_cost = _await_billed_input_cost( + client, registered.model_name, scoped_key, expected_rate=UPDATED_INPUT_RATE + ) + assert _approx_equal(input_cost, prompt_tokens * UPDATED_INPUT_RATE), ( + f"input_cost {input_cost} != {prompt_tokens} tokens * updated rate {UPDATED_INPUT_RATE} " + f"= {prompt_tokens * UPDATED_INPUT_RATE}; the partial update did not reach billing" + ) + + @pytest.mark.covers("mgmt.model.update.clear_persists") + def test_explicit_null_removes_the_key_from_the_stored_row( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + registered = _register(client, resources) + + stored = client.proxy.patch_model( + registered.model_id, + ModelPatchBody(litellm_params=LiteLLMParamsPatch(max_input_tokens=Clear(), input_cost_per_token=Clear())), + ) + stored_params = stored.litellm_params.model_fields_set + assert "max_input_tokens" not in stored_params, ( + f"stored litellm_params still carries max_input_tokens " + f"{stored.litellm_params.max_input_tokens} after an explicit null" + ) + assert "input_cost_per_token" not in stored_params, ( + f"stored litellm_params still carries input_cost_per_token " + f"{stored.litellm_params.input_cost_per_token} after an explicit null" + ) + assert "max_input_tokens" not in stored.model_info.model_fields_set, ( + f"stored model_info carries max_input_tokens {stored.model_info.max_input_tokens} after the clear" + ) + assert "input_cost_per_token" not in stored.model_info.model_fields_set, ( + f"stored model_info still mirrors input_cost_per_token {stored.model_info.input_cost_per_token}" + ) + + cost_map_input_rate = client.proxy.model_cost_map()[BACKEND_MODEL].input_cost_per_token + assert cost_map_input_rate is not None, f"cost map has no input rate for {BACKEND_MODEL}" + entry = _stored_entry( + client, + registered.model_name, + converged=lambda entry: "max_input_tokens" not in entry.litellm_params.model_fields_set, + ) + stored = "/model/info" + + _assert_untouched_keys_as_created(entry, stored) + served = entry.litellm_params.model_fields_set + assert "max_input_tokens" not in served, ( + f"{stored}: litellm_params still serves max_input_tokens {entry.litellm_params.max_input_tokens}" + ) + assert "input_cost_per_token" not in served, ( + f"{stored}: litellm_params still serves input_cost_per_token {entry.litellm_params.input_cost_per_token}" + ) + assert entry.model_info.input_cost_per_token == cost_map_input_rate, ( + f"{stored}: model_info.input_cost_per_token {entry.model_info.input_cost_per_token} is not the " + f"cost map's {cost_map_input_rate}; the cleared override {PINNED_INPUT_RATE} still resolves" + ) + + @pytest.mark.covers("mgmt.model.update.clear_persists") + def test_cleared_price_is_billed_at_the_cost_map_rate( + self, client: ManagementClient, resources: ResourceManager, scoped_key: str + ) -> None: + registered = _register(client, resources) + _ = client.proxy.patch_model( + registered.model_id, + ModelPatchBody(litellm_params=LiteLLMParamsPatch(max_input_tokens=Clear(), input_cost_per_token=Clear())), + ) + _ = _stored_entry( + client, + registered.model_name, + converged=lambda entry: "input_cost_per_token" not in entry.litellm_params.model_fields_set, + ) + cost_map_input_rate = client.proxy.model_cost_map()[BACKEND_MODEL].input_cost_per_token + assert cost_map_input_rate is not None, f"cost map has no input rate for {BACKEND_MODEL}" + + prompt_tokens, input_cost = _await_billed_input_cost( + client, registered.model_name, scoped_key, expected_rate=cost_map_input_rate + ) + + assert _approx_equal(input_cost, prompt_tokens * cost_map_input_rate), ( + f"input_cost {input_cost} != {prompt_tokens} tokens * cost map rate {cost_map_input_rate} " + f"= {prompt_tokens * cost_map_input_rate}" + ) + assert not _approx_equal(input_cost, prompt_tokens * PINNED_INPUT_RATE), ( + f"input_cost {input_cost} is still billed at the cleared override {PINNED_INPUT_RATE}" + ) + + @pytest.mark.covers("mgmt.model.delete.persists") + def test_delete_removes_the_deployment_everywhere( + self, client: ManagementClient, resources: ResourceManager, scoped_key: str + ) -> None: + registered = _register(client, resources) + _ = _stored_entry(client, registered.model_name, converged=lambda _entry: True) + + client.delete_model_strict(registered.model_id) + + _assert_absent_everywhere(client, registered.model_name) + refused = client.chat_status(scoped_key, registered.model_name, "hi this is a test") + assert refused.status_code == 400, ( + f"chat against the deleted model must be rejected 400, got {refused.status_code}: {refused.body[:300]}" + ) + envelope = _ErrorEnvelope.model_validate_json(refused.body) + assert envelope.error.message, f"400 body must carry an error message: {refused.body[:300]}" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 62810e6cfd9..8db37bd25a5 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -655,6 +655,11 @@ class OcrResponse(BaseModel): # ---------- spend logs ---------- +class CostBreakdown(BaseModel): + input_cost: float | None = None + output_cost: float | None = None + + class GuardrailEntityMatch(BaseModel): entity_type: str score: float @@ -672,6 +677,7 @@ class GuardrailRunRecord(BaseModel): class SpendLogMetadata(BaseModel): + cost_breakdown: CostBreakdown | None = None applied_guardrails: list[str] | None = None guardrail_information: list[GuardrailRunRecord] | None = None @@ -815,15 +821,42 @@ class CustomPricing(BaseModel): return prompt_tokens * self.input_cost_per_token + completion_tokens * self.output_cost_per_token +class DeploymentParams(CustomPricing): + """The litellm_params half of a /model/info row: the stored deployment as written, + credentials scrubbed. Unlike model_info it is never back-filled from the cost map, + so a key the store dropped is absent here (check `model_fields_set`).""" + + model: str | None = None + api_base: str | None = None + max_input_tokens: int | None = None + + +class DeploymentModelInfo(CustomPricing): + id: str | None = None + max_input_tokens: int | None = None + + class ModelInfoEntry(BaseModel): """One /model/info row. `litellm_params` is the configured deployment (carries any custom-pricing override); `model_info` is the price the proxy resolved for - it - the override merged over the cost-map defaults.""" + it - the override merged over the cost-map defaults, so a key cleared from the + stored blob reads as the cost-map default here.""" model_config = ConfigDict(protected_namespaces=()) model_name: str - litellm_params: CustomPricing = CustomPricing() - model_info: CustomPricing = CustomPricing() + litellm_params: DeploymentParams = DeploymentParams() + model_info: DeploymentModelInfo = DeploymentModelInfo() + + +class StoredDeployment(BaseModel): + """PATCH /model/{model_id}/update answers with the row as stored: both blobs raw, + nothing back-filled, so a cleared key is absent from `model_fields_set` of the + blob it was cleared from.""" + + model_config = ConfigDict(protected_namespaces=()) + model_name: str + litellm_params: DeploymentParams + model_info: DeploymentModelInfo class ModelInfoResponse(BaseModel): @@ -924,9 +957,10 @@ class LiteLLMParamsBody(BaseModel): timeout: float | None = None tpm: int | None = None weight: int | None = None + max_input_tokens: int | None = None -ModelMode = Literal["batch", "realtime", "image_generation"] +ModelMode = Literal["chat", "batch", "realtime", "image_generation"] class ModelInfoBody(BaseModel): @@ -936,6 +970,7 @@ class ModelInfoBody(BaseModel): # constraint when a prior run's teardown had not removed the row. id: str | None = None mode: ModelMode | None = None + max_input_tokens: int | None = None access_groups: list[str] | None = None team_id: str | None = None allowed_fails_policy: dict[str, int] | None = None @@ -964,6 +999,37 @@ class ModelUpdateBody(BaseModel): model_info: ModelInfoBody +class Clear(BaseModel): + """Serializes to JSON null. The transport dumps every body with exclude_none, so a + field set to this is how a patch carries the explicit null that removes a stored key.""" + + @model_serializer + def _as_null(self) -> None: + return None + + +class LiteLLMParamsPatch(BaseModel): + api_base: str | Clear | None = None + max_input_tokens: int | Clear | None = None + input_cost_per_token: float | Clear | None = None + output_cost_per_token: float | Clear | None = None + + +class ModelInfoPatch(BaseModel): + mode: ModelMode | Clear | None = None + max_input_tokens: int | Clear | None = None + + +class ModelPatchBody(BaseModel): + """PATCH /model/{model_id}/update body, JSON Merge Patch over the stored deployment: + a field left None is dropped from the body and unchanged, a field set to `Clear()` + is sent as null and removed, a field with a value is set.""" + + model_config = ConfigDict(protected_namespaces=()) + litellm_params: LiteLLMParamsPatch | None = None + model_info: ModelInfoPatch | None = None + + class ModelListEntry(BaseModel): id: str diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 1bac5116a9d..fa1b06fe7ed 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -10,10 +10,10 @@ from __future__ import annotations import time import warnings -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass -from functools import reduce from datetime import datetime +from functools import reduce from types import MappingProxyType from typing import Final @@ -62,6 +62,7 @@ from models import ( ModelMode, ModelNewBody, ModelNewResponse, + ModelPatchBody, ModelsListParams, ModelsListResponse, ModelUpdateBody, @@ -72,6 +73,7 @@ from models import ( SpendLogsPage, SpendLogsPageParams, SpendLogsParams, + StoredDeployment, ToolsetCreateBody, ToolsetRow, ToolsetUpdateBody, @@ -132,6 +134,103 @@ class NotServableOn: last_result: Result[ModelsListResponse] | None +type BodyReader[R: BaseModel] = Callable[[float], Result[R]] + + +@dataclass(frozen=True, slots=True) +class BodyNotConverged[R: BaseModel]: + """The deadline passed without a read the predicate accepted; `last_result` is the + final read, so the caller can tell a body that never matched from a read that + failed.""" + + last_result: Result[R] | None + + +@dataclass(frozen=True, slots=True) +class BodyConverged[R: BaseModel]: + """Every replica answered a body the predicate accepted; `bodies` is the last read + per replica.""" + + bodies: Mapping[str, R] + + +@dataclass(frozen=True, slots=True) +class BodyNeverConvergedOn[R: BaseModel]: + """`BodyNotConverged` labeled with the replica whose reads never satisfied the predicate.""" + + replica: str + last_result: Result[R] | None + + +def await_body_converged[R: BaseModel]( + read: BodyReader[R], + *, + predicate: Callable[[R], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> Success[R] | BodyNotConverged[R]: + """Poll `read` until it answers a body `predicate` accepts, or `timeout` passes. + + Each read's request timeout is clamped to the remaining budget, and the sleep + between reads to the time left, so the last read before the deadline is never + skipped. Clock and sleep are injected.""" + deadline: Final = now() + timeout + + def reads() -> Iterator[Result[R]]: + while (remaining := deadline - now()) > 0: + yield read(min(request_timeout, remaining)) + sleep(min(interval, max(deadline - now(), 0.0))) + + def attempts() -> Iterator[Success[R] | BodyNotConverged[R]]: + for result in reads(): + if isinstance(result, Success) and predicate(result.data): + yield result + return + yield BodyNotConverged(last_result=result) + + initial: Final[Success[R] | BodyNotConverged[R]] = BodyNotConverged(last_result=None) + return reduce(lambda _previous, result: result, attempts(), initial) + + +def await_body_converged_everywhere[R: BaseModel]( + readers: Mapping[str, BodyReader[R]], + *, + predicate: Callable[[R], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> BodyConverged[R] | BodyNeverConvergedOn[R]: + """`await_body_converged` against every replica in turn, each with the full budget, so a + write counts as landed only once every replica serves it.""" + def read_replica( + outcome: BodyConverged[R] | BodyNeverConvergedOn[R], + item: tuple[str, BodyReader[R]], + ) -> BodyConverged[R] | BodyNeverConvergedOn[R]: + if isinstance(outcome, BodyNeverConvergedOn): + return outcome + replica, read = item + match await_body_converged( + read, + predicate=predicate, + timeout=timeout, + interval=interval, + request_timeout=request_timeout, + now=now, + sleep=sleep, + ): + case Success(data=data): + return BodyConverged(bodies=MappingProxyType({**outcome.bodies, replica: data})) + case BodyNotConverged(last_result=last_result): + return BodyNeverConvergedOn(replica=replica, last_result=last_result) + initial: Final[BodyConverged[R] | BodyNeverConvergedOn[R]] = BodyConverged(bodies=MappingProxyType({})) + return reduce(read_replica, readers.items(), initial) + + def await_servable( list_models: Callable[[float], Result[ModelsListResponse]], *, @@ -658,6 +757,102 @@ class ProxyClient: ) ) + def patch_model(self, model_id: str, body: ModelPatchBody) -> StoredDeployment: + """JSON Merge Patch the deployment `model_id` via PATCH /model/{model_id}/update: + a field the body omits is unchanged, one sent as null is removed from the stored + row, one sent with a value is set. See ModelPatchBody for how a null is sent. + Returns the row as stored after the write.""" + return unwrap( + self.transport.patch( + f"/model/{model_id}/update", + headers=self.transport.master, + json=body, + response_type=StoredDeployment, + ) + ) + + def read_model_back_everywhere[R: BaseModel]( + self, path: str, response_type: type[R], *, predicate: Callable[[R], bool] + ) -> Mapping[str, R]: + """GET `path` on every replica until each answers a body `predicate` accepts, + polling to poll_timeout, and return the last body per replica. + + Fails naming the replica that never converged, so a write that reached one + gateway but not the others is caught instead of passing on whichever gateway + the balancer answered from. Falls back to the single proxy address when no + replica list is configured. + + `path` must be a data-plane route. The replicas are gateways, which serve only + the LLM surface, so a control-plane path answers on exactly one service and + 404s on every replica in a split deployment: asking each replica for one is + never the question the caller means. Read those through `self.transport` + instead, which routes them to the control plane.""" + if is_control_plane_path(path): + raise AssertionError( + f"read_model_back_everywhere({path!r}) asks every data-plane replica for a control-plane route. " + "The replicas are gateways and do not serve it; poll a data-plane path such as /v1/models " + "here, and read the control plane through the shared transport." + ) + readers: Final = { + url: self._body_reader(transport, path, response_type) + for url, transport in self._read_back_replicas().items() + } + outcome: Final = await_body_converged_everywhere( + readers, + predicate=predicate, + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case BodyConverged(bodies=bodies): + return bodies + case BodyNeverConvergedOn(replica=replica, last_result=last_result): + raise AssertionError( + f"GET {path} on {replica} never answered the expected body within " + f"{self.poll_timeout}s; last read: {last_result}" + ) + + def read_model_back[R: BaseModel](self, path: str, response_type: type[R], *, predicate: Callable[[R], bool]) -> R: + """GET `path` through the shared transport until the body satisfies `predicate`, + polling to poll_timeout, and return that body. + + The counterpart to `read_model_back_everywhere` for a control-plane route such as + /model/info: the stored row lives in one database behind one control plane, so + there is a single answer to converge on rather than one per gateway.""" + outcome: Final = await_body_converged_everywhere( + {CONTROL_PLANE_BASE_URL: self._body_reader(self.transport, path, response_type)}, + predicate=predicate, + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case BodyConverged(bodies=bodies): + return bodies[CONTROL_PLANE_BASE_URL] + case BodyNeverConvergedOn(last_result=last_result): + raise AssertionError( + f"GET {path} never answered the expected body within " + f"{self.poll_timeout}s; last read: {last_result}" + ) + + def _read_back_replicas(self) -> Mapping[str, Transport]: + return self.replicas or MappingProxyType({CONTROL_PLANE_BASE_URL: self.transport}) + + @staticmethod + def _body_reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> BodyReader[R]: + return lambda timeout: transport.get( + path, + headers=transport.master, + params=NoBody(), + response_type=response_type, + timeout=timeout, + ) + def delete_model(self, model_id: str) -> None: result = self.transport.post( "/model/delete", diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 3b84a47e3cc..cbf7f5648d4 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -20,8 +20,12 @@ from typing import Final, cast import pytest from e2e_config import parse_replica_urls from e2e_http import Result, Success -from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse +from models import KeyInfo, KeyInfoResponse, ModelInfoEntry, ModelInfoResponse, ModelListEntry, ModelsListResponse from proxy_client import ( + BodyReader, + BodyConverged, + BodyNeverConvergedOn, + await_body_converged_everywhere, ConvergeOutcome, Converged, EverywhereConverged, @@ -274,3 +278,54 @@ class TestReplicasFor: client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={}) with pytest.raises(AssertionError, match="no replica is configured"): _ = client.replicas_for("/v1/models") + + +def _info(*model_names: str) -> Success[ModelInfoResponse]: + entries: Final = [ModelInfoEntry(model_name=model_name) for model_name in model_names] + return Success(status_code=200, data=ModelInfoResponse(data=entries)) + + +def _reader(results: Iterable[Success[ModelInfoResponse]]) -> BodyReader[ModelInfoResponse]: + it: Final = iter(results) + return lambda _timeout: next(it) + + +def _lists_model(body: ModelInfoResponse) -> bool: + return any(entry.model_name == MODEL for entry in body.data) + + +def _read_back( + readers: Mapping[str, BodyReader[ModelInfoResponse]], +) -> tuple[BodyConverged[ModelInfoResponse] | BodyNeverConvergedOn[ModelInfoResponse], FakeClock]: + clock: Final = FakeClock() + outcome: Final = await_body_converged_everywhere( + readers, + predicate=_lists_model, + timeout=TIMEOUT, + interval=INTERVAL, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + return outcome, clock + + +class TestAwaitBodyConvergedEverywhere: + def test_waits_for_the_lagging_replica_and_returns_every_body(self) -> None: + readers: Final = { + "gateway-1": _reader(repeat(_info(MODEL))), + "gateway-2": _reader(chain(repeat(_info(), 2), repeat(_info(MODEL)))), + } + outcome, clock = _read_back(readers) + assert outcome == BodyConverged(bodies={"gateway-1": _info(MODEL).data, "gateway-2": _info(MODEL).data}) + assert clock.elapsed == 2 * INTERVAL + + @pytest.mark.parametrize("lagging", ["gateway-1", "gateway-2"]) + def test_fails_naming_the_replica_that_never_converges(self, lagging: str) -> None: + readers: Final = { + "gateway-1": _reader(repeat(_info(MODEL))), + "gateway-2": _reader(repeat(_info(MODEL))), + } | {lagging: _reader(repeat(_info()))} + outcome, clock = _read_back(readers) + assert outcome == BodyNeverConvergedOn(replica=lagging, last_result=_info()) + assert clock.elapsed >= TIMEOUT diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 44fdbaa3e41..804e073a4a0 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -306,6 +306,7 @@ CONTROL_PLANE_PREFIXES: tuple[str, ...] = ( "/config", "/guardrails", "/openapi.json", + "/public/", ) 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 c02f886fc31..1376727e296 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 @@ -3115,9 +3115,6 @@ class TestUpdateDBModelClearPricing: """Sending an explicit `null` for a pricing field must remove it from both `litellm_params` and `model_info` (SPECIAL_MODEL_INFO_PARAMS are mirrored between the two by Deployment.__init__). - - Restricted to SPECIAL_MODEL_INFO_PARAMS so non-pricing fields (e.g. team_id) - cannot be cleared via this path. """ def test_clear_input_cost_removes_from_both_blobs(self): @@ -3193,10 +3190,10 @@ class TestUpdateDBModelClearPricing: assert params["input_cost_per_token"] == 0.000001 assert params["output_cost_per_token"] == 0.000007 - def test_null_on_non_pricing_field_does_not_clear(self): - """Security guard: only SPECIAL_MODEL_INFO_PARAMS can be cleared via null. - Privileged or unrelated model_info fields (e.g. team_id) must be unaffected - by the null-clearing path so a team admin can't ungate a team-scoped model. + def test_null_on_one_field_leaves_other_fields_alone(self): + """A null clears only the key it names: pricing the patch never mentions and + the ownership key team_id stay put, so a team admin can't ungate a + team-scoped model through the clear path. """ from litellm.proxy.management_endpoints.model_management_endpoints import ( update_db_model, @@ -3217,8 +3214,6 @@ class TestUpdateDBModelClearPricing: model_info=ModelInfo(id="dep-pricing-1", team_id="team-keep-me"), ) - # Patch sends a null for api_base (non-SPECIAL field). Must NOT clear team_id - # or any other non-pricing field from the merged dict. result = update_db_model( db_model=db_model, updated_patch=updateDeployment( @@ -3394,6 +3389,171 @@ class TestUpdateDBModelClearPricing: assert info["cache_creation_input_token_cost"] == 0.000003 +_PROTECTED_MODEL_INFO_VALUES = { + "team_id": "team-keep-me", + "team_public_model_name": "team-facing-name", + "access_groups": ["group-a"], + "created_at": "2026-01-01T00:00:00+00:00", + "created_by": "creator", + "updated_at": "2026-01-02T00:00:00+00:00", + "updated_by": "updater", + "blocked": True, +} + + +def _build_db_model_with_pinned_model_info(): + """Deployment whose stored blobs pin non-pricing keys an earlier save wrote, next to a + pricing override, so a clear can be checked key by key.""" + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + return Deployment( + model_name="pinned-gpt-4o-mini", + litellm_params=LiteLLM_Params( + model="gpt-4o-mini", input_cost_per_token=0.000001, max_input_tokens=4096 + ), + model_info=ModelInfo( + id="dep-pinned-0", + max_input_tokens=4096, + mode="chat", + supports_vision=True, + **_PROTECTED_MODEL_INFO_VALUES, + ), + ) + + +class TestUpdateDBModelNullClearsAnyKey: + """JSON Merge Patch on PATCH /model/{id}/update: a key sent as null is removed from the + stored blob it was sent in, whatever the key, except the identity and ownership keys, + whose nulls are ignored.""" + + def test_model_info_nulls_remove_pinned_non_pricing_keys(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate( + {"model_info": {"max_input_tokens": None, "mode": None}} + ), + ) + + info = json.loads(result["model_info"]) + assert "max_input_tokens" not in info + assert "mode" not in info + assert info["supports_vision"] is True + assert info["input_cost_per_token"] == 0.000001 + + def test_litellm_params_null_removes_pinned_non_pricing_key(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate( + {"litellm_params": {"max_input_tokens": None}} + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "max_input_tokens" not in params + assert params["model"] == "gpt-4o-mini" + assert params["input_cost_per_token"] == 0.000001 + assert info["max_input_tokens"] == 4096 + + def test_omitted_key_is_untouched_by_a_null_elsewhere(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate( + {"model_info": {"mode": None, "supports_vision": False}} + ), + ) + + info = json.loads(result["model_info"]) + assert "mode" not in info + assert info["supports_vision"] is False + assert info["max_input_tokens"] == 4096 + + @pytest.mark.parametrize("field", sorted(_PROTECTED_MODEL_INFO_VALUES)) + def test_null_on_protected_key_is_ignored(self, field): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate({"model_info": {field: None}}), + ) + + info = json.loads(result["model_info"]) + assert info[field] == _PROTECTED_MODEL_INFO_VALUES[field] + assert info["max_input_tokens"] == 4096 + + def test_echoing_the_read_back_blob_preserves_every_stored_key(self): + """The Admin UI edit form submits the whole /model/info row back, and that read reports + every key the deployment never stored as an explicit null. Those nulls have to stay + no-ops: a write drops None before storing, so a null in the echoed blob always names a + key the stored row does not carry. + """ + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + db_model = _build_db_model_with_pinned_model_info() + echoed = { + "id": "dep-pinned-0", + "max_input_tokens": 4096, + "mode": "chat", + "supports_vision": True, + "input_cost_per_token": 0.000001, + "team_id": "team-keep-me", + "base_model": None, + "tier": None, + "max_output_tokens": None, + "supports_function_calling": None, + "cache_read_input_token_cost": None, + } + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment.model_validate({"model_info": echoed}), + ) + + info = json.loads(result["model_info"]) + assert info["max_input_tokens"] == 4096 + assert info["mode"] == "chat" + assert info["supports_vision"] is True + assert info["input_cost_per_token"] == 0.000001 + assert info["team_id"] == "team-keep-me" + for never_stored in ("base_model", "tier", "max_output_tokens", "supports_function_calling"): + assert never_stored not in info + + def test_null_on_pricing_key_still_clears_both_blobs(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate( + {"model_info": {"input_cost_per_token": None}} + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "input_cost_per_token" not in params + assert "input_cost_per_token" not in info + assert params["max_input_tokens"] == 4096 + assert info["max_input_tokens"] == 4096 + + class TestGetModelInfoWithIdBlocked: """`ProxyConfig.get_model_info_with_id` must propagate the DB-level `blocked` column into the in-memory `model_info` dict so the router filter can read it.""" diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 30b265905f3..d22ec60e61a 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -220,6 +220,198 @@ def test_should_store_full_pricing_under_deployment_model_id(): assert entry["output_cost_per_token"] == 0.0 +def test_should_drop_a_price_the_deployment_no_longer_carries(): + """Re-registering a deployment must replace its model_id entry, not merge onto it. + + A merge left the old rate in the cost map after an operator cleared the override, so + the deployment kept billing at a price its config no longer had. + """ + backend_model = "vertex_ai/gemini-2.5-flash" + model_id = "deployment-cleared-price" + original = {model_id: litellm.model_cost.get(model_id)} + + try: + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"input_cost_per_token": 0.005, "output_cost_per_token": 0.01}, + model=backend_model, + custom_llm_provider="vertex_ai", + ) + assert litellm.model_cost[model_id]["input_cost_per_token"] == 0.005 + + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"mode": "chat"}, + model=backend_model, + custom_llm_provider="vertex_ai", + ) + + entry = litellm.model_cost[model_id] + assert entry.get("input_cost_per_token") != 0.005, ( + "the cleared override survived re-registration, so the deployment still bills at it" + ) + assert entry.get("output_cost_per_token") != 0.01 + finally: + _restore_model_cost_entries(original) + + +def test_should_not_strip_a_builtin_entry_when_a_deployment_id_collides_with_it(): + """Deployments are keyed into the same cost map as the built-in catalog, so a deployment + whose id happens to name a real model must not evict that model's entry. + + Stripping it would take the pricing and capability flags every other deployment of that + model reads, process-wide, until the next price-map reload. Registering twice, because + the first registration is what would mark the entry as this deployment's own. + """ + colliding_id = "gpt-4o" + original = {colliding_id: litellm.model_cost.get(colliding_id)} + builtin_max_tokens = litellm.model_cost[colliding_id]["max_tokens"] + + try: + for _ in range(2): + Router._register_deployment_in_model_cost( + model_id=colliding_id, + model_info={"id": colliding_id, "db_model": True, "mode": "chat"}, + model="gpt-4o-mini", + custom_llm_provider="openai", + ) + + entry = litellm.model_cost[colliding_id] + assert entry["max_tokens"] == builtin_max_tokens, ( + "registering a deployment under a catalog model's name wiped that model's context window" + ) + assert entry["litellm_provider"] == "openai" + assert entry["supports_vision"] is True + finally: + _restore_model_cost_entries(original) + + +def test_should_drop_a_stale_price_even_when_the_deployment_declares_a_provider(): + """A deployment may carry `litellm_provider` in its own model_info, which must not be + read as "this is a catalog entry" and stop the stale price from being dropped.""" + model_id = "deployment-provider-tagged" + original = {model_id: litellm.model_cost.get(model_id)} + + try: + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"id": model_id, "litellm_provider": "openai", "input_cost_per_token": 0.005}, + model="gpt-4o-mini", + custom_llm_provider="openai", + ) + assert litellm.model_cost[model_id]["input_cost_per_token"] == 0.005 + + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"id": model_id, "litellm_provider": "openai", "mode": "chat"}, + model="gpt-4o-mini", + custom_llm_provider="openai", + ) + + assert litellm.model_cost[model_id].get("input_cost_per_token") != 0.005, ( + "a deployment that declares its provider kept billing at the price it no longer carries" + ) + finally: + _restore_model_cost_entries(original) + + +def test_should_give_a_cost_map_key_back_when_the_deployment_is_deleted(): + """Deleting a deployment releases its claim on the shared cost-map key. + + Held forever, a later catalog refresh that starts publishing a model under that same + name would be treated as the deleted deployment's own entry and evicted. + """ + from litellm.router import _DEPLOYMENT_COST_MAP_KEYS + + model_id = "deployment-to-delete" + original = {model_id: litellm.model_cost.get(model_id)} + router = Router( + model_list=[ + { + "model_name": "to-delete", + "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"id": model_id, "input_cost_per_token": 0.005}, + } + ] + ) + + try: + assert model_id in _DEPLOYMENT_COST_MAP_KEYS + + assert router.delete_deployment(id=model_id) is not None + + assert model_id not in _DEPLOYMENT_COST_MAP_KEYS, ( + "a deleted deployment kept its claim on the shared cost-map key" + ) + finally: + _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) + _restore_model_cost_entries(original) + + +def test_should_keep_the_cost_map_key_while_another_router_still_serves_it(): + """Two live routers can serve the same deployment id, and the claim is process-wide. + + Releasing it when only one of them drops the deployment would put the survivor back on + merging, so the price it just cleared would keep billing. + """ + from litellm.router import _DEPLOYMENT_COST_MAP_KEYS + + model_id = "deployment-served-twice" + original = {model_id: litellm.model_cost.get(model_id)} + entry = { + "model_name": "served-twice", + "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"id": model_id, "input_cost_per_token": 0.005}, + } + first = Router(model_list=[entry]) + second = Router(model_list=[entry]) + + try: + assert model_id in _DEPLOYMENT_COST_MAP_KEYS + + assert first.delete_deployment(id=model_id) is not None + + assert model_id in _DEPLOYMENT_COST_MAP_KEYS, ( + "the claim was released while another router still served the deployment" + ) + + assert second.delete_deployment(id=model_id) is not None + assert model_id not in _DEPLOYMENT_COST_MAP_KEYS + finally: + _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) + _restore_model_cost_entries(original) + + +def test_should_keep_the_cost_map_key_while_a_dynamically_built_router_serves_it(): + """A router built with no model_list still serves whatever add_deployment gives it, so it + counts when deciding whether the shared cost-map claim can be released.""" + from litellm.router import _DEPLOYMENT_COST_MAP_KEYS + + model_id = "deployment-added-dynamically" + original = {model_id: litellm.model_cost.get(model_id)} + entry = { + "model_name": "added-dynamically", + "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"id": model_id, "input_cost_per_token": 0.005}, + } + configured = Router(model_list=[entry]) + dynamic = Router() + dynamic.add_deployment(deployment=Deployment(**entry)) + + try: + assert configured.delete_deployment(id=model_id) is not None + + assert model_id in _DEPLOYMENT_COST_MAP_KEYS, ( + "the claim was released while a dynamically built router still served the deployment" + ) + + assert dynamic.delete_deployment(id=model_id) is not None + assert model_id not in _DEPLOYMENT_COST_MAP_KEYS + finally: + _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) + _restore_model_cost_entries(original) + + def test_should_preserve_builtin_pricing_regardless_of_deployment_order(): """ The built-in pricing should be preserved no matter which deployment diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 83b0d58f2b2..93aaf3ca58c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9064,8 +9064,9 @@ export interface paths { * Patch Model * @description PATCH Endpoint for partial model updates. * - * Only updates the fields specified in the request while preserving other existing values. - * Follows proper PATCH semantics by only modifying provided fields. + * JSON Merge Patch semantics over `litellm_params` and `model_info`: a key absent from the + * body is unchanged, a key sent as null is removed from the stored row, and a key sent with a + * value is set (identity and ownership keys such as `id` and `team_id` ignore a null). * * Args: * model_id: The ID of the model to update