feat(prometheus): add per-team litellm_team_members_metric gauge (#31506)

Emit litellm_team_members_metric on every team member add and delete,
labelled by team and team_alias and set to the team's authoritative
member count. Because it is set from the current membership rather than
incremented or decremented, it tracks the count up and down, never goes
negative, and self-corrects on the next change after a proxy restart.
Bulk member add is covered for free since it delegates to
team_member_add, and the helper no-ops when the Prometheus callback is
not registered.

Resolves LIT-3082
This commit is contained in:
Yassin Kortam 2026-06-27 22:19:50 +03:00 committed by GitHub
parent 1883f975e2
commit b2e708d5ae
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 190 additions and 0 deletions

View file

@ -259,6 +259,13 @@ class PrometheusLogger(CustomLogger):
),
)
# Number of members in a team
self.litellm_team_members_metric = self._gauge_factory(
"litellm_team_members_metric",
"Number of members in a team",
labelnames=self.get_labels_for_metric("litellm_team_members_metric"),
)
# Remaining Budget for Org
self.litellm_remaining_org_budget_metric = self._gauge_factory(
"litellm_remaining_org_budget_metric",
@ -3557,6 +3564,22 @@ class PrometheusLogger(CustomLogger):
)
)
def set_team_members_metric(self, team: LiteLLM_TeamTable) -> None:
"""Set the team members gauge to the team's current member count."""
enum_values = UserAPIKeyLabelValues(
team=team.team_id,
team_alias=team.team_alias or "",
)
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name="litellm_team_members_metric"
),
enum_values=enum_values,
)
self.litellm_team_members_metric.labels(**_labels).set(
len(team.members_with_roles)
)
async def _set_org_budget_metrics_after_api_request(
self,
org_id: Optional[str],

View file

@ -23,6 +23,7 @@ from pydantic import BaseModel
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.integrations.prometheus import PrometheusLogger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import (
BlockTeamRequest,
@ -2445,6 +2446,23 @@ async def _add_team_members_to_team(
return updated_team, updated_users, updated_team_memberships
def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None:
"""Update the Prometheus team members gauge after a membership change.
No-ops when the Prometheus callback is not registered, and never lets a
metric failure break the team add/delete request.
"""
prometheus_logger = PrometheusLogger.get_instance()
if prometheus_logger is None:
return
try:
prometheus_logger.set_team_members_metric(team)
except Exception as e:
verbose_proxy_logger.debug(
"Prometheus: failed to emit team members metric: %s", str(e)
)
async def _validate_and_populate_member_user_info(
member: Member,
prisma_client: PrismaClient,
@ -2665,6 +2683,9 @@ async def team_member_add(
raise HTTPException(
status_code=404, detail={"error": f"Team with id {data.team_id} not found"}
)
_emit_team_members_metric(complete_team_data)
return TeamAddMemberResponse(
**updated_team.model_dump(),
updated_users=updated_users,
@ -2792,6 +2813,8 @@ async def team_member_delete(
data={"members_with_roles": json.dumps(_db_new_team_members)}, # type: ignore
)
_emit_team_members_metric(existing_team_row)
## DELETE TEAM ID from USER ROW, IF EXISTS ##
# get user row
key_val = {}

View file

@ -215,6 +215,7 @@ DEFINED_PROMETHEUS_METRICS = Literal[
"litellm_remaining_team_budget_metric",
"litellm_team_max_budget_metric",
"litellm_team_budget_remaining_hours_metric",
"litellm_team_members_metric",
"litellm_remaining_org_budget_metric",
"litellm_org_max_budget_metric",
"litellm_org_budget_remaining_hours_metric",
@ -533,6 +534,11 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.TEAM_ALIAS.value,
]
litellm_team_members_metric = [
UserAPIKeyLabelNames.TEAM.value,
UserAPIKeyLabelNames.TEAM_ALIAS.value,
]
litellm_remaining_org_budget_metric = [
UserAPIKeyLabelNames.ORG_ID.value,
UserAPIKeyLabelNames.ORG_ALIAS.value,

View file

@ -927,3 +927,80 @@ def test_custom_latency_buckets():
REGISTRY.unregister(collector)
except Exception:
pass
class TestSetTeamMembersMetric:
"""litellm_team_members_metric tracks the current member count per team."""
def _gauge_value(self, team_id, team_alias):
return REGISTRY.get_sample_value(
"litellm_team_members_metric",
{"team": team_id, "team_alias": team_alias},
)
def test_metric_initialized(self, prometheus_logger):
assert hasattr(prometheus_logger, "litellm_team_members_metric")
assert prometheus_logger.litellm_team_members_metric is not None
@pytest.mark.parametrize("count", [0, 1, 3, 7])
def test_sets_gauge_to_member_count(self, prometheus_logger, count):
from litellm.proxy._types import LiteLLM_TeamTable, Member
team = LiteLLM_TeamTable(
team_id="team-a",
team_alias="Acme",
members_with_roles=[
Member(user_id=f"u{i}", role="user") for i in range(count)
],
)
prometheus_logger.set_team_members_metric(team)
assert self._gauge_value("team-a", "Acme") == float(count)
def test_gauge_reflects_latest_count_not_delta(self, prometheus_logger):
"""Re-emitting overwrites with the authoritative count (set, not inc/dec)."""
from litellm.proxy._types import LiteLLM_TeamTable, Member
members = [Member(user_id=f"u{i}", role="user") for i in range(4)]
team = LiteLLM_TeamTable(
team_id="team-b", team_alias="Beta", members_with_roles=members
)
prometheus_logger.set_team_members_metric(team)
assert self._gauge_value("team-b", "Beta") == 4.0
# Drop two members and re-emit: gauge must read 2, not 4 and not -2.
team.members_with_roles = members[:2]
prometheus_logger.set_team_members_metric(team)
assert self._gauge_value("team-b", "Beta") == 2.0
def test_none_alias_falls_back_to_empty_string(self, prometheus_logger):
from litellm.proxy._types import LiteLLM_TeamTable, Member
team = LiteLLM_TeamTable(
team_id="team-c",
team_alias=None,
members_with_roles=[Member(user_id="solo", role="admin")],
)
prometheus_logger.set_team_members_metric(team)
assert self._gauge_value("team-c", "") == 1.0
def test_teams_isolated_by_label(self, prometheus_logger):
from litellm.proxy._types import LiteLLM_TeamTable, Member
team_one = LiteLLM_TeamTable(
team_id="team-1",
team_alias="One",
members_with_roles=[Member(user_id="a", role="user")],
)
team_two = LiteLLM_TeamTable(
team_id="team-2",
team_alias="Two",
members_with_roles=[
Member(user_id="b", role="user"),
Member(user_id="c", role="user"),
Member(user_id="d", role="user"),
],
)
prometheus_logger.set_team_members_metric(team_one)
prometheus_logger.set_team_members_metric(team_two)
assert self._gauge_value("team-1", "One") == 1.0
assert self._gauge_value("team-2", "Two") == 3.0

View file

@ -9434,3 +9434,64 @@ async def test_team_info_forwards_key_limit_to_get_data():
)
assert mock_prisma.get_data.await_args.kwargs["limit"] == 7
class TestEmitTeamMembersMetric:
"""The _emit_team_members_metric seam between the team handlers and Prometheus."""
@pytest.fixture
def restore_callbacks(self):
import litellm
original = litellm.callbacks
yield
litellm.callbacks = original
def _team(self, member_count):
return LiteLLM_TeamTable(
team_id="team-x",
team_alias="X",
members_with_roles=[
Member(user_id=f"u{i}", role="user") for i in range(member_count)
],
)
def test_emits_with_team_when_logger_registered(self, restore_callbacks):
import litellm
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy.management_endpoints.team_endpoints import (
_emit_team_members_metric,
)
fake_logger = MagicMock(spec=PrometheusLogger)
litellm.callbacks = [fake_logger]
team = self._team(3)
_emit_team_members_metric(team)
fake_logger.set_team_members_metric.assert_called_once_with(team)
def test_noop_when_no_logger_registered(self, restore_callbacks):
import litellm
from litellm.proxy.management_endpoints.team_endpoints import (
_emit_team_members_metric,
)
litellm.callbacks = []
# Must not raise when Prometheus is not enabled.
_emit_team_members_metric(self._team(2))
def test_metric_failure_does_not_break_request(self, restore_callbacks):
import litellm
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy.management_endpoints.team_endpoints import (
_emit_team_members_metric,
)
fake_logger = MagicMock(spec=PrometheusLogger)
fake_logger.set_team_members_metric.side_effect = Exception("boom")
litellm.callbacks = [fake_logger]
# A metric failure must be swallowed, not propagated to the handler.
_emit_team_members_metric(self._team(1))
fake_logger.set_team_members_metric.assert_called_once()