From c09ac782e3a64d9f0224655b79459f446d9dd9af Mon Sep 17 00:00:00 2001 From: mynkyu Date: Thu, 27 Aug 2026 18:30:16 +0900 Subject: [PATCH 1/2] fix(router): stamp model_group when retrieving a batch Batch token usage is accounted on the retrieve call, not on create: a provider only reports token counts once the job finishes, so the usage arrives on aretrieve_batch and that is the spend log row the tokens land on. Router.acreate_batch stamps the requested model group into its metadata, but Router.aretrieve_batch never did. A batch is retrieved by id, so the request carries no model, and the router fans the lookup out over its deployments - leaving model_group unset on the one record that carries the tokens. /global/activity/model groups the spend logs by model_group, so every batch's tokens were bucketed under an empty group. Stamp the model group inside the per-deployment retrieve attempt, preferring an explicitly requested group and otherwise using the model_name of the deployment that answered, which is unambiguous even when the request named no model. An existing model_group in the metadata is left untouched, so nothing that already resolves a group changes. Scope is limited to aretrieve_batch: acompletion, aresponses and acreate_batch logging are untouched, and cost/spend attribution by model is unchanged. Signed-off-by: mynkyu --- litellm/router.py | 9 ++ .../test_router_batch_retrieve_model_group.py | 118 ++++++++++++++++++ 2 files changed, 127 insertions(+) create mode 100644 tests/test_litellm/test_router_batch_retrieve_model_group.py diff --git a/litellm/router.py b/litellm/router.py index f0ebb539bb7..d70e251f044 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5812,6 +5812,8 @@ class Router: """ try: parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) + requested_model_group: Final = model + metadata_variable_name: Final = _get_router_metadata_variable_name(function_name="aretrieve_batch") if model is not None: filtered_model_list: ( list[DeploymentTypedDict] | list[dict] | dict | None @@ -5848,6 +5850,13 @@ class Router: kwargs=new_kwargs, function_name="aretrieve_batch", ) + ## STAMP THE MODEL GROUP FOR SPEND TRACKING ## + # A batch is retrieved by id, so the request carries no model group of its + # own - only the deployment that answered knows it. Batch token usage lands + # on this retrieve call (the provider reports counts once the job finishes), + # so without this the tokens are logged under an empty model_group. + model_group: Final = requested_model_group or model_name["model_name"] + new_kwargs[metadata_variable_name].setdefault("model_group", model_group) new_kwargs.pop("custom_llm_provider", None) data.pop("custom_llm_provider", None) return await litellm.aretrieve_batch( diff --git a/tests/test_litellm/test_router_batch_retrieve_model_group.py b/tests/test_litellm/test_router_batch_retrieve_model_group.py new file mode 100644 index 00000000000..ef8a23e4917 --- /dev/null +++ b/tests/test_litellm/test_router_batch_retrieve_model_group.py @@ -0,0 +1,118 @@ +""" +model_group attribution on router batch retrieval. + +Batch token usage is accounted on the *retrieve* call, not on create: the +provider only knows the token counts once the job finishes, so +`LiteLLMBatch.usage` arrives on `aretrieve_batch` and that is the record the +spend log tokens land on. + +`aretrieve_batch` is addressed by batch_id, so the request carries no model, +and the router fans the lookup out across its deployments. These tests lock +that the winning deployment's model group is stamped on the emitted +StandardLoggingPayload, so `/global/activity/model` - which groups the spend +logs by `model_group` - can attribute those tokens instead of bucketing every +batch under "". +""" + +import asyncio +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +import litellm.batches.main as bm +from litellm import Router +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import LiteLLMBatch, Usage + +MODEL_GROUP = "vertex-gemini-2.5-flash-lite-dev" +DEPLOYMENT_MODEL = "vertex_ai/gemini-2.5-flash-lite" + + +class _PayloadCollector(CustomLogger): + def __init__(self): + super().__init__() + self.payloads = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.payloads.append(kwargs.get("standard_logging_object")) + + +@pytest.fixture +def router(): + return Router( + model_list=[ + { + "model_name": MODEL_GROUP, + "litellm_params": { + "model": DEPLOYMENT_MODEL, + "vertex_project": "fake-project", + "vertex_location": "us-central1", + "vertex_credentials": "fake-creds", + }, + } + ] + ) + + +@pytest.fixture +def collector(): + logger = _PayloadCollector() + previous = litellm.callbacks + litellm.callbacks = [logger] + try: + yield logger + finally: + litellm.callbacks = previous + + +@pytest.fixture +def vertex_retrieve(): + """Mock the vertex provider seam - the only real network boundary.""" + batch = LiteLLMBatch( + id="batch-1", + completion_window="24h", + created_at=0, + endpoint="/v1/chat/completions", + input_file_id="file-1", + object="batch", + status="completed", + usage=Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200), + ) + seam = MagicMock(name="vertex_ai_batches_instance") + seam.retrieve_batch.return_value = batch + with patch.object(bm, "vertex_ai_batches_instance", seam): + yield seam + + +async def _collected_payload(collector) -> dict: + for _ in range(50): # the success handler runs as a background task + payloads = [p for p in collector.payloads if p is not None] + if payloads: + return payloads[-1] + await asyncio.sleep(0.05) + raise AssertionError(f"no StandardLoggingPayload was emitted: {collector.payloads}") + + +@pytest.mark.asyncio +async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, vertex_retrieve): + """ + The proxy retrieves a managed batch by id only - no `model` in the request. + The router fans out over its deployments, so the model group is only known + from the deployment that answered. + """ + response = await router.aretrieve_batch(batch_id="batch-1") + + assert response.usage.total_tokens == 1200 + payload = await _collected_payload(collector) + assert payload["model"] == DEPLOYMENT_MODEL + assert payload["model_group"] == MODEL_GROUP + + +@pytest.mark.asyncio +async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, vertex_retrieve): + """An explicitly requested model group is what gets logged.""" + await router.aretrieve_batch(model=MODEL_GROUP, batch_id="batch-1") + + payload = await _collected_payload(collector) + assert payload["model_group"] == MODEL_GROUP From 7a1904f8c602b029471d464feaa8f5bba02708c2 Mon Sep 17 00:00:00 2001 From: mynkyu Date: Thu, 27 Aug 2026 19:01:18 +0900 Subject: [PATCH 2/2] test: fake the provider at the HTTP boundary in the batch model_group test The test-quality gate flagged the first version for patching an SDK internal (litellm.batches.main.vertex_ai_batches_instance) and for writing litellm.callbacks directly. Drive an openai-compatible deployment through respx instead, so the retrieve call and the usage accounting that reads the completed batch's output file both run for real, and install the collector with monkeypatch so nothing leaks into the next test. Signed-off-by: mynkyu --- .../test_router_batch_retrieve_model_group.py | 146 +++++++++++------- 1 file changed, 89 insertions(+), 57 deletions(-) diff --git a/tests/test_litellm/test_router_batch_retrieve_model_group.py b/tests/test_litellm/test_router_batch_retrieve_model_group.py index ef8a23e4917..b99ec50e041 100644 --- a/tests/test_litellm/test_router_batch_retrieve_model_group.py +++ b/tests/test_litellm/test_router_batch_retrieve_model_group.py @@ -1,35 +1,81 @@ """ model_group attribution on router batch retrieval. -Batch token usage is accounted on the *retrieve* call, not on create: the -provider only knows the token counts once the job finishes, so -`LiteLLMBatch.usage` arrives on `aretrieve_batch` and that is the record the -spend log tokens land on. +Batch token usage is accounted on the *retrieve* call, not on create: a provider +only reports token counts once the job finishes, so the usage is read off the +completed batch's output file during retrieve logging and that is the spend log +row the tokens land on. -`aretrieve_batch` is addressed by batch_id, so the request carries no model, -and the router fans the lookup out across its deployments. These tests lock -that the winning deployment's model group is stamped on the emitted -StandardLoggingPayload, so `/global/activity/model` - which groups the spend -logs by `model_group` - can attribute those tokens instead of bucketing every -batch under "". +A batch is retrieved by id, so the request carries no model and the router fans +the lookup out across its deployments. These tests lock that the answering +deployment's model group is stamped on the emitted StandardLoggingPayload, so +`/global/activity/model` - which groups the spend logs by `model_group` - can +attribute those tokens instead of bucketing every batch under "". + +The provider is faked at the HTTP boundary, so the whole retrieve + usage +accounting path runs for real. """ import asyncio -from unittest.mock import MagicMock, patch +import json +import httpx import pytest +import respx import litellm -import litellm.batches.main as bm from litellm import Router from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import LiteLLMBatch, Usage -MODEL_GROUP = "vertex-gemini-2.5-flash-lite-dev" -DEPLOYMENT_MODEL = "vertex_ai/gemini-2.5-flash-lite" +MODEL_GROUP = "gemini-batch-group" +DEPLOYMENT_MODEL = "openai/gpt-4o-mini" +API_BASE = "http://localhost:4001/v1" +BATCH_ID = "batch-1" +ROWS = 2 +TOKENS_PER_ROW = 600 + +COMPLETED_BATCH = { + "id": BATCH_ID, + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": "file-in-1", + "completion_window": "24h", + "status": "completed", + "output_file_id": "file-out-1", + "error_file_id": None, + "created_at": 0, + "completed_at": 1, + "request_counts": {"total": ROWS, "completed": ROWS, "failed": 0}, + "metadata": None, +} + +OUTPUT_JSONL = "\n".join( + json.dumps( + { + "id": f"req-{row}", + "custom_id": f"row-{row}", + "response": { + "status_code": 200, + "body": { + "id": f"chatcmpl-{row}", + "object": "chat.completion", + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 500, "completion_tokens": 100, "total_tokens": TOKENS_PER_ROW}, + }, + }, + } + ) + for row in range(ROWS) +) class _PayloadCollector(CustomLogger): + """Captures the StandardLoggingPayload the spend log is built from.""" + def __init__(self): super().__init__() self.payloads = [] @@ -37,6 +83,14 @@ class _PayloadCollector(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): self.payloads.append(kwargs.get("standard_logging_object")) + async def retrieve_batch_payload(self) -> dict: + for _ in range(100): # the success handler runs as a background task + for payload in self.payloads: + if payload and payload.get("call_type") == "aretrieve_batch": + return payload + await asyncio.sleep(0.05) + raise AssertionError(f"no aretrieve_batch payload was emitted: {self.payloads}") + @pytest.fixture def router(): @@ -46,9 +100,8 @@ def router(): "model_name": MODEL_GROUP, "litellm_params": { "model": DEPLOYMENT_MODEL, - "vertex_project": "fake-project", - "vertex_location": "us-central1", - "vertex_credentials": "fake-creds", + "api_base": API_BASE, + "api_key": "sk-fake", }, } ] @@ -56,63 +109,42 @@ def router(): @pytest.fixture -def collector(): +def collector(monkeypatch): logger = _PayloadCollector() - previous = litellm.callbacks - litellm.callbacks = [logger] - try: - yield logger - finally: - litellm.callbacks = previous + monkeypatch.setattr(litellm, "callbacks", [logger]) + return logger @pytest.fixture -def vertex_retrieve(): - """Mock the vertex provider seam - the only real network boundary.""" - batch = LiteLLMBatch( - id="batch-1", - completion_window="24h", - created_at=0, - endpoint="/v1/chat/completions", - input_file_id="file-1", - object="batch", - status="completed", - usage=Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200), - ) - seam = MagicMock(name="vertex_ai_batches_instance") - seam.retrieve_batch.return_value = batch - with patch.object(bm, "vertex_ai_batches_instance", seam): - yield seam - - -async def _collected_payload(collector) -> dict: - for _ in range(50): # the success handler runs as a background task - payloads = [p for p in collector.payloads if p is not None] - if payloads: - return payloads[-1] - await asyncio.sleep(0.05) - raise AssertionError(f"no StandardLoggingPayload was emitted: {collector.payloads}") +def provider(): + """Fake the provider at the HTTP boundary: the completed batch plus the + output file the usage accounting reads.""" + with respx.mock(assert_all_called=True) as respx_mock: + respx_mock.get(f"{API_BASE}/batches/{BATCH_ID}").mock(return_value=httpx.Response(200, json=COMPLETED_BATCH)) + respx_mock.get(f"{API_BASE}/files/file-out-1/content").mock(return_value=httpx.Response(200, text=OUTPUT_JSONL)) + yield respx_mock @pytest.mark.asyncio -async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, vertex_retrieve): +async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, provider): """ The proxy retrieves a managed batch by id only - no `model` in the request. The router fans out over its deployments, so the model group is only known from the deployment that answered. """ - response = await router.aretrieve_batch(batch_id="batch-1") + response = await router.aretrieve_batch(batch_id=BATCH_ID) - assert response.usage.total_tokens == 1200 - payload = await _collected_payload(collector) + assert response.id == BATCH_ID + payload = await collector.retrieve_batch_payload() + assert payload["total_tokens"] == ROWS * TOKENS_PER_ROW assert payload["model"] == DEPLOYMENT_MODEL assert payload["model_group"] == MODEL_GROUP @pytest.mark.asyncio -async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, vertex_retrieve): +async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, provider): """An explicitly requested model group is what gets logged.""" - await router.aretrieve_batch(model=MODEL_GROUP, batch_id="batch-1") + await router.aretrieve_batch(model=MODEL_GROUP, batch_id=BATCH_ID) - payload = await _collected_payload(collector) + payload = await collector.retrieve_batch_payload() assert payload["model_group"] == MODEL_GROUP