From a5e51de5c6d7afcae7124d07e39d4acf6dc22032 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 22 Sep 2026 20:56:29 +0000 Subject: [PATCH] test(integration): uncostable batches retire from the cost poll page so newer batches are costed (Pylon #7342) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/proxy_config.yaml | 1 + .../spend/test_batch_poll_starvation.py | 212 ++++++++++++++++++ 2 files changed, 213 insertions(+) create mode 100644 tests/integration/spend/test_batch_poll_starvation.py diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index a3b07f76d2f..94b0a3003fe 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -5,6 +5,7 @@ general_settings: store_model_in_db: true disable_spend_logs: false proxy_batch_write_at: 1 + proxy_batch_polling_interval: 1 litellm_settings: enable_redis_auth_cache: true cache: true diff --git a/tests/integration/spend/test_batch_poll_starvation.py b/tests/integration/spend/test_batch_poll_starvation.py new file mode 100644 index 00000000000..f68f10e3d8a --- /dev/null +++ b/tests/integration/spend/test_batch_poll_starvation.py @@ -0,0 +1,212 @@ +import json +import os +from hashlib import sha256 +from typing import Final + +import pytest +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.upstream import delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse +from pydantic import JsonValue + +from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE + +INPUT_COST_PER_TOKEN: Final = 0.001 +OUTPUT_COST_PER_TOKEN: Final = 0.002 +PROMPT_TOKENS: Final = 100 +COMPLETION_TOKENS: Final = 50 +BATCH_COST_SHARE: Final = 0.5 + +_INPUT_FILE: Final = JsonResponse( + content_type="application/json", + body={ + "id": "file-in-$REQUEST_ID", + "object": "file", + "purpose": "batch", + "bytes": 100, + "created_at": 1, + "filename": "in.jsonl", + "status": "processed", + }, +) + + +def _batch(status: str, output_file_id: str | None) -> dict[str, JsonValue]: + return { + "id": "batch-$REQUEST_ID", + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": "file-in-$REQUEST_ID", + "completion_window": "24h", + "status": status, + "output_file_id": output_file_id, + "error_file_id": None, + "created_at": 1, + "in_progress_at": 1, + "completed_at": 1 if status == "completed" else None, + "expires_at": 1, + "request_counts": {"total": 1, "completed": 1 if status == "completed" else 0, "failed": 0}, + "metadata": None, + } + + +def _accepting_routes() -> dict[str, JsonResponse | TextResponse]: + return { + "POST /files": _INPUT_FILE, + "POST /batches": JsonResponse(content_type="application/json", body=_batch("validating", None)), + } + + +def _gone_at_provider_routes() -> RoutedResponse: + return RoutedResponse( + content_type="application/x-routed", + routes={ + **_accepting_routes(), + "GET /batches/batch-$REQUEST_ID": JsonResponse( + content_type="application/json", + status=404, + body={ + "error": { + "message": "No batch found with id 'batch-$REQUEST_ID'.", + "type": "invalid_request_error", + "param": "id", + "code": "batch_not_found", + } + }, + ), + }, + ) + + +def _completed_routes() -> RoutedResponse: + output_line: Final = { + "id": "batch_req_1", + "custom_id": "r1", + "response": { + "status_code": 200, + "request_id": "$REQUEST_ID-1", + "body": { + "id": "chatcmpl-$REQUEST_ID-1", + "object": "chat.completion", + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + }, + }, + }, + "error": None, + } + return RoutedResponse( + content_type="application/x-routed", + routes={ + **_accepting_routes(), + "GET /batches/batch-$REQUEST_ID": JsonResponse( + content_type="application/json", body=_batch("completed", "file-out-$REQUEST_ID") + ), + "GET /files/file-out-$REQUEST_ID/content": TextResponse( + content_type="application/jsonl", body=json.dumps(output_line, separators=(",", ":")) + "\n" + ), + }, + ) + + +def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse) -> str: + scenario_id: Final = f"poll-{marker}-{sha256(os.urandom(16)).hexdigest()[:12]}" + handle: Final = register_scenario(scenario_id, routes) + scenario.cleanups.callback(delete_scenario, handle) + created: Final = scenario.gateway.post( + "/model/new", + { + "model_name": f"poll-{marker}-{sha256(scenario_id.encode()).hexdigest()[:12]}", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-scripted-provider", + "api_base": handle.api_base(), + "input_cost_per_token": INPUT_COST_PER_TOKEN, + "output_cost_per_token": OUTPUT_COST_PER_TOKEN, + }, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return string_value(created["model_name"]) + + +def _submitted_batch_id(gateway: Gateway, key: str, model_name: str) -> str: + request_line: Final = { + "custom_id": "r1", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": model_name, "messages": [{"role": "user", "content": "poll starvation"}]}, + } + file_response: Final = gateway.request_multipart( + "/v1/files", + {"purpose": "batch", "target_model_names": model_name}, + {"file": ("in.jsonl", (json.dumps(request_line) + "\n").encode(), "application/jsonl")}, + key=key, + ) + assert file_response.is_success, file_response.text + batch_response: Final = gateway.request( + "POST", + "/v1/batches", + { + "input_file_id": string_value(JSON_OBJECT.validate_json(file_response.content)["id"]), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": model_name, + }, + key=key, + ) + assert batch_response.is_success, batch_response.text + return string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) + + +def _managed_rows(batch_ids: tuple[str, ...]) -> list[dict[str, JsonValue]]: + placeholders: Final = ", ".join("%s" for _ in batch_ids) + return read_rows( + f'SELECT batch_processed FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id IN ({placeholders})', + batch_ids, + ) + + +@pytest.mark.timeout(180) +@pytest.mark.covers("quota_management.spend_tracking.batch_costs.uncostable_rows_retire_so_newer_batches_are_costed") +def test_batches_gone_at_provider_do_not_starve_a_newer_batch_out_of_cost_polling(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + gone_batch_ids: Final = tuple( + _submitted_batch_id( + gateway, key, _scripted_deployment(scenario, f"gone{index}", _gone_at_provider_routes()) + ) + for index in range(MAX_OBJECTS_PER_POLL_CYCLE) + ) + costable_batch_id: Final = _submitted_batch_id( + gateway, key, _scripted_deployment(scenario, "costable", _completed_routes()) + ) + spend_rows: Final = eventually( + lambda: read_rows( + 'SELECT call_type, status, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" ' + "WHERE api_key = %s AND call_type = %s", + (sha256(key.encode()).hexdigest(), "aretrieve_batch"), + ), + lambda rows: len(rows) == 1, + seconds=120, + ) + assert spend_rows == [ + { + "call_type": "aretrieve_batch", + "status": "success", + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "spend": pytest.approx( + BATCH_COST_SHARE + * (PROMPT_TOKENS * INPUT_COST_PER_TOKEN + COMPLETION_TOKENS * OUTPUT_COST_PER_TOKEN) + ), + } + ] + assert _managed_rows((costable_batch_id,)) == [{"batch_processed": True}] + assert _managed_rows(gone_batch_ids) == [{"batch_processed": True}] * MAX_OBJECTS_PER_POLL_CYCLE