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