mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
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:
parent
1883f975e2
commit
b2e708d5ae
5 changed files with 190 additions and 0 deletions
|
|
@ -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],
|
||||
|
|
|
|||
|
|
@ -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 = {}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue