From 6d796d0f1f441c577ee493483fe077eb3ca18a04 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 2 Jul 2026 21:19:54 +0530 Subject: [PATCH] feat(proxy): track cost for unmanaged Vertex AI batch jobs (#31442) * feat(proxy): track cost for unmanaged Vertex AI batch jobs CheckBatchCost previously skipped Vertex batches created via the raw GCS input_file_id path, since their unified_object_id is a raw provider job id that fails the base64 managed-id check. Behind the opt-in general_settings flag track_unmanaged_vertex_batch_cost, the poller now derives the model from the gs:// input_file_id, maps it to a configured vertex_ai deployment, polls the batch, computes cost, and marks batch_processed=True. * Update tracking for failed", "expired", "cancelled" * fix(proxy): apply ruff format to proxy_server.py * address greptile review feedback (greploop iteration 1) Filter unmanaged Vertex batch deployments by vertex_ai provider so a shared model group name can't route to a wrong-provider deployment. Move gs:// URI parsing into VertexAIBatchTransformation. Add test coverage for the failed/expired/cancelled terminal-status DB update. * fix: route unmanaged vertex batches to matching deployment --------- Co-authored-by: Cursor Agent --- .../proxy/common_utils/check_batch_cost.py | 233 ++++++++++-- .../llms/vertex_ai/batches/transformation.py | 18 +- litellm/proxy/dev_config.yaml | 3 + litellm/proxy/proxy_server.py | 1 + .../proxy_unit_tests/test_check_batch_cost.py | 340 +++++++++++++++++- 5 files changed, 565 insertions(+), 30 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 ee7745d0add..831a23ff3cd 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -13,6 +13,8 @@ from litellm.constants import ( ) if TYPE_CHECKING: + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy._types import LiteLLM_ManagedObjectTable from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -26,6 +28,7 @@ class CheckBatchCost: proxy_logging_obj: "ProxyLogging", prisma_client: "PrismaClient", llm_router: "Router", + track_unmanaged_vertex_batch_cost: bool = False, ): from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -33,6 +36,7 @@ class CheckBatchCost: self.proxy_logging_obj: ProxyLogging = proxy_logging_obj self.prisma_client: PrismaClient = prisma_client self.llm_router: Router = llm_router + self._track_unmanaged_vertex_batch_cost = track_unmanaged_vertex_batch_cost # Cached after the first poll cycle. Once we know the column is absent we skip # the guaranteed-failing primary query on every subsequent cycle. self._has_batch_processed_column: bool = True @@ -97,6 +101,182 @@ class CheckBatchCost: order={"created_at": "asc"}, ) + @staticmethod + def _record_error( + prom_logger: Optional["PrometheusLogger"], error_type: str + ) -> None: + if prom_logger is not None: + prom_logger.record_check_batch_cost_error(error_type) + + def _resolve_job_routing( + self, + job: "LiteLLM_ManagedObjectTable", + prom_logger: Optional["PrometheusLogger"], + ) -> Optional[Tuple[str, str]]: + """ + Resolve (model_id, batch_id) for a managed-object row, where model_id is a router + deployment id and batch_id is the raw provider batch id. + + Managed batches encode both in a base64 unified id. Unmanaged Vertex batches, created with + a raw gs:// input_file_id, store the raw provider job id as unified_object_id; when + track_unmanaged_vertex_batch_cost is enabled the model is derived from the gs:// path and + mapped to a configured vertex_ai deployment. Returns None (recording a metric) when the row + can't be routed. + """ + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + get_batch_id_from_unified_batch_id, + get_model_id_from_unified_batch_id, + ) + + unified_object_id = job.unified_object_id + decoded = _is_base64_encoded_unified_file_id(unified_object_id) + if decoded: + model_id = get_model_id_from_unified_batch_id(decoded) + if model_id is None: + verbose_proxy_logger.info( + f"Skipping job {unified_object_id} because it is not a valid model id" + ) + self._record_error(prom_logger, "invalid_model_id") + return None + return model_id, get_batch_id_from_unified_batch_id(decoded) + + if self._track_unmanaged_vertex_batch_cost: + return self._resolve_unmanaged_vertex_routing(job, prom_logger) + + verbose_proxy_logger.info( + f"Skipping job {unified_object_id} because it is not a valid unified object id" + ) + self._record_error(prom_logger, "invalid_unified_id") + return None + + def _resolve_unmanaged_vertex_routing( + self, + job: "LiteLLM_ManagedObjectTable", + prom_logger: Optional["PrometheusLogger"], + ) -> Optional[Tuple[str, str]]: + from litellm.llms.vertex_ai.batches.transformation import ( + VertexAIBatchTransformation, + ) + + input_file_id = self._get_input_file_id(job) + if not VertexAIBatchTransformation.is_unmanaged_gcs_batch_input_file_id( + input_file_id + ): + verbose_proxy_logger.info( + f"Skipping job {job.unified_object_id}: not an unmanaged vertex batch " + "(no gs:// input_file_id with a publishers/ model path)" + ) + self._record_error(prom_logger, "invalid_unified_id") + return None + assert input_file_id is not None # narrowed by is_unmanaged_gcs_batch_input_file_id + + bare_model_name = VertexAIBatchTransformation.get_bare_model_name_from_gcs_file( + input_file_id + ) + deployment_id = self._get_vertex_ai_deployment_id_for_bare_model( + bare_model_name + ) + if deployment_id is None: + verbose_proxy_logger.info( + f"Skipping unmanaged vertex batch {job.unified_object_id}: no vertex_ai " + f"deployment configured for model {bare_model_name}" + ) + self._record_error(prom_logger, "unmanaged_no_matching_deployment") + return None + + return deployment_id, job.unified_object_id + + def _get_vertex_ai_deployment_id_for_bare_model( + self, bare_model_name: str + ) -> Optional[str]: + model_group = self.llm_router.resolve_model_name_from_model_id(bare_model_name) + deployment_id = ( + self._get_vertex_ai_deployment_id(model_group) if model_group else None + ) + if deployment_id is not None: + return deployment_id + + return self._get_vertex_ai_deployment_id_from_matching_deployments( + bare_model_name + ) + + def _get_vertex_ai_deployment_id_from_matching_deployments( + self, bare_model_name: str + ) -> Optional[str]: + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + for deployment in self.llm_router.get_model_list(model_name=None) or []: + litellm_params = deployment.get("litellm_params") or {} + actual_model = litellm_params.get("model") + if not isinstance(actual_model, str): + continue + if not self._is_bare_model_match(actual_model, bare_model_name): + continue + try: + _, llm_provider, _, _ = get_llm_provider( + model=actual_model, + custom_llm_provider=litellm_params.get("custom_llm_provider"), + ) + except Exception: + continue + if llm_provider != "vertex_ai": + continue + model_info = deployment.get("model_info") or {} + deployment_id = model_info.get("id") + if isinstance(deployment_id, str): + return deployment_id + return None + + @staticmethod + def _is_bare_model_match(actual_model: str, bare_model_name: str) -> bool: + return ( + actual_model == bare_model_name + or actual_model.endswith(f"/{bare_model_name}") + or actual_model.endswith(f":{bare_model_name}") + ) + + def _get_vertex_ai_deployment_id(self, model_group: str) -> Optional[str]: + """ + Returns the first deployment id for `model_group` whose provider is vertex_ai, + skipping deployments from other providers that happen to share the model group name. + """ + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + for deployment_id in self.llm_router.get_model_ids(model_name=model_group): + deployment_info = self.llm_router.get_deployment(model_id=deployment_id) + if deployment_info is None: + continue + try: + _, llm_provider, _, _ = get_llm_provider( + model=deployment_info.litellm_params.model, + custom_llm_provider=deployment_info.litellm_params.custom_llm_provider, + ) + except Exception: + continue + if llm_provider == "vertex_ai": + return deployment_id + return None + + @staticmethod + def _get_input_file_id(job: "LiteLLM_ManagedObjectTable") -> Optional[str]: + import json + + from litellm.types.utils import LiteLLMBatch + + file_object = job.file_object + if isinstance(file_object, str): + try: + file_object = json.loads(file_object) + except (json.JSONDecodeError, ValueError): + return None + if not isinstance(file_object, dict): + return None + try: + return LiteLLMBatch.model_validate(file_object).input_file_id + except Exception: + return None + async def check_batch_cost(self): """ Check if the batch JOB has been tracked. @@ -114,8 +294,6 @@ class CheckBatchCost: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, - get_batch_id_from_unified_batch_id, - get_model_id_from_unified_batch_id, ) try: @@ -172,31 +350,10 @@ class CheckBatchCost: else: jobs = await self._fallback_find_jobs() for job in jobs: - # get the model from the job - unified_object_id = job.unified_object_id - decoded_unified_object_id = _is_base64_encoded_unified_file_id( - unified_object_id - ) - if not decoded_unified_object_id: - verbose_proxy_logger.info( - f"Skipping job {unified_object_id} because it is not a valid unified object id" - ) - if prom_logger: - prom_logger.record_check_batch_cost_error("invalid_unified_id") - continue - else: - unified_object_id = decoded_unified_object_id - - model_id = get_model_id_from_unified_batch_id(unified_object_id) - batch_id = get_batch_id_from_unified_batch_id(unified_object_id) - - if model_id is None: - verbose_proxy_logger.info( - f"Skipping job {unified_object_id} because it is not a valid model id" - ) - if prom_logger: - prom_logger.record_check_batch_cost_error("invalid_model_id") + routing = self._resolve_job_routing(job, prom_logger) + if routing is None: continue + model_id, batch_id = routing verbose_proxy_logger.info( f"Querying model ID: {model_id} for cost and usage of batch ID: {batch_id}" @@ -213,7 +370,7 @@ class CheckBatchCost: ) except Exception as e: verbose_proxy_logger.info( - f"Skipping job {unified_object_id} because of error querying model ID: {model_id} for cost and usage of batch ID: {batch_id}: {e}" + f"Skipping job {job.unified_object_id} because of error querying model ID: {model_id} for cost and usage of batch ID: {batch_id}: {e}" ) if prom_logger: prom_logger.record_check_batch_cost_error("provider_retrieval_error") @@ -287,7 +444,7 @@ class CheckBatchCost: deployment_info = self.llm_router.get_deployment(model_id=model_id) if deployment_info is None: verbose_proxy_logger.info( - f"Skipping job {unified_object_id} because it is not a valid deployment info" + f"Skipping job {job.unified_object_id} because it is not a valid deployment info" ) if prom_logger: prom_logger.record_check_batch_cost_error("deployment_not_found") @@ -413,6 +570,26 @@ class CheckBatchCost: f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" ) + elif response.status in ("failed", "expired", "cancelled"): + try: + update_data = { + "status": response.status, + "file_object": response.model_dump_json(), + } + if self._has_batch_processed_column: + update_data["batch_processed"] = True + await self.prisma_client.db.litellm_managedobjecttable.update( + where={"id": job.id}, + data=update_data, + ) + verbose_proxy_logger.info( + f"CheckBatchCost: marked job {job.id} as {response.status} in DB" + ) + except Exception as db_err: + verbose_proxy_logger.error( + f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}" + ) + # Record polling run metrics (always, even if nothing was processed) if prom_logger: prom_logger.record_check_batch_cost_run( diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index 4cac7212d62..6bbe8f75701 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -1,4 +1,4 @@ -from typing import Any, Dict +from typing import Any, Dict, Optional from litellm._uuid import uuid from litellm.llms.vertex_ai.common_utils import ( @@ -207,3 +207,19 @@ class VertexAIBatchTransformation: parts = model_path.split("/") model = f"publishers/{'/'.join(parts[:3])}" return model + + @classmethod + def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: Optional[str]) -> bool: + """ + Returns True if `input_file_id` is a raw gs:// Vertex batch input file (i.e. not a + LiteLLM-managed unified file id) with a `publishers/` model path that + `_get_model_from_gcs_file` can parse. + """ + return input_file_id is not None and input_file_id.startswith("gs://") and "publishers/" in input_file_id + + @classmethod + def get_bare_model_name_from_gcs_file(cls, gcs_file_uri: str) -> str: + """ + Extracts the bare model name (e.g. "gemini-1.5-flash-001") from a gcs file uri. + """ + return cls._get_model_from_gcs_file(gcs_file_uri).rsplit("/", 1)[-1] diff --git a/litellm/proxy/dev_config.yaml b/litellm/proxy/dev_config.yaml index 0dc90e658f2..65a4b8e7cbf 100644 --- a/litellm/proxy/dev_config.yaml +++ b/litellm/proxy/dev_config.yaml @@ -193,6 +193,9 @@ model_list: general_settings: master_key: sk-1234 + # Opt-in: let CheckBatchCost track cost for unmanaged Vertex batches created with a raw gs:// input_file_id. + # Requires a vertex_ai deployment configured for the batched model. Defaults to false. + # track_unmanaged_vertex_batch_cost: true sandbox_tools: - sandbox_tool_name: e2b_sandbox diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 04f0414f900..a29acf3b06f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7627,6 +7627,7 @@ class ProxyStartupEvent: proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client, llm_router=llm_router, + track_unmanaged_vertex_batch_cost=general_settings.get("track_unmanaged_vertex_batch_cost", False), ) scheduler.add_job( check_batch_cost_job.check_batch_cost, diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index e8acaf6fea6..63bbe147801 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1,13 +1,35 @@ """ Unit tests for CheckBatchCost class. Covers: stale-row cleanup (file_purpose scoping), paginated find_many, -and the batch_processed-column fallback query. +the batch_processed-column fallback query, and routing of unmanaged +Vertex batches (raw gs:// input_file_id, no managed unified id). """ from unittest.mock import AsyncMock, MagicMock, patch import pytest +_IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id" + + +def _unmanaged_vertex_file_object( + input_file_id="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash/abc.jsonl", + status="validating", +): + """A LiteLLMBatch JSON blob shaped like what the managed-files hook stores for an + unmanaged Vertex batch (raw gs:// input_file_id).""" + from litellm.types.utils import LiteLLMBatch + + return LiteLLMBatch( + id="8823717160934178816", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id=input_file_id, + object="batch", + status=status, + ).model_dump_json() + class TestCheckBatchCost: """Test suite for CheckBatchCost class""" @@ -375,6 +397,76 @@ class TestCheckBatchCost: ), "update() must include batch_processed=True when column is present" assert update_data["status"] == "complete" + @pytest.mark.asyncio + @pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"]) + async def test_terminal_status_marks_job_processed( + self, + check_batch_cost_instance, + mock_prisma_client, + mock_llm_router, + terminal_status, + ): + """When the provider reports a terminal status (failed/expired/cancelled), the row + must be written back with that status and batch_processed=True so it stops being + polled forever. + """ + from unittest.mock import patch + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) + + mock_job = MagicMock() + mock_job.id = "job-terminal-1" + mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.created_by = "user-1" + + assert check_batch_cost_instance._has_batch_processed_column is True + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + + mock_response = MagicMock() + mock_response.status = terminal_status + mock_response.model_dump_json.return_value = ( + f'{{"id":"batch-1","status":"{terminal_status}"}}' + ) + + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + + decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;" + + with ( + patch( + "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", + side_effect=[decoded_id, None], + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id", + return_value="model-123", + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id", + return_value="batch-456", + ), + ): + await check_batch_cost_instance.check_batch_cost() + + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), f"Expected update() to be called exactly once for a {terminal_status} job" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] + assert update_data["status"] == terminal_status + assert ( + update_data["batch_processed"] is True + ), "terminal-status update() must set batch_processed=True so polling stops" + @pytest.mark.asyncio async def test_raw_output_file_id_converted_to_managed_id( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router @@ -512,3 +604,249 @@ class TestCheckBatchCost: } assert mock_response.output_file_id == fake_managed_output_id assert mock_response.error_file_id == fake_managed_error_id + + +class TestUnmanagedVertexRouting: + """Routing of unmanaged Vertex batches whose unified_object_id is a raw provider job id.""" + + def _instance(self, track_unmanaged, router): + from litellm_enterprise.proxy.common_utils.check_batch_cost import ( + CheckBatchCost, + ) + + return CheckBatchCost( + proxy_logging_obj=MagicMock(), + prisma_client=MagicMock(), + llm_router=router, + track_unmanaged_vertex_batch_cost=track_unmanaged, + ) + + def _job(self, file_object=None): + job = MagicMock() + job.unified_object_id = "8823717160934178816" + job.file_object = ( + file_object if file_object is not None else _unmanaged_vertex_file_object() + ) + return job + + def test_flag_off_skips_unmanaged_id_unchanged(self): + """Default (flag off): a raw numeric unified_object_id is skipped exactly as before; + no model derivation or router lookup happens.""" + router = MagicMock() + instance = self._instance(track_unmanaged=False, router=router) + prom = MagicMock() + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(self._job(), prom) + + assert result is None + prom.record_check_batch_cost_error.assert_called_once_with("invalid_unified_id") + router.resolve_model_name_from_model_id.assert_not_called() + router.get_model_ids.assert_not_called() + + def _vertex_deployment(self): + deployment = MagicMock() + deployment.litellm_params.custom_llm_provider = "vertex_ai" + deployment.litellm_params.model = "vertex_ai/gemini-2.5-flash" + return deployment + + def test_flag_on_routes_to_vertex_deployment(self): + """Flag on: derive the bare model from the gs:// path, resolve it to a deployment id, + and use the raw unified_object_id as the provider batch id.""" + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = "gemini-2.5-flash" + router.get_model_ids.return_value = ["deploy-1"] + router.get_deployment = MagicMock(return_value=self._vertex_deployment()) + instance = self._instance(track_unmanaged=True, router=router) + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(self._job(), MagicMock()) + + assert result == ("deploy-1", "8823717160934178816") + # bare model name (trailing GCS segment), not the full publishers/.. path + router.resolve_model_name_from_model_id.assert_called_once_with( + "gemini-2.5-flash" + ) + router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") + + def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): + """Flag on, but the only deployment for the model group is a non-vertex_ai + provider: must not be selected, even though the model group name matches.""" + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = "gemini-2.5-flash" + router.get_model_ids.return_value = ["deploy-openai"] + non_vertex_deployment = MagicMock() + non_vertex_deployment.litellm_params.custom_llm_provider = "openai" + non_vertex_deployment.litellm_params.model = "gpt-4o" + router.get_deployment = MagicMock(return_value=non_vertex_deployment) + instance = self._instance(track_unmanaged=True, router=router) + prom = MagicMock() + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(self._job(), prom) + + assert result is None + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) + + def test_flag_on_uses_later_vertex_deployment_with_matching_suffix(self): + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = "azure-gemini" + router.get_model_ids.return_value = ["deploy-azure"] + non_vertex_deployment = MagicMock() + non_vertex_deployment.litellm_params.custom_llm_provider = "azure" + non_vertex_deployment.litellm_params.model = "azure/gemini-2.5-flash" + router.get_deployment = MagicMock(return_value=non_vertex_deployment) + router.get_model_list.return_value = [ + { + "model_name": "azure-gemini", + "litellm_params": { + "model": "azure/gemini-2.5-flash", + "custom_llm_provider": "azure", + }, + "model_info": {"id": "deploy-azure"}, + }, + { + "model_name": "vertex-gemini", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "custom_llm_provider": "vertex_ai", + }, + "model_info": {"id": "deploy-vertex"}, + }, + ] + instance = self._instance(track_unmanaged=True, router=router) + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(self._job(), MagicMock()) + + assert result == ("deploy-vertex", "8823717160934178816") + router.get_model_ids.assert_called_once_with(model_name="azure-gemini") + + def test_flag_on_no_matching_deployment_records_metric(self): + """Flag on but no vertex_ai deployment for the model: skip with a distinct metric.""" + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = None + router.get_model_ids.return_value = [] + instance = self._instance(track_unmanaged=True, router=router) + prom = MagicMock() + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(self._job(), prom) + + assert result is None + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) + + def test_flag_on_non_gcs_input_is_not_unmanaged_vertex(self): + """Flag on, but input_file_id is not a gs:// publishers path: treat as unroutable, + do not attempt model derivation.""" + router = MagicMock() + instance = self._instance(track_unmanaged=True, router=router) + prom = MagicMock() + job = self._job( + file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123") + ) + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(job, prom) + + assert result is None + prom.record_check_batch_cost_error.assert_called_once_with("invalid_unified_id") + router.resolve_model_name_from_model_id.assert_not_called() + + @pytest.mark.asyncio + async def test_end_to_end_costs_unmanaged_batch(self): + """Flag on, completed unmanaged batch: the poller polls Vertex with the raw job id, + computes cost, and marks batch_processed=True. Fails before this change (the row is + skipped at the unified-id gate).""" + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = "gemini-2.5-flash" + router.get_model_ids.return_value = ["deploy-1"] + + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.output_file_id = "gs://bucket/out/predictions.jsonl" + mock_response.error_file_id = None + mock_response.completed_at = None + mock_response.created_at = None + mock_response.model_dump_json.return_value = ( + '{"id":"8823717160934178816","status":"completed"}' + ) + router.aretrieve_batch = AsyncMock(return_value=mock_response) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"vertex_project": "p", "vertex_location": "us-central1"} + ) + + deployment = MagicMock() + deployment.litellm_params.custom_llm_provider = "vertex_ai" + deployment.litellm_params.model = "vertex_ai/gemini-2.5-flash" + deployment.model_name = "gemini-2.5-flash" + deployment.model_info.model_dump.return_value = {} + router.get_deployment = MagicMock(return_value=deployment) + + instance = self._instance(track_unmanaged=True, router=router) + instance.proxy_logging_obj.get_proxy_hook.return_value = None + instance._has_batch_processed_column = True + + prisma = instance.prisma_client + prisma.db = MagicMock() + prisma.db.litellm_managedobjecttable = MagicMock() + prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + prisma.db.litellm_managedobjecttable.update = AsyncMock() + prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[self._job()] + ) + prisma.db.litellm_usertable = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + mock_file_content = MagicMock() + mock_file_content.content = b'{"id":"req-1"}' + + with ( + patch(_IS_B64, side_effect=[False, None]), + patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + return_value=mock_file_content, + ), + patch( + "litellm.batches.batch_utils._get_file_content_as_dictionary", + return_value=[{"id": "req-1"}], + ), + patch( + "litellm.batches.batch_utils.calculate_batch_cost_and_usage", + new_callable=AsyncMock, + return_value=( + 0.01, + {"prompt_tokens": 10, "completion_tokens": 5}, + ["gemini-2.5-flash"], + ), + ), + patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", + return_value=("gemini-2.5-flash", "vertex_ai", None, None), + ), + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, + ): + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + mock_logging_cls.return_value = mock_logging_obj + + await instance.check_batch_cost() + + router.aretrieve_batch.assert_awaited_once() + assert router.aretrieve_batch.call_args[1]["model"] == "deploy-1" + assert router.aretrieve_batch.call_args[1]["batch_id"] == "8823717160934178816" + + mock_logging_obj.async_success_handler.assert_awaited_once() + assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.01 + + assert prisma.db.litellm_managedobjecttable.update.call_count == 1 + update_data = prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert update_data["batch_processed"] is True + assert update_data["status"] == "complete"