From 2bb55035bf1ae23aae2058c4cb2c48b2752e9697 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sun, 6 Sep 2026 09:47:37 +0000 Subject: [PATCH] fix(model-management): honor an explicit null as a clear on model update PATCH /model/{model_id}/update merged the patch with exclude_none and then popped explicit nulls only for the mirrored pricing fields, so a null sent for max_input_tokens, mode, supports_vision or any other key was dropped and a value pinned by an earlier save could never be removed. The route now follows JSON Merge Patch over both blobs: 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. Ownership and identity keys keep ignoring a null, as do the fields the stored models require, since clearing one writes a row no reload can rebuild. Mirrored pricing keys still clear from both blobs. Clearing a price also needed the router to stop merging a deployment's cost-map entry onto its previous registration, which left the old rate in place and kept billing at a price the deployment no longer carried. Adds a create, read, partial-update, clear, enforce, delete lifecycle e2e that reads back on every replica, and a harness helper for that read-back. --- .../model_management_endpoints.py | 71 ++-- litellm/router.py | 6 +- tests/e2e/coverage_registry/mgmt.yaml | 2 + .../management/test_model_lifecycle_e2e.py | 317 ++++++++++++++++++ tests/e2e/models.py | 80 ++++- tests/e2e/proxy_client.py | 146 ++++++++ tests/e2e/test_proxy_client.py | 69 +++- .../test_model_management_endpoints.py | 139 +++++++- .../test_router_model_cost_isolation.py | 35 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +- 10 files changed, 829 insertions(+), 41 deletions(-) create mode 100644 tests/e2e/management/test_model_lifecycle_e2e.py diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 0f19e9ce149..113477dbbdd 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,45 @@ 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", + } +) + +# Clearing a required field writes a row no reload can rebuild: both blobs load through +# LiteLLM_Params / ModelInfo, which reject it. +_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 and the keys the stored models require + are left alone, and 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 +788,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 +846,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 3f450661946..987c231bffb 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9745,7 +9745,10 @@ 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. The + deployment's own ``model_id`` entry 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. 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. @@ -9762,6 +9765,7 @@ class Router: } if model_id is not None: + litellm.model_cost.pop(model_id, None) litellm.register_model( model_cost={model_id: model_info}, persist_across_reloads=False, diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index d571fb36546..55123decabc 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -74,3 +74,5 @@ - {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 on every replica"} +- {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"} 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..00b995553bb --- /dev/null +++ b/tests/e2e/management/test_model_lifecycle_e2e.py @@ -0,0 +1,317 @@ +"""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. + +Every read-back goes through ProxyClient.read_back_everywhere, which polls /model/info +on every URL in PROXY_REPLICA_URLS, so a write that reached only one gateway fails +naming the gateway that never converged. +""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Mapping +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, + 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 + + +@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 _entry_everywhere( + client: ManagementClient, + model_name: str, + *, + converged: Callable[[ModelInfoEntry], bool], +) -> Mapping[str, ModelInfoEntry]: + """The /model/info row for `model_name` from every replica, once each replica's + row satisfies `converged`.""" + + def has_converged(body: ModelInfoResponse) -> bool: + entry: Final = _entry(body, model_name) + return entry is not None and converged(entry) + + bodies: Final = client.proxy.read_back_everywhere("/model/info", ModelInfoResponse, predicate=has_converged) + return {replica: entry for replica, body in bodies.items() if (entry := _entry(body, model_name)) is not None} + + +def _assert_absent_everywhere(client: ManagementClient, model_name: str) -> None: + def gone(body: ModelInfoResponse) -> bool: + return _entry(body, model_name) is None + + _ = client.proxy.read_back_everywhere("/model/info", ModelInfoResponse, predicate=gone) + + +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 + + +class TestModelLifecycle: + @pytest.mark.covers("mgmt.model.add.persists") + def test_create_reads_back_every_field_on_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + registered = _register(client, resources) + + entries = _entry_everywhere(client, registered.model_name, converged=lambda _entry: True) + + for replica, entry in entries.items(): + _assert_untouched_keys_as_created(entry, replica) + assert entry.litellm_params.input_cost_per_token == PINNED_INPUT_RATE, ( + f"{replica}: 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"{replica}: max_input_tokens {entry.litellm_params.max_input_tokens} != {PINNED_MAX_INPUT_TOKENS}" + ) + assert entry.model_info.id == registered.model_id, ( + f"{replica}: model_info.id {entry.model_info.id!r} != {registered.model_id!r}" + ) + + @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}" + ) + + entries = _entry_everywhere( + client, + registered.model_name, + converged=lambda entry: entry.litellm_params.input_cost_per_token == UPDATED_INPUT_RATE, + ) + + for replica, entry in entries.items(): + _assert_untouched_keys_as_created(entry, replica) + assert entry.litellm_params.max_input_tokens == PINNED_MAX_INPUT_TOKENS, ( + f"{replica}: 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"{replica}: 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 = _billed_input_cost(client, registered.model_name, scoped_key) + 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}" + entries = _entry_everywhere( + client, + registered.model_name, + converged=lambda entry: "max_input_tokens" not in entry.litellm_params.model_fields_set, + ) + + for replica, entry in entries.items(): + _assert_untouched_keys_as_created(entry, replica) + served = entry.litellm_params.model_fields_set + assert "max_input_tokens" not in served, ( + f"{replica}: litellm_params still serves max_input_tokens {entry.litellm_params.max_input_tokens}" + ) + assert "input_cost_per_token" not in served, ( + f"{replica}: litellm_params still serves input_cost_per_token " + f"{entry.litellm_params.input_cost_per_token}" + ) + assert entry.model_info.input_cost_per_token == cost_map_input_rate, ( + f"{replica}: 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())), + ) + _ = _entry_everywhere( + 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 = _billed_input_cost(client, registered.model_name, scoped_key) + + 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) + _ = _entry_everywhere(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 2c6c0e9bbd4..13ae18c7be0 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -10,7 +10,7 @@ from collections.abc import Sequence from datetime import datetime from typing import Literal -from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_validator +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_serializer, model_validator # ---------- keys ---------- @@ -573,8 +573,18 @@ class OcrResponse(BaseModel): # ---------- spend logs ---------- +class CostBreakdown(BaseModel): + input_cost: float | None = None + output_cost: float | None = None + + +class SpendLogMetadata(BaseModel): + cost_breakdown: CostBreakdown | None = None + + class SpendLogRow(BaseModel): request_id: str | None = None + metadata: SpendLogMetadata | None = None api_key: str | None = None model: str | None = None spend: float | None = None @@ -711,15 +721,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): @@ -820,9 +857,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): @@ -832,6 +870,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 @@ -860,6 +899,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 3b15e57dfcf..a201c4c7be9 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -16,6 +16,8 @@ from datetime import datetime from types import MappingProxyType from typing import Final +from pydantic import BaseModel + from e2e_http import ( AnthropicHeaders, AuthHeaders, @@ -58,6 +60,7 @@ from models import ( ModelMode, ModelNewBody, ModelNewResponse, + ModelPatchBody, ModelsListParams, ModelsListResponse, ModelUpdateBody, @@ -68,6 +71,7 @@ from models import ( SpendLogsPage, SpendLogsPageParams, SpendLogsParams, + StoredDeployment, ) from e2e_config import ( CONTROL_PLANE_BASE_URL, @@ -125,6 +129,89 @@ class NotServableOn: last_result: Result[ModelsListResponse] | None +type BodyReader[R: BaseModel] = Callable[[float], Result[R]] + + +@dataclass(frozen=True, slots=True) +class NotConverged[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 Converged[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 NeverConvergedOn[R: BaseModel]: + """`NotConverged` labeled with the replica whose reads never satisfied the predicate.""" + + replica: str + last_result: Result[R] | None + + +def await_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] | NotConverged[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 + last_result: Result[R] | None = None + while (remaining := deadline - now()) > 0: + last_result = read(min(request_timeout, remaining)) + if isinstance(last_result, Success) and predicate(last_result.data): + return last_result + sleep(min(interval, max(deadline - now(), 0.0))) + return NotConverged(last_result=last_result) + + +def await_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], +) -> Converged[R] | NeverConvergedOn[R]: + """`await_converged` against every replica in turn, each with the full budget, so a + write counts as landed only once every replica serves it.""" + bodies: dict[str, R] = {} + for replica, read in readers.items(): + match await_converged( + read, + predicate=predicate, + timeout=timeout, + interval=interval, + request_timeout=request_timeout, + now=now, + sleep=sleep, + ): + case Success(data=data): + bodies[replica] = data + case NotConverged(last_result=last_result): + return NeverConvergedOn(replica=replica, last_result=last_result) + return Converged(bodies=MappingProxyType(bodies)) + + def await_servable( list_models: Callable[[float], Result[ModelsListResponse]], *, @@ -433,6 +520,65 @@ 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_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.""" + readers: Final = { + url: self._body_reader(transport, path, response_type) + for url, transport in self._read_back_replicas().items() + } + outcome: Final = await_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 Converged(bodies=bodies): + return bodies + case NeverConvergedOn(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_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 a508c97b9fb..e902d58b349 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -3,8 +3,9 @@ No proxy needed and no ``e2e`` marker: this pins that a model registered through the control plane only counts as servable once every configured replica lists it on /v1/models, which is what keeps a two-gateway stack from handing a test a -model that one gateway has not reloaded yet. The fakes are plain pollers and an -injected clock, so nothing here monkeypatches anything. +model that one gateway has not reloaded yet, and that a read-back after a write +converges only once every replica serves the written state. The fakes are plain +pollers and an injected clock, so nothing here monkeypatches anything. """ from __future__ import annotations @@ -18,8 +19,17 @@ import pytest from e2e_config import parse_replica_urls from e2e_http import Success -from models import ModelListEntry, ModelsListResponse -from proxy_client import ModelsPoller, NotServableOn, Servable, await_servable_everywhere +from models import ModelInfoEntry, ModelInfoResponse, ModelListEntry, ModelsListResponse +from proxy_client import ( + BodyReader, + Converged, + ModelsPoller, + NeverConvergedOn, + NotServableOn, + Servable, + await_converged_everywhere, + await_servable_everywhere, +) MODEL: Final = "gpt-under-test" TIMEOUT: Final = 10.0 @@ -85,3 +95,54 @@ class TestParseReplicaUrls: def test_falls_back_to_the_data_plane_address_when_unset(self) -> None: assert parse_replica_urls("", "http://lb") == ("http://lb",) + + +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[Converged[ModelInfoResponse] | NeverConvergedOn[ModelInfoResponse], FakeClock]: + clock: Final = FakeClock() + outcome: Final = await_converged_everywhere( + readers, + predicate=_lists_model, + timeout=TIMEOUT, + interval=INTERVAL, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + return outcome, clock + + +class TestAwaitConvergedEverywhere: + 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 == Converged(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 == NeverConvergedOn(replica=lagging, last_result=_info()) + assert clock.elapsed >= TIMEOUT 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..f9dbf6b7e25 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 @@ -3090,9 +3090,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): @@ -3168,10 +3165,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, @@ -3192,8 +3189,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( @@ -3369,6 +3364,132 @@ 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_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..5cd67a0f9b3 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -220,6 +220,41 @@ 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_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 b1534c19670..8a69d8755ca 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8995,8 +8995,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