fix(dynamic_rate_limiter_v3): publish a team's share gauges for requests served while the model is idle (#51709)

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

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)

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)