From ece1484bc7428da8f66683d4897c9c1d7d504e19 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 11 Nov 2026 22:55:59 -0800 Subject: [PATCH] feat(rate-limit): team_model_limit_max_borrow_multiple caps how far a team borrows idle model capacity (#51359) * feat(rate-limit): team_model_limit_max_borrow_multiple caps how far a team borrows idle model capacity With team_model_limit_borrows_idle_capacity_from_other_teams on, a team could use everything up to the model's own tpm and rpm while the model was idle. Setting team_model_limit_max_borrow_multiple in the team's metadata caps the team at its model_tpm_limit and model_rpm_limit times that multiple, with a 429 whose rate_limited_by header reads model_per_team_max_borrow_multiple. Unset keeps the previous behavior. Values below 1 or non-numeric are logged and ignored. * fix(rate-limit): reject infinite or oversized team_model_limit_max_borrow_multiple values instead of failing the rate check * fix(team): make team_model_limit_max_borrow_multiple proxy-controlled team metadata so team admins cannot lift their own cap --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/dynamic_rate_limiter_v3.py | 116 ++++++++++++---- .../hooks/parallel_request_limiter_v3.py | 44 +++++- .../team_admin_field_permissions.py | 1 + ..._team_model_limit_borrows_idle_capacity.py | 32 ++++- .../hooks/test_dynamic_rate_limiter_v3.py | 129 ++++++++++++++++++ .../test_team_admin_field_permissions.py | 1 + 6 files changed, 296 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 5a6d6fc4cdf..3654d0b9176 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -4,6 +4,7 @@ Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting import os from collections.abc import Callable, Mapping +from dataclasses import dataclass from datetime import datetime from typing import TYPE_CHECKING, Final, Literal @@ -27,6 +28,7 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ( map_v3_rate_limit_type, ) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( # noqa: F401 # legacy module exports + BorrowingTeamModelLimits, PROXY_MaxParallelRequestsHandler_v3, RateLimitDescriptor, RateLimitDescriptorRateLimitObject, @@ -72,6 +74,19 @@ def _usage_of(response_obj: object) -> Usage | None: TEAM_MODEL_LIMIT_AT_SATURATION: Final = "model_per_team_saturated" +TEAM_MODEL_LIMIT_MAX_BORROW_MULTIPLE: Final = "model_per_team_max_borrow_multiple" + + +@dataclass(frozen=True, slots=True) +class _BorrowingTeamShare: + """The team's share and, when the team caps how far it borrows, its ceiling, both on the share's counter.""" + + limits: BorrowingTeamModelLimits + share: RateLimitDescriptor + ceiling: RateLimitDescriptor | None + + def enforced(self, saturated: bool) -> RateLimitDescriptor | None: + return self.share if saturated else self.ceiling def _team_limit_name(rate_limit_type: str) -> str: @@ -392,9 +407,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return descriptors - def _create_borrowing_team_model_limit_descriptor( - self, model: str, user_api_key_dict: UserAPIKeyAuth - ) -> RateLimitDescriptor | None: + def _borrowing_team_share(self, model: str, user_api_key_dict: UserAPIKeyAuth) -> _BorrowingTeamShare | None: """One counter per team and resolved model group, so an alias and its target share the team's share.""" rate_limited: Final = self.v3_limiter.rate_limited_model(model) or RateLimitedModel( requested=model, group=model @@ -402,9 +415,18 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): limits: Final = borrowing_team_model_limits(user_api_key_dict, rate_limited) if limits.tpm is None and limits.rpm is None: return None + ceiling: Final = limits.borrow_ceiling() + value: Final = f"{user_api_key_dict.team_id}:{rate_limited.group}" + return _BorrowingTeamShare( + limits=limits, + share=self._team_share_descriptor(value, limits), + ceiling=None if ceiling is None else self._team_share_descriptor(value, ceiling), + ) + + def _team_share_descriptor(self, value: str, limits: BorrowingTeamModelLimits) -> RateLimitDescriptor: return RateLimitDescriptor( key=TEAM_MODEL_LIMIT_AT_SATURATION, - value=f"{user_api_key_dict.team_id}:{rate_limited.group}", + value=value, rate_limit={ "tokens_per_unit": limits.tpm, "requests_per_unit": limits.rpm, @@ -412,6 +434,55 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): }, ) + def _team_share_error( + self, + model: str, + status: RateLimitStatus, + saturation: float, + priority: str | None, + user_api_key_dict: UserAPIKeyAuth, + team_share: _BorrowingTeamShare, + ) -> ProxyRateLimitError: + if saturation >= _get_priority_settings().saturation_threshold: + return self._team_limit_error(model, status, saturation, priority, user_api_key_dict) + return self._team_borrow_ceiling_error( + model, status, saturation, priority, user_api_key_dict, team_share.limits + ) + + def _team_borrow_ceiling_error( + self, + model: str, + status: RateLimitStatus, + saturation: float, + priority: str | None, + user_api_key_dict: UserAPIKeyAuth, + limits: BorrowingTeamModelLimits, + ) -> ProxyRateLimitError: + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(model) + limit_name: Final = _team_limit_name(status["rate_limit_type"]) + return ProxyRateLimitError( + detail={ + "error": f"Team {limit_name} borrow cap reached for {model}: the team may use at most " + f"{limits.max_borrow_multiple}x its {limit_name} while borrowing idle capacity. " + f"Team: {user_api_key_dict.team_id}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Team {limit_name}: {limits.share(status['rate_limit_type'])}, " + f"Borrow cap: {status['current_limit']}, " + f"Remaining: {status['limit_remaining']}, " + f"Model saturation: {saturation:.1%}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "rate_limited_by": TEAM_MODEL_LIMIT_MAX_BORROW_MULTIPLE, + "x-litellm-priority": priority or "default", + "x-litellm-saturation": f"{saturation:.2%}", + }, + rate_limit_type=map_v3_rate_limit_type(status["rate_limit_type"]), + model=resolved_model, + llm_provider=llm_provider, + ) + def _team_limit_error( self, model: str, @@ -449,28 +520,28 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): """Hold a Responses WebSocket turn to the team's share once the model saturates, like the regular limiter does.""" if requested_model is None: return - team_limit_descriptor: Final = self._create_borrowing_team_model_limit_descriptor( - model=requested_model, user_api_key_dict=user_api_key_dict - ) + team_share: Final = self._borrowing_team_share(model=requested_model, user_api_key_dict=user_api_key_dict) model_group_info: Final = self.llm_router.registry().group(requested_model) - if team_limit_descriptor is None or model_group_info is None: + if team_share is None or model_group_info is None: return saturation: Final = await self._check_model_saturation(requested_model, model_group_info) - if saturation < _get_priority_settings().saturation_threshold: + enforced: Final = team_share.enforced(saturation >= _get_priority_settings().saturation_threshold) + if enforced is None: return - tpm_windows: Final = _socket_turn_tpm_windows((team_limit_descriptor,)) + tpm_windows: Final = _socket_turn_tpm_windows((enforced,)) response: Final = await self.v3_limiter.should_rate_limit( tpm_windows, parent_otel_span=user_api_key_dict.parent_otel_span, read_only=True, skip_parallel_check=True ) over_limit: Final = next((s for s in response["statuses"] if s["code"] == "OVER_LIMIT"), None) if over_limit is None: return - raise self._team_limit_error( + raise self._team_share_error( requested_model, over_limit, saturation, self._get_priority_from_user_api_key_dict(user_api_key_dict), user_api_key_dict, + team_share, ) def _create_model_tracking_descriptor( @@ -569,18 +640,17 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # asyncio.Lock + in-memory fallback for single-process deployments. # All-or-nothing: if any enforced descriptor would exceed its limit, # no counter is modified and the response carries "OVER_LIMIT". - team_limit_descriptor: Final = self._create_borrowing_team_model_limit_descriptor( - model=model, user_api_key_dict=user_api_key_dict - ) - get_or_create_request_stash().borrowing_team_model_limit = team_limit_descriptor - if team_limit_descriptor is not None: - descriptors_to_check.append(team_limit_descriptor) + team_share: Final = self._borrowing_team_share(model=model, user_api_key_dict=user_api_key_dict) + get_or_create_request_stash().borrowing_team_model_limit = None if team_share is None else team_share.share + if team_share is not None: + descriptors_to_check.append(team_share.share) + enforced_team_share: Final = None if team_share is None else team_share.enforced(should_enforce_priority) enforced_descriptors: Final[list[RateLimitDescriptor]] = [model_wide_descriptor] if priority_descriptors and should_enforce_priority: enforced_descriptors.extend(priority_descriptors) - if team_limit_descriptor is not None and should_enforce_priority: - enforced_descriptors.append(team_limit_descriptor) + if enforced_team_share is not None: + enforced_descriptors.append(enforced_team_share) per_request_increment: Final[dict[Literal["requests", "tokens"], int]] = { "requests": 1, @@ -645,8 +715,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): llm_provider=llm_provider, ) - if descriptor_key == TEAM_MODEL_LIMIT_AT_SATURATION: - raise self._team_limit_error(model, status, saturation, priority, user_api_key_dict) + if descriptor_key == TEAM_MODEL_LIMIT_AT_SATURATION and team_share is not None: + raise self._team_share_error(model, status, saturation, priority, user_api_key_dict, team_share) # Fail-closed guard: overall_code says OVER_LIMIT but no status # matched a descriptor key we know how to translate into a 429. @@ -687,9 +757,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # priority allocation has accurate usage when it later becomes # enforced. The increment-then-check semantics of should_rate_limit # are fine because we ignore the OVER_LIMIT response. - if team_limit_descriptor is not None and not should_enforce_priority: + if team_share is not None and enforced_team_share is None: await self.v3_limiter.atomic_check_and_increment_by_n( - descriptors=[team_limit_descriptor], + descriptors=[team_share.share], increments=[per_request_increment], parent_otel_span=user_api_key_dict.parent_otel_span, ) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 7f266e67b41..6b53e952f98 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -18,6 +18,7 @@ from datetime import datetime, timezone from types import MappingProxyType from typing import ( TYPE_CHECKING, + Annotated, Any, Final, Literal, @@ -27,7 +28,7 @@ from typing import ( ) from fastapi import HTTPException -from pydantic import TypeAdapter, ValidationError +from pydantic import Field, TypeAdapter, ValidationError from starlette.status import HTTP_400_BAD_REQUEST, HTTP_503_SERVICE_UNAVAILABLE from typing_extensions import NotRequired, ReadOnly @@ -199,17 +200,54 @@ def inherited_team_model_limit( return team_limit +TEAM_MODEL_LIMIT_MAX_BORROW_MULTIPLE: Final = "team_model_limit_max_borrow_multiple" +MAX_BORROW_MULTIPLE_CEILING: Final = 1_000_000 +_MAX_BORROW_MULTIPLE: Final[TypeAdapter[float]] = TypeAdapter( + Annotated[float, Field(ge=1, le=MAX_BORROW_MULTIPLE_CEILING, allow_inf_nan=False)] +) + + @dataclass(frozen=True, slots=True) class BorrowingTeamModelLimits: """Inherited team per-model limits the dynamic rate limiter enforces at saturation instead of the regular limiter.""" tpm: int | None rpm: int | None + max_borrow_multiple: float | None = None + + def share(self, rate_limit_type: str) -> int | None: + return self.rpm if rate_limit_type == "requests" else self.tpm + + def borrow_ceiling(self) -> "BorrowingTeamModelLimits | None": + """The most the team may use while borrowing, or None when it may borrow up to the model's own limit.""" + if self.max_borrow_multiple is None: + return None + return BorrowingTeamModelLimits( + tpm=None if self.tpm is None else int(self.tpm * self.max_borrow_multiple), + rpm=None if self.rpm is None else int(self.rpm * self.max_borrow_multiple), + ) NO_BORROWING: Final = BorrowingTeamModelLimits(tpm=None, rpm=None) +def team_model_limit_max_borrow_multiple(user_api_key_dict: UserAPIKeyAuth) -> float | None: + raw: Final[object] = (user_api_key_dict.team_metadata or {}).get(TEAM_MODEL_LIMIT_MAX_BORROW_MULTIPLE) + if raw is None: + return None + try: + return _MAX_BORROW_MULTIPLE.validate_python(raw) + except ValidationError: + verbose_proxy_logger.warning( + "Ignoring %s=%r on team %s: expected a number from 1 to %d, so the team borrows up to the model limit", + TEAM_MODEL_LIMIT_MAX_BORROW_MULTIPLE, + raw, + user_api_key_dict.team_id, + MAX_BORROW_MULTIPLE_CEILING, + ) + return None + + def borrowing_team_model_limits(user_api_key_dict: UserAPIKeyAuth, model: RateLimitedModel) -> BorrowingTeamModelLimits: settings: Final = litellm.priority_reservation_settings if ( @@ -224,7 +262,9 @@ def borrowing_team_model_limits(user_api_key_dict: UserAPIKeyAuth, model: RateLi ) if limits == NO_BORROWING or not dynamic_rate_limiter_covers(model.group): return NO_BORROWING - return limits + return BorrowingTeamModelLimits( + tpm=limits.tpm, rpm=limits.rpm, max_borrow_multiple=team_model_limit_max_borrow_multiple(user_api_key_dict) + ) def _resolve_media_models(router: "Router | None", model: str, team_id: str | None) -> tuple[str, ...]: diff --git a/litellm/proxy/management_endpoints/team_admin_field_permissions.py b/litellm/proxy/management_endpoints/team_admin_field_permissions.py index 08a84ca91c3..f80aa2a2e29 100644 --- a/litellm/proxy/management_endpoints/team_admin_field_permissions.py +++ b/litellm/proxy/management_endpoints/team_admin_field_permissions.py @@ -108,6 +108,7 @@ PROXY_CONTROLLED_TEAM_METADATA_KEYS: Final[frozenset[str]] = frozenset( "phoenix_project_name_override", "phoenix_project_name", "opik", + "team_model_limit_max_borrow_multiple", ) ) _NOT_COLUMNS: Final[frozenset[str]] = frozenset({"team_id", "metadata"}) diff --git a/tests/integration/routing/test_team_model_limit_borrows_idle_capacity.py b/tests/integration/routing/test_team_model_limit_borrows_idle_capacity.py index 03951dfca21..b9f853d7da8 100644 --- a/tests/integration/routing/test_team_model_limit_borrows_idle_capacity.py +++ b/tests/integration/routing/test_team_model_limit_borrows_idle_capacity.py @@ -16,6 +16,7 @@ MODEL_TPM: Final = 1000 MODEL_RPM: Final = 1000 TEAM_TPM: Final = 100 TEAM_RPM: Final = 3 +MAX_BORROW_MULTIPLE: Final = 2 UPSTREAM_REPLY: Final = json.dumps( { "id": "chatcmpl_team_model_limit_borrowing", @@ -68,7 +69,7 @@ def saturated_model(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[ def _team_calls( - owned: Gateway, upstream: str, calls: int, limit: str = "model_tpm_limit" + owned: Gateway, upstream: str, calls: int, limit: str = "model_tpm_limit", max_borrow_multiple: int | None = None ) -> tuple[httpx.Response, ...]: with owned.scenario() as scenario: model: Final = scenario.model( @@ -78,8 +79,13 @@ def _team_calls( tpm=MODEL_TPM, rpm=MODEL_RPM, ) + cap: Final = ( + {} if max_borrow_multiple is None else {"team_model_limit_max_borrow_multiple": max_borrow_multiple} + ) key: Final = scenario.key( - team_id=scenario.team(metadata={limit: {model: TEAM_TPM if limit == "model_tpm_limit" else TEAM_RPM}}) + team_id=scenario.team( + metadata={limit: {model: TEAM_TPM if limit == "model_tpm_limit" else TEAM_RPM}, **cap} + ) ) return tuple( owned.request( @@ -115,3 +121,25 @@ def test_team_is_held_to_its_model_limit_once_the_model_saturates( assert first.headers.get("rate_limit_type") == rate_limit_type, dict(first.headers) assert f"Team {limit} reached" in first.json()["error"]["message"], first.text assert "no idle capacity to borrow" in first.json()["error"]["message"], first.text + + +@pytest.mark.parametrize( + "limit, rate_limit_type", [("model_tpm_limit", "tokens"), ("model_rpm_limit", "requests")], ids=["tpm", "rpm"] +) +def test_team_borrowing_stops_at_its_max_borrow_multiple_while_the_model_is_idle( + idle_model: tuple[Gateway, str], limit: str, rate_limit_type: str +) -> None: + owned, upstream = idle_model + responses: Final = _team_calls(owned, upstream, calls=12, limit=limit, max_borrow_multiple=MAX_BORROW_MULTIPLE) + admitted: Final = tuple(response for response in responses if response.status_code == 200) + blocked: Final = tuple(response for response in responses if response.status_code == 429) + assert blocked, [r.text for r in responses] + assert responses[0].status_code == 200, responses[0].text + share: Final = TEAM_TPM if limit == "model_tpm_limit" else TEAM_RPM + tokens_per_call: Final = 30 if limit == "model_tpm_limit" else 1 + assert len(admitted) * tokens_per_call <= share * MAX_BORROW_MULTIPLE + tokens_per_call, [r.text for r in responses] + first: Final = blocked[0] + assert first.headers.get("rate_limited_by") == "model_per_team_max_borrow_multiple", dict(first.headers) + assert first.headers.get("rate_limit_type") == rate_limit_type, dict(first.headers) + assert f"Team {limit} borrow cap reached" in first.json()["error"]["message"], first.text + assert f"{MAX_BORROW_MULTIPLE}.0x its {limit}" in first.json()["error"]["message"], first.text diff --git a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py index 4977e8d5737..a8265680856 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -2080,3 +2080,132 @@ async def test_a_websocket_turn_still_borrows_while_the_model_is_idle(monkeypatc handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.5) await _admit_then_settle(handler, _team_key(), 150) await handler.enforce_token_limits_for_frame(user_api_key_dict=_team_key(), requested_model="shared-model") + + +def _capped_team_key(multiple: object, limit: str = "model_tpm_limit", share: int = 100) -> UserAPIKeyAuth: + return _team_key(team_limits={limit: {"shared-model": share}, "team_model_limit_max_borrow_multiple": multiple}) + + +def _held_to_borrow_cap(exc: ProxyRateLimitError, rate_limit_type: str, limit_sentence: str) -> None: + assert exc.status_code == 429 + assert exc.headers is not None + assert exc.headers["rate_limited_by"] == "model_per_team_max_borrow_multiple" + assert exc.headers["rate_limit_type"] == rate_limit_type + assert str(exc.detail).startswith( + "{'error': 'Team " + limit_sentence + " borrow cap reached for shared-model: the team may use at most " + ), exc.detail + + +@pytest.mark.asyncio +async def test_team_borrows_past_its_share_while_it_stays_under_its_max_borrow_multiple( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _admit_then_settle(handler, _capped_team_key(2), 190) + await _check(handler, _capped_team_key(2), saturation=0.49) + + +@pytest.mark.asyncio +async def test_team_at_its_max_borrow_multiple_is_refused_while_the_model_is_idle(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _admit_then_settle(handler, _capped_team_key(2), 200) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, _capped_team_key(2), saturation=0.0) + _held_to_borrow_cap(exc.value, "tokens", "model_tpm_limit") + assert ( + "2.0x its model_tpm_limit while borrowing idle capacity. Team: team-a, Rate limit type: tokens, " + "Team model_tpm_limit: 100, Borrow cap: 200, Remaining: 0, Model saturation: 0.0%" + ) in str(exc.value.detail), exc.value.detail + + +@pytest.mark.asyncio +async def test_a_fractional_max_borrow_multiple_caps_the_team_at_that_fraction_of_its_share( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _admit_then_settle(handler, _capped_team_key(1.5), 149) + await _check(handler, _capped_team_key(1.5), saturation=0.0) + await _settle_tokens(handler, 1) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, _capped_team_key(1.5), saturation=0.0) + _held_to_borrow_cap(exc.value, "tokens", "model_tpm_limit") + assert "Borrow cap: 150," in str(exc.value.detail), exc.value.detail + + +@pytest.mark.asyncio +async def test_the_max_borrow_multiple_caps_requests_as_well_as_tokens(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + auth: Final = _capped_team_key(2, limit="model_rpm_limit", share=2) + for _ in range(4): + await _check(handler, auth, saturation=0.0) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, auth, saturation=0.0) + _held_to_borrow_cap(exc.value, "requests", "model_rpm_limit") + assert "Team model_rpm_limit: 2, Borrow cap: 4, Remaining: 0," in str(exc.value.detail), exc.value.detail + + +@pytest.mark.asyncio +async def test_a_team_under_its_max_borrow_multiple_is_still_held_to_its_share_once_the_model_saturates( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _admit_then_settle(handler, _capped_team_key(5), 150) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, _capped_team_key(5), saturation=0.5) + _held_to_team_limit(exc.value, "tokens", "model_tpm_limit") + + +@pytest.mark.asyncio +async def test_a_team_at_its_max_borrow_multiple_is_held_to_its_share_not_its_cap_once_the_model_saturates( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _admit_then_settle(handler, _capped_team_key(2), 100) + await _check(handler, _capped_team_key(2), saturation=0.49) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, _capped_team_key(2), saturation=0.5) + _held_to_team_limit(exc.value, "tokens", "model_tpm_limit") + + +@pytest.mark.asyncio +async def test_another_teams_borrowing_does_not_count_against_a_teams_max_borrow_multiple( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _admit_then_settle(handler, _team_key("team-b"), 150) + await _admit_then_settle(handler, _capped_team_key(2), 150) + await _check(handler, _capped_team_key(2), saturation=0.0) + + +@pytest.mark.parametrize( + "multiple", + [0.5, 0, -2, "two", None, "Infinity", float("inf"), 1e308], + ids=["below-one", "zero", "negative", "text", "null", "infinity-text", "infinity", "overflows-the-share"], +) +@pytest.mark.asyncio +async def test_an_unusable_max_borrow_multiple_is_ignored_so_the_team_borrows_up_to_the_model_limit( + monkeypatch: pytest.MonkeyPatch, multiple: object +): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _admit_then_settle(handler, _capped_team_key(multiple), 300) + await _check(handler, _capped_team_key(multiple), saturation=0.49) + + +@pytest.mark.asyncio +async def test_a_websocket_turn_is_refused_at_the_teams_max_borrow_multiple_while_the_model_is_idle( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.5) + await _admit_then_settle(handler, _capped_team_key(2), 200) + with pytest.raises(ProxyRateLimitError) as exc: + await handler.enforce_token_limits_for_frame( + user_api_key_dict=_capped_team_key(2), requested_model="shared-model" + ) + _held_to_borrow_cap(exc.value, "tokens", "model_tpm_limit") + + +@pytest.mark.asyncio +async def test_a_websocket_turn_under_the_teams_max_borrow_multiple_still_borrows(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.5) + await _admit_then_settle(handler, _capped_team_key(2), 190) + await handler.enforce_token_limits_for_frame(user_api_key_dict=_capped_team_key(2), requested_model="shared-model") diff --git a/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py b/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py index 88b02ed1aa1..c3e7ee8ad0d 100644 --- a/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py +++ b/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py @@ -344,6 +344,7 @@ _PROXY_READ_TEAM_METADATA_KEYS: Final = ( "phoenix_project_name_override", "phoenix_project_name", "opik", + "team_model_limit_max_borrow_multiple", ) _STORED_METADATA: Final = {"cost_center": "eng", "region": "eu", "priority": "high", "guardrails": ["pii"]}