From 1a9899091afa21ce5b5818bd594bc1f61ea3bdff Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Mon, 13 Jul 2026 17:40:20 -0700 Subject: [PATCH] fix(ptu): backfill CLI bootstraps its own Prisma and forces flag The script imported prisma_client and general_settings from proxy_server, but both are only populated during the async proxy startup path, so a standalone invocation exited early with an empty result. Rewritten to construct its own PrismaClient from DATABASE_URL, connect it, and pass force=True to run_ptu_reservation_rollup so backfill runs regardless of whether the config flag is on. Adds a force kwarg on the rollup entry point (default False keeps the scheduler contract intact). 8 new tests: 6 CLI-side, 2 rollup-side for the force branch. --- .../spend_tracking/ptu_reservation_rollup.py | 29 ++--- scripts/ptu_reservation_backfill.py | 48 +++++--- .../test_ptu_reservation_rollup.py | 28 +++++ .../test_ptu_reservation_backfill.py | 106 ++++++++++++++++++ 4 files changed, 181 insertions(+), 30 deletions(-) create mode 100644 tests/test_scripts/test_ptu_reservation_backfill.py diff --git a/litellm/proxy/spend_tracking/ptu_reservation_rollup.py b/litellm/proxy/spend_tracking/ptu_reservation_rollup.py index 2d80ef80fef..f5a2061c969 100644 --- a/litellm/proxy/spend_tracking/ptu_reservation_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_reservation_rollup.py @@ -88,23 +88,26 @@ async def _upsert_ptu_daily_row( async def run_ptu_reservation_rollup( prisma_client: Any, target_date: date | None = None, + *, + force: bool = False, ) -> RollupResult: """Rollup one UTC day of flat PTU cost across all active reservations. - Defaults to yesterday UTC. Callable from the scheduler and from the CLI - backfill helper; both paths are idempotent under the ``LiteLLM_DailyTeamSpend`` - unique constraint. + Defaults to yesterday UTC. ``force=True`` bypasses the feature-flag check + so the CLI backfill can run when the scheduler is off. Idempotent under + the LiteLLM_DailyTeamSpend unique constraint on every invocation path. """ - from litellm.proxy.proxy_server import general_settings + if not force: + from litellm.proxy.proxy_server import general_settings - if not general_settings.get("enable_ptu_cost_attribution", False): - verbose_proxy_logger.debug("PTU rollup: feature flag off, skipping") - return RollupResult( - day=target_date or (datetime.now(timezone.utc).date() - timedelta(days=1)), - reservations_processed=0, - rows_written=0, - skipped_flag_off=True, - ) + if not general_settings.get("enable_ptu_cost_attribution", False): + verbose_proxy_logger.debug("PTU rollup: feature flag off, skipping") + return RollupResult( + day=target_date or (datetime.now(timezone.utc).date() - timedelta(days=1)), + reservations_processed=0, + rows_written=0, + skipped_flag_off=True, + ) if prisma_client is None: verbose_proxy_logger.warning("PTU rollup: prisma_client is None, skipping") @@ -136,7 +139,7 @@ async def run_ptu_reservation_rollup( flat_cost=flat_cost, ) rows_written += 1 - except Exception as exc: + except Exception as exc: # noqa: BLE001 # one bad reservation must not stop the batch; logged and continued verbose_proxy_logger.error( "PTU rollup: upsert failed for reservation=%s day=%s: %s", reservation.id, diff --git a/scripts/ptu_reservation_backfill.py b/scripts/ptu_reservation_backfill.py index 9e95a767e41..e73c54b8b25 100644 --- a/scripts/ptu_reservation_backfill.py +++ b/scripts/ptu_reservation_backfill.py @@ -1,6 +1,9 @@ """Re-run the PTU reservation daily rollup for a specific UTC date. -Idempotent under the LiteLLM_DailyTeamSpend unique constraint. +Idempotent under the LiteLLM_DailyTeamSpend unique constraint. Reads +DATABASE_URL from the environment, connects Prisma directly, and bypasses +the ``enable_ptu_cost_attribution`` config flag so operators can backfill +without turning the feature on for live traffic. Usage: python scripts/ptu_reservation_backfill.py --date 2026-07-12 @@ -42,26 +45,37 @@ def _dates_from_args(args: argparse.Namespace) -> list[date]: async def _run(dates: list[date]) -> int: - from litellm.proxy.proxy_server import prisma_client + database_url = os.environ.get("DATABASE_URL") + if not database_url: + print("ERROR: DATABASE_URL is not set", file=sys.stderr) + return 2 + + from litellm._logging import verbose_proxy_logger + from litellm.caching.dual_cache import DualCache from litellm.proxy.spend_tracking.ptu_reservation_rollup import ( run_ptu_reservation_rollup, ) + from litellm.proxy.utils import PrismaClient, ProxyLogging - if prisma_client is None: - print("ERROR: prisma_client is None; is DATABASE_URL set?", file=sys.stderr) - return 2 - - total_rows = 0 - for target in dates: - result = await run_ptu_reservation_rollup(prisma_client, target_date=target) - print( - f"[{result.day.isoformat()}] " - f"reservations={result.reservations_processed} rows_written={result.rows_written}" - f"{' (flag off, skipped)' if result.skipped_flag_off else ''}" - ) - total_rows += result.rows_written - print(f"total rows written: {total_rows}") - return 0 + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + prisma_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) + await prisma_client.connect() + try: + total_rows = 0 + for target in dates: + result = await run_ptu_reservation_rollup(prisma_client, target_date=target, force=True) + print( + f"[{result.day.isoformat()}] " + f"reservations={result.reservations_processed} rows_written={result.rows_written}" + ) + total_rows += result.rows_written + print(f"total rows written: {total_rows}") + return 0 + finally: + try: + await prisma_client.db.disconnect() + except Exception as exc: + verbose_proxy_logger.debug("prisma disconnect failed: %s", exc) def main() -> int: diff --git a/tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py index 6a74a674c74..f20f7cbceb7 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py @@ -309,3 +309,31 @@ async def test_rollup_boundary_effective_from_at_day_start_is_active(mock_prisma result = await run_ptu_reservation_rollup(prisma, target_date=date(2026, 7, 12)) assert result.rows_written == 1 + + +@pytest.mark.asyncio +async def test_rollup_force_bypasses_flag_off(mock_prisma, monkeypatch): + prisma, mock_daily, mock_reservation = mock_prisma + monkeypatch.setitem(ps.general_settings, "enable_ptu_cost_attribution", False) + mock_reservation.find_many = AsyncMock(return_value=[_r(id="res_1", ptu_count=1, cost_per_ptu=200.0)]) + + result = await run_ptu_reservation_rollup(prisma, target_date=date(2026, 7, 12), force=True) + + assert result.skipped_flag_off is False + assert result.rows_written == 1 + mock_reservation.find_many.assert_awaited_once() + mock_daily.upsert.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_rollup_force_default_false_still_honors_flag(mock_prisma, monkeypatch): + prisma, mock_daily, mock_reservation = mock_prisma + monkeypatch.setitem(ps.general_settings, "enable_ptu_cost_attribution", False) + mock_reservation.find_many = AsyncMock(return_value=[_r(id="res_1")]) + + result = await run_ptu_reservation_rollup(prisma, target_date=date(2026, 7, 12)) + + assert result.skipped_flag_off is True + assert result.rows_written == 0 + mock_reservation.find_many.assert_not_awaited() + mock_daily.upsert.assert_not_awaited() diff --git a/tests/test_scripts/test_ptu_reservation_backfill.py b/tests/test_scripts/test_ptu_reservation_backfill.py new file mode 100644 index 00000000000..ca760a5a907 --- /dev/null +++ b/tests/test_scripts/test_ptu_reservation_backfill.py @@ -0,0 +1,106 @@ +import argparse +import importlib.util +import os +from datetime import date +from pathlib import Path + +import pytest + + +_SCRIPT_PATH = Path(__file__).parent.parent.parent / "scripts" / "ptu_reservation_backfill.py" + + +def _load_script(): + spec = importlib.util.spec_from_file_location("_ptu_backfill", _SCRIPT_PATH) + assert spec is not None and spec.loader is not None + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def test_dates_from_args_single_date(): + mod = _load_script() + args = argparse.Namespace(date=date(2026, 7, 12), date_range=None) + assert mod._dates_from_args(args) == [date(2026, 7, 12)] + + +def test_dates_from_args_inclusive_range(): + mod = _load_script() + args = argparse.Namespace(date=None, date_range="2026-07-10:2026-07-12") + assert mod._dates_from_args(args) == [ + date(2026, 7, 10), + date(2026, 7, 11), + date(2026, 7, 12), + ] + + +def test_dates_from_args_single_day_range(): + mod = _load_script() + args = argparse.Namespace(date=None, date_range="2026-07-12:2026-07-12") + assert mod._dates_from_args(args) == [date(2026, 7, 12)] + + +def test_dates_from_args_rejects_reversed_range(): + mod = _load_script() + args = argparse.Namespace(date=None, date_range="2026-07-15:2026-07-01") + with pytest.raises(ValueError): + mod._dates_from_args(args) + + +@pytest.mark.asyncio +async def test_run_exits_when_database_url_missing(monkeypatch, capsys): + mod = _load_script() + monkeypatch.delenv("DATABASE_URL", raising=False) + code = await mod._run([date(2026, 7, 12)]) + assert code == 2 + err = capsys.readouterr().err + assert "DATABASE_URL" in err + + +@pytest.mark.asyncio +async def test_run_calls_rollup_with_force_true(monkeypatch, capsys): + mod = _load_script() + + class _FakePrismaClient: + def __init__(self, *_args, **_kwargs): + self.connected = False + self.disconnected = False + self.db = self + + async def connect(self): + self.connected = True + + async def disconnect(self): + self.disconnected = True + + fake_client_holder = {} + + def _fake_ctor(*args, **kwargs): + client = _FakePrismaClient(*args, **kwargs) + fake_client_holder["client"] = client + return client + + from litellm.proxy import utils as proxy_utils + monkeypatch.setattr(proxy_utils, "PrismaClient", _fake_ctor) + + from litellm.proxy.spend_tracking import ptu_reservation_rollup as rollup_mod + from litellm.proxy.spend_tracking.ptu_reservation_rollup import RollupResult + + calls: list[dict] = [] + + async def _fake_rollup(prisma, *, target_date, force=False): + calls.append({"target_date": target_date, "force": force}) + return RollupResult(day=target_date, reservations_processed=1, rows_written=1) + + monkeypatch.setattr(rollup_mod, "run_ptu_reservation_rollup", _fake_rollup) + + monkeypatch.setenv("DATABASE_URL", "postgresql://fake@localhost/test") + + code = await mod._run([date(2026, 7, 10), date(2026, 7, 11)]) + assert code == 0 + assert [c["target_date"] for c in calls] == [date(2026, 7, 10), date(2026, 7, 11)] + assert all(c["force"] is True for c in calls) + out = capsys.readouterr().out + assert "total rows written: 2" in out + assert fake_client_holder["client"].connected is True + assert fake_client_holder["client"].disconnected is True