feat(ptu): fetch azure_billing flat cost via Cost Management API (LIT-4077 POC) - #34048
feat(ptu): fetch azure_billing flat cost via Cost Management API (LIT-4077 POC)#34048yucheng-berri wants to merge 4 commits into
Conversation
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
|
@greptileai please review |
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Greptile SummaryThis PR adds Azure Cost Management API integration for PTU reservation cost attribution (LIT-4077). A thin async client (
Confidence Score: 4/5Safe to merge for public-cloud deployments; sovereign cloud operators will see 401s from Azure until the token scope is fixed. The new client's Entra ID token scope is hardcoded to Files Needing Attention:
|
| Filename | Overview |
|---|---|
| litellm/integrations/azure_cost_management/azure_cost_management_client.py | New async client for the Azure Cost Management API. Correctly wraps HTTP errors and network errors, protects client_secret from repr, and handles the zero-rows case — but the Entra ID token scope is hardcoded to the public cloud audience, breaking sovereign cloud deployments when management_base_url is overridden. |
| litellm/proxy/spend_tracking/ptu_reservation_rollup.py | Extends rollup to support azure_billing reservations via injected AzureCostFetcher. Non-USD currency guard is correct; Protocol definition omits last_currency, which could silently suppress the guard for future custom implementations. |
| litellm/proxy/proxy_server.py | Adds _DemoAzureCostFetcher stub and _build_azure_cost_fetcher_if_enabled helper; scheduler now wraps rollup in a thin async callable that builds the fetcher just-in-time. Clean and well-guarded. |
| litellm/proxy/management_endpoints/ptu_reservation_endpoints.py | Relaxes the blanket azure_billing rejection to require azure_resource_id instead, matching the Pydantic domain validator. Change is minimal and correct. |
| scripts/ptu_reservation_backfill.py | Backfill CLI now reads AZURE_SUBSCRIPTION_ID and builds an Azure fetcher matching the scheduler's code path. Logging setup added at entry point. Clean. |
| tests/test_litellm/integrations/azure_cost_management/test_azure_cost_management_client.py | Good unit test coverage for the new client: HTTP errors, network errors, zero-row response, multi-row sum, repr safety, and sovereign cloud URL override. All tests use mocks with no real network calls. |
| tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py | Existing tests correctly converted to async; new tests cover azure_billing fetch, non-USD skip, fetcher error, and the end-to-end rollup write path. Coverage is thorough. |
| litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py | Description field updated to document the azure_billing auto-pull; enable_azure_ptu_billing_pull was intentionally removed and regression-guarded by a new test. |
| litellm/types/llms/custom_http.py | Adds AzureCostManagement enum value to httpxSpecialProvider so the client gets its own HTTP handler pool. Straightforward addition. |
Reviews (5): Last reviewed commit: "feat(ptu): sovereign-cloud base URL + de..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Greptile SummaryThis POC adds an
Confidence Score: 3/5The Azure billing fetch path silently writes non-USD costs to ptu_flat_cost without conversion; any customer on a non-USD Azure subscription will see incorrect spend data from the first rollup run. The non-USD currency issue is a present silent data-corruption path: no error, no rejection, just wrong amounts in the DB for non-USD subscriptions. Combined with the unreachable exception branch and the incomplete error wrapping in the new client, the integration layer needs hardening before merge. litellm/integrations/azure_cost_management/azure_cost_management_client.py and litellm/proxy/spend_tracking/ptu_reservation_rollup.py — the currency-handling gap spans both files.
|
| Filename | Overview |
|---|---|
| litellm/integrations/azure_cost_management/azure_cost_management_client.py | New thin Azure Cost Management API client; non-USD currency amounts are returned as-is without conversion or skip, and the exception-handling contract has a gap for non-HTTP network errors. |
| litellm/proxy/spend_tracking/ptu_reservation_rollup.py | Rollup now async-branches on cost_source; azure_billing path fetches from Azure but writes non-USD amounts to ptu_flat_cost without any currency guard or skip. |
| litellm/proxy/proxy_server.py | Adds _build_azure_cost_fetcher_if_enabled helper and wraps PTU rollup in a closure; AzureCostManagementError imported but never used inside the function. |
| litellm/proxy/management_endpoints/ptu_reservation_endpoints.py | CRUD gate relaxed correctly: azure_billing now allowed when azure_resource_id is provided; error message updated to reflect new requirement. |
| litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py | New enable_azure_ptu_billing_pull flag added to UISettings, ALLOWED_UI_SETTINGS_FIELDS, and _RUNTIME_GENERAL_SETTINGS_FLAGS; follows the established pattern correctly. |
| scripts/ptu_reservation_backfill.py | Backfill CLI gains _build_backfill_azure_fetcher and passes it to the rollup; mirrors the proxy path correctly. |
| tests/test_litellm/integrations/azure_cost_management/test_azure_cost_management_client.py | New mock-only tests for the Azure client covering normal, zero-rows, non-USD, HTTP error, and request shape cases; no real network calls. |
| tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py | Existing sync tests correctly converted to async, new azure_billing coverage added; test for non-USD rollup behaviour is absent despite the silent write risk. |
| litellm/proxy/dev_config.yaml | Dev config updated with azure_ptu_billing.subscription_id env var reference; follows existing litellm os.environ/ convention. |
Reviews (2): Last reviewed commit: "feat(ptu): fetch azure_billing flat cost..." | Re-trigger Greptile
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
|
@greptileai please re-review — all 5 findings from the prior pass are addressed in ac86a11 (currency guard, repr redaction, network error wrap, unused import, unreachable branch). Individual replies posted on each thread with the commit hash. CI also picks up an async-clients check fix and an obsolete stage-1 test update. |
|
@greptileai re-review please. c6b5c3b collapses the two PTU flags into one: |
| except httpx.RequestError as exc: | ||
| raise AzureCostManagementError(f"Azure Cost Management network error: {exc}") from exc |
There was a problem hiding this comment.
🟡 Azure cost lookups that time out are not converted to the client's documented error type
A timed-out Azure cost lookup is not wrapped into the client's own error (the except httpx.RequestError at litellm/integrations/azure_cost_management/azure_cost_management_client.py:149) because the shared HTTP layer already converts timeouts into a different, unrelated error type before this code sees them, so the promised error type never surfaces for timeouts.
Impact: A timeout during an Azure cost fetch raises an unexpected error type instead of the documented one, and the test that claims to cover this scenario passes without exercising real behavior.
How the HTTP handler pre-empts the client's except clauses
AzureCostManagementClient.get_daily_cost calls self._http.post(...) where self._http is an AsyncHTTPHandler (litellm/integrations/azure_cost_management/azure_cost_management_client.py:92). Inside AsyncHTTPHandler.post (litellm/llms/custom_httpx/http_handler.py), httpx.TimeoutException is caught and re-raised as litellm.Timeout, which is NOT a subclass of httpx.RequestError. Therefore the except httpx.RequestError branch at line 149 never runs for timeouts; litellm.Timeout propagates instead of AzureCostManagementError, contradicting the docstring at lines 103-104 ("Raises AzureCostManagementError on non-2xx responses or ... network failure").
The unit test test_get_daily_cost_wraps_network_errors mocks the handler to raise httpx.ReadTimeout directly, so it never sees the real handler's conversion and gives false confidence. In practice the impact is limited because the rollup wraps the call in a broad except Exception (litellm/proxy/spend_tracking/ptu_reservation_rollup.py:68), so a bad fetch still degrades to 0.0.
Note: HTTP status errors ARE handled correctly, since AsyncHTTPHandler raises MaskedHTTPStatusError, a subclass of httpx.HTTPStatusError.
Prompt for agents
In litellm/integrations/azure_cost_management/azure_cost_management_client.py, get_daily_cost posts via AsyncHTTPHandler.post, which (see litellm/llms/custom_httpx/http_handler.py) converts httpx.TimeoutException into litellm.Timeout before it can reach the except httpx.RequestError branch. As a result, timeouts are not wrapped into AzureCostManagementError as the docstring promises, and test_get_daily_cost_wraps_network_errors passes only because it mocks the handler to raise httpx.ReadTimeout directly (bypassing the real conversion). Consider either catching litellm.Timeout (and other litellm exception types the handler may raise) in get_daily_cost and wrapping them in AzureCostManagementError, or updating the test to exercise the real handler behavior so it would catch this gap. Confirm the exact exception types AsyncHTTPHandler.post can raise before deciding what to catch.
Was this helpful? React with 👍 or 👎 to provide feedback.
Stage 6 POC on top of the LIT-1697 stack. Adds an azure_billing branch to the daily rollup so a reservation with cost_source=azure_billing pulls accrued cost from the Azure Cost Management API instead of computing the manual (ptu_count * cost_per_ptu) / days-in-month formula. The client is a thin async wrapper around one endpoint the rollup needs; Entra ID auth is inherited from get_azure_ad_token_from_entra_id (no new dependency). Deliberately narrow: no retry, no reconciliation for Azure's 24-72h reporting lag, no UI toggle for cost_source. Those are follow-ups. Gated behind a second UI settings flag enable_azure_ptu_billing_pull that composes below enable_ptu_cost_attribution. The scheduler builds a fetcher just-in-time so runtime flag flips take effect without a proxy restart; manual reservations keep working when either flag is off. Behavior changes: azure_billing cost_source is now accepted by /ptu_reservation/new when azure_resource_id is present. The rollup uses the injected fetcher for those; when no fetcher is available (flag off or env creds missing) the reservation is skipped with a warning and no sentinel row is written.
Currency guard (P1): the rollup now inspects azure_fetcher.last_currency after each call and skips writes when the response reports a non-USD currency, so a customer on a non-USD Azure subscription cannot land a raw foreign-currency amount in ptu_flat_cost as if it were USD. The warning names the reservation, day, and currency so operators can spot the skip and file a follow-up for currency conversion. Client hardening (P2): - client_secret is field(repr=False) so repr(config) no longer exposes it - httpx.RequestError family (ConnectError, ReadTimeout, RemoteProtocolError) is wrapped into AzureCostManagementError so the client honors its documented interface for every failure path, not just HTTPStatusError - Removed the unreachable AttributeError branch in _parse_cost_and_currency - Removed the unused AzureCostManagementError import in proxy_server.py CI (ensure_async_clients_test): the client no longer instantiates AsyncHTTPHandler directly; it goes through get_async_httpx_client with a new httpxSpecialProvider.AzureCostManagement enum value so it shares the same cached connection pool contract as the rest of the integrations. Behavior changes: azure_billing reservations that receive a non-USD Cost Management response are now skipped with a warning instead of writing the raw amount; POST /ptu_reservation/new with cost_source=azure_billing and no azure_resource_id continues to return 400 (stage 1 test updated to assert the new error message rather than the old blanket rejection).
…_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.
0ffa82f to
1b85426
Compare
ca26fd3 to
74145dd
Compare
…ughs Two small additions to make LIT-4077 easier to test and demo without a real PTU. AzureCostManagementConfig now accepts management_base_url, defaulting to https://management.azure.com and reading AZURE_MANAGEMENT_BASE_URL from env. That enables sovereign clouds (Azure Government uses management.usgovcloudapi.net, Azure China uses management.chinacloudapi.cn) and, incidentally, lets a local HTTP fake stand in for Azure during end-to-end tests. The client builds URLs from the config field instead of the previously hardcoded host. _build_azure_cost_fetcher_if_enabled now returns a small _DemoAzureCostFetcher when AZURE_PTU_DEMO_USD is set. The stub returns that USD amount for every reservation/day so a customer walkthrough can show plausible azure_billing figures without provisioning a PTU or waiting for Azure's 24-72h reporting lag. A WARNING fires every rollup so the mode is impossible to miss in production logs; unsetting the env var restores the real client path. Tests cover the new config field (default, env override, base URL propagation into request URLs).
| result = await run_ptu_reservation_rollup( | ||
| prisma_client, target_date=target, force=True, azure_fetcher=azure_fetcher | ||
| ) |
There was a problem hiding this comment.
🟡 Backfill tool passes an argument the existing backfill test's stand-in does not accept, breaking that test
The backfill run now always forwards an extra Azure-fetcher argument to the rollup (azure_fetcher=azure_fetcher at scripts/ptu_reservation_backfill.py:92-94), but the existing backfill test replaces the rollup with a stand-in that doesn't accept it, so that test errors out.
Impact: The already-checked-in backfill test fails, blocking CI for this change.
Signature mismatch between new call site and existing test double
scripts/ptu_reservation_backfill.py:87 builds a fetcher and scripts/ptu_reservation_backfill.py:92-94 calls run_ptu_reservation_rollup(prisma_client, target_date=target, force=True, azure_fetcher=azure_fetcher). The pre-existing (unmodified) test tests/test_scripts/test_ptu_reservation_backfill.py:91 defines async def _fake_rollup(prisma, *, target_date, force=False) and monkeypatches it in at :95. Since _fake_rollup has no azure_fetcher parameter and no **kwargs, invoking mod._run(...) at :99 raises TypeError: _fake_rollup() got an unexpected keyword argument 'azure_fetcher', failing test_run_calls_rollup_with_force_true. The PR added the new keyword to the caller but did not update this mapped test's fake to match.
Prompt for agents
The backfill script now calls run_ptu_reservation_rollup with a new keyword argument azure_fetcher (scripts/ptu_reservation_backfill.py lines 92-94). The existing test tests/test_scripts/test_ptu_reservation_backfill.py defines a fake rollup _fake_rollup(prisma, *, target_date, force=False) at line 91 that does not accept azure_fetcher, so test_run_calls_rollup_with_force_true will raise TypeError when _run forwards azure_fetcher. Update that test's _fake_rollup signature to accept azure_fetcher (e.g. add azure_fetcher=None) and optionally assert its value, so the existing test passes with the new call site.
Was this helpful? React with 👍 or 👎 to provide feedback.
| def _default_token_provider_factory(config: AzureCostManagementConfig) -> TokenProvider: | ||
| from litellm.llms.azure.common_utils import get_azure_ad_token_from_entra_id | ||
|
|
||
| return get_azure_ad_token_from_entra_id( | ||
| tenant_id=config.tenant_id, | ||
| client_id=config.client_id, | ||
| client_secret=config.client_secret, | ||
| scope="https://management.azure.com/.default", | ||
| ) |
There was a problem hiding this comment.
🟡 Azure cost lookups against sovereign clouds request a token for the wrong cloud and will fail authentication
The access token is always requested for the public Azure cloud (scope="https://management.azure.com/.default" at litellm/integrations/azure_cost_management/azure_cost_management_client.py:70) even when the operator points the client at a different cloud's endpoint, so cost lookups against that cloud are rejected as unauthorized.
Impact: Operators on sovereign/government Azure clouds cannot fetch PTU billing costs; every call is denied even though a different endpoint was configured.
Base URL is configurable but the token audience is hardcoded to public cloud
AzureCostManagementConfig supports overriding management_base_url via AZURE_MANAGEMENT_BASE_URL (azure_cost_management_client.py:56), and tests test_config_management_base_url_reads_env_override / test_get_daily_cost_hits_management_base_url_from_config explicitly frame this as sovereign-cloud support (e.g. https://management.usgovcloudapi.net). However _default_token_provider_factory (azure_cost_management_client.py:63-71) always passes scope="https://management.azure.com/.default". When the client runs in production with no injected token_provider, the default factory acquires a token whose audience is the public cloud, which the sovereign-cloud Cost Management endpoint will reject with a 401. The scope should be derived from config.management_base_url (e.g. f"{config.management_base_url}/.default").
| def _default_token_provider_factory(config: AzureCostManagementConfig) -> TokenProvider: | |
| from litellm.llms.azure.common_utils import get_azure_ad_token_from_entra_id | |
| return get_azure_ad_token_from_entra_id( | |
| tenant_id=config.tenant_id, | |
| client_id=config.client_id, | |
| client_secret=config.client_secret, | |
| scope="https://management.azure.com/.default", | |
| ) | |
| def _default_token_provider_factory(config: AzureCostManagementConfig) -> TokenProvider: | |
| from litellm.llms.azure.common_utils import get_azure_ad_token_from_entra_id | |
| return get_azure_ad_token_from_entra_id( | |
| tenant_id=config.tenant_id, | |
| client_id=config.client_id, | |
| client_secret=config.client_secret, | |
| scope=f"{config.management_base_url}/.default", | |
| ) |
Was this helpful? React with 👍 or 👎 to provide feedback.
|
@greptileai review latest head |
| def _default_token_provider_factory(config: AzureCostManagementConfig) -> TokenProvider: | ||
| from litellm.llms.azure.common_utils import get_azure_ad_token_from_entra_id | ||
|
|
||
| return get_azure_ad_token_from_entra_id( | ||
| tenant_id=config.tenant_id, | ||
| client_id=config.client_id, | ||
| client_secret=config.client_secret, | ||
| scope="https://management.azure.com/.default", | ||
| ) |
There was a problem hiding this comment.
_default_token_provider_factory always requests a token with scope "https://management.azure.com/.default", ignoring config.management_base_url. When an operator sets AZURE_MANAGEMENT_BASE_URL=https://management.usgovcloudapi.net (a deliberate feature — there is even a test for the URL override), the API requests go to the sovereign cloud endpoint but carry a token issued for the public cloud audience. Azure will reject these with 401 Unauthorized on every rollup run for that operator.
The scope should be derived from config.management_base_url, e.g. f"{config.management_base_url.rstrip('/')}/.default".
|
Closing as superseded. This is stage 6 of the v1 PTU design, which put PTU config on a separate reservation table. The shipped design puts it on the model deployment instead (#35341, #35343, #35391, #35393, #36829), so this branch is built on a table that no longer exists. The idea here (read real accrued cost from the Azure Cost Management API instead of computing ptu_count x cost_per_ptu_per_hour) never shipped and is still open as a follow-up; it would need rewriting onto the model. Branch litellm_lit4077_stage6_azure_billing_poc is kept. |
2 participants
Relevant issues
Part of the LIT-1697 Azure PTU cost attribution stack; stacked on #33439 (stage 5)
Linear ticket
Resolves LIT-4077
Pre-Submission checklist
Opened as a draft POC to align on the customer-facing story before a real review pass. Full merge review follows after v1 (stages 1-5) ships and the customer walkthrough lands
Delays in PR merge?
Not applicable; POC, not requesting review yet
Screenshots / Proof of Fix
Live end-to-end against a real Azure subscription. Both flags enabled at runtime, a reservation with
cost_source=azure_billingcreated via API, the backfill CLI acquires a real Entra ID token, calls Cost Management, and writes the responseBackfill log:
The $0.00 is a real Azure response for a resource id with no billing history in the dev sub; a customer's own PTU deployment returns its accrued cost on the same code path. Currency is USD, verified against the sub via the Cost Management response
Validator rejection when
azure_resource_idis missing:Restart persistence (both flags survive
_sync_ui_settings_to_general_settings):Type
New Feature
Changes
A thin async client at
litellm/integrations/azure_cost_management/calls one Cost Management endpoint the rollup needs. Entra ID auth is inherited fromget_azure_ad_token_from_entra_idso no new dependency lands. The client is deliberately narrow (one public method) and follows thelitellm/integrations/pattern (seedatadog/,langfuse/, etc.)A new UI settings flag
enable_azure_ptu_billing_pullcomposes belowenable_ptu_cost_attributionand reuses the stage 5 machinery (UISettings model,ALLOWED_UI_SETTINGS_FIELDS,_RUNTIME_GENERAL_SETTINGS_FLAGS). Persistence and startup sync are inherited without touching those code paths_compute_daily_flat_costin the rollup branches oncost_source: manual keeps the existing (ptu_count * cost_per_ptu) / days-in-month formula,azure_billingcalls the injected fetcher, unknown values return 0 with a warning. The function is now async since it awaits the client, and existing sync tests were converted accordinglyThe scheduler wraps the rollup in a small callable that builds a fetcher just-in-time via
_build_azure_cost_fetcher_if_enabledso runtime flag flips take effect without a proxy restart. The helper readsgeneral_settings.enable_azure_ptu_billing_pull, the configuredsubscription_id, and env vars for creds; on any missing input it returns None and the rollup no-ops forazure_billingrowsThe CRUD gate is relaxed: the earlier blanket reject on
cost_source=azure_billingbecomes a check thatazure_resource_idis present. The pydantic domain validator from stage 1 already enforces the rest (ptu_count/cost_per_ptumust be null,azure_resource_idrequired)The backfill CLI grows a matching helper (
_build_backfill_azure_fetcher) so operators runningscripts/ptu_reservation_backfill.pyexercise the same code path as the scheduler whenAZURE_SUBSCRIPTION_IDand Entra ID env creds are setDeferred to follow-up tickets: retry/backoff on 429/503, sliding-window reconciliation for Azure's 24-72h reporting lag, UI toggle for
cost_sourceon the reservation create form, non-USD currency handling, Prometheus counter for fetch failuresBehavior changes
cost_source=azure_billingonPOST /ptu_reservation/newis now accepted (previously rejected with "not supported in this release") whenazure_resource_idis presentFor reservations with
cost_source=azure_billingthat reach the daily rollup: withenable_azure_ptu_billing_pullon and Entra ID env creds resolvable,flat_costis fetched from Azure; with the flag off or creds missing, the reservation is skipped with a warning and no sentinel row is writtenNo change to manual reservations or to the read path (
/team/daily/activity, Usage page)QA runbook
Prereqs: dev proxy on 4097 with
STORE_MODEL_IN_DB=True; env varsAZURE_TENANT_ID,AZURE_CLIENT_ID,AZURE_CLIENT_SECRET,AZURE_SUBSCRIPTION_IDfrom a service principal withCost Management Readeron the target subscription;config.yamlhasgeneral_settings.azure_ptu_billing.subscription_id: os.environ/AZURE_SUBSCRIPTION_IDcurl -X PATCH http://localhost:4097/update/ui_settings -H 'Authorization: Bearer sk-1234' -H 'Content-Type: application/json' -d '{"enable_ptu_cost_attribution": true, "enable_azure_ptu_billing_pull": true}'curl -X POST http://localhost:4097/ptu_reservation/newwith a body that setscost_source=azure_billingand a real Azureazure_resource_id.venv/bin/python scripts/ptu_reservation_backfill.py --date $(date -u +%F); watch for the log linePTU rollup: azure_billing reservation=... returned $...curl 'http://localhost:4097/team/daily/activity?team_ids=<team>&start_date=<today>&end_date=<today>';flat_costis populated when the Azure resource has billing historySynced UI settings to general_settings on startup: ['enable_ptu_cost_attribution', 'enable_azure_ptu_billing_pull']/update/ui_settingsand re-run backfill; azure_billing reservation is skipped with a warning, manual reservations continue to accrueFinal Attestation