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:
Yucheng Zhu 2026-07-21 11:34:03 -07:00
parent 978d1ece6b
commit 74145dd1ad
4 changed files with 31 additions and 19 deletions

View file

@ -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

View file

@ -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.
"""

View file

@ -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

View file

@ -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
):