From 74145dd1adb6389b290afcc781ba89863c04e2ae Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Tue, 21 Jul 2026 11:34:03 -0700 Subject: [PATCH] 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. --- litellm/proxy/proxy_server.py | 16 ++++++++-------- .../spend_tracking/ptu_reservation_rollup.py | 9 ++++++--- .../ui_crud_endpoints/proxy_setting_endpoints.py | 9 +-------- .../test_proxy_setting_endpoints.py | 16 ++++++++++++++++ 4 files changed, 31 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 63ab2085b20..bb6c4747d02 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 diff --git a/litellm/proxy/spend_tracking/ptu_reservation_rollup.py b/litellm/proxy/spend_tracking/ptu_reservation_rollup.py index e9f91ef1d00..f0cfb8d0161 100644 --- a/litellm/proxy/spend_tracking/ptu_reservation_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_reservation_rollup.py @@ -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. """ diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 656d2f5e86f..c2ed350a2ea 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 6b3ab154fe7..99bc68840b6 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -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 ):