diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 81b9c80a24b..e201ccbe5c1 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -128,6 +128,10 @@ def _tightest_of(*shares: tuple[RateLimitDescriptor, int] | None) -> tuple[RateL return min((share for share in shares if share is not None), key=lambda share: share[1], default=None) +def _share_on_model(share: RateLimitDescriptor | None, team: str, model_group: str) -> RateLimitDescriptor | None: + return share if share is not None and share["value"] == f"{team}:{model_group}" else None + + def _capped_idle_share( share: RateLimitDescriptor | None, cap: RateLimitDescriptor | None ) -> tuple[RateLimitDescriptor, int] | None: @@ -689,6 +693,19 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): _tightest_team_share((stash.borrowing_team_model_limit, stash.priority_class_share)), _capped_idle_share(stash.borrowing_team_model_share, stash.borrowing_team_model_cap), ) + ) or _tightest_team_share( + ( + hold.share + if hold is not None + else _share_on_model( + stash.borrowing_team_model_share, + settlement.user_api_key_team_id, + ( + self.v3_limiter.rate_limited_model(model_group) + or RateLimitedModel(requested=model_group, group=model_group) + ).group, + ), + ) ) if share is None: return diff --git a/tests/integration/routing/test_dynamic_limiter_prometheus_metrics.py b/tests/integration/routing/test_dynamic_limiter_prometheus_metrics.py index 480fef6b5a7..afdca1716db 100644 --- a/tests/integration/routing/test_dynamic_limiter_prometheus_metrics.py +++ b/tests/integration/routing/test_dynamic_limiter_prometheus_metrics.py @@ -172,3 +172,26 @@ def test_a_team_under_its_share_borrows_nothing(saturated_model: tuple[Gateway, assert tuple(_team_series(samples, "litellm_dynamic_limiter_share_tokens", team).values()) == (TEAM_TPM,) assert tuple(_team_series(samples, "litellm_dynamic_limiter_borrowed_tokens", team).values()) == (0,) assert not _team_series(samples, "litellm_dynamic_limiter_throttled_total", team) + + +def test_a_team_served_while_the_model_is_idle_publishes_its_share_with_nothing_borrowed( + saturated_model: tuple[Gateway, str], +) -> None: + owned, upstream = saturated_model + with owned.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{OPENAI_MODEL}", api_base=f"{upstream}/v1", api_key="synthetic-openai-key", tpm=MODEL_TPM + ) + team: Final = scenario.team(metadata={"model_tpm_limit": {model: TEAM_TPM}}) + _settled_calls(owned, model, team, scenario.key(team_id=team), IDLE_CALLS) + samples: Final = eventually( + lambda: scrape(owned), + lambda scraped: bool(_team_series(scraped, "litellm_dynamic_limiter_share_tokens", team)), + seconds=30, + return_last_on_timeout=True, + ) + assert IDLE_CALLS * REPLY_TOKENS > TEAM_TPM + assert IDLE_CALLS * REPLY_TOKENS < MODEL_TPM * SATURATION_THRESHOLD + assert tuple(_team_series(samples, "litellm_dynamic_limiter_share_tokens", team).values()) == (TEAM_TPM,) + assert tuple(_team_series(samples, "litellm_dynamic_limiter_borrowed_tokens", team).values()) == (0,) + assert not _team_series(samples, "litellm_dynamic_limiter_throttled_total", team) 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 8a385b4a3fa..b24608239e2 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -3566,6 +3566,60 @@ async def test_a_capped_idle_teams_borrowing_is_read_against_its_share_not_the_w assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **TEAM_LABELS) == 50 +@pytest.mark.asyncio +async def test_an_uncapped_team_served_while_the_model_is_idle_publishes_its_share_with_nothing_borrowed( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + registry = _prometheus_sink(monkeypatch) + await _admit_at_the_hook(handler, _team_key(), tokens=150) + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **TEAM_LABELS) == 100 + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **TEAM_LABELS) == 0 + + +@pytest.mark.asyncio +async def test_an_idle_fallback_served_for_an_uncapped_team_publishes_its_share_on_the_fallback_model( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + registry = _prometheus_sink(monkeypatch) + team: Final = _team_key(team_limits={"model_tpm_limit": {"other-model": 40, "shared-model": 100}}) + assert await handler.async_pre_call_hook(team, DualCache(), {"model": "other-model"}, "completion") is None + assert await handler.async_fallback_target_allowed(team, "shared-model", {}) is True + await _settle_tokens(handler, 150, model_group="shared-model") + shared: Final = {"model_group": "shared-model", "team": "team-a", "priority": "default"} + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **shared) == 100 + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **shared) == 0 + + +@pytest.mark.asyncio +async def test_an_uncapped_alias_request_served_while_the_model_is_idle_publishes_its_share( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True, model_group_alias={"alias-model": "shared-model"}) + registry = _prometheus_sink(monkeypatch) + assert await handler.async_pre_call_hook(_team_key(), DualCache(), {"model": "alias-model"}, "completion") is None + await _settle_tokens(handler, 150, model_group="alias-model") + alias: Final = {"model_group": "alias-model", "team": "team-a", "priority": "default"} + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **alias) == 100 + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **alias) == 0 + + +@pytest.mark.asyncio +async def test_an_idle_fallback_to_a_model_the_team_has_no_share_on_publishes_no_share_there( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + registry = _prometheus_sink(monkeypatch) + team: Final = _team_key(team_limits={"model_tpm_limit": {"other-model": 40}}) + assert await handler.async_pre_call_hook(team, DualCache(), {"model": "other-model"}, "completion") is None + assert await handler.async_fallback_target_allowed(team, "shared-model", {}) is True + await _settle_tokens(handler, 150, model_group="shared-model") + shared: Final = {"model_group": "shared-model", "team": "team-a"} + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **shared) is None + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **shared) is None + + @pytest.mark.asyncio async def test_an_idle_fallback_served_for_a_capped_team_still_publishes_its_borrowing(monkeypatch: pytest.MonkeyPatch): handler = _borrowing_handler(monkeypatch, borrows=True)