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

@@ -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

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