From a6b9f4d894fcb4dc3c44d764a24c996eb7bbc2a8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 20:49:17 -0800 Subject: [PATCH] fix(batches): set the created endpoint on the batch object the hook saves and the logger reads Retrieve, cancel and the cost poller now write the endpoint saved at create time onto the provider's batch object itself instead of a copy. The managed files hook persists that object and swaps its raw provider id for the unified id in place, and the spend logger reads the same object, so a copy either let the hook save the provider's endpoint back or kept the raw id on the inline cost row. --- .../proxy/common_utils/check_batch_cost.py | 6 +- litellm/proxy/batches_endpoints/endpoints.py | 15 +-- .../openai_files_endpoints/common_utils.py | 6 +- .../spend/test_batch_observability.py | 106 +++++++++++------- .../proxy/batches_endpoints/test_endpoints.py | 36 +++++- 5 files changed, 115 insertions(+), 54 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index c2f1da3df4f..e7b12c36d94 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -1163,7 +1163,7 @@ class CheckBatchCost: - if not, return False - if so, return True """ - from litellm.proxy.openai_files_endpoints.common_utils import with_created_batch_endpoint + from litellm.proxy.openai_files_endpoints.common_utils import created_batch_endpoint try: from litellm.integrations.prometheus import PrometheusLogger @@ -1226,7 +1226,7 @@ class CheckBatchCost: ) try: - provider_response = await self.llm_router.aretrieve_batch( + response = await self.llm_router.aretrieve_batch( model=model_id, batch_id=batch_id, litellm_metadata={ @@ -1245,7 +1245,7 @@ class CheckBatchCost: else: await self._expire_if_old_and_unfinished(job) continue - response = with_created_batch_endpoint(provider_response, self._stored_batch(job)) + response.endpoint = created_batch_endpoint(response, self._stored_batch(job)) if self._is_completed_with_output_id_pending(response): verbose_proxy_logger.info( f"CheckBatchCost: batch {batch_id} is completed but its output file id " diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index b3408f2d6e7..89d9c9326bc 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -52,6 +52,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 # apply_team_provider_credentials, authorize_model_for_key, batch_cost_poller_is_active, + created_batch_endpoint, decode_model_from_file_id, encode_batch_response_ids, encode_file_id_with_model, @@ -67,7 +68,6 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 # prepare_data_with_credentials, update_batch_in_database, validate_managed_id_requirement, - with_created_batch_endpoint, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import request_tags_from_request_data from litellm.proxy.route_llm_request import raise_if_required_body_param_missing @@ -790,12 +790,13 @@ async def retrieve_batch( **data, ) - retrieved_batch: Final = with_created_batch_endpoint(response, stored_batch) + response.endpoint = created_batch_endpoint(response, stored_batch) + # FIX: Update the database with the latest state from provider await update_batch_in_database( batch_id=batch_id, unified_batch_id=unified_batch_id, - response=retrieved_batch, + response=response, managed_files_obj=managed_files_obj, prisma_client=prisma_client, verbose_proxy_logger=verbose_proxy_logger, @@ -806,7 +807,7 @@ async def retrieve_batch( ### CALL HOOKS ### - modify outgoing data response = await proxy_logging_obj.post_call_success_hook( - data=data, user_api_key_dict=user_api_key_dict, response=retrieved_batch + data=data, user_api_key_dict=user_api_key_dict, response=response ) # Fix: bug_feb14_batch_retrieve_returns_raw_input_file_id @@ -1240,11 +1241,11 @@ async def cancel_batch( prisma_client=prisma_client, verbose_proxy_logger=verbose_proxy_logger, ) - cancelled_batch: Final = with_created_batch_endpoint(response, stored_batch) + response.endpoint = created_batch_endpoint(response, stored_batch) await update_batch_in_database( batch_id=batch_id, unified_batch_id=unified_batch_id, - response=cancelled_batch, + response=response, managed_files_obj=managed_files_obj, prisma_client=prisma_client, verbose_proxy_logger=verbose_proxy_logger, @@ -1255,7 +1256,7 @@ async def cancel_batch( ### CALL HOOKS ### - modify outgoing data response = await proxy_logging_obj.post_call_success_hook( - data=data, user_api_key_dict=user_api_key_dict, response=cancelled_batch + data=data, user_api_key_dict=user_api_key_dict, response=response ) ### ALERTING ### diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 1a057f1e516..1222b7004a5 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1607,10 +1607,8 @@ def batch_cost_poller_is_active() -> bool: return False -def with_created_batch_endpoint(provider_batch: "LiteLLMBatch", stored_batch: "LiteLLMBatch | None") -> "LiteLLMBatch": - if stored_batch is None or not stored_batch.endpoint or stored_batch.endpoint == provider_batch.endpoint: - return provider_batch - return provider_batch.model_copy(update={"endpoint": stored_batch.endpoint}) +def created_batch_endpoint(provider_batch: "LiteLLMBatch", stored_batch: "LiteLLMBatch | None") -> str: + return stored_batch.endpoint if stored_batch is not None and stored_batch.endpoint else provider_batch.endpoint def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool: diff --git a/tests/integration/spend/test_batch_observability.py b/tests/integration/spend/test_batch_observability.py index ca9c26ff513..bd3462febd0 100644 --- a/tests/integration/spend/test_batch_observability.py +++ b/tests/integration/spend/test_batch_observability.py @@ -2,12 +2,16 @@ from __future__ import annotations import json import uuid +from dataclasses import dataclass from hashlib import sha256 +from pathlib import Path from typing import Final +import httpx import pytest -from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows +from integration._support.process import owned_proxy from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse from pydantic import JsonValue @@ -59,11 +63,11 @@ def _failed_line(index: int) -> str: ) -def _batch(status: str, *, files_ready: bool) -> dict[str, JsonValue]: +def _batch(status: str, *, files_ready: bool, endpoint: str = "/v1/chat/completions") -> dict[str, JsonValue]: return { "id": "batch-$REQUEST_ID", "object": "batch", - "endpoint": "/v1/chat/completions", + "endpoint": endpoint, "errors": None, "input_file_id": "file-in-$REQUEST_ID", "completion_window": "24h", @@ -79,7 +83,7 @@ def _batch(status: str, *, files_ready: bool) -> dict[str, JsonValue]: } -def _provider_routes() -> RoutedResponse: +def _provider_routes(created_endpoint: str = "/v1/chat/completions") -> RoutedResponse: output_lines: Final = ( _successful_line(1, REASONING_TOKENS[0]), _failed_line(2), @@ -102,7 +106,8 @@ def _provider_routes() -> RoutedResponse: }, ), "POST /batches": JsonResponse( - content_type="application/json", body=_batch("validating", files_ready=False) + content_type="application/json", + body=_batch("validating", files_ready=False, endpoint=created_endpoint), ), "GET /batches/batch-$REQUEST_ID": JsonResponse( content_type="application/json", body=_batch("completed", files_ready=True) @@ -150,44 +155,53 @@ def _metadata(row: dict[str, JsonValue]) -> dict[str, JsonValue]: return object_value(JSON_OBJECT.validate_json(value) if isinstance(value, str) else value) +@dataclass(frozen=True, slots=True) +class CompletedBatch: + key: str + batch_id: str + retrieval: httpx.Response + + +def _complete_a_batch(gateway: Gateway, scenario: Scenario, created_endpoint: str, *, managed: bool) -> CompletedBatch: + scenario_id: Final = f"batch-observability-{uuid.uuid4().hex[:12]}" + handle: Final = register_scenario(scenario_id, _provider_routes(created_endpoint)) + scenario.cleanups.callback(delete_scenario, handle) + model_name: Final = scenario.model(api_base=handle.api_base()) + key: Final = scenario.key(models=[model_name]) + file_response: Final = gateway.request_multipart( + "/v1/files", + {"purpose": "batch", "target_model_names" if managed else "model": model_name}, + {"file": ("in.jsonl", _input_file(model_name), "application/jsonl")}, + key=key, + ) + assert file_response.status_code == 200, file_response.text + input_file_id: Final = string_value(JSON_OBJECT.validate_json(file_response.content)["id"]) + batch_response: Final = gateway.request( + "POST", + "/v1/batches", + {"input_file_id": input_file_id, "endpoint": created_endpoint, "completion_window": "24h"} + | ({} if managed else {"model": model_name}), + key=key, + ) + assert batch_response.status_code == 200, batch_response.text + batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) + retrieval: Final = eventually( + lambda: gateway.request("GET", f"/v1/batches/{batch_id}", key=key), + lambda response: response.status_code == 200 and response.json()["status"] == "completed", + seconds=30, + ) + assert retrieval.status_code == 200, retrieval.text + return CompletedBatch(key=key, batch_id=batch_id, retrieval=retrieval) + + @pytest.mark.covers("spend.batches.retrieval_row_aggregates_reasoning_tokens_and_per_request_counts") def test_batch_retrieval_row_sums_reasoning_tokens_and_counts_output_and_error_file_failures( gateway: Gateway, ) -> None: with gateway.scenario() as scenario: - scenario_id: Final = f"batch-observability-{uuid.uuid4().hex[:12]}" - handle: Final = register_scenario(scenario_id, _provider_routes()) - scenario.cleanups.callback(delete_scenario, handle) - model_name: Final = scenario.model(api_base=handle.api_base()) - key: Final = scenario.key(models=[model_name]) - file_response: Final = gateway.request_multipart( - "/v1/files", - {"purpose": "batch", "model": model_name}, - {"file": ("in.jsonl", _input_file(model_name), "application/jsonl")}, - key=key, - ) - assert file_response.status_code == 200, file_response.text - input_file_id: Final = string_value(JSON_OBJECT.validate_json(file_response.content)["id"]) - batch_response: Final = gateway.request( - "POST", - "/v1/batches", - { - "input_file_id": input_file_id, - "endpoint": "/v1/chat/completions", - "completion_window": "24h", - "model": model_name, - }, - key=key, - ) - assert batch_response.status_code == 200, batch_response.text - batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) - retrieval: Final = eventually( - lambda: gateway.request("GET", f"/v1/batches/{batch_id}", key=key), - lambda response: response.status_code == 200 and response.json()["status"] == "completed", - seconds=30, - ) - assert retrieval.status_code == 200, retrieval.text - rows: Final = eventually(lambda: _retrieval_rows(key), lambda values: len(values) == 1, seconds=70) + batch: Final = _complete_a_batch(gateway, scenario, "/v1/chat/completions", managed=False) + retrieval: Final = batch.retrieval + rows: Final = eventually(lambda: _retrieval_rows(batch.key), lambda values: len(values) == 1, seconds=70) row: Final = rows[0] metadata: Final = _metadata(row) usage: Final = object_value(metadata["usage_object"]) @@ -199,3 +213,19 @@ def test_batch_retrieval_row_sums_reasoning_tokens_and_counts_output_and_error_f ) assert metadata["batch_successful_requests"] == 2, (retrieval.text, metadata) assert metadata["batch_failed_requests"] == 1 + ERROR_FILE_FAILURES, (retrieval.text, metadata) + + +@pytest.mark.timeout(240) +def test_managed_batch_keeps_its_created_endpoint_when_the_provider_reports_another( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + owned_proxy(gateway, tmp_path, {"PROXY_BATCH_POLLING_ENABLED": "false"}) as candidate, + candidate.scenario() as scenario, + ): + batch: Final = _complete_a_batch(candidate, scenario, "/v1/embeddings", managed=True) + stored: Final = candidate.request("GET", f"/v1/batches/{batch.batch_id}", key=batch.key) + assert (batch.retrieval.json()["endpoint"], stored.json()["endpoint"]) == ( + "/v1/embeddings", + "/v1/embeddings", + ), stored.text diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index 28643d82a46..00ff9161062 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/tests/unit/proxy/batches_endpoints/test_endpoints.py @@ -1909,20 +1909,48 @@ async def test_retrieve__db_non_terminal_state_syncs_with_provider(retrieve_harn assert retrieve_harness.update_batch_in_db.call_count == 1 +def _hook_rewriting_the_batch_id_to(batch_id): + def hook(**kwargs): + kwargs["response"].id = batch_id + return kwargs["response"] + + return hook + + @pytest.mark.asyncio async def test_retrieve__provider_sync_keeps_the_endpoint_the_batch_was_created_for(retrieve_harness): stored = make_batch(id="batch-unified-blob", status="validating").model_copy(update={"endpoint": "/v1/embeddings"}) + provider_batch = make_batch(id="batch-raw-xyz", status="in_progress") retrieve_harness.get_batch_from_db.return_value = (MagicMock(status="validating"), stored) - retrieve_harness.router.aretrieve_batch.return_value = make_batch(id="batch-raw-xyz", status="in_progress") + retrieve_harness.router.aretrieve_batch.return_value = provider_batch + retrieve_harness.logging.post_call_success_hook.side_effect = _hook_rewriting_the_batch_id_to("batch-unified-blob") with patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID): resp = await call_retrieve(retrieve_harness, "batch-unified-blob") assert (resp.endpoint, resp.status) == ("/v1/embeddings", "in_progress") assert retrieve_harness.update_batch_in_db.call_args.kwargs["response"].endpoint == "/v1/embeddings" + assert retrieve_harness.logging.post_call_success_hook.call_args.kwargs["response"].endpoint == "/v1/embeddings" + assert (provider_batch.id, provider_batch.endpoint) == ("batch-unified-blob", "/v1/embeddings") assert resp._hidden_params["unified_batch_id"] == UNIFIED_BATCH_ID +@pytest.mark.asyncio +async def test_retrieve__a_saved_batch_without_an_endpoint_keeps_the_providers(retrieve_harness): + stored = make_batch(id="batch-unified-blob", status="validating").model_copy(update={"endpoint": ""}) + provider_batch = make_batch(id="batch-raw-xyz", status="in_progress").model_copy( + update={"endpoint": "/v1/embeddings"} + ) + retrieve_harness.get_batch_from_db.return_value = (MagicMock(status="validating"), stored) + retrieve_harness.router.aretrieve_batch.return_value = provider_batch + + with patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID): + resp = await call_retrieve(retrieve_harness, "batch-unified-blob") + + assert resp.endpoint == "/v1/embeddings" + assert retrieve_harness.update_batch_in_db.call_args.kwargs["response"].endpoint == "/v1/embeddings" + + @pytest.mark.asyncio @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"]) async def test_retrieve__executed_batch_served_from_db_in_every_status(retrieve_harness, status): @@ -2695,7 +2723,9 @@ async def test_cancel__db_write_receives_caller_auth(cancel_harness): async def test_cancel__keeps_the_endpoint_the_batch_was_created_for(cancel_harness): stored_row = MagicMock(status="in_progress") stored = make_batch(id="batch-unified-blob", status="in_progress").model_copy(update={"endpoint": "/v1/embeddings"}) - cancel_harness.router.acancel_batch.return_value = make_batch(id="batch-raw-xyz", status="cancelling") + provider_batch = make_batch(id="batch-raw-xyz", status="cancelling") + cancel_harness.router.acancel_batch.return_value = provider_batch + cancel_harness.logging.post_call_success_hook.side_effect = _hook_rewriting_the_batch_id_to("batch-unified-blob") with ( patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID), @@ -2706,6 +2736,8 @@ async def test_cancel__keeps_the_endpoint_the_batch_was_created_for(cancel_harne assert (resp.endpoint, resp.status) == ("/v1/embeddings", "cancelling") written = cancel_harness.update_batch_in_db.call_args.kwargs assert (written["response"].endpoint, written["db_batch_object"]) == ("/v1/embeddings", stored_row) + assert cancel_harness.logging.post_call_success_hook.call_args.kwargs["response"].endpoint == "/v1/embeddings" + assert (provider_batch.id, provider_batch.endpoint) == ("batch-unified-blob", "/v1/embeddings") @pytest.mark.asyncio