diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ebcfab090b5..8e51e250319 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2840,6 +2840,11 @@ async def update_key_fn( """ Update an existing API key's parameters. + The body is a merge patch: a field left out keeps its stored value, and on the key's own columns + an explicit null clears it. The metadata-backed fields below are the exception, merging into the + stored metadata instead: passing one as null leaves it unchanged, while `metadata` itself + replaces the stored metadata wholesale. + Parameters: - key: Optional[str] - The key to update. Either key or key_alias must be provided. - key_alias: Optional[str] - User-friendly key alias. If key is omitted, also identifies the key to update (must match exactly one key, same as /key/delete's key_aliases) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index d571fb36546..860d96a50b4 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -5,6 +5,8 @@ - {id: mgmt.key.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:2462", rationale: "Budget/model changes persist"} - {id: mgmt.key.update.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:2462", rationale: "Non-admin cannot escalate perms"} - {id: mgmt.key.update.happy_path, module: mgmt, tier: P1, surface: ui, assertions: [happy_path], source: "key_management_endpoints.py:2462", rationale: "Key edit through the dashboard"} +- {id: mgmt.key.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "key_management_endpoints.py:2829", rationale: "A partial /key/update changes only the field it names; alias, models, limits, budget window, team and metadata read back unchanged on every gateway replica"} +- {id: mgmt.key.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "key_management_endpoints.py:2829", rationale: "An explicit null on /key/update clears max_budget and budget_duration, and the derived budget_reset_at with it, on every gateway replica"} - {id: mgmt.key.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3122", rationale: "Deletion revokes future calls"} - {id: mgmt.key.delete.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:3122", rationale: "Non-owner cannot delete"} - {id: mgmt.key.info.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3380", rationale: "Info reflects all writes"} diff --git a/tests/e2e/management/test_key_lifecycle_e2e.py b/tests/e2e/management/test_key_lifecycle_e2e.py new file mode 100644 index 00000000000..4c8effc4d24 --- /dev/null +++ b/tests/e2e/management/test_key_lifecycle_e2e.py @@ -0,0 +1,299 @@ +"""Live e2e: one virtual key walked through its whole lifecycle, read back on every +gateway replica. + +Create, read, partial update, clear, enforce, delete: one method per step, and every +step creates its own team and key (both deleted on teardown) so a step reruns or skips +on its own. Writes go through the control plane; read-backs poll every URL in +PROXY_REPLICA_URLS until each replica converges, because a write that is visible on the +gateway that took it and stale on its neighbour is exactly the failure this file exists +to catch. Revocation is the slowest of those: a deleted key stays usable on the other +replicas until their auth cache entry expires, so the delete step polls each of them +rather than asserting once. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +import pytest + +from e2e_config import unique_marker +from e2e_http import Result, StreamingResponse, Success, UnknownApiError, unwrap +from lifecycle import ResourceManager +from management_client import MODEL_ACCESS_DENIED_MARKER, ManagementClient +from models import ( + CLEAR, + ChatBody, + ChatMessage, + KeyGenerateBody, + KeyGenerateResponse, + KeyInfo, + KeyInfoParams, + KeyInfoResponse, + KeyMetadata, + KeyUpdateBody, + LiteLLMParamsBody, + TeamNewBody, +) +from proxy_client import Converged, NotConverged, Poller, await_converged, await_converged_everywhere +from transport import Transport + +pytestmark = pytest.mark.e2e + +BACKING_MODEL: Final = "gpt-4o-mini" +DENIED_MODEL: Final = "gpt-5.5" +MAX_BUDGET: Final = 25.0 +TPM_LIMIT: Final = 313131 +RPM_LIMIT: Final = 323232 +UPDATED_RPM_LIMIT: Final = 424242 +BUDGET_DURATION: Final = "30d" + + +@dataclass(frozen=True, slots=True) +class CreatedKey: + written: KeyGenerateBody + response: KeyGenerateResponse + + @property + def key(self) -> str: + return self.response.key + + +@pytest.fixture(scope="module") +def mock_deployment(client: ManagementClient) -> Iterator[str]: + """A deployment that answers from a canned response, so the enforcement step needs no + provider key. The alias carries a unique marker, like every other model this suite + registers, so concurrent runs never share one model group.""" + model_name: Final = f"e2e-key-lifecycle-{unique_marker()}" + model_id: Final = client.proxy.create_model(model_name, LiteLLMParamsBody(model=BACKING_MODEL, mock_response="ok")) + try: + yield model_name + finally: + client.proxy.delete_model(model_id) + + +def _await[T](client: ManagementClient, poller: Poller[T], converged: Callable[[T], bool], failure: str) -> T: + outcome: Final = await_converged( + poller, + converged=converged, + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case Converged(result=result): + return result + case NotConverged(last_result=last): + pytest.fail(f"{failure}; last outcome: {last}") + + +def _chat_poller(transport: Transport, key: str, model: str) -> Poller[StreamingResponse]: + return lambda: transport.send( + "/chat/completions", + headers=transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"say hi {unique_marker()}")], + max_tokens=16, + ), + ) + + +def _create_key(client: ManagementClient, resources: ResourceManager, model: str) -> CreatedKey: + marker: Final = unique_marker() + team_id: Final = client.create_team(TeamNewBody(team_alias=f"e2e-key-lifecycle-team-{marker}")) + resources.defer(lambda: client.delete_team(team_id)) + written: Final = KeyGenerateBody( + key_alias=f"e2e-key-lifecycle-{marker}", + models=[model], + max_budget=MAX_BUDGET, + tpm_limit=TPM_LIMIT, + rpm_limit=RPM_LIMIT, + budget_duration=BUDGET_DURATION, + metadata=KeyMetadata(tag=marker), + team_id=team_id, + ) + response: Final = unwrap(client.generate_key(written)) + resources.defer(lambda: client.proxy.delete_key(response.key)) + return CreatedKey(written=written, response=response) + + +def _key_info_everywhere( + client: ManagementClient, key: str, settled: Callable[[KeyInfo], bool] +) -> Mapping[str, KeyInfo]: + def converged(result: Result[KeyInfoResponse]) -> bool: + return isinstance(result, Success) and settled(result.data.info) + + reads: Final = client.proxy.read_back_everywhere( + "/key/info", params=KeyInfoParams(key=key), response_type=KeyInfoResponse, converged=converged + ) + return MappingProxyType({replica: unwrap(read).info for replica, read in reads.items()}) + + +def _is_key_not_found(result: Result[KeyInfoResponse]) -> bool: + return isinstance(result, UnknownApiError) and result.status_code == 404 + + +def _assert_reads_back(info: KeyInfo, expected: KeyGenerateBody, replica: str) -> None: + for field, observed, wanted in ( + ("key_alias", info.key_alias, expected.key_alias), + ("models", info.models, expected.models), + ("max_budget", info.max_budget, expected.max_budget), + ("tpm_limit", info.tpm_limit, expected.tpm_limit), + ("rpm_limit", info.rpm_limit, expected.rpm_limit), + ("budget_duration", info.budget_duration, expected.budget_duration), + ("team_id", info.team_id, expected.team_id), + ("metadata", info.metadata, expected.metadata), + ): + assert observed == wanted, f"{replica}: /key/info reports {field}={observed!r}, expected {wanted!r}" + + +def _poll_chat_ok(client: ManagementClient, key: str, model: str) -> None: + _ = _await( + client, + _chat_poller(client.proxy.transport, key, model), + lambda outcome: outcome.ok, + f"chat on {model} never succeeded for the key before the deadline", + ) + + +def _warm_every_replica(client: ManagementClient, key: str, model: str) -> None: + """Serve one call from every replica, so each has the key in its auth cache. Without + this the revocation check below would only prove a replica rejects a key it never + knew, which is true of any random string.""" + for replica, transport in client.proxy.replicas.items(): + _ = _await( + client, + _chat_poller(transport, key, model), + lambda outcome: outcome.ok, + f"{replica}: chat on {model} never succeeded for the key before the deadline", + ) + + +def _assert_chat_rejected_everywhere(client: ManagementClient, key: str, model: str) -> None: + outcomes: Final = await_converged_everywhere( + {replica: _chat_poller(transport, key, model) for replica, transport in client.proxy.replicas.items()}, + converged=lambda outcome: outcome.status_code == 401, + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + for replica, outcome in outcomes.items(): + assert isinstance(outcome, Converged), ( + f"{replica}: the deleted key was still accepted on chat after " + f"{client.proxy.poll_timeout}s, last status {outcome.last_result.status_code}" + ) + + +class TestKeyLifecycle: + def test_create_echoes_every_field_written( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + + response: Final = created.response + for field, observed, wanted in ( + ("key_alias", response.key_alias, created.written.key_alias), + ("models", response.models, created.written.models), + ("max_budget", response.max_budget, created.written.max_budget), + ("tpm_limit", response.tpm_limit, created.written.tpm_limit), + ("rpm_limit", response.rpm_limit, created.written.rpm_limit), + ("budget_duration", response.budget_duration, created.written.budget_duration), + ("team_id", response.team_id, created.written.team_id), + ("metadata", response.metadata, created.written.metadata), + ): + assert observed == wanted, f"/key/generate echoed {field}={observed!r}, sent {wanted!r}" + + def test_read_reflects_the_create_on_every_replica( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + + infos: Final = _key_info_everywhere( + client, created.key, lambda info: info.key_alias == created.written.key_alias + ) + for replica, info in infos.items(): + _assert_reads_back(info, created.written, replica) + assert info.budget_reset_at is not None, ( + f"{replica}: /key/info reports no budget_reset_at for budget_duration={BUDGET_DURATION!r}" + ) + + @pytest.mark.covers("mgmt.key.update.preserves_unrelated_fields") + def test_partial_update_changes_only_the_named_field( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + before: Final = _key_info_everywhere(client, created.key, lambda info: info.rpm_limit == RPM_LIMIT) + + _ = unwrap(client.update_key(KeyUpdateBody(key=created.key, rpm_limit=UPDATED_RPM_LIMIT))) + + after: Final = _key_info_everywhere(client, created.key, lambda info: info.rpm_limit == UPDATED_RPM_LIMIT) + for replica, info in after.items(): + _assert_reads_back(info, created.written.model_copy(update={"rpm_limit": UPDATED_RPM_LIMIT}), replica) + assert info.budget_reset_at == before[replica].budget_reset_at, ( + f"{replica}: budget_reset_at moved from {before[replica].budget_reset_at!r} to " + f"{info.budget_reset_at!r} on a /key/update that did not name budget_duration" + ) + + @pytest.mark.covers("mgmt.key.update.clear_persists") + def test_explicit_null_clears_the_budget_and_its_reset_time( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + _ = _key_info_everywhere(client, created.key, lambda info: info.max_budget == MAX_BUDGET) + + _ = unwrap(client.update_key(KeyUpdateBody(key=created.key, max_budget=CLEAR, budget_duration=CLEAR))) + + cleared: Final = _key_info_everywhere(client, created.key, lambda info: info.max_budget is None) + for replica, info in cleared.items(): + assert info.budget_duration is None, ( + f"{replica}: budget_duration={info.budget_duration!r} survived an explicit null" + ) + assert info.budget_reset_at is None, ( + f"{replica}: clearing budget_duration left budget_reset_at={info.budget_reset_at!r}" + ) + _assert_reads_back( + info, created.written.model_copy(update={"max_budget": None, "budget_duration": None}), replica + ) + + def test_key_serves_its_model_and_is_denied_others( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + + _poll_chat_ok(client, created.key, mock_deployment) + + denied: Final = client.chat_status(created.key, DENIED_MODEL, f"say hi {unique_marker()}") + assert denied.status_code == 403, ( + f"chat on {DENIED_MODEL!r} outside the key's model list must be denied 403, got " + f"{denied.status_code}: {denied.body[:300]}" + ) + assert MODEL_ACCESS_DENIED_MARKER in denied.body, ( + f"403 body must be a model-access denial, got: {denied.body[:300]}" + ) + + def test_delete_revokes_info_and_chat_on_every_replica( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + """The teardown's deferred delete fires again on the already-deleted key by + design: the deferred cleanup must survive this test failing before the + in-body delete, and a repeat /key/delete is a cheap no-op the warn-only + teardown absorbs.""" + created: Final = _create_key(client, resources, mock_deployment) + _warm_every_replica(client, created.key, mock_deployment) + + client.delete_key_strict(created.key) + + _ = client.proxy.read_back_everywhere( + "/key/info", + params=KeyInfoParams(key=created.key), + response_type=KeyInfoResponse, + converged=_is_key_not_found, + ) + _assert_chat_rejected_everywhere(client, created.key, mock_deployment) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index c01687d0b31..fa124bdf6bd 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -8,9 +8,9 @@ from __future__ import annotations from collections.abc import Sequence from datetime import datetime -from typing import Literal +from typing import Final, Literal -from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_validator +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_serializer, model_validator # ---------- keys ---------- @@ -49,6 +49,7 @@ class KeyMetadata(BaseModel): logging: list[KeyLoggingCallback] | None = None priority: str | None = None batch_enqueued_token_limit: int | None = None + tag: str | None = None class ObjectPermission(BaseModel): @@ -81,6 +82,14 @@ class KeyGenerateBody(BaseModel): class KeyGenerateResponse(BaseModel): key: str + key_alias: str | None = None + models: list[str] = [] + max_budget: float | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None + budget_duration: str | None = None + team_id: str | None = None + metadata: KeyMetadata | None = None class KeyRegenerateBody(BaseModel): @@ -122,6 +131,7 @@ class KeyInfo(BaseModel): blocked: bool | None = None spend: float | None = None max_budget: float | None = None + budget_duration: str | None = None budget_reset_at: str | None = None budget_id: str | None = None litellm_budget_table: LiteLLMBudgetTable | None = None @@ -919,12 +929,34 @@ class CredentialCreateResponse(BaseModel): # ---------- key / team / user / organization management ---------- +class Cleared(BaseModel): + """An explicit JSON null in a merge-patch body. The transport drops `None` fields + before sending (`exclude_none`), so `None` means "leave the stored value alone"; a + field set to `CLEAR` reaches the wire as `null`, which tells the proxy to clear it.""" + + model_config = ConfigDict(frozen=True) + + @model_serializer + def _as_null(self) -> None: + return None + + +CLEAR: Final = Cleared() + + class KeyUpdateBody(BaseModel): + """POST /key/update is a merge patch: a field left `None` is dropped from the body and + keeps its stored value, `CLEAR` sends an explicit null that clears it (`budget_duration` + clears `budget_reset_at` with it), and `metadata` replaces the stored metadata wholesale.""" + key: str models: list[str] | None = None key_alias: str | None = None tpm_limit: int | None = None rpm_limit: int | None = None + max_budget: float | Cleared | None = None + budget_duration: str | Cleared | None = None + metadata: KeyMetadata | None = None class KeyBlockBody(BaseModel): diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 3b15e57dfcf..520cbfde5a9 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, @@ -149,9 +151,7 @@ def await_servable( last_result: Result[ModelsListResponse] | None = None while True: t = now() - phase_deadline = ( - started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds - ) + phase_deadline = started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds remaining = phase_deadline - t if remaining <= 0: if ( @@ -164,9 +164,7 @@ def await_servable( poll_timeout = min(request_timeout, remaining) last_result = list_models(poll_timeout) - listed = isinstance(last_result, Success) and any( - entry.id == model_name for entry in last_result.data.data - ) + listed = isinstance(last_result, Success) and any(entry.id == model_name for entry in last_result.data.data) t = now() if not listed: first_seen_at = None @@ -179,9 +177,7 @@ def await_servable( elif t - first_seen_at >= db_sync_seconds: return Servable() - phase_deadline = ( - started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds - ) + phase_deadline = started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds wait = min(interval, phase_deadline - now()) if wait > 0: sleep(wait) @@ -239,6 +235,88 @@ def servable_timeout_message( ) +type Poller[T] = Callable[[], T] + + +@dataclass(frozen=True, slots=True) +class Converged[T]: + result: T + + +@dataclass(frozen=True, slots=True) +class NotConverged[T]: + """The deadline passed without a read satisfying the predicate; `last_result` is + the final read, so the caller can tell a stale body from a failed request.""" + + last_result: T + + +type ConvergeOutcome[T] = Converged[T] | NotConverged[T] + + +def await_converged[T]( + poll: Poller[T], + *, + converged: Callable[[T], bool], + timeout: float, + interval: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> ConvergeOutcome[T]: + """Poll until a read satisfies `converged` or `timeout` elapses. + + Polls before testing the deadline, so a zero or already-spent budget still gets one + attempt, and sleeps only min(interval, time left), so the attempt that lands exactly + on the deadline is taken rather than skipped. Clock and sleep are injected.""" + deadline: Final = now() + timeout + while True: + result = poll() + if converged(result): + return Converged(result=result) + remaining = deadline - now() + if remaining <= 0: + return NotConverged(last_result=result) + sleep(min(interval, remaining)) + + +def await_converged_everywhere[T]( + pollers: Mapping[str, Poller[T]], + *, + converged: Callable[[T], bool], + timeout: float, + interval: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> Mapping[str, ConvergeOutcome[T]]: + """`await_converged` against every replica in turn, each with the full budget, so a + replica that lags behind the one a write landed on is polled until it catches up + rather than failing on its first stale read.""" + return MappingProxyType( + { + replica: await_converged( + poll, converged=converged, timeout=timeout, interval=interval, now=now, sleep=sleep + ) + for replica, poll in pollers.items() + } + ) + + +def first_lagging_replica[T]( + outcomes: Mapping[str, ConvergeOutcome[T]], +) -> tuple[str, NotConverged[T]] | None: + return next( + ((replica, outcome) for replica, outcome in outcomes.items() if isinstance(outcome, NotConverged)), + None, + ) + + +def converge_timeout_message(*, what: str, replica: str, timeout: float, last_result: object) -> str: + return ( + f"{what} on {replica} never converged within {timeout}s of the write " + f"(control/data-plane propagation issue); last read: {last_result}" + ) + + @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport @@ -290,6 +368,52 @@ class ProxyClient: ) ).info + def read_back_everywhere[R: BaseModel]( + self, + path: str, + *, + params: BaseModel, + response_type: type[R], + converged: Callable[[Result[R]], bool], + ) -> Mapping[str, Result[R]]: + """GET `path` under the master key on every replica in PROXY_REPLICA_URLS (the + data-plane URL alone when the stack exports no per-gateway addresses), polling + each to poll_timeout until its read satisfies `converged`. Returns that read per + replica, or fails naming the first replica that never converged and its last + read. Behind a load balancer the single address proves one replica converged, + not all of them; only per-gateway addresses make this a fleet-wide proof.""" + outcomes: Final = await_converged_everywhere( + { + url: self._body_poller(transport, path, params, response_type) + for url, transport in self.replicas.items() + }, + converged=converged, + timeout=self.poll_timeout, + interval=self.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + lagging: Final = first_lagging_replica(outcomes) + if lagging is not None: + replica, outcome = lagging + raise AssertionError( + converge_timeout_message( + what=f"GET {path}", + replica=replica, + timeout=self.poll_timeout, + last_result=outcome.last_result, + ) + ) + return MappingProxyType( + {replica: outcome.result for replica, outcome in outcomes.items() if isinstance(outcome, Converged)} + ) + + @staticmethod + def _body_poller[R: BaseModel]( + transport: Transport, path: str, params: BaseModel, response_type: type[R] + ) -> Poller[Result[R]]: + return lambda: transport.get(path, headers=transport.master, params=params, response_type=response_type) + def model_info(self) -> list[ModelInfoEntry]: """Every configured deployment with the price the proxy resolved for it (config override merged over cost-map defaults).""" @@ -320,9 +444,7 @@ class ProxyClient: response_type=FileListResponse, ) - def list_fine_tuning_jobs( - self, key: str, params: FineTuningJobsParams - ) -> Result[FineTuningJobsResponse]: + def list_fine_tuning_jobs(self, key: str, params: FineTuningJobsParams) -> Result[FineTuningJobsResponse]: return self.transport.get( "/v1/fine_tuning/jobs", headers=self.transport.bearer(key), @@ -375,7 +497,11 @@ class ProxyClient: ) ).model_id written_at = time.monotonic() - self._await_model_servable(body.model_name, listed_for) + try: + self._await_model_servable(body.model_name, listed_for) + except BaseException: + self.delete_model(model_id) + raise settle_propagation(written_at) return model_id diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index a508c97b9fb..2caac58333f 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -1,9 +1,11 @@ -"""Harness coverage for the model barrier that gates on every replica. +"""Harness coverage for the barriers that gate on every replica. 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 +on /v1/models, and that a management write only counts as read back once every +replica's read satisfies the caller's predicate, which is what keeps a two-gateway +stack from handing a test a model or a key that one gateway has not caught up on +yet. The fakes are plain pollers standing in for each replica's transport plus an injected clock, so nothing here monkeypatches anything. """ @@ -12,18 +14,33 @@ from __future__ import annotations from collections.abc import Iterable, Mapping from dataclasses import dataclass from itertools import chain, repeat +from types import MappingProxyType from typing import Final 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 e2e_http import Result, Success +from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse +from proxy_client import ( + Poller, + ConvergeOutcome, + Converged, + ModelsPoller, + NotConverged, + NotServableOn, + Servable, + await_converged_everywhere, + await_servable_everywhere, + first_lagging_replica, + converge_timeout_message, +) MODEL: Final = "gpt-under-test" TIMEOUT: Final = 10.0 INTERVAL: Final = 2.0 +RPM_BEFORE_UPDATE: Final = 100 +RPM_AFTER_UPDATE: Final = 200 @dataclass @@ -78,6 +95,91 @@ class TestAwaitServableEverywhere: assert _await(pollers) == Servable() +def _key_info(rpm_limit: int) -> Success[KeyInfoResponse]: + return Success(status_code=200, data=KeyInfoResponse(info=KeyInfo(rpm_limit=rpm_limit))) + + +def _reads(results: Iterable[Result[KeyInfoResponse]]) -> Poller[Result[KeyInfoResponse]]: + it: Final = iter(results) + return lambda: next(it) + + +def _updated(result: Result[KeyInfoResponse]) -> bool: + return isinstance(result, Success) and result.data.info.rpm_limit == RPM_AFTER_UPDATE + + +def _converge( + pollers: Mapping[str, Poller[Result[KeyInfoResponse]]], clock: FakeClock +) -> Mapping[str, ConvergeOutcome[Result[KeyInfoResponse]]]: + return await_converged_everywhere( + pollers, + converged=_updated, + timeout=TIMEOUT, + interval=INTERVAL, + now=clock.now, + sleep=clock.sleep, + ) + + +class TestAwaitConvergedEverywhere: + def test_waits_for_the_replica_that_lags_behind_the_write(self) -> None: + clock: Final = FakeClock() + pollers: Final = MappingProxyType( + { + "gateway-1": _reads(repeat(_key_info(RPM_AFTER_UPDATE))), + "gateway-2": _reads( + chain(repeat(_key_info(RPM_BEFORE_UPDATE), 2), repeat(_key_info(RPM_AFTER_UPDATE))) + ), + } + ) + outcomes: Final = _converge(pollers, clock) + assert outcomes == { + "gateway-1": Converged(result=_key_info(RPM_AFTER_UPDATE)), + "gateway-2": Converged(result=_key_info(RPM_AFTER_UPDATE)), + } + assert first_lagging_replica(outcomes) is None + assert clock.elapsed == 2 * INTERVAL + + def test_names_the_replica_that_never_converges_with_its_last_read(self) -> None: + clock: Final = FakeClock() + pollers: Final = MappingProxyType( + { + "gateway-1": _reads(repeat(_key_info(RPM_AFTER_UPDATE))), + "gateway-2": _reads(repeat(_key_info(RPM_BEFORE_UPDATE))), + } + ) + outcomes: Final = _converge(pollers, clock) + assert first_lagging_replica(outcomes) == ( + "gateway-2", + NotConverged(last_result=_key_info(RPM_BEFORE_UPDATE)), + ) + assert clock.elapsed == TIMEOUT + message: Final = converge_timeout_message( + what="GET /key/info", + replica="gateway-2", + timeout=TIMEOUT, + last_result=_key_info(RPM_BEFORE_UPDATE), + ) + assert "gateway-2" in message and "/key/info" in message and str(RPM_BEFORE_UPDATE) in message + + def test_each_replica_gets_its_own_full_budget(self) -> None: + """A replica that converges late must not eat into the next replica's budget: both + need most of the timeout here, so one shared deadline would starve the second.""" + clock: Final = FakeClock() + slow: Final = chain(repeat(_key_info(RPM_BEFORE_UPDATE), 3), repeat(_key_info(RPM_AFTER_UPDATE))) + pollers: Final = MappingProxyType( + { + "gateway-1": _reads(slow), + "gateway-2": _reads( + chain(repeat(_key_info(RPM_BEFORE_UPDATE), 3), repeat(_key_info(RPM_AFTER_UPDATE))) + ), + } + ) + outcomes: Final = _converge(pollers, clock) + assert first_lagging_replica(outcomes) is None + assert clock.elapsed == 2 * 3 * INTERVAL + + class TestParseReplicaUrls: def test_splits_and_trims_the_gateway_addresses(self) -> None: raw: Final = " http://127.0.0.1:4010/, http://127.0.0.1:4011 " diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 108e50b52a1..36fd744efc0 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8017,6 +8017,11 @@ export interface paths { * Update Key Fn * @description Update an existing API key's parameters. * + * The body is a merge patch: a field left out keeps its stored value, and on the key's own columns + * an explicit null clears it. The metadata-backed fields below are the exception, merging into the + * stored metadata instead: passing one as null leaves it unchanged, while `metadata` itself + * replaces the stored metadata wholesale. + * * Parameters: * - key: Optional[str] - The key to update. Either key or key_alias must be provided. * - key_alias: Optional[str] - User-friendly key alias. If key is omitted, also identifies the key to update (must match exactly one key, same as /key/delete's key_aliases)