mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
refactor(ptu): collapse enable_azure_ptu_billing_pull into enable_ptu_cost_attribution
Two flags for the same feature confused reviewers with no benefit. Presence of general_settings.azure_ptu_billing.subscription_id together with the Entra ID env vars is now the sole signal that an operator wants the automated Cost Management pull; if either is missing, azure_billing reservations no-op with a warning and manual reservations continue to accrue via the formula. Removes the enable_azure_ptu_billing_pull field from UISettings, its ALLOWED and _RUNTIME_GENERAL_SETTINGS_FLAGS entries, and the corresponding gate inside _build_azure_cost_fetcher_if_enabled. Extends the existing runtime-flag test with a regression assertion that the removed flag is not reintroduced. Behavior changes: the UI Settings panel now shows one PTU toggle instead of two. An operator who previously had the outer flag on but the pull flag off will, after this change, see azure_billing reservations start pulling from Azure as soon as their creds and subscription_id are configured; the fallback path when either is missing is unchanged.
This commit is contained in:
parent
978d1ece6b
commit
74145dd1ad
4 changed files with 31 additions and 19 deletions
|
|
@ -7465,19 +7465,19 @@ def giveup(e):
|
|||
|
||||
|
||||
def _build_azure_cost_fetcher_if_enabled() -> Optional[Any]:
|
||||
"""Return an AzureCostManagementClient when the pull flag + config + creds are all present, else None.
|
||||
"""Return an AzureCostManagementClient when config + creds are present, else None.
|
||||
|
||||
Evaluated at rollup call time so runtime flag toggles take effect without
|
||||
a proxy restart.
|
||||
Presence of general_settings.azure_ptu_billing.subscription_id together with the
|
||||
Entra ID env vars is the sole signal for "operator wants the auto-pull"; if any
|
||||
piece is missing, azure_billing reservations no-op with a warning inside the
|
||||
rollup and manual reservations continue to accrue via the formula.
|
||||
|
||||
Evaluated at rollup call time so runtime config changes take effect without a
|
||||
proxy restart.
|
||||
"""
|
||||
if not general_settings.get("enable_azure_ptu_billing_pull", False):
|
||||
return None
|
||||
azure_ptu_billing = general_settings.get("azure_ptu_billing") or {}
|
||||
subscription_id = azure_ptu_billing.get("subscription_id")
|
||||
if not subscription_id:
|
||||
verbose_proxy_logger.warning(
|
||||
"enable_azure_ptu_billing_pull is true but general_settings.azure_ptu_billing.subscription_id is not set"
|
||||
)
|
||||
return None
|
||||
try:
|
||||
from litellm.integrations.azure_cost_management import AzureCostManagementClient
|
||||
|
|
|
|||
|
|
@ -58,7 +58,8 @@ async def _compute_daily_flat_cost(
|
|||
if azure_fetcher is None or reservation.azure_resource_id is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"PTU rollup: reservation=%s cost_source=azure_billing skipped "
|
||||
"(enable_azure_ptu_billing_pull off or azure_resource_id missing)",
|
||||
"(azure_ptu_billing.subscription_id / Entra ID env vars not configured, "
|
||||
"or azure_resource_id missing on the reservation)",
|
||||
getattr(reservation, "id", "?"),
|
||||
)
|
||||
return 0.0
|
||||
|
|
@ -155,8 +156,10 @@ async def run_ptu_reservation_rollup(
|
|||
|
||||
Defaults to yesterday UTC. ``force=True`` bypasses the feature-flag check
|
||||
so the CLI backfill can run when the scheduler is off. ``azure_fetcher``
|
||||
is injected by the proxy startup wiring when ``enable_azure_ptu_billing_pull``
|
||||
is on; azure_billing reservations no-op with a warning when it is None.
|
||||
is injected by the proxy startup wiring when
|
||||
``general_settings.azure_ptu_billing.subscription_id`` and the Entra ID env
|
||||
vars are configured; azure_billing reservations no-op with a warning when
|
||||
it is None.
|
||||
Idempotent under the LiteLLM_DailyTeamSpend unique constraint on every
|
||||
invocation path.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -184,12 +184,7 @@ class UISettings(BaseModel):
|
|||
|
||||
enable_ptu_cost_attribution: bool = Field(
|
||||
default=False,
|
||||
description="If true, enables admin-registered PTU reservations and daily flat-cost attribution on team daily spend. Governs the /ptu_reservation CRUD endpoints, the daily rollup job, the PTU Reservations UI page, and the Flat Cost column on the Usage page.",
|
||||
)
|
||||
|
||||
enable_azure_ptu_billing_pull: bool = Field(
|
||||
default=False,
|
||||
description="If true and enable_ptu_cost_attribution is also true, reservations with cost_source='azure_billing' fetch daily flat cost from the Azure Cost Management API instead of the manual PTU * cost/day formula. Requires general_settings.azure_ptu_billing.subscription_id and Entra ID env vars (AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_CLIENT_SECRET).",
|
||||
description="If true, enables admin-registered PTU reservations and daily flat-cost attribution on team daily spend. Governs the /ptu_reservation CRUD endpoints, the daily rollup job, the PTU Reservations UI page, and the Flat Cost column on the Usage page. When general_settings.azure_ptu_billing.subscription_id and the Entra ID env vars (AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_CLIENT_SECRET) are also set, reservations with cost_source='azure_billing' auto-fetch daily flat cost from the Azure Cost Management API; otherwise only manual reservations accrue.",
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -217,7 +212,6 @@ ALLOWED_UI_SETTINGS_FIELDS = {
|
|||
"disable_key_generate_for_org_admin",
|
||||
"enable_chat_ui",
|
||||
"enable_ptu_cost_attribution",
|
||||
"enable_azure_ptu_billing_pull",
|
||||
}
|
||||
|
||||
# Flags that must be synced from the persisted UISettings into
|
||||
|
|
@ -232,7 +226,6 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS = [
|
|||
"allow_vector_stores_for_team_admins",
|
||||
"disable_key_generate_for_org_admin",
|
||||
"enable_ptu_cost_attribution",
|
||||
"enable_azure_ptu_billing_pull",
|
||||
]
|
||||
|
||||
# Extension point: packages outside OSS (e.g. litellm_enterprise) can
|
||||
|
|
|
|||
|
|
@ -1391,6 +1391,22 @@ class TestProxySettingEndpoints:
|
|||
assert "enable_ptu_cost_attribution" in _RUNTIME_GENERAL_SETTINGS_FLAGS
|
||||
assert "enable_ptu_cost_attribution" in ALLOWED_UI_SETTINGS_FIELDS
|
||||
|
||||
def test_azure_pull_is_not_a_separate_flag(self):
|
||||
"""Regression: presence of azure_ptu_billing.subscription_id + Entra ID env vars is the sole signal for the auto-pull.
|
||||
|
||||
Adding a second UI flag was confusing and served no purpose; this asserts we don't
|
||||
reintroduce one under this exact name.
|
||||
"""
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
_RUNTIME_GENERAL_SETTINGS_FLAGS,
|
||||
ALLOWED_UI_SETTINGS_FIELDS,
|
||||
UISettings,
|
||||
)
|
||||
|
||||
assert "enable_azure_ptu_billing_pull" not in _RUNTIME_GENERAL_SETTINGS_FLAGS
|
||||
assert "enable_azure_ptu_billing_pull" not in ALLOWED_UI_SETTINGS_FIELDS
|
||||
assert "enable_azure_ptu_billing_pull" not in UISettings.model_fields
|
||||
|
||||
def test_get_sso_settings_from_database(
|
||||
self, mock_proxy_config, mock_auth, monkeypatch
|
||||
):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue