From 643cdfc4f00a1649c5e07721456608e6f764c972 Mon Sep 17 00:00:00 2001 From: shrey-berri Date: Thu, 29 Oct 2026 09:06:43 -0700 Subject: [PATCH] fix(spend): read a failure filter's sessions once on the session page With status_filter=failure, the session page inlined the session-outcomes subquery in its read, so each parallel worker re-sorted the whole window for it: on 2M rows, 2.2x slower than base with 3x the blocks and temp. The failure filter now reads those sessions once, as a `status_sessions AS MATERIALIZED` CTE, in the page read and the count. A success filter keeps most sessions and stays inlined, where it is faster. Other routes keep the inlined subquery. Results are unchanged. --- .../spend_management_endpoints.py | 58 ++++++++---- .../test_spend_management_endpoints.py | 91 ++++++++++++++++++- 2 files changed, 127 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 34f5bf31f05..fe71dc06d57 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2950,16 +2950,26 @@ async def ui_view_spend_logs( sql_params.append(json.dumps(used_client_oauth_token)) p += 1 - if status_filter is not None and group_by_session is True and not is_search_lookup: - # The window goes in as literals (parse_date's UTC datetimes, checked by - # _utc_timestamp_literal). Prisma reuses its prepared statements, and from the 6th - # run Postgres plans a $N window as 0.5% of the table, then probes every - # session by api_key: ~2 min instead of ~2.5 s over 700k rows. - session_filter_conditions: Final = ( - " AND ".join([*window_literal_conditions, *sql_conditions[len(window_literal_conditions) :]]) or "TRUE" - ) - sql_conditions.append( - f"""({_SESSION_GROUP_KEY_SQL}) IN ( + session_page_route: Final = ( + group_by_session is True + and not is_v2 + and not is_request_id_lookup + and not is_search_lookup + and sort_by == "startTime" + ) + # The session page reads a failure filter's sessions once, as a MATERIALIZED CTE: inlined in the read, + # each parallel worker re-sorts the whole window for it (3x the blocks and temp on 2M rows). A success + # filter keeps most sessions, and there the inlined form is faster. + materialize_status_sessions: Final = session_page_route and status_filter == "failure" + # The window goes in as literals (parse_date's UTC datetimes, checked by + # _utc_timestamp_literal). Prisma reuses its prepared statements, and from the 6th + # run Postgres plans a $N window as 0.5% of the table, then probes every + # session by api_key: ~2 min instead of ~2.5 s over 700k rows. + session_filter_conditions: Final = ( + " AND ".join([*window_literal_conditions, *sql_conditions[len(window_literal_conditions) :]]) or "TRUE" + ) + status_sessions_sql: Final = ( + f""" SELECT session_key, api_key FROM ( SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL}) {_SESSION_KEY_EXPR} AS session_key, api_key, status @@ -2968,18 +2978,20 @@ async def ui_view_spend_logs( ORDER BY {_SESSION_GROUP_KEY_SQL}, {_SESSION_REPRESENTATIVE_ORDER_SQL} ) AS session_outcomes WHERE COALESCE(status, 'success') = ${p} - )""" + """ + if status_filter is not None and group_by_session is True and not is_search_lookup + else None + ) + if status_sessions_sql is not None: + sql_conditions.append( + f"({_SESSION_GROUP_KEY_SQL}) IN (SELECT session_key, api_key FROM status_sessions)" + if materialize_status_sessions + else f"({_SESSION_GROUP_KEY_SQL}) IN ({status_sessions_sql})" ) sql_params.append(status_filter) p += 1 - if ( - group_by_session is True - and not is_v2 - and not is_request_id_lookup - and not is_search_lookup - and sort_by == "startTime" - ): + if session_page_route: return await _ui_session_grouped_spend_logs( prisma_client=prisma_client, sql_conditions=sql_conditions, @@ -2990,6 +3002,7 @@ async def ui_view_spend_logs( sort_desc=order_direction != "asc", session_cursor=session_cursor, status_filtered=status_filter is not None and not is_search_lookup, + status_sessions_sql=status_sessions_sql if materialize_status_sessions else None, ) # Build the ORDER BY expression. ttft_ms is computed from @@ -3197,6 +3210,7 @@ async def _ui_session_grouped_spend_logs( sort_desc: bool, session_cursor: str | None, status_filtered: bool = False, + status_sessions_sql: str | None = None, ) -> Mapping[str, object]: """ One row per session, keyset-paginated by session last activity. @@ -3231,6 +3245,10 @@ async def _ui_session_grouped_spend_logs( row_where_clause: Final = " AND ".join(row_conditions) if row_conditions else "TRUE" cmp_op: Final = "<" if sort_desc else ">" direction: Final = "DESC" if sort_desc else "ASC" + # ``status_sessions_sql`` (a failure filter's sessions, which the last condition reads) is computed once per read. + status_cte: Final = ( + f"status_sessions AS MATERIALIZED ({status_sessions_sql})," if status_sessions_sql is not None else "" + ) cursor: Final = _parse_session_cursor(session_cursor) cursor_clause: Final = ( @@ -3298,7 +3316,7 @@ async def _ui_session_grouped_spend_logs( ) carried_keys: str = f"${carry_index}::text[], ${carry_index + 1}::text[]" page_query: str = f""" - WITH scanned AS MATERIALIZED ( + WITH {status_cte} scanned AS MATERIALIZED ( {_newest_session_rows_sql(where_clause, read_rows, not_after)} ), sessions AS ( @@ -3376,7 +3394,7 @@ async def _ui_session_grouped_spend_logs( if _read_falls_short(count_rows, session_count, rows_read, cap + 1): continue count_query: str = f""" - WITH scanned AS MATERIALIZED ({_newest_session_rows_sql(where_clause, count_rows, None)}) + WITH {status_cte} scanned AS MATERIALIZED ({_newest_session_rows_sql(where_clause, count_rows, None)}) SELECT COUNT(*) AS rows_scanned, (SELECT COUNT(*) FROM (SELECT 1 FROM scanned GROUP BY session_key, api_key) AS s) AS session_count diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index e309c7eb00a..ef7b6ac1b3f 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -8013,8 +8013,9 @@ async def test_ui_view_spend_logs_group_by_session_status_inlines_the_window(cli ) assert response.status_code == 200, response.text read = next(call for call in mock_prisma.db.query_raw.await_args_list if "rows_scanned" in call.args[0]) - outer, outcomes = read.args[0].split("SELECT DISTINCT ON", 1) - outcomes = outcomes.split("AS session_outcomes", 1)[0] + before, outcomes = read.args[0].split("SELECT DISTINCT ON", 1) + outcomes, after = outcomes.split("AS session_outcomes", 1) + outer = before + after for bound in (start_date, end_date): literal = datetime.datetime.strptime(bound, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc).isoformat() assert f"('{literal}'::timestamptz AT TIME ZONE 'UTC')" in outcomes @@ -8028,6 +8029,92 @@ async def test_ui_view_spend_logs_group_by_session_status_inlines_the_window(cli app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.asyncio +@pytest.mark.parametrize("status_filter", ["failure", "success"]) +async def test_ui_view_spend_logs_group_by_session_materializes_only_failed_sessions( + client, monkeypatch, status_filter +): + """A failure filter's sessions are read once as a MATERIALIZED CTE (inlined, each parallel worker re-sorts the window); success stays inlined.""" + page_rows = [_session_page_row("sess-1", "2026-08-29 08:00:00")] + mock_prisma = _session_grouped_mock_prisma(page_rows, 1, [_session_representative_row("req-1", "sess-1")]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "status_filter": status_filter, + "model": "gpt-4o", + "session_cursor": "2026-08-29 09:00:00|hashed-key|req-solo", + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + assert [row["session_id"] for row in response.json()["data"]] == ["sess-1"] + reads = [call.args for call in mock_prisma.db.query_raw.await_args_list if "rows_scanned" in call.args[0]] + assert len(reads) == 2, "the cursor page's read and the window count both carry the status filter" + for read in reads: + assert read[0].count("AS session_outcomes") == 1 + assert "COALESCE(status, 'success') = $4" in read[0] and read[4] == status_filter + if status_filter == "failure": + assert read[0].lstrip().startswith("WITH status_sessions AS MATERIALIZED (") + assert "IN (SELECT session_key, api_key FROM status_sessions)" in read[0] + else: + assert "status_sessions" not in read[0] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_failure_off_the_session_page_stays_inlined(client, monkeypatch): + """Only the session page declares the status_sessions CTE, so other grouped sorts keep the inlined failure subquery.""" + queries = [] + + async def mock_query_raw(sql_query, *params): + queries.append(sql_query) + return [{"total_count": 0}] if "total_count" in sql_query else [] + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "status_filter": "failure", + "sort_by": "spend", + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + assert queries and all("status_sessions" not in query for query in queries) + assert any("AS session_outcomes" in query for query in queries) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_group_by_session_jumps_to_page_without_cursor(client, monkeypatch): """page > 1 with no session_cursor (the UI's last-page jump) serves the sessions that page starts at."""