mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(batches): keep the endpoint a batch was created for on retrieve, cancel and cost polling
Backport of #50784 to lts/1.108.x. Applied fromde995b48e8,2184695446,a6b9f4d894,7f51cb1ce7,d54a31f1a8,2617ad780e,c11902afe1(litellm_bedrock_batch_retrieve_endpoint).
This commit is contained in:
mateo-berri
parent
2a61f2114d
commit
6bd520f1a3
12 changed files with 295 additions and 43 deletions
|
|
@ -627,17 +627,22 @@ class CheckBatchCost:
|
|||
return None
|
||||
|
||||
@classmethod
|
||||
def _get_input_file_id(cls, job: "ManagedObjectRow") -> Optional[str]:
|
||||
def _stored_batch(cls, job: "ManagedObjectRow") -> Optional["LiteLLMBatch"]:
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
parsed_file_object: Final = cls._parse_file_object(job)
|
||||
if parsed_file_object is None:
|
||||
return None
|
||||
try:
|
||||
return LiteLLMBatch.model_validate(parsed_file_object).input_file_id
|
||||
return LiteLLMBatch.model_validate(parsed_file_object)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _get_input_file_id(cls, job: "ManagedObjectRow") -> Optional[str]:
|
||||
stored_batch: Final = cls._stored_batch(job)
|
||||
return stored_batch.input_file_id if stored_batch is not None else None
|
||||
|
||||
@classmethod
|
||||
def _get_stored_model_id(cls, job: "ManagedObjectRow") -> Optional[str]:
|
||||
"""The deployment that created the batch, written into file_object by the managed
|
||||
|
|
@ -1105,6 +1110,8 @@ class CheckBatchCost:
|
|||
- if not, return False
|
||||
- if so, return True
|
||||
"""
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import created_batch_endpoint
|
||||
|
||||
try:
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
|
|
@ -1188,6 +1195,7 @@ class CheckBatchCost:
|
|||
if self._is_batch_gone_at_provider(e, batch_id) and self._batch_deployment_exists(model_id):
|
||||
await self._retire_job(job, f"batch {batch_id} no longer exists at the provider")
|
||||
continue
|
||||
response.endpoint = created_batch_endpoint(response, self._stored_batch(job))
|
||||
|
||||
## RETRIEVE THE BATCH JOB OUTPUT FILE
|
||||
if response.status in PROVIDER_TERMINAL_BATCH_STATUSES and response.output_file_id is not None:
|
||||
|
|
|
|||
|
|
@ -315,7 +315,10 @@ class BedrockBatchesHandler:
|
|||
# provider uses, so model_list / env / role-assumption configs are
|
||||
# honored. We instantiate BedrockBatchesConfig (which extends
|
||||
# BaseAWSLLM) lazily to avoid a circular import at module load.
|
||||
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
|
||||
from litellm.llms.bedrock.batches.transformation import (
|
||||
BedrockBatchesConfig,
|
||||
openai_endpoint_for_bedrock_job_model,
|
||||
)
|
||||
|
||||
creds: Final = BedrockBatchesConfig().resolve_credentials(AwsAuthParams.model_validate(kwargs), region)
|
||||
|
||||
|
|
@ -404,7 +407,7 @@ class BedrockBatchesHandler:
|
|||
request_counts=_record_counts_from_response(response),
|
||||
metadata=openai_batch_metadata,
|
||||
completion_window="24h",
|
||||
endpoint="/v1/chat/completions",
|
||||
endpoint=openai_endpoint_for_bedrock_job_model(response.get("modelId")),
|
||||
input_file_id=input_uri,
|
||||
output_file_id=output_file_uri if openai_status == "completed" else None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final, Literal, cast
|
|||
from httpx import Headers, Response
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.exceptions import ModelNotMappedError
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix, is_bedrock_arn
|
||||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
BEDROCK_MANAGED_S3_BATCH_PREFIX,
|
||||
|
|
@ -32,6 +33,7 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders, Usage
|
|||
from ..base_aws_llm import BaseAWSLLM
|
||||
from ..common_utils import (
|
||||
CommonBatchFilesUtils,
|
||||
get_cached_model_info,
|
||||
merge_bedrock_aws_request_params,
|
||||
resolve_s3_bucket_owner,
|
||||
resolve_s3_encryption_key_id,
|
||||
|
|
@ -82,6 +84,21 @@ def _validate_bedrock_tags(raw_tags: object) -> list[BedrockTag]:
|
|||
) from e
|
||||
|
||||
|
||||
def openai_endpoint_for_bedrock_job_model(
|
||||
model_id: str | None,
|
||||
) -> Literal["/v1/chat/completions", "/v1/embeddings"]:
|
||||
if not model_id:
|
||||
return "/v1/chat/completions"
|
||||
try:
|
||||
model_info: Final = get_cached_model_info()(model=model_id, custom_llm_provider="bedrock")
|
||||
except ModelNotMappedError:
|
||||
return "/v1/chat/completions"
|
||||
is_embedding_model: Final = (
|
||||
model_info.get("mode") == "embedding" or model_info.get("output_vector_size") is not None
|
||||
)
|
||||
return "/v1/embeddings" if is_embedding_model else "/v1/chat/completions"
|
||||
|
||||
|
||||
def titan_embedding_usage_from_batch_output(model_output: Mapping[str, object]) -> Usage | None:
|
||||
"""Titan embedding batch lines report usage as a top-level inputTextTokenCount, not a usage block."""
|
||||
if "embedding" not in model_output and "embeddingsByType" not in model_output:
|
||||
|
|
@ -639,7 +656,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
|||
return LiteLLMBatch(
|
||||
id=job_arn,
|
||||
object="batch",
|
||||
endpoint="/v1/chat/completions",
|
||||
endpoint=openai_endpoint_for_bedrock_job_model(response_data.get("modelId")),
|
||||
errors=errors,
|
||||
input_file_id=input_file_id,
|
||||
completion_window="24h",
|
||||
|
|
|
|||
|
|
@ -16994,6 +16994,7 @@
|
|||
"max_tokens": 512,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 1024,
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/",
|
||||
"supports_embedding_image_input": true
|
||||
},
|
||||
|
|
@ -17004,6 +17005,7 @@
|
|||
"max_tokens": 512,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 1024,
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/",
|
||||
"supports_embedding_image_input": true
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -635,6 +636,7 @@ async def retrieve_batch(
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
stored_batch: Final = response
|
||||
|
||||
if executed_batch and response is None:
|
||||
raise batch_error(404, f"No batch found with id '{batch_id}'.")
|
||||
|
|
@ -779,6 +781,8 @@ async def retrieve_batch(
|
|||
**data,
|
||||
)
|
||||
|
||||
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,
|
||||
|
|
@ -1221,6 +1225,14 @@ async def cancel_batch(
|
|||
managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files")
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
db_batch_object, stored_batch = await get_batch_from_database(
|
||||
batch_id=batch_id,
|
||||
unified_batch_id=unified_batch_id,
|
||||
managed_files_obj=managed_files_obj,
|
||||
prisma_client=prisma_client,
|
||||
verbose_proxy_logger=verbose_proxy_logger,
|
||||
)
|
||||
response.endpoint = created_batch_endpoint(response, stored_batch)
|
||||
await update_batch_in_database(
|
||||
batch_id=batch_id,
|
||||
unified_batch_id=unified_batch_id,
|
||||
|
|
@ -1228,6 +1240,7 @@ async def cancel_batch(
|
|||
managed_files_obj=managed_files_obj,
|
||||
prisma_client=prisma_client,
|
||||
verbose_proxy_logger=verbose_proxy_logger,
|
||||
db_batch_object=db_batch_object,
|
||||
operation="cancel",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1480,6 +1480,13 @@ def batch_cost_poller_is_active() -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def created_batch_endpoint(provider_batch: "LiteLLMBatch", stored_batch: "LiteLLMBatch | None") -> str:
|
||||
if stored_batch is None or not stored_batch.endpoint:
|
||||
return provider_batch.endpoint
|
||||
provider_named_the_endpoint: Final = provider_batch.endpoint not in ("", "/v1/chat/completions")
|
||||
return provider_batch.endpoint if provider_named_the_endpoint else stored_batch.endpoint
|
||||
|
||||
|
||||
def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool:
|
||||
"""Whether a "completed" batch may be retired from cost recovery.
|
||||
|
||||
|
|
|
|||
|
|
@ -16994,6 +16994,7 @@
|
|||
"max_tokens": 512,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 1024,
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/",
|
||||
"supports_embedding_image_input": true
|
||||
},
|
||||
|
|
@ -17004,6 +17005,7 @@
|
|||
"max_tokens": 512,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 1024,
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/",
|
||||
"supports_embedding_image_input": true
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -165,6 +165,24 @@ def test_handle_model_invocation_job_status_completed(patched_boto3):
|
|||
assert batch.metadata["output_s3_uri"] == OUTPUT_PREFIX
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"job_model_id,expected_endpoint",
|
||||
[
|
||||
("amazon.titan-embed-text-v2:0", "/v1/embeddings"),
|
||||
("anthropic.claude-3-5-sonnet-20240620-v1:0", "/v1/chat/completions"),
|
||||
],
|
||||
)
|
||||
def test_job_status_endpoint_follows_the_job_models_mode(
|
||||
patched_boto3: tuple[MagicMock, MagicMock], job_model_id: str, expected_endpoint: str
|
||||
) -> None:
|
||||
fake_client, _ = patched_boto3
|
||||
fake_client.get_model_invocation_job.return_value = {**_fake_boto3_response(), "modelId": job_model_id}
|
||||
|
||||
batch: Final = BedrockBatchesHandler.handle_model_invocation_job_status(batch_id=JOB_ARN)
|
||||
|
||||
assert batch.endpoint == expected_endpoint
|
||||
|
||||
|
||||
@pytest.mark.parametrize("success_count,error_count", [(100, 0), (86, 14)])
|
||||
def test_completed_job_maps_provider_record_counts(patched_boto3, success_count, error_count):
|
||||
fake_client, _ = patched_boto3
|
||||
|
|
|
|||
|
|
@ -14,11 +14,14 @@ URL/ARN handling, and the error class. AWS auth/sigv4 is the only external seam
|
|||
we mock; everything else runs for real.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
|
|
@ -805,6 +808,49 @@ def test_retrieve_response_no_errors_when_no_message(config):
|
|||
assert out.errors is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"job_model_id,expected_endpoint",
|
||||
[
|
||||
("amazon.titan-embed-text-v2:0", "/v1/embeddings"),
|
||||
("arn:aws:bedrock:us-west-2::foundation-model/amazon.titan-embed-text-v2:0", "/v1/embeddings"),
|
||||
("anthropic.claude-3-5-sonnet-20240620-v1:0", "/v1/chat/completions"),
|
||||
("a-model-the-cost-map-does-not-list", "/v1/chat/completions"),
|
||||
(None, "/v1/chat/completions"),
|
||||
],
|
||||
)
|
||||
def test_retrieve_response_endpoint_follows_the_job_models_mode(
|
||||
config: BedrockBatchesConfig, job_model_id: str | None, expected_endpoint: str
|
||||
) -> None:
|
||||
out: Final = config.transform_retrieve_batch_response(
|
||||
model=None,
|
||||
raw_response=_raw({"jobArn": ARN, "status": "InProgress", "modelId": job_model_id}),
|
||||
logging_obj=MagicMock(),
|
||||
litellm_params={},
|
||||
)
|
||||
assert out.endpoint == expected_endpoint
|
||||
|
||||
|
||||
def _is_bedrock_embedding_model(cost_map_entry: Mapping[str, object]) -> bool:
|
||||
return cost_map_entry.get("litellm_provider") == "bedrock" and cost_map_entry.get("mode") == "embedding"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"job_model_id",
|
||||
sorted(model_id for model_id, entry in litellm.model_cost.items() if _is_bedrock_embedding_model(entry)),
|
||||
)
|
||||
def test_retrieve_response_endpoint_survives_a_batch_deployment_overriding_the_models_mode(
|
||||
config: BedrockBatchesConfig, monkeypatch: pytest.MonkeyPatch, job_model_id: str
|
||||
) -> None:
|
||||
monkeypatch.setitem(litellm.model_cost, job_model_id, {**litellm.model_cost[job_model_id], "mode": "batch"})
|
||||
out: Final = config.transform_retrieve_batch_response(
|
||||
model=None,
|
||||
raw_response=_raw({"jobArn": ARN, "status": "InProgress", "modelId": job_model_id}),
|
||||
logging_obj=MagicMock(),
|
||||
litellm_params={},
|
||||
)
|
||||
assert out.endpoint == "/v1/embeddings"
|
||||
|
||||
|
||||
def test_retrieve_response_enriches_metadata(config):
|
||||
body = {
|
||||
"jobArn": ARN,
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ from contextlib import ExitStack
|
|||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Dict, Final, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -1841,6 +1842,66 @@ 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: str) -> Callable[..., LiteLLMBatch]:
|
||||
def hook(*, response: LiteLLMBatch, **_: object) -> LiteLLMBatch:
|
||||
response.id = batch_id
|
||||
return response
|
||||
|
||||
return hook
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve__provider_sync_keeps_the_endpoint_the_batch_was_created_for(retrieve_harness):
|
||||
stored: Final = make_batch(id="batch-unified-blob", status="validating").model_copy(update={"endpoint": "/v1/embeddings"})
|
||||
provider_batch: Final = 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 = 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: Final = 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: Final = make_batch(id="batch-unified-blob", status="validating").model_copy(update={"endpoint": ""})
|
||||
provider_batch: Final = 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: Final = 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
|
||||
async def test_retrieve__an_embeddings_endpoint_from_the_provider_corrects_a_saved_chat_endpoint(retrieve_harness):
|
||||
stored: Final = make_batch(id="batch-unified-blob", status="validating").model_copy(
|
||||
update={"endpoint": "/v1/chat/completions"}
|
||||
)
|
||||
provider_batch: Final = 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: Final = 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):
|
||||
|
|
@ -2595,6 +2656,27 @@ async def test_cancel__db_write_receives_caller_auth(cancel_harness):
|
|||
assert cancel_harness.update_batch_in_db.call_args.kwargs["user_api_key_dict"] is caller
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel__keeps_the_endpoint_the_batch_was_created_for(cancel_harness):
|
||||
stored_row: Final = MagicMock(status="in_progress")
|
||||
stored: Final = make_batch(id="batch-unified-blob", status="in_progress").model_copy(update={"endpoint": "/v1/embeddings"})
|
||||
provider_batch: Final = 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),
|
||||
patch.object(endpoints, "get_batch_from_database", AsyncMock(return_value=(stored_row, stored))),
|
||||
):
|
||||
resp: Final = await call_cancel(cancel_harness, "batch-unified-blob")
|
||||
|
||||
assert (resp.endpoint, resp.status) == ("/v1/embeddings", "cancelling")
|
||||
written: Final = 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
|
||||
async def test_cancel__unified_missing_model_id_400(cancel_harness):
|
||||
# unified id with no model_id segment -> get_model_id returns None -> 400.
|
||||
|
|
|
|||
|
|
@ -3662,6 +3662,30 @@ class TestMultiPodBatchCostClaim:
|
|||
assert row.batch_processed is True
|
||||
await self._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("has_batch_processed_column", (True, False))
|
||||
async def test_billing_a_batch_keeps_the_endpoint_it_was_created_for(self, has_batch_processed_column: bool) -> None:
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
row: Final = _FakeManagedObjectRow()
|
||||
row.file_object = LiteLLMBatch(
|
||||
id="batch-456",
|
||||
status="in_progress",
|
||||
input_file_id="file-input-1",
|
||||
endpoint="/v1/embeddings",
|
||||
completion_window="24h",
|
||||
created_at=1,
|
||||
object="batch",
|
||||
).model_dump_json()
|
||||
instance: Final = self._instance(self._prisma(row, []), self._router())
|
||||
instance._has_batch_processed_column = has_batch_processed_column
|
||||
|
||||
with self._billing_patches([]):
|
||||
await instance.check_batch_cost()
|
||||
|
||||
assert row.status == "complete"
|
||||
assert json.loads(row.file_object)["endpoint"] == "/v1/embeddings"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_schema_without_batch_processed_still_bills(self):
|
||||
"""Older schemas have no column to claim, so they keep the pre-fix behavior instead
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue