From dc54c09246c60b0846ef0cde9bde99f6b1294163 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 13 Oct 2026 21:03:54 +0000 Subject: [PATCH] feat(control_plane): fan-out primitive and global usage and fleet aggregates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/control_plane/app.py | 13 +- .../proxy/control_plane/common/exceptions.py | 16 +- litellm/proxy/control_plane/common/fanout.py | 90 +++++++ litellm/proxy/control_plane/container.py | 3 + .../control_plane/dataplanes/dependencies.py | 24 +- .../control_plane/dataplanes/identity.py | 35 +++ .../proxy/control_plane/dataplanes/router.py | 31 +-- litellm/proxy/control_plane/fleet/__init__.py | 0 litellm/proxy/control_plane/fleet/router.py | 74 ++++++ litellm/proxy/control_plane/fleet/service.py | 137 ++++++++++ litellm/proxy/control_plane/settings.py | 3 +- litellm/proxy/control_plane/usage/__init__.py | 0 litellm/proxy/control_plane/usage/router.py | 170 +++++++++++++ litellm/proxy/control_plane/usage/service.py | 163 ++++++++++++ .../types/proxy/control_plane/dataplanes.py | 9 + litellm/types/proxy/control_plane/fleet.py | 55 ++++ litellm/types/proxy/control_plane/usage.py | 75 ++++++ .../proxy/control_plane/common/test_fanout.py | 104 ++++++++ .../proxy/control_plane/fleet/__init__.py | 0 .../proxy/control_plane/fleet/test_router.py | 239 ++++++++++++++++++ tests/unit/proxy/control_plane/support.py | 93 +++++++ .../unit/proxy/control_plane/test_settings.py | 19 ++ .../proxy/control_plane/usage/__init__.py | 0 .../proxy/control_plane/usage/test_router.py | 185 ++++++++++++++ .../proxy/control_plane/usage/test_service.py | 113 +++++++++ 25 files changed, 1615 insertions(+), 36 deletions(-) create mode 100644 litellm/proxy/control_plane/common/fanout.py create mode 100644 litellm/proxy/control_plane/dataplanes/identity.py create mode 100644 litellm/proxy/control_plane/fleet/__init__.py create mode 100644 litellm/proxy/control_plane/fleet/router.py create mode 100644 litellm/proxy/control_plane/fleet/service.py create mode 100644 litellm/proxy/control_plane/usage/__init__.py create mode 100644 litellm/proxy/control_plane/usage/router.py create mode 100644 litellm/proxy/control_plane/usage/service.py create mode 100644 litellm/types/proxy/control_plane/fleet.py create mode 100644 litellm/types/proxy/control_plane/usage.py create mode 100644 tests/unit/proxy/control_plane/common/test_fanout.py create mode 100644 tests/unit/proxy/control_plane/fleet/__init__.py create mode 100644 tests/unit/proxy/control_plane/fleet/test_router.py create mode 100644 tests/unit/proxy/control_plane/support.py create mode 100644 tests/unit/proxy/control_plane/usage/__init__.py create mode 100644 tests/unit/proxy/control_plane/usage/test_router.py create mode 100644 tests/unit/proxy/control_plane/usage/test_service.py diff --git a/litellm/proxy/control_plane/app.py b/litellm/proxy/control_plane/app.py index d8e5fc093b3..a01ba31b6b8 100644 --- a/litellm/proxy/control_plane/app.py +++ b/litellm/proxy/control_plane/app.py @@ -15,9 +15,16 @@ from litellm.proxy.control_plane.common.exceptions import ( install_exception_handlers, ) from litellm.proxy.control_plane.container import Container -from litellm.proxy.control_plane.dataplanes.dependencies import get_dataplane_client, get_registry +from litellm.proxy.control_plane.dataplanes.dependencies import ( + get_control_plane_settings, + get_dataplane_client, + get_fanout, + get_registry, +) from litellm.proxy.control_plane.dataplanes.router import router as dataplanes_router +from litellm.proxy.control_plane.fleet.router import router as fleet_router from litellm.proxy.control_plane.health.router import router as health_router +from litellm.proxy.control_plane.usage.router import router as usage_router from litellm.proxy.middleware.request_size_limit_middleware import RequestSizeLimitMiddleware TITLE: Final = "LiteLLM Control Plane" @@ -82,9 +89,13 @@ def create_app(container: Container, license_check: License) -> FastAPI: app.dependency_overrides[get_token_verifier] = lambda: container.token_verifier app.dependency_overrides[get_login_service] = lambda: container.login_service app.dependency_overrides[get_registry] = lambda: container.registry + app.dependency_overrides[get_control_plane_settings] = lambda: container.settings app.dependency_overrides[get_dataplane_client] = lambda: container.dataplane_client + app.dependency_overrides[get_fanout] = lambda: container.fanout install_exception_handlers(app) app.include_router(health_router) app.include_router(auth_router, prefix=container.settings.api_prefix) app.include_router(dataplanes_router, prefix=container.settings.api_prefix) + app.include_router(usage_router, prefix=container.settings.api_prefix) + app.include_router(fleet_router, prefix=container.settings.api_prefix) return app diff --git a/litellm/proxy/control_plane/common/exceptions.py b/litellm/proxy/control_plane/common/exceptions.py index 097a3f74bd2..3a47a0f9dc4 100644 --- a/litellm/proxy/control_plane/common/exceptions.py +++ b/litellm/proxy/control_plane/common/exceptions.py @@ -25,6 +25,10 @@ def error_body(slug: str, title: str, status: int, detail: str) -> ErrorBody: return ErrorBody(type=f"{ERROR_TYPE_BASE}{slug}", title=title, status=status, detail=detail) +def bad_request(detail: str) -> ErrorBody: + return error_body("bad-request", "Bad Request", HTTPStatus.BAD_REQUEST, detail) + + def unauthorized(detail: str) -> ErrorBody: return error_body("unauthorized", "Unauthorized", 401, detail) @@ -41,6 +45,14 @@ def too_many_requests(detail: str) -> ErrorBody: return error_body("too-many-requests", "Too Many Requests", 429, detail) +def bad_gateway(detail: str) -> ErrorBody: + return error_body("bad-gateway", "Bad Gateway", 502, detail) + + +def service_unavailable(detail: str) -> ErrorBody: + return error_body("service-unavailable", "Service Unavailable", 503, detail) + + class ControlPlaneException(Exception): """An error raised from a dependency or handler, rendered by the app-level handler below.""" @@ -50,10 +62,6 @@ class ControlPlaneException(Exception): super().__init__(body.detail) -def bad_gateway(detail: str) -> ErrorBody: - return error_body("bad-gateway", "Bad Gateway", 502, detail) - - def _render(body: ErrorBody) -> JSONResponse: response: Final = error_response(body) if body.status == 401: diff --git a/litellm/proxy/control_plane/common/fanout.py b/litellm/proxy/control_plane/common/fanout.py new file mode 100644 index 00000000000..bde3624a201 --- /dev/null +++ b/litellm/proxy/control_plane/common/fanout.py @@ -0,0 +1,90 @@ +"""Bounded scatter/gather over dataplanes. A dataplane that fails is a Missed value in the result, never an exception.""" + +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import Final, Generic, TypeVar + +from pydantic import TypeAdapter, ValidationError + +from litellm.proxy.control_plane.dataplanes.client import CallResult, DataplaneReply, DataplaneUnreachable +from litellm.proxy.control_plane.dataplanes.credentials import CredentialUnavailable +from litellm.proxy.control_plane.dataplanes.registry import Dataplane + +T: Final = TypeVar("T") +M: Final = TypeVar("M") + + +@dataclass(frozen=True, slots=True) +class Failure: + reason: str + + +@dataclass(frozen=True, slots=True) +class Missed: + dataplane: Dataplane + reason: str + + +@dataclass(frozen=True, slots=True) +class Reached(Generic[T]): + dataplane: Dataplane + value: T + + +@dataclass(frozen=True, slots=True) +class Gathered(Generic[T]): + reached: tuple[Reached[T], ...] + missed: tuple[Missed, ...] + + @property + def partial(self) -> bool: + return bool(self.missed) + + @property + def nothing_answered(self) -> bool: + return bool(self.missed) and not self.reached + + +@dataclass(frozen=True, slots=True) +class Fanout: + concurrency: int + _gate: asyncio.Semaphore = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + object.__setattr__(self, "_gate", asyncio.Semaphore(self.concurrency)) + + async def gather( + self, dataplanes: tuple[Dataplane, ...], fetch: Callable[[Dataplane], Awaitable[T | Failure]] + ) -> Gathered[T]: + async def bounded(dataplane: Dataplane) -> Reached[T] | Missed: + async with self._gate: + outcome: Final = await fetch(dataplane) + if isinstance(outcome, Failure): + return Missed(dataplane, outcome.reason) + return Reached(dataplane, outcome) + + outcomes: Final = await asyncio.gather(*(bounded(dataplane) for dataplane in dataplanes)) + return Gathered( + reached=tuple(outcome for outcome in outcomes if isinstance(outcome, Reached)), + missed=tuple(outcome for outcome in outcomes if isinstance(outcome, Missed)), + ) + + +def parse_reply(result: CallResult, shape: TypeAdapter[M]) -> M | Failure: + match result: + case DataplaneReply(status=status, body=body) if status < 400: + try: + return shape.validate_python(body) + except ValidationError as error: + return Failure(f"unexpected response shape: {error.error_count()} validation error(s)") + case DataplaneReply(status=status, credential=credential): + return Failure(f"HTTP {status} with the {credential} credential") + case DataplaneUnreachable(detail=detail, attempts=attempts): + return Failure(f"unreachable after {attempts} attempt(s): {detail}") + case CredentialUnavailable(reason=reason): + return Failure(f"no usable credential: {reason}") + + +def describe_missed(missed: tuple[Missed, ...]) -> str: + return ", ".join(f"{miss.dataplane.id}: {miss.reason}" for miss in missed) diff --git a/litellm/proxy/control_plane/container.py b/litellm/proxy/control_plane/container.py index a45268e9d9c..e73212c2b19 100644 --- a/litellm/proxy/control_plane/container.py +++ b/litellm/proxy/control_plane/container.py @@ -24,6 +24,7 @@ from litellm.proxy.control_plane.auth.rbac import ( build_policy, parse_role_definition, ) +from litellm.proxy.control_plane.common.fanout import Fanout from litellm.proxy.control_plane.dataplanes.client import DataplaneClient, HttpTransport, RetryPolicy from litellm.proxy.control_plane.dataplanes.credentials import TokenExchanger, http_token_exchanger from litellm.proxy.control_plane.dataplanes.registry import DataplaneRegistry, build_registry @@ -39,6 +40,7 @@ class Container: login_service: LoginService registry: DataplaneRegistry dataplane_client: DataplaneClient + fanout: Fanout class InvalidRolePolicy(ValueError): @@ -143,4 +145,5 @@ def build_container( timeout_seconds=settings.dataplane_timeout_seconds, retry=RetryPolicy(attempts=settings.dataplane_retry_attempts), ), + fanout=Fanout(concurrency=settings.fanout_concurrency), ) diff --git a/litellm/proxy/control_plane/dataplanes/dependencies.py b/litellm/proxy/control_plane/dataplanes/dependencies.py index bf081c8c719..22928c90012 100644 --- a/litellm/proxy/control_plane/dataplanes/dependencies.py +++ b/litellm/proxy/control_plane/dataplanes/dependencies.py @@ -1,14 +1,20 @@ from typing import Annotated, Final -from fastapi import Depends +from fastapi import Depends, Query from fastapi.security import HTTPAuthorizationCredentials from litellm.proxy.control_plane.auth.dependencies import bearer_scheme, get_principal from litellm.proxy.control_plane.auth.rbac import Principal from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, forbidden, not_found, unauthorized +from litellm.proxy.control_plane.common.fanout import Fanout from litellm.proxy.control_plane.dataplanes.client import DataplaneClient from litellm.proxy.control_plane.dataplanes.credentials import Caller -from litellm.proxy.control_plane.dataplanes.registry import Dataplane, DataplaneRegistry +from litellm.proxy.control_plane.dataplanes.registry import Dataplane, DataplaneRegistry, visible_to +from litellm.proxy.control_plane.settings import ControlPlaneSettings + + +def get_control_plane_settings() -> ControlPlaneSettings: + raise RuntimeError("get_control_plane_settings is not bound, build the app with create_app") def get_registry() -> DataplaneRegistry: @@ -19,6 +25,10 @@ def get_dataplane_client() -> DataplaneClient: raise RuntimeError("get_dataplane_client is not bound, build the app with create_app") +def get_fanout() -> Fanout: + raise RuntimeError("get_fanout is not bound, build the app with create_app") + + async def get_caller( credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(bearer_scheme)], principal: Annotated[Principal, Depends(get_principal)], @@ -39,3 +49,13 @@ async def dataplane_access( if not principal.dataplanes.allows(dataplane_id): raise ControlPlaneException(forbidden(f"Not bound to dataplane {dataplane_id!r}.")) return dataplane + + +async def selected_dataplanes( + principal: Annotated[Principal, Depends(get_principal)], + registry: Annotated[DataplaneRegistry, Depends(get_registry)], + dataplane: Annotated[str | None, Query(description="Restrict a fleet-wide call to one dataplane id.")] = None, +) -> tuple[Dataplane, ...]: + if dataplane is None: + return visible_to(registry, principal) + return (await dataplane_access(dataplane, principal, registry),) diff --git a/litellm/proxy/control_plane/dataplanes/identity.py b/litellm/proxy/control_plane/dataplanes/identity.py new file mode 100644 index 00000000000..5b92b4aa298 --- /dev/null +++ b/litellm/proxy/control_plane/dataplanes/identity.py @@ -0,0 +1,35 @@ +from typing import Final + +from pydantic import BaseModel, ConfigDict, ValidationError + +USER_INFO_PATH: Final = "/user/info" + + +class UserInfo(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + user_id: str | None = None + user_email: str | None = None + user_role: str | None = None + + +class UserInfoBody(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + user_id: str | None = None + user_info: UserInfo | None = None + + +def parse_user_info(body: object) -> UserInfoBody | None: + if not isinstance(body, dict): + return None + try: + return UserInfoBody.model_validate(body) + except ValidationError: + return None + + +def resolved_user_id(info: UserInfoBody) -> str | None: + if info.user_id is not None: + return info.user_id + return None if info.user_info is None else info.user_info.user_id diff --git a/litellm/proxy/control_plane/dataplanes/router.py b/litellm/proxy/control_plane/dataplanes/router.py index 9522ae056ba..d23928cac64 100644 --- a/litellm/proxy/control_plane/dataplanes/router.py +++ b/litellm/proxy/control_plane/dataplanes/router.py @@ -2,7 +2,6 @@ from http import HTTPStatus from typing import Annotated, Final from fastapi import APIRouter, Depends, Security -from pydantic import BaseModel, ConfigDict, ValidationError from litellm.proxy.control_plane.auth.dependencies import require_scopes from litellm.proxy.control_plane.auth.rbac import Principal, Scope @@ -15,6 +14,7 @@ from litellm.proxy.control_plane.dataplanes.dependencies import ( get_dataplane_client, get_registry, ) +from litellm.proxy.control_plane.dataplanes.identity import USER_INFO_PATH, UserInfo, parse_user_info from litellm.proxy.control_plane.dataplanes.registry import Dataplane, DataplaneRegistry, visible_to from litellm.types.proxy.control_plane.dataplanes import ( DataplaneHealthResponse, @@ -26,22 +26,6 @@ from litellm.types.proxy.control_plane.dataplanes import ( router: Final = APIRouter(prefix="/dataplanes", tags=["dataplanes"]) LIVENESS_PATH: Final = "/health/liveliness" -USER_INFO_PATH: Final = "/user/info" - - -class _UserInfo(BaseModel): - model_config = ConfigDict(frozen=True, extra="ignore") - - user_id: str | None = None - user_email: str | None = None - user_role: str | None = None - - -class _UserInfoBody(BaseModel): - model_config = ConfigDict(frozen=True, extra="ignore") - - user_id: str | None = None - user_info: _UserInfo | None = None def _to_response(dataplane: Dataplane) -> DataplaneResponse: @@ -96,10 +80,10 @@ async def dataplane_whoami( result: Final = await client.get(dataplane, USER_INFO_PATH, caller) match result: case DataplaneReply(status=status, body=body, credential=credential) if status < 400: - info: Final = _identity_body(body) + info: Final = parse_user_info(body) if info is None: raise ControlPlaneException(bad_gateway("Dataplane returned a malformed /user/info body.")) - user: Final = info.user_info or _UserInfo() + user: Final = info.user_info or UserInfo() return DataplaneIdentityResponse( id=dataplane.id, credential=credential, @@ -119,12 +103,3 @@ async def dataplane_whoami( raise ControlPlaneException(bad_gateway(f"Dataplane unreachable after {attempts} attempt(s): {detail}")) case CredentialUnavailable(reason=reason): raise ControlPlaneException(bad_gateway(f"No usable credential for this dataplane: {reason}")) - - -def _identity_body(body: object) -> _UserInfoBody | None: - if not isinstance(body, dict): - return None - try: - return _UserInfoBody.model_validate(body) - except ValidationError: - return None diff --git a/litellm/proxy/control_plane/fleet/__init__.py b/litellm/proxy/control_plane/fleet/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/control_plane/fleet/router.py b/litellm/proxy/control_plane/fleet/router.py new file mode 100644 index 00000000000..4a5642a7212 --- /dev/null +++ b/litellm/proxy/control_plane/fleet/router.py @@ -0,0 +1,74 @@ +from typing import Annotated, Final, TypeAlias, TypeVar + +from fastapi import APIRouter, Depends, Security + +from litellm.proxy.control_plane.auth.dependencies import require_scopes +from litellm.proxy.control_plane.auth.rbac import Principal, Scope +from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, service_unavailable +from litellm.proxy.control_plane.common.fanout import Fanout, Gathered, describe_missed +from litellm.proxy.control_plane.dataplanes.client import DataplaneClient +from litellm.proxy.control_plane.dataplanes.credentials import Caller +from litellm.proxy.control_plane.dataplanes.dependencies import ( + get_caller, + get_dataplane_client, + get_fanout, + selected_dataplanes, +) +from litellm.proxy.control_plane.dataplanes.registry import Dataplane +from litellm.proxy.control_plane.fleet.service import fetch_models, fetch_teams, group_models +from litellm.types.proxy.control_plane.dataplanes import UnreachableDataplane +from litellm.types.proxy.control_plane.fleet import FleetModel, FleetModelsResponse, FleetTeam, FleetTeamsResponse + +T: Final = TypeVar("T") + +router: Final = APIRouter(prefix="/fleet", tags=["fleet"]) + +FleetRead: TypeAlias = Annotated[Principal, Security(require_scopes, scopes=[Scope.FLEET_READ.value])] +Selected: TypeAlias = Annotated[tuple[Dataplane, ...], Depends(selected_dataplanes)] + + +def _or_503(gathered: Gathered[T]) -> Gathered[T]: + if gathered.nothing_answered: + raise ControlPlaneException(service_unavailable(f"No dataplane answered: {describe_missed(gathered.missed)}")) + return gathered + + +def _flatten(gathered: Gathered[tuple[T, ...]]) -> tuple[T, ...]: + return tuple(item for reached in gathered.reached for item in reached.value) # comprehension-ok: flatten + + +def _unreachable(gathered: Gathered[T]) -> tuple[UnreachableDataplane, ...]: + return tuple(UnreachableDataplane(id=miss.dataplane.id, reason=miss.reason) for miss in gathered.missed) + + +@router.get("/models", response_model=FleetModelsResponse) +async def fleet_models( + _principal: FleetRead, + dataplanes: Selected, + caller: Annotated[Caller, Depends(get_caller)], + client: Annotated[DataplaneClient, Depends(get_dataplane_client)], + fanout: Annotated[Fanout, Depends(get_fanout)], +) -> FleetModelsResponse: + """Every deployment across the fleet, plus which dataplanes serve each public model name.""" + gathered: Final[Gathered[tuple[FleetModel, ...]]] = _or_503( + await fanout.gather(dataplanes, lambda dataplane: fetch_models(client, dataplane, caller)) + ) + items: Final = _flatten(gathered) + return FleetModelsResponse( + items=items, groups=group_models(items), partial=gathered.partial, unreachable=_unreachable(gathered) + ) + + +@router.get("/teams", response_model=FleetTeamsResponse) +async def fleet_teams( + _principal: FleetRead, + dataplanes: Selected, + caller: Annotated[Caller, Depends(get_caller)], + client: Annotated[DataplaneClient, Depends(get_dataplane_client)], + fanout: Annotated[Fanout, Depends(get_fanout)], +) -> FleetTeamsResponse: + """Every team the caller may see on each dataplane.""" + gathered: Final[Gathered[tuple[FleetTeam, ...]]] = _or_503( + await fanout.gather(dataplanes, lambda dataplane: fetch_teams(client, dataplane, caller)) + ) + return FleetTeamsResponse(items=_flatten(gathered), partial=gathered.partial, unreachable=_unreachable(gathered)) diff --git a/litellm/proxy/control_plane/fleet/service.py b/litellm/proxy/control_plane/fleet/service.py new file mode 100644 index 00000000000..110f70b7d3c --- /dev/null +++ b/litellm/proxy/control_plane/fleet/service.py @@ -0,0 +1,137 @@ +"""Fleet inventory: the union of what every dataplane serves, read through the caller's own identity.""" + +from http import HTTPStatus +from itertools import groupby +from types import MappingProxyType +from typing import Final + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from litellm.proxy.control_plane.common.fanout import Failure, parse_reply +from litellm.proxy.control_plane.dataplanes.client import CallResult, DataplaneClient, DataplaneReply +from litellm.proxy.control_plane.dataplanes.credentials import Caller +from litellm.proxy.control_plane.dataplanes.identity import USER_INFO_PATH, parse_user_info, resolved_user_id +from litellm.proxy.control_plane.dataplanes.registry import Dataplane +from litellm.types.proxy.control_plane.fleet import FleetModel, FleetModelGroup, FleetTeam + +MODEL_INFO_PATH: Final = "/model/info" +TEAM_LIST_PATH: Final = "/team/list" + + +class _Lenient(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + +class _LitellmParams(_Lenient): + model: str | None = None + custom_llm_provider: str | None = None + + +class _ModelInfo(_Lenient): + id: str | None = None + litellm_provider: str | None = None + mode: str | None = None + + +class _ModelEntry(_Lenient): + model_name: str + litellm_params: _LitellmParams = _LitellmParams() + model_info: _ModelInfo = _ModelInfo() + + +class _ModelInfoBody(_Lenient): + data: tuple[_ModelEntry, ...] | _ModelEntry = () + + +class _TeamEntry(_Lenient): + team_id: str + team_alias: str | None = None + organization_id: str | None = None + spend: float | None = None + max_budget: float | None = None + members_with_roles: tuple[object, ...] = () + models: tuple[str, ...] = () + blocked: bool = False + + +_MODELS: Final = TypeAdapter(_ModelInfoBody) +_TEAMS: Final = TypeAdapter(tuple[_TeamEntry, ...]) + + +async def fetch_models( + client: DataplaneClient, dataplane: Dataplane, caller: Caller +) -> tuple[FleetModel, ...] | Failure: + body: Final = parse_reply(await client.get(dataplane, MODEL_INFO_PATH, caller), _MODELS) + if isinstance(body, Failure): + return body + entries: Final = body.data if isinstance(body.data, tuple) else (body.data,) + return tuple( + FleetModel( + dataplane_id=dataplane.id, + model_name=entry.model_name, + litellm_model=entry.litellm_params.model, + provider=entry.litellm_params.custom_llm_provider or entry.model_info.litellm_provider, + model_id=entry.model_info.id, + mode=entry.model_info.mode, + ) + for entry in entries + ) + + +NOT_ALLOWED_TO_LIST_ALL: Final = frozenset({HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN}) + + +def _refused(result: CallResult) -> bool: + return isinstance(result, DataplaneReply) and result.status in NOT_ALLOWED_TO_LIST_ALL + + +async def _list_teams(client: DataplaneClient, dataplane: Dataplane, caller: Caller) -> CallResult: + """Every team when the dataplane lets the caller see them all, else the teams the caller belongs to.""" + everything: Final = await client.get(dataplane, TEAM_LIST_PATH, caller) + if not _refused(everything): + return everything + identity: Final = await client.get(dataplane, USER_INFO_PATH, caller) + match identity: + case DataplaneReply(status=status, body=body) if status < HTTPStatus.BAD_REQUEST: + info: Final = parse_user_info(body) + user_id: Final = None if info is None else resolved_user_id(info) + if user_id is None: + return everything + own: Final = MappingProxyType({"user_id": user_id}) + return await client.get(dataplane, TEAM_LIST_PATH, caller, params=tuple(own.items())) + case DataplaneReply(status=HTTPStatus.NOT_FOUND, credential=credential): + return DataplaneReply(HTTPStatus.OK, (), credential) + case _: + return everything + + +async def fetch_teams(client: DataplaneClient, dataplane: Dataplane, caller: Caller) -> tuple[FleetTeam, ...] | Failure: + teams: Final = parse_reply(await _list_teams(client, dataplane, caller), _TEAMS) + if isinstance(teams, Failure): + return teams + return tuple( + FleetTeam( + dataplane_id=dataplane.id, + team_id=team.team_id, + team_alias=team.team_alias, + organization_id=team.organization_id, + spend=team.spend or 0.0, + max_budget=team.max_budget, + members=len(team.members_with_roles), + models=team.models, + blocked=team.blocked, + ) + for team in teams + ) + + +def group_models(models: tuple[FleetModel, ...]) -> tuple[FleetModelGroup, ...]: + by_name: Final = sorted(models, key=_model_name) + return tuple( + FleetModelGroup(model_name=name, dataplanes=tuple(sorted({model.dataplane_id for model in served}))) + for name, served in groupby(by_name, key=_model_name) + ) + + +def _model_name(model: FleetModel) -> str: + return model.model_name diff --git a/litellm/proxy/control_plane/settings.py b/litellm/proxy/control_plane/settings.py index 90490960a68..030c602b153 100644 --- a/litellm/proxy/control_plane/settings.py +++ b/litellm/proxy/control_plane/settings.py @@ -110,7 +110,8 @@ class ControlPlaneSettings(BaseSettings): dataplanes_file: str | None = None dataplane_timeout_seconds: float = DEFAULT_DATAPLANE_TIMEOUT_SECONDS dataplane_retry_attempts: int = 3 - fanout_concurrency: int = 16 + fanout_concurrency: int = Field(default=16, ge=1) + usage_max_window_days: int = Field(default=366, ge=1) api_prefix: str = "/api/v1" @field_validator("master_key") diff --git a/litellm/proxy/control_plane/usage/__init__.py b/litellm/proxy/control_plane/usage/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/control_plane/usage/router.py b/litellm/proxy/control_plane/usage/router.py new file mode 100644 index 00000000000..4961825d813 --- /dev/null +++ b/litellm/proxy/control_plane/usage/router.py @@ -0,0 +1,170 @@ +from datetime import UTC, date, datetime, timedelta +from typing import Annotated, Final, TypeAlias + +from fastapi import APIRouter, Depends, Query, Security + +from litellm.proxy.control_plane.auth.dependencies import require_scopes +from litellm.proxy.control_plane.auth.rbac import Principal, Scope +from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, bad_request, service_unavailable +from litellm.proxy.control_plane.common.fanout import Fanout, Gathered, describe_missed +from litellm.proxy.control_plane.dataplanes.client import DataplaneClient +from litellm.proxy.control_plane.dataplanes.credentials import Caller +from litellm.proxy.control_plane.dataplanes.dependencies import ( + get_caller, + get_control_plane_settings, + get_dataplane_client, + get_fanout, + selected_dataplanes, +) +from litellm.proxy.control_plane.dataplanes.registry import Dataplane +from litellm.proxy.control_plane.settings import ControlPlaneSettings +from litellm.proxy.control_plane.usage.service import ( + DataplaneActivity, + DataplaneDays, + UsageWindow, + daily_rows, + dataplane_totals, + fetch_daily_activity, + ranked, + totals, +) +from litellm.types.proxy.control_plane.dataplanes import UnreachableDataplane +from litellm.types.proxy.control_plane.usage import ( + DailyUsageResponse, + DailyUsageRow, + DataplaneUsage, + TopUsageResponse, + UsageGroupBy, + UsageRank, + UsageSummaryResponse, + UsageWindowResponse, +) + +DEFAULT_WINDOW_DAYS: Final = 30 + +router: Final = APIRouter(prefix="/usage", tags=["usage"]) + +UsageRead: TypeAlias = Annotated[Principal, Security(require_scopes, scopes=[Scope.USAGE_READ.value])] +Selected: TypeAlias = Annotated[tuple[Dataplane, ...], Depends(selected_dataplanes)] +GroupBy: TypeAlias = Annotated[UsageGroupBy, Query(description="Dimension to split the fleet's usage by.")] + + +def usage_window( + settings: Annotated[ControlPlaneSettings, Depends(get_control_plane_settings)], + start_date: Annotated[date | None, Query(description="First day, inclusive. Defaults to 30 days ago.")] = None, + end_date: Annotated[date | None, Query(description="Last day, inclusive. Defaults to today (UTC).")] = None, +) -> UsageWindow: + end: Final = end_date if end_date is not None else datetime.now(UTC).date() + default_window_days: Final = min(DEFAULT_WINDOW_DAYS, settings.usage_max_window_days) + start: Final = start_date if start_date is not None else end - timedelta(days=default_window_days - 1) + if start > end: + raise ControlPlaneException(bad_request(f"start_date {start} is after end_date {end}.")) + if (end - start).days + 1 > settings.usage_max_window_days: + raise ControlPlaneException(bad_request(f"Usage windows may not exceed {settings.usage_max_window_days} days.")) + return UsageWindow(start=start, end=end) + + +Window: TypeAlias = Annotated[UsageWindow, Depends(usage_window)] + + +async def _gather_days( + fanout: Fanout, client: DataplaneClient, caller: Caller, dataplanes: tuple[Dataplane, ...], window: UsageWindow +) -> Gathered[DataplaneActivity]: + gathered: Final = await fanout.gather( + dataplanes, lambda dataplane: fetch_daily_activity(client, dataplane, caller, window) + ) + if gathered.nothing_answered: + raise ControlPlaneException(service_unavailable(f"No dataplane answered: {describe_missed(gathered.missed)}")) + return gathered + + +def _fleet(gathered: Gathered[DataplaneActivity]) -> tuple[DataplaneDays, ...]: + return tuple(DataplaneDays(reached.dataplane.id, reached.value.days) for reached in gathered.reached) + + +def _unreachable(gathered: Gathered[DataplaneActivity]) -> tuple[UnreachableDataplane, ...]: + return tuple(UnreachableDataplane(id=miss.dataplane.id, reason=miss.reason) for miss in gathered.missed) + + +def _truncated_dataplanes(gathered: Gathered[DataplaneActivity]) -> tuple[str, ...]: + return tuple(reached.dataplane.id for reached in gathered.reached if reached.value.keys_truncated) + + +def _window(window: UsageWindow) -> UsageWindowResponse: + return UsageWindowResponse(start_date=window.start, end_date=window.end) + + +@router.get("/summary", response_model=UsageSummaryResponse) +async def usage_summary( + _principal: UsageRead, + dataplanes: Selected, + window: Window, + caller: Annotated[Caller, Depends(get_caller)], + client: Annotated[DataplaneClient, Depends(get_dataplane_client)], + fanout: Annotated[Fanout, Depends(get_fanout)], +) -> UsageSummaryResponse: + """Spend and token totals across every dataplane the caller is bound to.""" + gathered: Final = await _gather_days(fanout, client, caller, dataplanes, window) + fleet: Final = _fleet(gathered) + truncated: Final = _truncated_dataplanes(gathered) + return UsageSummaryResponse( + window=_window(window), + totals=totals(fleet), + by_dataplane=tuple(DataplaneUsage(id=days.dataplane_id, metrics=dataplane_totals(days)) for days in fleet), + partial=gathered.partial or bool(truncated), + unreachable=_unreachable(gathered), + truncated_dataplanes=truncated, + ) + + +@router.get("/daily", response_model=DailyUsageResponse) +async def usage_daily( + _principal: UsageRead, + dataplanes: Selected, + window: Window, + caller: Annotated[Caller, Depends(get_caller)], + client: Annotated[DataplaneClient, Depends(get_dataplane_client)], + fanout: Annotated[Fanout, Depends(get_fanout)], + group_by: GroupBy = "dataplane", +) -> DailyUsageResponse: + """Per-day usage across the fleet, split by dataplane, model, provider, key, team or endpoint.""" + gathered: Final = await _gather_days(fanout, client, caller, dataplanes, window) + truncated: Final = _truncated_dataplanes(gathered) + return DailyUsageResponse( + window=_window(window), + group_by=group_by, + rows=tuple( + DailyUsageRow(date=day, group=group, metrics=metrics) + for (day, group), metrics in daily_rows(_fleet(gathered), group_by) + ), + partial=gathered.partial or bool(truncated), + unreachable=_unreachable(gathered), + truncated_dataplanes=truncated, + ) + + +@router.get("/top", response_model=TopUsageResponse) +async def usage_top( + _principal: UsageRead, + dataplanes: Selected, + window: Window, + caller: Annotated[Caller, Depends(get_caller)], + client: Annotated[DataplaneClient, Depends(get_dataplane_client)], + fanout: Annotated[Fanout, Depends(get_fanout)], + group_by: GroupBy = "model", + limit: Annotated[int, Query(ge=1, le=100)] = 10, +) -> TopUsageResponse: + """The biggest spenders in the window along one dimension, summed across the fleet.""" + gathered: Final = await _gather_days(fanout, client, caller, dataplanes, window) + truncated: Final = _truncated_dataplanes(gathered) + return TopUsageResponse( + window=_window(window), + group_by=group_by, + limit=limit, + items=tuple( + UsageRank(group=group, metrics=metrics) for group, metrics in ranked(_fleet(gathered), group_by, limit) + ), + partial=gathered.partial or bool(truncated), + unreachable=_unreachable(gathered), + truncated_dataplanes=truncated, + ) diff --git a/litellm/proxy/control_plane/usage/service.py b/litellm/proxy/control_plane/usage/service.py new file mode 100644 index 00000000000..e156988574c --- /dev/null +++ b/litellm/proxy/control_plane/usage/service.py @@ -0,0 +1,163 @@ +"""Fleet-wide usage: fetch each dataplane's daily activity, then pure reducers over SpendMetrics.""" + +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from datetime import date +from functools import reduce +from itertools import groupby +from types import MappingProxyType +from typing import Final, TypeAlias + +from pydantic import TypeAdapter + +from litellm.constants import USAGE_TOP_API_KEYS_MAX +from litellm.proxy.control_plane.common.fanout import Failure, parse_reply +from litellm.proxy.control_plane.dataplanes.client import DataplaneClient +from litellm.proxy.control_plane.dataplanes.credentials import Caller +from litellm.proxy.control_plane.dataplanes.registry import Dataplane +from litellm.types.proxy.control_plane.usage import UsageGroupBy +from litellm.types.proxy.management_endpoints.common_daily_activity import ( + DailySpendData, + SpendAnalyticsPaginatedResponse, + SpendMetrics, +) + +DAILY_ACTIVITY_PATH: Final = "/user/daily/activity/aggregated" +NO_TEAM: Final = "(no team)" + +_DAILY: Final = TypeAdapter(SpendAnalyticsPaginatedResponse) +_NUMBERS: Final = TypeAdapter(Mapping[str, float]) + +DatedGroup: TypeAlias = tuple[date, str] +Grouped: TypeAlias = tuple[tuple[DatedGroup, SpendMetrics], ...] + + +@dataclass(frozen=True, slots=True) +class UsageWindow: + start: date + end: date + + def params(self) -> Mapping[str, str]: + return MappingProxyType( + { + "start_date": self.start.isoformat(), + "end_date": self.end.isoformat(), + } + ) + + +@dataclass(frozen=True, slots=True) +class DataplaneDays: + dataplane_id: str + days: tuple[DailySpendData, ...] + + +@dataclass(frozen=True, slots=True) +class DataplaneActivity: + days: tuple[DailySpendData, ...] + keys_truncated: bool + + +async def fetch_daily_activity( + client: DataplaneClient, dataplane: Dataplane, caller: Caller, window: UsageWindow +) -> DataplaneActivity | Failure: + parsed: Final = parse_reply( + await client.get( + dataplane, + DAILY_ACTIVITY_PATH, + caller, + params=tuple(window.params().items()) + (("api_key_limit", str(USAGE_TOP_API_KEYS_MAX)),), + ), + _DAILY, + ) + if isinstance(parsed, Failure): + return parsed + metadata: Final = parsed.metadata + keys_truncated: Final = ( + metadata.api_key_limit is not None + and metadata.total_api_keys is not None + and metadata.total_api_keys > metadata.api_key_limit + ) + return DataplaneActivity(days=tuple(parsed.results), keys_truncated=keys_truncated) + + +def add_metrics(left: SpendMetrics, right: SpendMetrics) -> SpendMetrics: + first: Final = _NUMBERS.validate_python(left.model_dump()) + second: Final = _NUMBERS.validate_python(right.model_dump()) + return SpendMetrics.model_validate({name: value + second[name] for name, value in first.items()}) + + +def sum_metrics(metrics: Iterable[SpendMetrics]) -> SpendMetrics: + return reduce(add_metrics, metrics, SpendMetrics()) + + +def _key_label(api_key: str, alias: str | None) -> str: + return alias if alias else api_key + + +def group_day(day: DailySpendData, group_by: UsageGroupBy, dataplane_id: str) -> Mapping[str, SpendMetrics]: + breakdown: Final = day.breakdown + match group_by: + case "dataplane": + return MappingProxyType({dataplane_id: day.metrics}) + case "model": + return MappingProxyType({name: entry.metrics for name, entry in breakdown.models.items()}) + case "model_group": + return MappingProxyType({name: entry.metrics for name, entry in breakdown.model_groups.items()}) + case "provider": + return MappingProxyType({name: entry.metrics for name, entry in breakdown.providers.items()}) + case "endpoint": + return MappingProxyType({name: entry.metrics for name, entry in breakdown.endpoints.items()}) + case "api_key": + by_label: Final = tuple( + (_key_label(key, entry.metadata.key_alias), entry.metrics) for key, entry in breakdown.api_keys.items() + ) + return MappingProxyType({label: metrics for (_, label), metrics in _merge_by_key(by_label, day.date)}) + case "team": + by_team: Final = tuple( + (entry.metadata.team_id or NO_TEAM, entry.metrics) for entry in breakdown.api_keys.values() + ) + return MappingProxyType({team: metrics for (_, team), metrics in _merge_by_key(by_team, day.date)}) + + +def _merge_by_key(entries: Iterable[tuple[str, SpendMetrics]], day: date) -> Grouped: + return merge_grouped(((day, group), metrics) for group, metrics in entries) + + +def merge_grouped(entries: Iterable[tuple[DatedGroup, SpendMetrics]]) -> Grouped: + ordered: Final = sorted(entries, key=_dated_group) + return tuple( + (key, sum_metrics(metrics for _, metrics in bucket)) for key, bucket in groupby(ordered, key=_dated_group) + ) + + +def _dated_group(entry: tuple[DatedGroup, SpendMetrics]) -> DatedGroup: + return entry[0] + + +def _dated_groups(days: DataplaneDays, group_by: UsageGroupBy) -> Iterable[tuple[DatedGroup, SpendMetrics]]: + for day in days.days: + for group, metrics in group_day(day, group_by, days.dataplane_id).items(): + yield (day.date, group), metrics + + +def daily_rows(fleet: Iterable[DataplaneDays], group_by: UsageGroupBy) -> Grouped: + return merge_grouped( + entry for days in fleet for entry in _dated_groups(days, group_by) + ) # comprehension-ok: flatten + + +def ranked(fleet: Iterable[DataplaneDays], group_by: UsageGroupBy, limit: int) -> tuple[tuple[str, SpendMetrics], ...]: + per_group: Final = merge_grouped( + ((date.min, group), metrics) for (_, group), metrics in daily_rows(fleet, group_by) + ) + by_spend: Final = sorted(per_group, key=lambda item: item[1].spend, reverse=True) + return tuple((group, metrics) for (_, group), metrics in by_spend[:limit]) + + +def totals(fleet: Iterable[DataplaneDays]) -> SpendMetrics: + return sum_metrics(day.metrics for days in fleet for day in days.days) # comprehension-ok: flatten + + +def dataplane_totals(days: DataplaneDays) -> SpendMetrics: + return sum_metrics(day.metrics for day in days.days) diff --git a/litellm/types/proxy/control_plane/dataplanes.py b/litellm/types/proxy/control_plane/dataplanes.py index a62d103878d..9948642f7c0 100644 --- a/litellm/types/proxy/control_plane/dataplanes.py +++ b/litellm/types/proxy/control_plane/dataplanes.py @@ -42,3 +42,12 @@ class DataplaneIdentityResponse(LiteLLMBaseModel): user_id: str | None user_email: str | None user_role: str | None + + +class UnreachableDataplane(LiteLLMBaseModel): + """A dataplane that did not answer a fleet-wide call, named so the caller knows what the numbers miss.""" + + model_config = ConfigDict(frozen=True) + + id: str + reason: str diff --git a/litellm/types/proxy/control_plane/fleet.py b/litellm/types/proxy/control_plane/fleet.py new file mode 100644 index 00000000000..90bb88f0a31 --- /dev/null +++ b/litellm/types/proxy/control_plane/fleet.py @@ -0,0 +1,55 @@ +from pydantic import ConfigDict + +from litellm.types.proxy.control_plane.dataplanes import UnreachableDataplane +from litellm.types.utils import LiteLLMBaseModel + + +class FleetModel(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + dataplane_id: str + model_name: str + litellm_model: str | None + provider: str | None + model_id: str | None + mode: str | None + + +class FleetModelGroup(LiteLLMBaseModel): + """One public model name and every dataplane that serves it.""" + + model_config = ConfigDict(frozen=True) + + model_name: str + dataplanes: tuple[str, ...] + + +class FleetModelsResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + items: tuple[FleetModel, ...] + groups: tuple[FleetModelGroup, ...] + partial: bool + unreachable: tuple[UnreachableDataplane, ...] + + +class FleetTeam(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + dataplane_id: str + team_id: str + team_alias: str | None + organization_id: str | None + spend: float + max_budget: float | None + members: int + models: tuple[str, ...] + blocked: bool + + +class FleetTeamsResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + items: tuple[FleetTeam, ...] + partial: bool + unreachable: tuple[UnreachableDataplane, ...] diff --git a/litellm/types/proxy/control_plane/usage.py b/litellm/types/proxy/control_plane/usage.py new file mode 100644 index 00000000000..d70b394d5b3 --- /dev/null +++ b/litellm/types/proxy/control_plane/usage.py @@ -0,0 +1,75 @@ +from datetime import date +from typing import Literal, TypeAlias + +from pydantic import ConfigDict + +from litellm.types.proxy.control_plane.dataplanes import UnreachableDataplane +from litellm.types.proxy.management_endpoints.common_daily_activity import SpendMetrics +from litellm.types.utils import LiteLLMBaseModel + +UsageGroupBy: TypeAlias = Literal["dataplane", "model", "model_group", "provider", "api_key", "team", "endpoint"] + + +class UsageWindowResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + start_date: date + end_date: date + + +class DataplaneUsage(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + id: str + metrics: SpendMetrics + + +class UsageSummaryResponse(LiteLLMBaseModel): + """Fleet-wide totals for the window, with the per-dataplane split and whoever did not answer.""" + + model_config = ConfigDict(frozen=True) + + window: UsageWindowResponse + totals: SpendMetrics + by_dataplane: tuple[DataplaneUsage, ...] + partial: bool + unreachable: tuple[UnreachableDataplane, ...] + truncated_dataplanes: tuple[str, ...] + + +class DailyUsageRow(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + date: date + group: str + metrics: SpendMetrics + + +class DailyUsageResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + window: UsageWindowResponse + group_by: UsageGroupBy + rows: tuple[DailyUsageRow, ...] + partial: bool + unreachable: tuple[UnreachableDataplane, ...] + truncated_dataplanes: tuple[str, ...] + + +class UsageRank(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + group: str + metrics: SpendMetrics + + +class TopUsageResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + window: UsageWindowResponse + group_by: UsageGroupBy + limit: int + items: tuple[UsageRank, ...] + partial: bool + unreachable: tuple[UnreachableDataplane, ...] + truncated_dataplanes: tuple[str, ...] diff --git a/tests/unit/proxy/control_plane/common/test_fanout.py b/tests/unit/proxy/control_plane/common/test_fanout.py new file mode 100644 index 00000000000..a0090e5c594 --- /dev/null +++ b/tests/unit/proxy/control_plane/common/test_fanout.py @@ -0,0 +1,104 @@ +import asyncio +from dataclasses import dataclass, field +from typing import Final + +import pytest +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from litellm.proxy.control_plane.common.fanout import Failure, Fanout, Missed, parse_reply +from litellm.proxy.control_plane.dataplanes.client import CallResult, DataplaneReply, DataplaneUnreachable +from litellm.proxy.control_plane.dataplanes.credentials import Caller, CredentialUnavailable, IssuedCredential +from litellm.proxy.control_plane.dataplanes.registry import Dataplane + + +class _NoCredential: + async def headers(self, caller: Caller) -> IssuedCredential: + raise AssertionError("fanout never issues credentials itself") + + +def _dataplane(dataplane_id: str) -> Dataplane: + return Dataplane(dataplane_id, dataplane_id, f"https://{dataplane_id}.test", _NoCredential(), ()) + + +@dataclass +class _InFlight: # mutable-ok: test probe counting concurrent fetches + now: int = 0 + peak: int = 0 + seen: list[str] = field(default_factory=list) + + +async def test_gather_bounds_concurrency_and_keeps_dataplane_order() -> None: + probe: Final = _InFlight() + + async def fetch(dataplane: Dataplane) -> str | Failure: + probe.now += 1 + probe.peak = max(probe.peak, probe.now) + await asyncio.sleep(0) + probe.now -= 1 + probe.seen.append(dataplane.id) + return dataplane.id.upper() + + dataplanes: Final = tuple(_dataplane(name) for name in ("a", "b", "c", "d", "e")) + gathered: Final = await Fanout(concurrency=2).gather(dataplanes, fetch) + assert probe.peak == 2, f"semaphore of 2 let {probe.peak} fetches run at once" + assert [reached.value for reached in gathered.reached] == ["A", "B", "C", "D", "E"] + assert gathered.missed == () and not gathered.partial + + +async def test_concurrent_gathers_share_the_configured_concurrency_limit() -> None: + probe: Final = _InFlight() + fanout: Final = Fanout(concurrency=2) + dataplanes: Final = tuple(_dataplane(name) for name in ("a", "b", "c")) + + async def fetch(dataplane: Dataplane) -> str: + probe.now += 1 + probe.peak = max(probe.peak, probe.now) + await asyncio.sleep(0) + probe.now -= 1 + return dataplane.id + + await asyncio.gather(fanout.gather(dataplanes, fetch), fanout.gather(dataplanes, fetch)) + assert probe.peak == 2, f"two concurrent requests ran {probe.peak} fetches at once" + + +async def test_a_failing_dataplane_is_a_missed_value_and_marks_the_result_partial() -> None: + async def fetch(dataplane: Dataplane) -> int | Failure: + return Failure("boom") if dataplane.id == "b" else 1 + + gathered: Final = await Fanout(concurrency=4).gather(tuple(_dataplane(name) for name in "abc"), fetch) + assert [reached.dataplane.id for reached in gathered.reached] == ["a", "c"] + assert gathered.missed == (Missed(gathered.missed[0].dataplane, "boom"),) + assert gathered.missed[0].dataplane.id == "b" + assert gathered.partial and not gathered.nothing_answered + + +async def test_nothing_answered_only_when_every_dataplane_missed() -> None: + async def fetch(dataplane: Dataplane) -> int | Failure: + return Failure("down") + + all_down: Final = await Fanout(concurrency=1).gather((_dataplane("a"),), fetch) + none_asked: Final = await Fanout(concurrency=1).gather((), fetch) + assert all_down.nothing_answered and not none_asked.nothing_answered + + +class _Shape(BaseModel): + model_config = ConfigDict(frozen=True) + + n: int + + +_SHAPE: Final = TypeAdapter(_Shape) + + +@pytest.mark.parametrize( + ("result", "expected"), + [ + (DataplaneReply(200, {"n": 3}, "service_token"), _Shape(n=3)), + (DataplaneReply(200, {"n": "x"}, "service_token"), Failure("unexpected response shape: 1 validation error(s)")), + (DataplaneReply(403, {"error": "no"}, "oidc_exchange"), Failure("HTTP 403 with the oidc_exchange credential")), + (DataplaneUnreachable("ConnectError", 3), Failure("unreachable after 3 attempt(s): ConnectError")), + (CredentialUnavailable("service_token", "no token"), Failure("no usable credential: no token")), + ], +) +def test_parse_reply_turns_every_bad_outcome_into_a_failure(result: CallResult, expected: _Shape | Failure) -> None: + assert parse_reply(result, _SHAPE) == expected diff --git a/tests/unit/proxy/control_plane/fleet/__init__.py b/tests/unit/proxy/control_plane/fleet/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/control_plane/fleet/test_router.py b/tests/unit/proxy/control_plane/fleet/test_router.py new file mode 100644 index 00000000000..b307e6e8c8e --- /dev/null +++ b/tests/unit/proxy/control_plane/fleet/test_router.py @@ -0,0 +1,239 @@ +from typing import Final + +import httpx +import pytest + +from litellm.litellm_core_utils.oauth_token_exchange import OAuthTokenExchangeConfig +from litellm.proxy.control_plane.auth.rbac import ( + USAGE_VIEWER_ROLE, + VIEWER_ROLE, + build_principal, + default_policy, +) +from litellm.proxy.control_plane.dataplanes.client import DataplaneClient, RetryPolicy +from litellm.proxy.control_plane.dataplanes.credentials import Caller, OidcTokenExchange +from litellm.proxy.control_plane.dataplanes.identity import USER_INFO_PATH +from litellm.proxy.control_plane.dataplanes.registry import Dataplane +from litellm.proxy.control_plane.fleet.service import fetch_teams +from tests.unit.proxy.control_plane.support import RoutedTransport, response, rig + +EU_MODELS: Final = "https://eu.test/model/info" +US_MODELS: Final = "https://us.test/model/info" +EU_TEAMS: Final = "https://eu.test/team/list" +US_TEAMS: Final = "https://us.test/team/list" +EU_USER_INFO: Final = f"https://eu.test{USER_INFO_PATH}" + + +def _model(name: str, provider: str, model_id: str, upstream: str | None = None) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {"model": upstream or f"{provider}/{name}", "custom_llm_provider": provider}, + "model_info": {"id": model_id, "mode": "chat", "db_model": True}, + } + + +async def _exchange(config: OAuthTokenExchangeConfig, subject_token: str) -> str: + assert config.audience == "dataplane" + assert subject_token == "idp-token" + return "dataplane-token" + + +def test_fleet_models_unions_deployments_and_groups_them_by_public_name() -> None: + client, _, minter = rig( + { + EU_MODELS: [ + response(200, {"data": [_model("gpt-4o", "azure", "m1"), _model("claude", "anthropic", "m2")]}) + ], + US_MODELS: [response(200, {"data": [_model("gpt-4o", "openai", "m3")]})], + } + ) + body: Final = client.get("/api/v1/fleet/models", headers=minter.headers()).json() + assert [(m["dataplane_id"], m["model_name"], m["provider"], m["model_id"]) for m in body["items"]] == [ + ("eu", "gpt-4o", "azure", "m1"), + ("eu", "claude", "anthropic", "m2"), + ("us", "gpt-4o", "openai", "m3"), + ] + assert body["groups"] == [ + {"model_name": "claude", "dataplanes": ["eu"]}, + {"model_name": "gpt-4o", "dataplanes": ["eu", "us"]}, + ] + assert body["partial"] is False + + +def test_fleet_models_accepts_a_single_deployment_object_from_a_dataplane() -> None: + client, _, minter = rig( + { + EU_MODELS: [response(200, {"data": _model("gpt-4o", "azure", "m1")})], + US_MODELS: [response(200, {"data": []})], + } + ) + + body: Final = client.get("/api/v1/fleet/models", headers=minter.headers()).json() + + assert body["partial"] is False + assert body["unreachable"] == [] + assert body["items"] == [ + { + "dataplane_id": "eu", + "model_name": "gpt-4o", + "litellm_model": "azure/gpt-4o", + "provider": "azure", + "model_id": "m1", + "mode": "chat", + } + ] + + +def test_fleet_models_tolerates_an_unparseable_dataplane_and_reports_it() -> None: + client, _, minter = rig( + { + EU_MODELS: [response(200, {"data": [_model("gpt-4o", "azure", "m1")]})], + US_MODELS: [response(200, {"data": "x"})], + } + ) + body: Final = client.get("/api/v1/fleet/models", headers=minter.headers()).json() + assert [m["model_id"] for m in body["items"]] == ["m1"] and body["partial"] is True + assert body["unreachable"][0]["id"] == "us" and "validation error" in body["unreachable"][0]["reason"] + + +def test_fleet_teams_counts_members_and_keeps_budget_fields() -> None: + team: Final = { + "team_id": "t1", + "team_alias": "ml", + "organization_id": None, + "spend": 12.5, + "max_budget": 100.0, + "members_with_roles": [{"user_id": "a", "role": "admin"}, {"user_id": "b", "role": "user"}], + "models": ["gpt-4o"], + "blocked": False, + "metadata": {"ignored": True}, + } + client, _, minter = rig({EU_TEAMS: [response(200, [team])], US_TEAMS: [response(200, [])]}) + body: Final = client.get("/api/v1/fleet/teams", headers=minter.headers()).json() + assert body["items"] == [ + { + "dataplane_id": "eu", + "team_id": "t1", + "team_alias": "ml", + "organization_id": None, + "spend": 12.5, + "max_budget": 100.0, + "members": 2, + "models": ["gpt-4o"], + "blocked": False, + } + ] + + +def test_fleet_teams_falls_back_to_the_callers_own_teams_when_the_dataplane_refuses_to_list_all() -> None: + own: Final = {"team_id": "t-own", "members_with_roles": [{"user_id": "alice"}]} + client, transport, minter = rig( + { + EU_TEAMS: [ + response(401, {"detail": {"error": "Only admin users can query all teams"}}), + response(200, [own]), + ], + EU_USER_INFO: [response(200, {"user_id": "alice"})], + US_TEAMS: [response(200, [])], + } + ) + body: Final = client.get("/api/v1/fleet/teams", headers=minter.headers()).json() + assert [(t["dataplane_id"], t["team_id"], t["members"]) for t in body["items"]] == [("eu", "t-own", 1)] + assert body["partial"] is False + eu_calls: Final = [(url, params) for url, params, _ in transport.calls if url in {EU_TEAMS, EU_USER_INFO}] + assert eu_calls == [ + (EU_TEAMS, None), + (EU_USER_INFO, None), + (EU_TEAMS, (("user_id", "alice"),)), + ] + + +def test_fleet_teams_treats_a_user_info_not_found_as_an_empty_reached_dataplane() -> None: + client, _, minter = rig( + { + EU_TEAMS: [response(403)], + EU_USER_INFO: [response(404)], + US_TEAMS: [response(200, [])], + } + ) + + body: Final = client.get("/api/v1/fleet/teams", headers=minter.headers()).json() + + assert body["items"] == [] + assert body["partial"] is False + assert body["unreachable"] == [] + + +@pytest.mark.asyncio +async def test_fleet_teams_uses_dataplane_identity_for_oidc_exchange_callers() -> None: + team: Final = {"team_id": "t-own", "members_with_roles": [{"user_id": "dataplane-user"}]} + routes: Final = { + EU_TEAMS: iter( + [ + response(401, {"detail": {"error": "Only admin users can query all teams"}}), + response(200, [team]), + ] + ), + EU_USER_INFO: iter([response(200, {"user_id": "dataplane-user"})]), + } + transport: Final = RoutedTransport(routes) + client: Final = DataplaneClient(transport=transport, timeout_seconds=1, retry=RetryPolicy(attempts=1)) + credential: Final = OidcTokenExchange( + OAuthTokenExchangeConfig( + token_endpoint="https://identity.test/token", + client_id="control-plane", + client_secret="secret", + profile="rfc8693", + audience="dataplane", + scope=None, + ), + _exchange, + ) + dataplane: Final = Dataplane( + id="eu", + name="EU", + url="https://eu.test", + credential=credential, + credential_kinds=("oidc_exchange",), + ) + principal: Final = build_principal( + default_policy(), + "idp-subject", + "idp@example.com", + frozenset({VIEWER_ROLE}), + None, + frozenset({"eu"}), + "oidc", + ) + caller: Final = Caller(principal=principal, bearer="idp-token") + + teams: Final = await fetch_teams(client, dataplane, caller) + + assert isinstance(teams, tuple) + assert tuple(team.team_id for team in teams) == ("t-own",) + assert [(url, params) for url, params, _ in transport.calls] == [ + (EU_TEAMS, None), + (EU_USER_INFO, None), + (EU_TEAMS, (("user_id", "dataplane-user"),)), + ] + assert all( + headers is not None and headers["Authorization"] == "Bearer dataplane-token" + for _, _, headers in transport.calls + ) + + +def test_fleet_routes_need_fleet_read_and_503_when_nobody_answers() -> None: + client, _, minter = rig({EU_TEAMS: [httpx.ConnectError("down")] * 3, US_TEAMS: [response(502, "bad")] * 3}) + denied: Final = client.get("/api/v1/fleet/teams", headers=minter.headers(role=USAGE_VIEWER_ROLE)) + assert denied.status_code == 403 and "fleet:read" in denied.json()["detail"] + assert client.get("/api/v1/fleet/teams", headers=minter.headers()).status_code == 503 + + +def test_a_team_with_null_spend_lists_as_zero_instead_of_failing_the_fleet() -> None: + fresh: Final = {"team_id": "new", "spend": None} + client, _, minter = rig({EU_TEAMS: [response(200, [fresh])], US_TEAMS: [response(200, [])]}) + reply: Final = client.get("/api/v1/fleet/teams", headers=minter.headers()) + assert reply.status_code == 200, reply.text + body: Final = reply.json() + assert [(item["team_id"], item["spend"]) for item in body["items"]] == [("new", 0.0)] + assert body["partial"] is False diff --git a/tests/unit/proxy/control_plane/support.py b/tests/unit/proxy/control_plane/support.py new file mode 100644 index 00000000000..0a199ab2f34 --- /dev/null +++ b/tests/unit/proxy/control_plane/support.py @@ -0,0 +1,93 @@ +"""Shared rig for control plane route tests: a routed fake transport and a licensed app over two dataplanes.""" + +from collections import defaultdict +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass, field +from datetime import timedelta +from typing import Final + +import httpx +from fastapi.testclient import TestClient + +from litellm.litellm_core_utils.oauth_token_exchange import OAuthTokenExchangeConfig +from litellm.proxy.control_plane.app import create_app +from litellm.proxy.control_plane.auth.jwt import SessionTokenCodec +from litellm.proxy.control_plane.auth.rbac import VIEWER_ROLE, RolePolicy, build_principal +from litellm.proxy.control_plane.container import build_container +from litellm.proxy.control_plane.settings import ControlPlaneSettings + +MASTER_KEY: Final = "01234567890123456789012345678901" +DATAPLANES: Final = { + "eu": {"url": "https://eu.test", "service_token": {"token": "sk-eu"}}, + "us": {"url": "https://us.test", "name": "US East", "service_token": {"token": "sk-us"}}, +} + + +class Licensed: + def is_premium(self) -> bool: + return True + + +@dataclass +class RoutedTransport: # mutable-ok: records calls for assertions + """Answers each URL from its own queue of outcomes, so fan-out order does not matter to a test.""" + + routes: Mapping[str, Iterator[httpx.Response | Exception]] + calls: list[tuple[str, tuple[tuple[str, str], ...] | None, dict[str, str] | None]] = field(default_factory=list) + + async def request( + self, + method: str, + url: str, + *, + params: Sequence[tuple[str, str]] | None = None, + json: dict[str, object] | None = None, + headers: dict[str, str] | None = None, + timeout: float, + ) -> httpx.Response: + assert method == "GET", "fleet reads never write" + self.calls.append((url, None if params is None else tuple(params), headers)) + outcome: Final = next(self.routes[url], None) + if outcome is None: + raise AssertionError(f"no scripted outcome left for {url}") + if isinstance(outcome, Exception): + raise outcome + return outcome + + +async def no_exchange(config: OAuthTokenExchangeConfig, subject_token: str) -> str: + raise AssertionError("no oidc dataplane configured") + + +def response(status: int, body: object = None) -> httpx.Response: + return httpx.Response(status, json=body, request=httpx.Request("GET", "https://x")) + + +@dataclass(frozen=True) +class Minter: + codec: SessionTokenCodec + policy: RolePolicy + + def headers(self, *dataplanes: str, role: str = VIEWER_ROLE) -> dict[str, str]: + principal: Final = build_principal( + self.policy, "alice", "alice@x.test", frozenset({role}), None, frozenset(dataplanes), "session" + ) + return {"Authorization": f"Bearer {self.codec.mint(principal)}"} + + +def rig( + routes: Mapping[str, list[httpx.Response | Exception]], usage_max_window_days: int = 366 +) -> tuple[TestClient, RoutedTransport, Minter]: + settings: Final = ControlPlaneSettings.model_validate( + { + "LITELLM_MASTER_KEY": MASTER_KEY, + "dataplanes": DATAPLANES, + "usage_max_window_days": usage_max_window_days, + } + ) + queues: defaultdict[str, Iterator[httpx.Response | Exception]] = defaultdict(lambda: iter(())) + queues.update({url: iter(outcomes) for url, outcomes in routes.items()}) + transport: Final = RoutedTransport(queues) + container: Final = build_container(settings, transport=transport, exchanger=no_exchange) + codec: Final = SessionTokenCodec(MASTER_KEY, settings.session_issuer, timedelta(hours=1), container.policy) + return TestClient(create_app(container, Licensed())), transport, Minter(codec, container.policy) diff --git a/tests/unit/proxy/control_plane/test_settings.py b/tests/unit/proxy/control_plane/test_settings.py index 3f441f605bb..b963d1409b2 100644 --- a/tests/unit/proxy/control_plane/test_settings.py +++ b/tests/unit/proxy/control_plane/test_settings.py @@ -60,3 +60,22 @@ def test_dataplanes_environment_and_file_cannot_both_be_configured(monkeypatch: with pytest.raises(ValueError, match="CONTROL_PLANE_DATAPLANES and CONTROL_PLANE_DATAPLANES_FILE"): load_settings() + + +def test_fanout_concurrency_must_be_positive() -> None: + with pytest.raises(ValidationError): + ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY, "fanout_concurrency": 0}) + + +def test_master_key_minimum_length_uses_utf8_bytes() -> None: + with pytest.raises(ValidationError): + ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": "é" * 15 + "x"}) + settings: Final = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": "é" * 16}) + assert len(settings.master_key.get_secret_value().encode("utf-8")) == 32 + + +def test_usage_window_maximum_is_configurable() -> None: + settings: Final = ControlPlaneSettings.model_validate( + {"LITELLM_MASTER_KEY": MASTER_KEY, "usage_max_window_days": 12} + ) + assert settings.usage_max_window_days == 12 diff --git a/tests/unit/proxy/control_plane/usage/__init__.py b/tests/unit/proxy/control_plane/usage/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/control_plane/usage/test_router.py b/tests/unit/proxy/control_plane/usage/test_router.py new file mode 100644 index 00000000000..70a9fd1e340 --- /dev/null +++ b/tests/unit/proxy/control_plane/usage/test_router.py @@ -0,0 +1,185 @@ +from datetime import date +from typing import Final + +import httpx + +from litellm.constants import USAGE_TOP_API_KEYS_MAX +from litellm.proxy.control_plane.auth.rbac import USAGE_VIEWER_ROLE +from tests.unit.proxy.control_plane.support import MASTER_KEY, response, rig + +EU: Final = "https://eu.test/user/daily/activity/aggregated" +US: Final = "https://us.test/user/daily/activity/aggregated" +WINDOW: Final = {"start_date": "2026-10-01", "end_date": "2026-10-02"} + + +def _page( + rows: list[dict[str, object]], api_key_limit: int | None = None, total_api_keys: int | None = None +) -> dict[str, object]: + return { + "results": rows, + "metadata": { + "has_more": False, + "page": 1, + "total_pages": 1, + "api_key_limit": api_key_limit, + "total_api_keys": total_api_keys, + }, + } + + +def _row(day: str, spend: float, models: dict[str, float] | None = None) -> dict[str, object]: + return { + "date": day, + "metrics": {"spend": spend, "api_requests": 1}, + "breakdown": {"models": {name: {"metrics": {"spend": cost}} for name, cost in (models or {}).items()}}, + } + + +def test_summary_adds_both_dataplanes_and_forwards_the_window_and_caller() -> None: + client, transport, minter = rig( + { + EU: [response(200, _page([_row("2026-10-01", 1.5), _row("2026-10-02", 0.5)]))], + US: [response(200, _page([_row("2026-10-01", 2.0)]))], + } + ) + body: Final = client.get( + "/api/v1/usage/summary", params=WINDOW, headers=minter.headers(role=USAGE_VIEWER_ROLE) + ).json() + assert body["window"] == WINDOW + assert (body["totals"]["spend"], body["totals"]["api_requests"]) == (4.0, 3) + assert [(dp["id"], dp["metrics"]["spend"]) for dp in body["by_dataplane"]] == [("eu", 2.0), ("us", 2.0)] + assert body["partial"] is False and body["unreachable"] == [] + url, params, headers = transport.calls[0] + assert params == ( + ("start_date", "2026-10-01"), + ("end_date", "2026-10-02"), + ("api_key_limit", str(USAGE_TOP_API_KEYS_MAX)), + ) + assert headers is not None and headers["X-LiteLLM-Impersonate-User"] == "alice" + + +def test_summary_does_not_report_a_dataplane_with_many_aggregated_rows_as_missed() -> None: + rows: Final = [_row("2026-10-01", 1.0) for _ in range(1001)] + client, transport, minter = rig({EU: [response(200, _page(rows))], US: [response(200, _page([]))]}) + body: Final = client.get("/api/v1/usage/summary", params=WINDOW, headers=minter.headers()).json() + assert body["totals"]["spend"] == 1001.0 + assert body["partial"] is False and body["unreachable"] == [] + assert sum(1 for url, _, _ in transport.calls if url == EU) == 1 + + +def test_one_dataplane_down_gives_partial_numbers_and_names_it() -> None: + client, _, minter = rig( + {EU: [response(200, _page([_row("2026-10-01", 1.0)]))], US: [httpx.ConnectError("refused")] * 3} + ) + body: Final = client.get("/api/v1/usage/summary", params=WINDOW, headers=minter.headers()).json() + assert body["totals"]["spend"] == 1.0 and body["partial"] is True + assert [u["id"] for u in body["unreachable"]] == ["us"] + assert "ConnectError" in body["unreachable"][0]["reason"] + + +def test_every_dataplane_down_is_a_503_error() -> None: + client, _, minter = rig({EU: [response(503, "maintenance")] * 3, US: [httpx.ReadTimeout("slow")] * 3}) + down: Final = client.get("/api/v1/usage/summary", params=WINDOW, headers=minter.headers()) + assert down.status_code == 503 + assert down.json()["type"] == "urn:litellm:error:service-unavailable" + assert "eu:" in down.json()["detail"] and "us:" in down.json()["detail"] + + +def test_binding_and_dataplane_filter_narrow_the_fan_out() -> None: + client, transport, minter = rig( + { + EU: [response(200, _page([_row("2026-10-01", 1.0)]))] * 2, + US: [response(200, _page([_row("2026-10-01", 9.0)]))], + } + ) + bound: Final = client.get("/api/v1/usage/summary", params=WINDOW, headers=minter.headers("eu")).json() + assert [dp["id"] for dp in bound["by_dataplane"]] == ["eu"] + filtered: Final = client.get( + "/api/v1/usage/summary", params={**WINDOW, "dataplane": "eu"}, headers={"Authorization": f"Bearer {MASTER_KEY}"} + ).json() + assert [dp["id"] for dp in filtered["by_dataplane"]] == ["eu"] + assert all(url == EU for url, _, _ in transport.calls) + assert ( + client.get( + "/api/v1/usage/summary", params={**WINDOW, "dataplane": "us"}, headers=minter.headers("eu") + ).status_code + == 403 + ) + + +def test_daily_groups_by_model_across_dataplanes_and_top_ranks_them() -> None: + client, _, minter = rig( + { + EU: [response(200, _page([_row("2026-10-01", 3.0, {"gpt-4o": 2.0, "claude": 1.0})]))] * 2, + US: [response(200, _page([_row("2026-10-01", 5.0, {"gpt-4o": 5.0})]))] * 2, + } + ) + daily: Final = client.get( + "/api/v1/usage/daily", params={**WINDOW, "group_by": "model"}, headers=minter.headers() + ).json() + assert [(row["date"], row["group"], row["metrics"]["spend"]) for row in daily["rows"]] == [ + ("2026-10-01", "claude", 1.0), + ("2026-10-01", "gpt-4o", 7.0), + ] + top: Final = client.get("/api/v1/usage/top", params={**WINDOW, "limit": 1}, headers=minter.headers()).json() + assert top["group_by"] == "model" and top["limit"] == 1 + assert [(item["group"], item["metrics"]["spend"]) for item in top["items"]] == [("gpt-4o", 7.0)] + + +def test_usage_routes_need_the_usage_read_scope_and_default_the_window_to_a_month() -> None: + client, transport, minter = rig({EU: [response(200, _page([]))], US: [response(200, _page([]))]}) + assert client.get("/api/v1/usage/daily", params=WINDOW).status_code == 401 + body: Final = client.get("/api/v1/usage/summary", headers=minter.headers()).json() + start, end = date.fromisoformat(body["window"]["start_date"]), date.fromisoformat(body["window"]["end_date"]) + assert (end - start).days == 29 + + +def test_default_usage_window_respects_a_smaller_configured_maximum() -> None: + client, _, minter = rig( + {EU: [response(200, _page([]))], US: [response(200, _page([]))]}, usage_max_window_days=7 + ) + result: Final = client.get("/api/v1/usage/summary", headers=minter.headers()) + assert result.status_code == 200 + body: Final = result.json() + start, end = date.fromisoformat(body["window"]["start_date"]), date.fromisoformat(body["window"]["end_date"]) + assert (end - start).days + 1 == 7 + + +def test_reversed_window_is_a_400_before_any_dataplane_is_asked() -> None: + client, transport, minter = rig({EU: [], US: []}) + reversed_window: Final = client.get( + "/api/v1/usage/summary", + params={"start_date": "2026-10-03", "end_date": "2026-10-01"}, + headers=minter.headers(), + ) + assert reversed_window.status_code == 400 and "after end_date" in reversed_window.json()["detail"] + assert transport.calls == [] + + +def test_window_over_the_configured_maximum_is_rejected_before_fanout() -> None: + client, transport, minter = rig({EU: [], US: []}) + too_wide: Final = client.get( + "/api/v1/usage/summary", + params={"start_date": "2025-09-30", "end_date": "2026-10-01"}, + headers=minter.headers(), + ) + assert too_wide.status_code == 400 + assert transport.calls == [] + + +def test_truncated_api_key_aggregation_marks_the_dataplane_partial() -> None: + client, transport, minter = rig( + { + EU: [ + response( + 200, + _page([_row("2026-10-01", 1.0)], api_key_limit=USAGE_TOP_API_KEYS_MAX, total_api_keys=1001), + ) + ], + US: [response(200, _page([]))], + } + ) + body: Final = client.get("/api/v1/usage/summary", params=WINDOW, headers=minter.headers()).json() + assert body["partial"] is True + assert body["truncated_dataplanes"] == ["eu"] + assert any(("api_key_limit", str(USAGE_TOP_API_KEYS_MAX)) in params for _, params, _ in transport.calls) diff --git a/tests/unit/proxy/control_plane/usage/test_service.py b/tests/unit/proxy/control_plane/usage/test_service.py new file mode 100644 index 00000000000..a538dc93f00 --- /dev/null +++ b/tests/unit/proxy/control_plane/usage/test_service.py @@ -0,0 +1,113 @@ +from datetime import date +from typing import Final + +import pytest + +from litellm.proxy.control_plane.usage.service import ( + NO_TEAM, + DataplaneDays, + add_metrics, + daily_rows, + group_day, + ranked, + totals, +) +from litellm.types.proxy.control_plane.usage import UsageGroupBy +from litellm.types.proxy.management_endpoints.common_daily_activity import ( + BreakdownMetrics, + DailySpendData, + KeyMetadata, + KeyMetricWithMetadata, + MetricWithMetadata, + SpendMetrics, +) + +D1: Final = date(2026, 10, 1) +D2: Final = date(2026, 10, 2) + + +def _m(spend: float, requests: int = 1) -> SpendMetrics: + return SpendMetrics(spend=spend, api_requests=requests, successful_requests=requests, total_tokens=requests * 10) + + +def _day( + day: date, + spend: float, + models: dict[str, float] | None = None, + keys: dict[str, tuple[float, str | None, str | None]] | None = None, +) -> DailySpendData: + return DailySpendData( + date=day, + metrics=_m(spend), + breakdown=BreakdownMetrics( + models={name: MetricWithMetadata(metrics=_m(cost)) for name, cost in (models or {}).items()}, + api_keys={ + key: KeyMetricWithMetadata(metrics=_m(cost), metadata=KeyMetadata(key_alias=alias, team_id=team)) + for key, (cost, alias, team) in (keys or {}).items() + }, + ), + ) + + +def test_add_metrics_sums_every_numeric_field() -> None: + left: Final = SpendMetrics.model_validate({name: 1 for name in SpendMetrics.model_fields}) + right: Final = SpendMetrics.model_validate({name: 2 for name in SpendMetrics.model_fields}) + total: Final = add_metrics(left, right) + assert total.model_dump() == {name: 3 for name in SpendMetrics.model_fields} + + +def test_totals_add_up_across_dataplanes_and_days() -> None: + fleet: Final = ( + DataplaneDays("eu", (_day(D1, 1.0), _day(D2, 2.0))), + DataplaneDays("us", (_day(D1, 4.0),)), + ) + assert totals(fleet) == _m(7.0, requests=3) + + +def test_daily_rows_merge_the_same_model_on_the_same_day_from_two_dataplanes() -> None: + fleet: Final = ( + DataplaneDays( + "eu", (_day(D1, 3.0, models={"gpt-4o": 2.0, "claude": 1.0}), _day(D2, 1.0, models={"claude": 1.0})) + ), + DataplaneDays("us", (_day(D1, 5.0, models={"gpt-4o": 5.0}),)), + ) + rows: Final = daily_rows(fleet, "model") + assert [(day, group, metrics.spend) for (day, group), metrics in rows] == [ + (D1, "claude", 1.0), + (D1, "gpt-4o", 7.0), + (D2, "claude", 1.0), + ] + assert daily_rows(fleet, "dataplane") == ( + ((D1, "eu"), _m(3.0)), + ((D1, "us"), _m(5.0)), + ((D2, "eu"), _m(1.0)), + ) + + +def test_team_grouping_folds_keys_into_their_team_and_labels_keys_by_alias() -> None: + day: Final = _day( + D1, + 6.0, + keys={"hash-a": (1.0, "alice-key", "t1"), "hash-b": (2.0, None, "t1"), "hash-c": (3.0, "svc", None)}, + ) + assert {team: metrics.spend for team, metrics in group_day(day, "team", "eu").items()} == {"t1": 3.0, NO_TEAM: 3.0} + assert set(group_day(day, "api_key", "eu")) == {"alice-key", "hash-b", "svc"} + + +def test_ranked_orders_by_fleet_spend_and_honours_the_limit() -> None: + fleet: Final = ( + DataplaneDays("eu", (_day(D1, 0, models={"a": 1.0, "b": 5.0}), _day(D2, 0, models={"a": 1.0}))), + DataplaneDays("us", (_day(D1, 0, models={"c": 3.0, "a": 2.0}),)), + ) + assert [(group, metrics.spend) for group, metrics in ranked(fleet, "model", limit=2)] == [("b", 5.0), ("a", 4.0)] + assert len(ranked(fleet, "model", limit=10)) == 3 + + +@pytest.mark.parametrize("dimension", ["model_group", "provider", "endpoint"]) +def test_empty_breakdowns_group_to_nothing(dimension: UsageGroupBy) -> None: + assert dict(group_day(_day(D1, 1.0), dimension, "eu")) == {} + + +def test_two_keys_sharing_an_alias_add_up_instead_of_overwriting() -> None: + day: Final = _day(D1, 8.0, keys={"hash-a": (3.0, "shared", None), "hash-b": (5.0, "shared", None)}) + assert {label: metrics.spend for label, metrics in group_day(day, "api_key", "eu").items()} == {"shared": 8.0}