mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
feat(batches): store each batch line's request and response in S3 and GCS (#41691)
* feat(batches): emit per-line JSONL batch records to callbacks
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): only emit line items for final batches
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* chore(deps): bump soupsieve to 2.9.2 for GHSA-gjv8-xp57-g29c and GHSA-j934-xhv5-fg8f
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* style(batches): satisfy ruff format for line item logging
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* chore(ui): regenerate schema.d.ts for store_batch_line_items_in_callbacks
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(batches): cover line-item edge shapes and general settings propagation
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): keep deployment credentials out of callback-visible params
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(batches): declare credential transport attribute on Logging
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(batches): integration contract for per-line batch callback events
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): resolve store_batch_line_items_in_callbacks through the settings store
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(batches): assert per-line events beside the aggregate and create events
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): drop resurrected use_legacy_interactions_schema shim from merge
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(batches): assert line item events carry no credentials
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): reset batch line-item flag when the setting is absent
Removing store_batch_line_items_in_callbacks from the DB general settings left the module flag at its previous value because the effect early-returned when the key was neither in the DB nor owned by YAML. The proxy owns the global, so an absent resolved value now writes False.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): map anthropic batch lines through the anthropic message transformer
Anthropic result lines carry a message-shaped body, so ModelResponse(**body) produced empty choices and a usage key that blocked the batch-rate usage attach. Extract the existing AnthropicResponse -> ModelResponse conversion used for /v1/messages logging into anthropic_message_to_model_response and reuse it for anthropic line items.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* style: tag the anthropic result cast with a cast-ok reason
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* style: shorten cast-ok reason so ruff format passes
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): skip line-item callbacks for providers without a line reconstruction
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): reconstruct bedrock and mistral batch lines through their provider transformations
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* style: wrap widened provider literals so ruff format passes
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): pair bedrock lines by recordId and type them from the reconstructed response
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): keep line item events out of proxy limiters and strip credentials from child params
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): keep cast-ok tags on the flagged cast lines for the discipline gate
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): skip failed line item events in parallel request limiter failure hooks
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(batches): tolerate lifecycle spend rows when asserting no line item rows
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): place mutable-ok tags on the lines the discipline gate flags
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): skip router budget accounting for batch line item events
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): keep anthropic line errors and skip native vertex batch rows
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(batches): restore test-quality tags on line item tests
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test: move batch line-item unit test to tests/unit after merging main
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): emit batch line-item callbacks once per batch
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): release line-item claim when fan-out emitted nothing
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): make the line-item claim release ownership-safe
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): keep bodyless failed batch lines reaching callbacks
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(logging): seed batch line hidden params in get_hidden_params defaults
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* perf(batches): reuse fetched batch result files for line-item logging
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): release line-item claim when nothing was emitted
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): type batch line-item callback sites to satisfy basedpyright budget
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* style: wrap batch hidden-params cast for ruff format
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): reason suppressions for batch line-item lint deltas
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* chore: drop stale mutable-ok suppressions after LIT002 removal
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(batches): skip batch line-item events in ClickHouse spend sink and router quota counters (#44486)
* fix(batches): skip batch line-item events in ClickHouse spend sink and router quota counters
Batch line-item callbacks (call_type=acompletion + litellm_params.batch_parent_id)
bypassed the call-type-only batch guards:
- ClickHouseSpendLogger wrote one spend row per line on top of the aggregate
aretrieve_batch row, billing every batch twice with the flag on.
- Router deployment callbacks and the routing-strategy usage handlers charged the
historical batch's tokens/rpm to the live TPM/RPM counters, so after
Router.aretrieve_batch the next chat raised RouterRateLimitError.
Skip these events wherever the aggregate batch retrieve is already skipped.
* fix(logging): keep internal litellm_params authoritative in model_call_details
A request body key named litellm_params lands in optional_params, and
update_environment_variables merged **optional_params after the internal
entry, so a caller could overwrite the logging kwargs' litellm_params and
forge markers like batch_parent_id. Live-reproduced: a Gemini chat whose
body carried {"litellm_params": {"batch_parent_id": ...}} returned 200
with no ClickHouse spend row. Re-assert the internal litellm_params after
the merge so every sink (ClickHouse, Postgres, router quota counters)
reads the same trusted value, and assert the aggregate aretrieve_batch
row still writes in the ClickHouse line-item test (review finding).
* fix(clickhouse): skip failed batch line items in the spend sink too
A failed JSONL line is part of the aggregate aretrieve_batch row's
request counts; writing it its own spend_logs row inflated request
counts with the line-item flag on. Extend the batch line-item guard to
the failure handler (review finding) and pin both directions: a failed
line item writes no row, a genuine live failure still does.
* fix(batches): declare batch_line_items service family on claim calls
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(integrations): keep batch line items out of the built-in metering sinks
With store_batch_line_items_in_callbacks on, per-line batch events carry a
real response_cost and reached six metering sinks that lacked the
is_batch_line_item_event guard, so a completed batch metered aggregate +
per-line (roughly 2x true spend): prometheus (litellm_spend_metric and
request counters), openmeter, lago, datadog cost management (FOCUS
BilledCost), newrelic metrics, and the OTel v1/v2 gen_ai.usage.cost
metrics. Skip line-item events in each sink's success/failure metering
entry points, matching the guards already on the Postgres and ClickHouse
spend sinks. OTel spans for line items are still emitted; only their
cost/token metrics are skipped. Tests drive real sink instances on the
real dispatch lists through a completed 2-line batch and assert each
sink meters exactly the aggregate.
* style: format metering-sinks test
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(otel): skip synthetic line-item durations in v2 metrics; isolate prometheus test
Follow-ups from review on the sink guards: record_failure() now skips
batch line items too (a failed line's interval measures retrieval and
emission, not its model call), and the prometheus guard test snapshots
and restores the process-wide registry collectors instead of leaving
them unregistered. The OTel v2 test is now a flag on/off A/B over the
whole metric snapshot: cost and usage values must match exactly and no
metric may gain datapoints with the flag on.
* chore: finish main merge conflict resolution for batch line items
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* style: sort imports in litellm_logging after main merge
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(batches): use the public hook class names main now exports
* fix(batches): only the batch writer can mark a callback event as a line item
amoderation, transcription and speech build litellm_params from every request kwarg,
so a body field named batch_parent_id reached Logging.litellm_params and every
spend, budget and rate-limit hook skipped the request. update_environment_variables
now drops an incoming batch_parent_id; the line-item writer sets it through
Logging.mark_batch_line_item. Batch line hidden params are emitted only when set.
* chore(batches): keep the lint gates within main's ceilings after the merge
Typed the helpers main now types (Mapping litellm_params, Any Span), replaced the
banned cast with a literal match, dropped isinstance checks main made redundant,
and moved test-quality reasons onto the lines the gate reads.
* fix(anthropic): read the message id off the validated message
The anthropic_messages logging path can hand over an already parsed message object,
which has no .get; main read its id through getattr, so take it from the validated
AnthropicResponse instead.
* fix(batches): an unreadable error file no longer drops the output line events
Costing counts an error file it could not fetch as no failed lines; line logging
refetched it and the raise aborted the whole fan-out, so the successful lines were
never emitted. The error file now reads as empty with a warning.
* test(batches): pin billing for batch_parent_id on every endpoint and value shape
Chat completions, embeddings, responses and messages join moderations, speech and
transcriptions, and speech is driven with an empty, int, list, object and 5 KB value.
* test(batches): bill every request of a concurrent burst mixing forged and plain bodies
* revert: keep the line-item fan-out all-or-nothing on an unreadable error file
85fc4498cc emitted the output lines and kept the batch claim when the error file
could not be read, so its failed lines were never retried. Raising releases the
claim and the next retrieve retries the whole fan-out.
* test(batches): pin the all-or-nothing retry and warm the burst model first
An unreadable error file emits no lines and the next retrieve emits all of them.
The burst cell drives one speech model after a warm-up request so a cold worker
cannot fail it, and waits for the batched key spend.
* fix(batches): run the line fan-out outside the aggregate event's deadline
The logging worker cancels a callback coroutine after its per-callback
deadline, and the line fan-out ran inside the aggregate aretrieve_batch event
before its callbacks, so a large batch or a slow sink could cancel the
aggregate billing event. The fan-out now runs in its own task, built from a
copy of the parent's litellm_params taken before the aggregate callbacks
write their request-local markers.
Anthropic keeps no input file and reports input_file_id "None", which the
fan-out fetched and failed on, so no Anthropic line ever arrived. That id is
now treated as no input file.
* fix(batches): let running line fan-outs finish before proxy shutdown closes their clients
A started fan-out keeps its batch claim, so a shutdown that cut it short
skipped the unsent lines until the claim expired. proxy_shutdown_event now
waits up to MAX_TIME_TO_CLEAR_QUEUE for them first. Also replace a test's
type: ignore with an isinstance check.
* fix(batches): keep batch line fields off callback kwargs and skip fan-outs at shutdown
Line bodies were spread into the callback kwargs, so a line's metadata could settle router rate limit reservations; they now stay in optional_params, and the router rate limit and prompt caching hooks skip line events.
A retrieve logged after the shutdown drain, or by the logging worker's exit flush, no longer starts a fan-out its loop would destroy while it holds the batch claim.
The output file is fetched before the error file again.
* fix(batches): log batch lines under the retrieve's callback settings and keep them out of alerting and shadow eval
Line events inherit the retrieve request's key and team callback settings, so turn_off_message_logging and per-team logging credentials apply to them, and they keep the system prompt in the logged messages.
Shadow eval, Datadog metrics, Slack daily reports and PagerDuty skip line events, since the batch API served them and the aggregate event already counts the batch.
* test(batches): line events follow the key's message logging opt-out
* test(batches): line events follow team logging settings, keep the line system prompt and survive a restart once
* fix(batches): keep in-memory line claims for 10,000 batches and pin managed batch emission across poller and retrieves
* Keep batch line unit tests in-process and drop repeated per-sink comments
* Share one batch provider alias across batch_utils and line item logging
* Report Bedrock line status codes, skip unknown Bedrock lines without raising, fetch the input file after the native Vertex check
* Read a swallowed Redis claim failure as unavailable, price unpriceable lines at zero, keep the OTel failure docstring
* Call handle_completed_batch when line items are off, keep a cache test from leaking Redis into the line claim cache, pin an unpriceable line's zero cost on the callback kwargs
* Hand the cost poller's credential snapshot to line item file reads, so a managed Bedrock batch's lines read its input from the deployment bucket
* Log batch lines through the sync callbacks a direct async call runs, keep failed lines out of router cooldown counts, and emit an error line whose body is text
* fix(batches): log failed Bedrock embedding lines as embeddings with their input
A Bedrock batch input record carries no url, so a failed Titan v2 line fell back to acompletion and its
inputText never reached the event messages. The failed line now takes aembedding from the deployment
model and the record's inputText is the line input
* fix(batches): keep a Bedrock line's system prompt in its logged messages
Bedrock batch input records hold the system prompt as content blocks, which logging only reads off the kwargs as text, so Claude and Nova lines dropped it
* feat(batches): store per-line batch requests and responses in S3/GCS via batch_line_item_callbacks
Replace the per-line logging path with a dedicated hook. When the hourly batch cost
poller finishes a managed batch, it pairs every output and error line with its input
line, prices each line, and hands one payload to every logger listed in
litellm_settings.batch_line_item_callbacks. s3_v2 and gcs_bucket write it as one
NDJSON object per batch under batch_line_items/<day>/<batch id>.jsonl.
Bodies are dropped and the header kept when the creator key or team, the global
setting, or the logger itself turns off message logging, and when the creator can
no longer be found. A failed upload is logged and never blocks spend tracking.
* fix(batches): store batch line items after the batch is billed and retired
Send a billed batch's lines only after its row is marked complete, so a slow or
failing sink never holds a claimed batch open. Skip reading the input file when
the bodies are dropped anyway, still send the error lines of a batch whose output
file is gone, and end every NDJSON record with a newline.
* test(batches): cover team-level and deleted-key redaction of batch line items
* fix(batches): store batch line items in the background so a hung bucket never holds up billing
The poller now retires a batch first and hands its line items to a per-worker background store that runs one batch at a time. A slow, hung or failing bucket no longer delays billing or retiring the other batches in the poll cycle. While a poll cycle's worth of batches is still waiting to store, the next batch's lines are dropped with an error.
* fix(batches): store an unbilled batch's line items only once it is retired
When retiring an expired, cancelled or all-failed batch fails, the batch stays on the poll page, so its lines are no longer read and stored again on every poll cycle until it is retired. The poller also releases a billed batch's output bytes before reading the next batch's output, as it did before line items existed.
* fix(batches): price and write batch line items off the event loop
Pricing every line of a large batch and serializing its object ran on the
proxy's event loop, so a 50k-line batch stalled the worker for seconds.
Both now run in a worker thread. An object the bucket refuses for good is
no longer kept in s3_v2's drop marks, which only a log-queue flush clears.
Vertex AI input rows carry the custom_id in their key or request labels, so
their lines now get the request too.
* fix(batches): pair no request for a shared custom_id and wait for an unclaimed batch's retry
A Vertex AI embeddings entry whose input is a list is split into one row per
item, all with the same custom_id, so its line now gets no request instead of
the last row. On a database without batch_processed nothing claims a billed
batch, so when marking it complete fails it stays in the poll page and its
lines now wait for the next poll instead of being written on every cycle.
* fix(batches): store batch line items only once the poll has billed every batch
Pricing a large batch's lines on a worker thread still competes with the poll for the GIL, so the next batches
in the same poll were billed and retired seconds later. The background stores now wait until the poll cycle that
queued them has finished.
* fix(batches): hash line-item objects off the event loop and keep lines with a lone surrogate
A batch's whole line-items object was hashed for its S3 PUT on the event loop, and a line
holding a lone surrogate from a provider's JSON escapes made the object impossible to
encode as UTF-8, so the batch's object was never written. The object is now prepared in a
worker thread, and only the line that cannot be encoded is written with escaped characters.
* fix(batches): keep request and response for batches created without a virtual key
A batch created with a JWT has no API key on its row, and its line items were stored without
request and response as if its creator had been deleted. Such a batch now follows its team's
and the proxy's message logging settings, the same as a master key batch. A deleted key or team
or a failed lookup still drops the bodies.
* fix(batches): pair line items with the right request on every deployment and skip empty objects
A batch retired with nothing billable now reads its input file by the id the provider reported, so a batch that
ran on the second deployment of a model group keeps each line's request. A Vertex AI line read as custom_id
"unknown" pairs with no input row, and a batch whose files hold no readable line writes no object
* fix(batches): let OAuth2 batches follow team and proxy message logging settings
The line-item redaction check looked up every saved creator credential other than the master key or a CLI session as a virtual key. An OAuth2 batch saves the raw token, so it read like a deleted key and lost its request and response
It now looks up only a saved value shaped like a virtual key's hash, reusing is_valid_sha256_hash, so OAuth2 batches follow their team's logging settings and the proxy-wide setting while deleted keys still fail closed
* fix(batches): hash the creator credential and write batch line items in parts
The creator credential reaches line metadata only as a sha256 hash, so an
OAuth2 token or the master key alias is never written to the bucket. A queued
store keeps only whether the batch was billed and downloads the output file
again at its turn, logging an error naming the batch if that fails. Lines are
built and uploaded one part at a time, each closing once its raw request and
result bytes reach BATCH_LINE_ITEMS_MAX_PART_BYTES (16 MiB by default), with
the input file indexed by byte span instead of parsed up front.
* docs(batches): say async_log_batch_line_items is called once per part
* fix(batches): log an error when a batch's deployment is gone before its line items are stored
---------
Co-authored-by: shivam <shivam@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: yucheng <yucheng@berri.ai>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
1f041a3265
commit
30ada9d26a
17 changed files with 2850 additions and 14 deletions
|
|
@ -2,7 +2,8 @@
|
|||
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
import asyncio
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import replace as dataclasses_replace
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
|
|
@ -13,15 +14,22 @@ from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
|
|||
ManagedObjectRow,
|
||||
)
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import (
|
||||
BATCH_LINE_ITEMS_MAX_PART_BYTES,
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
|
||||
MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
)
|
||||
from litellm.repositories.table_repositories import ManagedObjectRepository
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.verification_token_repository import VerificationTokenRepository
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.batches.batch_line_items import BatchFileProvider
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
|
|
@ -67,6 +75,19 @@ def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[Manage
|
|||
return ManagedObjectRepository(prisma_client).table
|
||||
|
||||
|
||||
def _token_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_VerificationToken]":
|
||||
return VerificationTokenRepository(prisma_client).table
|
||||
|
||||
|
||||
def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_TeamTable]":
|
||||
return TeamRepository(prisma_client).table
|
||||
|
||||
|
||||
def _row_metadata(row: object) -> dict[str, object]:
|
||||
metadata: Final = getattr(row, "metadata", None)
|
||||
return cast(dict[str, object], metadata) if isinstance(metadata, dict) else {} # cast-ok: prisma Json object column
|
||||
|
||||
|
||||
class CheckBatchCost:
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -87,6 +108,10 @@ class CheckBatchCost:
|
|||
# the guaranteed-failing primary query on every subsequent cycle.
|
||||
self._has_batch_processed_column: bool = True
|
||||
self.batch_processed_support_confirmed: bool = False
|
||||
self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
|
||||
self._line_item_store_turn: Final = asyncio.Lock()
|
||||
self._poll_idle: Final = asyncio.Event()
|
||||
self._poll_idle.set()
|
||||
|
||||
@staticmethod
|
||||
def _is_missing_batch_processed_column_error(err: Exception) -> bool:
|
||||
|
|
@ -299,9 +324,9 @@ class CheckBatchCost:
|
|||
|
||||
async def _finalize_unbilled_terminal_job(
|
||||
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
|
||||
) -> None:
|
||||
) -> bool:
|
||||
"""Persist a terminal batch that has nothing billable, converting any raw
|
||||
provider file ids to managed ids, and take it out of the poll page."""
|
||||
provider file ids to managed ids, and take it out of the poll page. Returns whether it was persisted."""
|
||||
try:
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_is_base64_encoded_unified_file_id,
|
||||
|
|
@ -329,10 +354,12 @@ class CheckBatchCost:
|
|||
verbose_proxy_logger.info(
|
||||
f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
|
||||
)
|
||||
return True
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}"
|
||||
)
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _record_error(
|
||||
|
|
@ -794,7 +821,191 @@ class CheckBatchCost:
|
|||
|
||||
return model_name, str(llm_provider) if llm_provider else None
|
||||
|
||||
async def check_batch_cost(self):
|
||||
async def _retire_unbilled_terminal_job(
|
||||
self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str
|
||||
) -> None:
|
||||
"""Persist a terminal batch that has nothing billable, then store its lines, reading its input file by the id
|
||||
the provider reported rather than the managed id that was persisted."""
|
||||
provider_input_file_id: Final = response.input_file_id
|
||||
if await self._finalize_unbilled_terminal_job(job, response):
|
||||
response.input_file_id = provider_input_file_id
|
||||
self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=False)
|
||||
|
||||
def _store_batch_line_items_in_background(
|
||||
self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
|
||||
) -> None:
|
||||
"""Store a retired batch's lines off the poll loop, one batch at a time and only between polls, so neither
|
||||
pricing the lines nor a slow or hung bucket holds up billing or retiring the other batches. While a poll
|
||||
cycle's worth of batches is still waiting, the lines of the next one are dropped with an error."""
|
||||
if not litellm.batch_line_item_callbacks:
|
||||
return
|
||||
if len(self._line_item_stores) >= MAX_OBJECTS_PER_POLL_CYCLE:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: {len(self._line_item_stores)} batches are still waiting to store their line items, "
|
||||
f"so the line items of batch {batch_id} (job {job.id}) were not stored"
|
||||
)
|
||||
return
|
||||
store: Final = asyncio.create_task(
|
||||
self._store_batch_line_items_in_turn(job, response, model_id, batch_id, billed=billed)
|
||||
)
|
||||
self._line_item_stores.add(store)
|
||||
store.add_done_callback(self._line_item_stores.discard)
|
||||
|
||||
async def _store_batch_line_items_in_turn(
|
||||
self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
|
||||
) -> None:
|
||||
async with self._line_item_store_turn:
|
||||
await self._poll_idle.wait()
|
||||
await self._store_batch_line_items(job, response, model_id, batch_id, billed=billed)
|
||||
|
||||
async def _store_batch_line_items(
|
||||
self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
|
||||
) -> None:
|
||||
"""Send every line of a finished batch to batch_line_item_callbacks, one part at a time. A billed batch's
|
||||
output file is read again here, so nothing waiting for its turn holds a batch's files. Failures, including an
|
||||
output file that can no longer be read, are logged and never raised, so the batch stays billed and retired
|
||||
whether or not its lines were stored."""
|
||||
if not litellm.batch_line_item_callbacks:
|
||||
return
|
||||
from litellm.batches.batch_line_items import (
|
||||
BatchFileProvider,
|
||||
BatchLineItemSource,
|
||||
batch_line_item_parts,
|
||||
dispatch_batch_line_items,
|
||||
line_item_metadata,
|
||||
)
|
||||
from litellm.batches.batch_utils import fetch_batch_managed_file_content
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info
|
||||
|
||||
try:
|
||||
deployment_info: Final = self.llm_router.get_deployment(model_id=model_id)
|
||||
if deployment_info is None:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: deployment {model_id} is no longer in the router, so the line items of batch "
|
||||
f"{batch_id} (job {job.id}) were not stored"
|
||||
)
|
||||
return
|
||||
litellm_model_name: Final = deployment_info.litellm_params.model
|
||||
model_name, llm_provider, _, _ = get_llm_provider(
|
||||
model=litellm_model_name,
|
||||
custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
|
||||
)
|
||||
cost_provider: Final = cast(BatchFileProvider, llm_provider)
|
||||
credentials: Final = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
|
||||
file_access: Final = {
|
||||
**credentials,
|
||||
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
|
||||
}
|
||||
output_content: Final = (
|
||||
await fetch_batch_managed_file_content(
|
||||
response.output_file_id, custom_llm_provider=cost_provider, litellm_params=file_access
|
||||
)
|
||||
if billed and response.output_file_id
|
||||
else b""
|
||||
)
|
||||
error_content: Final = await self._batch_file_content(response.error_file_id, cost_provider, file_access)
|
||||
if not output_content and not error_content:
|
||||
return
|
||||
redact_bodies: Final = await self._creator_turned_off_message_logging(job)
|
||||
source: Final = BatchLineItemSource(
|
||||
batch_id=job.unified_object_id,
|
||||
finished_at=response.completed_at
|
||||
or response.failed_at
|
||||
or response.expired_at
|
||||
or response.cancelled_at
|
||||
or response.created_at,
|
||||
model=model_name,
|
||||
custom_llm_provider=llm_provider,
|
||||
cost_provider=cost_provider,
|
||||
model_info=deployment_pricing_model_info(model_id=model_id, deployment_model=litellm_model_name),
|
||||
metadata=line_item_metadata(await self._creator_attribution.batch_metadata(job, batch_id)),
|
||||
)
|
||||
input_content: Final = (
|
||||
b""
|
||||
if redact_bodies
|
||||
else await self._batch_file_content(response.input_file_id, cost_provider, file_access)
|
||||
)
|
||||
parts: Final = batch_line_item_parts(
|
||||
source, input_content, (output_content, error_content), BATCH_LINE_ITEMS_MAX_PART_BYTES
|
||||
)
|
||||
handed = 0
|
||||
# pricing a part's lines is CPU-bound, so each part is built off the event loop
|
||||
while (payload := await asyncio.to_thread(next, parts, None)) is not None:
|
||||
await dispatch_batch_line_items(payload, redact_bodies=redact_bodies)
|
||||
handed += len(payload["lines"])
|
||||
if handed:
|
||||
verbose_proxy_logger.info(
|
||||
f"CheckBatchCost: handed the {handed} line items of batch {batch_id} to batch_line_item_callbacks"
|
||||
)
|
||||
except Exception as error: # noqa: BLE001 # storing the lines is best effort, the batch is already billed
|
||||
verbose_proxy_logger.exception(
|
||||
f"CheckBatchCost: the line items of batch {batch_id} (job {job.id}) were not stored: {error}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _batch_file_content(
|
||||
file_id: str | None, custom_llm_provider: "BatchFileProvider", file_access: Mapping[str, object]
|
||||
) -> bytes:
|
||||
"""A batch file's raw JSONL, or nothing when the batch has no such file or the provider will not serve it."""
|
||||
from litellm.batches.batch_utils import fetch_batch_managed_file_content
|
||||
|
||||
if not file_id:
|
||||
return b""
|
||||
try:
|
||||
return await fetch_batch_managed_file_content(
|
||||
file_id, custom_llm_provider=custom_llm_provider, litellm_params=file_access
|
||||
)
|
||||
except Exception as error: # noqa: BLE001 # a missing input or error file leaves those lines out, not the rest
|
||||
verbose_proxy_logger.warning(f"CheckBatchCost: could not read batch file {file_id}: {error}")
|
||||
return b""
|
||||
|
||||
async def _creator_turned_off_message_logging(self, job: "ManagedObjectRow") -> bool:
|
||||
"""Whether the creating key's or team's logging settings turn off message logging, falling back to the
|
||||
proxy-wide setting. A batch whose saved credential is not a virtual key's hash, such as one created under JWT
|
||||
or OAuth2 auth or with the master key, has only its team's and the proxy's settings. A key or team that can no
|
||||
longer be read counts as turned off."""
|
||||
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import get_dynamic_logging_metadata
|
||||
from litellm.proxy.proxy_server import proxy_config
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
api_key: Final = getattr(job, "api_key", None)
|
||||
team_id: Final = getattr(job, "team_id", None)
|
||||
has_key_row: Final = api_key is not None and is_valid_sha256_hash(api_key)
|
||||
try:
|
||||
key_row: Final = (
|
||||
await _token_table(self.prisma_client).find_unique(where={"token": api_key}) if has_key_row else None
|
||||
)
|
||||
team_row: Final = (
|
||||
await _team_table(self.prisma_client).find_unique(where={"team_id": team_id}) if team_id else None
|
||||
)
|
||||
except Exception as error:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: could not read the logging settings of batch job {job.id}'s creator, "
|
||||
f"so its line items are stored without request and response: {error}"
|
||||
)
|
||||
return True
|
||||
if (has_key_row and key_row is None) or (team_id and team_row is None):
|
||||
return True
|
||||
settings: Final = get_dynamic_logging_metadata(
|
||||
UserAPIKeyAuth(metadata=_row_metadata(key_row), team_metadata=_row_metadata(team_row), team_id=team_id),
|
||||
proxy_config,
|
||||
)
|
||||
callback_vars: Final = (settings.callback_vars if settings is not None else None) or {}
|
||||
creator_setting: Final = str_to_bool(callback_vars.get("turn_off_message_logging"))
|
||||
return creator_setting if creator_setting is not None else litellm.turn_off_message_logging is True
|
||||
|
||||
async def check_batch_cost(self) -> None:
|
||||
"""Run one poll cycle. Batch line items wait until it has billed and retired every batch it found."""
|
||||
self._poll_idle.clear()
|
||||
try:
|
||||
await self._poll_batches()
|
||||
finally:
|
||||
self._poll_idle.set()
|
||||
|
||||
async def _poll_batches(self):
|
||||
"""
|
||||
Check if the batch JOB has been tracked.
|
||||
- get all status="validating" and file_purpose="batch" jobs
|
||||
|
|
@ -908,7 +1119,7 @@ class CheckBatchCost:
|
|||
f"CheckBatchCost: output file {response.output_file_id} of batch {batch_id} "
|
||||
f"does not exist at the provider; retiring job {job.id} unbilled"
|
||||
)
|
||||
await self._finalize_unbilled_terminal_job(job, response)
|
||||
await self._retire_unbilled_terminal_job(job, response, model_id, batch_id)
|
||||
continue
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to track cost for batch {batch_id} "
|
||||
|
|
@ -938,6 +1149,10 @@ class CheckBatchCost:
|
|||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
|
||||
)
|
||||
if not self._has_batch_processed_column:
|
||||
# nothing claimed the row, so it stays in the poll page and its lines wait for that retry
|
||||
continue
|
||||
self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True)
|
||||
|
||||
elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
|
|
@ -950,7 +1165,7 @@ class CheckBatchCost:
|
|||
f"has not appeared yet; leaving job {job.id} for the next poll cycle"
|
||||
)
|
||||
continue
|
||||
await self._finalize_unbilled_terminal_job(job, response)
|
||||
await self._retire_unbilled_terminal_job(job, response, model_id, batch_id)
|
||||
|
||||
# Record polling run metrics (always, even if nothing was processed)
|
||||
if prom_logger:
|
||||
|
|
|
|||
|
|
@ -118,6 +118,7 @@ success_callback: List[CALLBACK_TYPES] = []
|
|||
failure_callback: List[CALLBACK_TYPES] = []
|
||||
service_callback: List[CALLBACK_TYPES] = []
|
||||
audit_log_callbacks: List[CALLBACK_TYPES] = []
|
||||
batch_line_item_callbacks: Sequence[CALLBACK_TYPES] = ()
|
||||
# logging_callback_manager is lazy-loaded via __getattr__
|
||||
_custom_logger_compatible_callbacks_literal = Literal[
|
||||
"lago",
|
||||
|
|
|
|||
250
litellm/batches/batch_line_items.py
Normal file
250
litellm/batches/batch_line_items.py
Normal file
|
|
@ -0,0 +1,250 @@
|
|||
"""Per-line records of a finished batch for the loggers listed in batch_line_item_callbacks."""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import re
|
||||
from collections import Counter
|
||||
from collections.abc import Iterable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from itertools import accumulate, chain, count
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.batches.batch_utils import batch_cost_and_usage, iter_batch_input_lines, parse_batch_output_line
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import init_custom_logger_compatible_class, is_valid_sha256_hash
|
||||
from litellm.llms.vertex_ai.files.transformation import vertex_batch_row_custom_id
|
||||
from litellm.types.utils import (
|
||||
ModelInfo,
|
||||
StandardBatchLineItem,
|
||||
StandardBatchLineItemMetadata,
|
||||
StandardBatchLineItemsPayload,
|
||||
StandardBatchLineItemUsage,
|
||||
)
|
||||
|
||||
BatchFileProvider: TypeAlias = Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"]
|
||||
|
||||
_VERTEX_AI_UNKNOWN_CUSTOM_ID: Final = "unknown" # what a Vertex AI batch row that carries no custom_id is read as
|
||||
_LINE: Final = re.compile(rb"[^\n]+")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BatchLineItemSource:
|
||||
batch_id: str
|
||||
finished_at: int
|
||||
model: str
|
||||
custom_llm_provider: str | None
|
||||
cost_provider: BatchFileProvider
|
||||
model_info: ModelInfo | None
|
||||
metadata: StandardBatchLineItemMetadata
|
||||
|
||||
|
||||
def line_item_metadata(attribution: Mapping[str, object]) -> StandardBatchLineItemMetadata:
|
||||
return StandardBatchLineItemMetadata(
|
||||
user_api_key_hash=_creator_key_hash(attribution),
|
||||
user_api_key_alias=_text(attribution, "user_api_key_alias"),
|
||||
user_api_key_team_id=_text(attribution, "user_api_key_team_id"),
|
||||
user_api_key_team_alias=_text(attribution, "user_api_key_team_alias"),
|
||||
user_api_key_user_id=_text(attribution, "user_api_key_user_id"),
|
||||
user_api_key_org_id=_text(attribution, "user_api_key_org_id"),
|
||||
)
|
||||
|
||||
|
||||
def _text(attribution: Mapping[str, object], name: str) -> str | None:
|
||||
value: Final = attribution.get(name)
|
||||
return value if isinstance(value, str) else None
|
||||
|
||||
|
||||
def _creator_key_hash(attribution: Mapping[str, object]) -> str | None:
|
||||
"""The creating credential as a SHA-256 hash. A virtual key is saved hashed already, while a credential saved as
|
||||
is, such as an OAuth2 access token, is hashed here so the object never holds it."""
|
||||
credential: Final = _text(attribution, "user_api_key_hash")
|
||||
if credential is None or is_valid_sha256_hash(credential):
|
||||
return credential
|
||||
return hashlib.sha256(credential.encode()).hexdigest()
|
||||
|
||||
|
||||
def _line_key(entry: Mapping[str, object]) -> str | None:
|
||||
"""A line's custom_id. Bedrock calls it recordId, and native Vertex rows carry it in the row key or request labels."""
|
||||
key: Final = entry.get("custom_id", entry.get("recordId"))
|
||||
return key if isinstance(key, str) else vertex_batch_row_custom_id(entry)
|
||||
|
||||
|
||||
def batch_line_item_parts(
|
||||
source: BatchLineItemSource, input_content: bytes, result_contents: Iterable[bytes], max_part_bytes: int
|
||||
) -> Iterator[StandardBatchLineItemsPayload]:
|
||||
"""Pair every output and error line with the input line that shares its custom_id, priced one line at a time, in
|
||||
parts that each close once their raw request and result lines reach max_part_bytes. Input lines stay unparsed until
|
||||
their result line is reached, so only one part's records are held at a time."""
|
||||
sized: Final = _sized_line_items(source, input_content, result_contents)
|
||||
for part in count():
|
||||
lines = tuple(_until_full(sized, max_part_bytes))
|
||||
if not lines:
|
||||
return
|
||||
yield StandardBatchLineItemsPayload(
|
||||
batch_id=source.batch_id, finished_at=source.finished_at, part=part, lines=lines
|
||||
)
|
||||
|
||||
|
||||
def _sized_line_items(
|
||||
source: BatchLineItemSource, input_content: bytes, result_contents: Iterable[bytes]
|
||||
) -> Iterator[tuple[StandardBatchLineItem, int]]:
|
||||
"""Each output and error line's record, with the raw bytes of that line and of the input line it pairs with."""
|
||||
request_spans: Final = _request_spans_by_custom_id(input_content)
|
||||
for raw_result in chain.from_iterable(map(iter_batch_input_lines, result_contents)):
|
||||
entry = parse_batch_output_line(raw_result)
|
||||
if entry is None:
|
||||
continue
|
||||
custom_id = _line_key(entry)
|
||||
span = _request_span_for(source, request_spans, custom_id)
|
||||
request = None if span is None else parse_batch_output_line(input_content[span[0] : span[1]])
|
||||
yield (
|
||||
_line_item(source, custom_id, request, entry),
|
||||
len(raw_result) + (0 if span is None else span[1] - span[0]),
|
||||
)
|
||||
|
||||
|
||||
def _until_full(
|
||||
sized: Iterator[tuple[StandardBatchLineItem, int]], max_part_bytes: int
|
||||
) -> Iterator[StandardBatchLineItem]:
|
||||
"""Take records until their raw bytes reach max_part_bytes, leaving the rest in sized for the next part."""
|
||||
for line, part_bytes in accumulate(sized, _with_bytes_so_far):
|
||||
yield line
|
||||
if part_bytes >= max_part_bytes:
|
||||
return
|
||||
|
||||
|
||||
def _with_bytes_so_far(
|
||||
so_far: tuple[StandardBatchLineItem, int], line: tuple[StandardBatchLineItem, int]
|
||||
) -> tuple[StandardBatchLineItem, int]:
|
||||
return line[0], so_far[1] + line[1]
|
||||
|
||||
|
||||
def _request_spans_by_custom_id(input_content: bytes) -> Mapping[str | None, tuple[int, int]]:
|
||||
"""Where each input line starts and ends, by its custom_id. A custom_id that several input lines share, as when a
|
||||
Vertex AI embeddings entry is split into one row per input, is paired with none of them."""
|
||||
keyed: Final = tuple(_keyed_spans(input_content))
|
||||
counts: Final = Counter(key for key, _ in keyed)
|
||||
return {key: span for key, span in keyed if counts[key] == 1}
|
||||
|
||||
|
||||
def _keyed_spans(input_content: bytes) -> Iterator[tuple[str | None, tuple[int, int]]]:
|
||||
for line in _LINE.finditer(input_content):
|
||||
entry = parse_batch_output_line(line[0]) if line[0].strip() else None
|
||||
if entry is not None:
|
||||
yield _line_key(entry), line.span()
|
||||
|
||||
|
||||
def _line_item(
|
||||
source: BatchLineItemSource, custom_id: str | None, request: Mapping[str, object] | None, entry: dict[str, object]
|
||||
) -> StandardBatchLineItem:
|
||||
priced: Final = batch_cost_and_usage(
|
||||
(entry,), source.cost_provider, model_name=source.model, model_info=source.model_info
|
||||
)
|
||||
succeeded: Final = priced.successful_requests > 0
|
||||
return StandardBatchLineItem(
|
||||
batch_id=source.batch_id,
|
||||
custom_id=custom_id,
|
||||
status="success" if succeeded else "failure",
|
||||
request=request,
|
||||
response=entry,
|
||||
usage=StandardBatchLineItemUsage(
|
||||
prompt_tokens=priced.usage.prompt_tokens,
|
||||
completion_tokens=priced.usage.completion_tokens,
|
||||
total_tokens=priced.usage.total_tokens,
|
||||
)
|
||||
if succeeded
|
||||
else None,
|
||||
response_cost=priced.cost,
|
||||
model=source.model,
|
||||
custom_llm_provider=source.custom_llm_provider,
|
||||
metadata=source.metadata,
|
||||
)
|
||||
|
||||
|
||||
def _request_span_for(
|
||||
source: BatchLineItemSource, request_spans: Mapping[str | None, tuple[int, int]], custom_id: str | None
|
||||
) -> tuple[int, int] | None:
|
||||
"""Where the input line that shares a result line's custom_id is. A Vertex AI line read as "unknown" pairs with
|
||||
none."""
|
||||
if custom_id is None or (source.cost_provider == "vertex_ai" and custom_id == _VERTEX_AI_UNKNOWN_CUSTOM_ID):
|
||||
return None
|
||||
return request_spans.get(custom_id)
|
||||
|
||||
|
||||
def without_request_and_response(payload: StandardBatchLineItemsPayload) -> StandardBatchLineItemsPayload:
|
||||
return StandardBatchLineItemsPayload(
|
||||
batch_id=payload["batch_id"],
|
||||
finished_at=payload["finished_at"],
|
||||
part=payload["part"],
|
||||
lines=tuple(
|
||||
StandardBatchLineItem(
|
||||
batch_id=line["batch_id"],
|
||||
custom_id=line["custom_id"],
|
||||
status=line["status"],
|
||||
request=None,
|
||||
response=None,
|
||||
usage=line["usage"],
|
||||
response_cost=line["response_cost"],
|
||||
model=line["model"],
|
||||
custom_llm_provider=line["custom_llm_provider"],
|
||||
metadata=line["metadata"],
|
||||
)
|
||||
for line in payload["lines"]
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def dispatch_batch_line_items(payload: StandardBatchLineItemsPayload, redact_bodies: bool) -> None:
|
||||
"""Hand one part's lines to every batch_line_item_callbacks logger. A failing logger is logged and never stops the
|
||||
rest."""
|
||||
await asyncio.gather(
|
||||
*(_deliver(callback, payload, redact_bodies) for callback in litellm.batch_line_item_callbacks)
|
||||
)
|
||||
|
||||
|
||||
async def _deliver(
|
||||
callback: litellm.CALLBACK_TYPES, payload: StandardBatchLineItemsPayload, redact_bodies: bool
|
||||
) -> None:
|
||||
try:
|
||||
logger: Final = _resolve_logger(callback)
|
||||
if logger is None:
|
||||
verbose_logger.error(
|
||||
"batch_line_item_callbacks: %r cannot store batch line items, part %d of batch %s was not stored",
|
||||
callback,
|
||||
payload["part"],
|
||||
payload["batch_id"],
|
||||
)
|
||||
return
|
||||
await logger.async_log_batch_line_items(
|
||||
without_request_and_response(payload)
|
||||
if redact_bodies or _logger_turned_off_message_logging(logger)
|
||||
else payload
|
||||
)
|
||||
except Exception as error: # noqa: BLE001 # one failing sink must not stop the other sinks or the cost poller
|
||||
verbose_logger.exception(
|
||||
"batch_line_item_callbacks: %r failed to store the %d lines of part %d of batch %s: %s",
|
||||
callback,
|
||||
len(payload["lines"]),
|
||||
payload["part"],
|
||||
payload["batch_id"],
|
||||
error,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_logger(callback: litellm.CALLBACK_TYPES) -> CustomLogger | None:
|
||||
match callback:
|
||||
case CustomLogger():
|
||||
return callback
|
||||
case "s3_v2" | "gcs_bucket":
|
||||
return init_custom_logger_compatible_class(callback, internal_usage_cache=None, llm_router=None)
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _logger_turned_off_message_logging(logger: CustomLogger) -> bool:
|
||||
return (
|
||||
getattr(logger, "turn_off_message_logging", False) is True
|
||||
or getattr(logger, "message_logging", True) is not True
|
||||
)
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
import json
|
||||
from collections.abc import Iterable, Iterator, Mapping
|
||||
from collections.abc import Iterable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import replace as dataclasses_replace
|
||||
from enum import Enum
|
||||
|
|
@ -70,7 +70,17 @@ def batch_cost_is_final(batch: Batch) -> bool:
|
|||
|
||||
|
||||
async def calculate_batch_cost_and_usage(
|
||||
file_content_dictionary: list[dict[str, object]],
|
||||
file_content_dictionary: Sequence[dict[str, object]],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
|
||||
model_name: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
) -> BatchCostUsageResult:
|
||||
"""Calculate the cost and usage of a batch, see batch_cost_and_usage."""
|
||||
return batch_cost_and_usage(file_content_dictionary, custom_llm_provider, model_name, model_info)
|
||||
|
||||
|
||||
def batch_cost_and_usage(
|
||||
file_content_dictionary: Sequence[dict[str, object]], # mutable-ok: the rows calculate_batch_cost_and_usage takes
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
|
||||
model_name: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
|
|
@ -436,7 +446,7 @@ def _provider_output_file_id(output_file_id: str) -> str:
|
|||
return extracted
|
||||
|
||||
|
||||
async def _fetch_batch_managed_file_content(
|
||||
async def fetch_batch_managed_file_content(
|
||||
file_id: str,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai",
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
|
|
@ -467,6 +477,9 @@ async def _fetch_batch_managed_file_content(
|
|||
return _file_content.content
|
||||
|
||||
|
||||
_fetch_batch_managed_file_content: Final = fetch_batch_managed_file_content
|
||||
|
||||
|
||||
async def _fetch_batch_output_file_content(
|
||||
batch: Batch,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai",
|
||||
|
|
@ -604,12 +617,12 @@ def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict[str, object
|
|||
whole batch's cost accounting.
|
||||
"""
|
||||
for line in iter_batch_input_lines(file_content):
|
||||
entry = _parse_batch_output_line(line)
|
||||
entry = parse_batch_output_line(line)
|
||||
if entry is not None:
|
||||
yield entry
|
||||
|
||||
|
||||
def _parse_batch_output_line(line: bytes) -> dict[str, object] | None:
|
||||
def parse_batch_output_line(line: bytes) -> dict[str, object] | None:
|
||||
try:
|
||||
parsed: Final[object] = json.loads(line)
|
||||
except ValueError as e:
|
||||
|
|
|
|||
|
|
@ -1862,6 +1862,10 @@ RESET_BUDGET_JOB_NAME: Final = "reset_budget_job"
|
|||
RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900
|
||||
PROXY_BATCH_POLLING_INTERVAL: Final = get_env_int("PROXY_BATCH_POLLING_INTERVAL", 3600)
|
||||
MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, get_env_int("MAX_OBJECTS_PER_POLL_CYCLE", 50))
|
||||
# raw request and result bytes per batch line-item object, so a large batch is built and stored one part at a time
|
||||
BATCH_LINE_ITEMS_MAX_PART_BYTES: Final = get_env_int_in_range(
|
||||
"BATCH_LINE_ITEMS_MAX_PART_BYTES", 16 * 1024 * 1024, 1, 1024 * 1024 * 1024
|
||||
)
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, get_env_int("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))
|
||||
PENDING_RESPONSE_STATUSES: Final = ("queued", "in_progress")
|
||||
TERMINAL_RESPONSE_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "incomplete"})
|
||||
|
|
|
|||
29
litellm/integrations/batch_line_item_objects.py
Normal file
29
litellm/integrations/batch_line_item_objects.py
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
"""How the line items of one finished batch are laid out as a single object in a storage bucket."""
|
||||
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
|
||||
from litellm.types.utils import StandardBatchLineItemsPayload
|
||||
|
||||
|
||||
def batch_line_items_ndjson(payload: StandardBatchLineItemsPayload) -> str:
|
||||
return "".join(f"{_ndjson_line(line)}\n" for line in payload["lines"])
|
||||
|
||||
|
||||
def _ndjson_line(line: object) -> str:
|
||||
"""One line as JSON, with non-ASCII text kept readable. A line holding a lone surrogate, which a provider's
|
||||
JSON escapes can carry, has no UTF-8 form, so that line alone is written with every non-ASCII character escaped."""
|
||||
text: Final = json.dumps(line, default=str, ensure_ascii=False)
|
||||
try:
|
||||
text.encode("utf-8")
|
||||
except UnicodeEncodeError:
|
||||
return json.dumps(line, default=str)
|
||||
return text
|
||||
|
||||
|
||||
def batch_line_items_object_name(payload: StandardBatchLineItemsPayload) -> str:
|
||||
"""batch_line_items/<finished day, UTC>/<batch id>/part-<part>.jsonl, the same name on every retry so a rewrite
|
||||
replaces it."""
|
||||
day: Final = datetime.fromtimestamp(payload["finished_at"], tz=timezone.utc).strftime("%Y-%m-%d")
|
||||
return f"batch_line_items/{day}/{payload['batch_id'].replace('/', '_')}/part-{payload['part']:05d}.jsonl"
|
||||
|
|
@ -22,6 +22,7 @@ from litellm.types.utils import (
|
|||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
StandardAuditLogPayload,
|
||||
StandardBatchLineItemsPayload,
|
||||
StandardCallbackDynamicParams,
|
||||
StandardLoggingPayload,
|
||||
)
|
||||
|
|
@ -195,6 +196,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
async def async_log_audit_log_event(self, audit_log: "StandardAuditLogPayload"):
|
||||
"""Called when an audit log is created. Override in subclasses to handle."""
|
||||
|
||||
async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None:
|
||||
"""Called once per part of a finished batch's lines, when this logger is listed in batch_line_item_callbacks."""
|
||||
|
||||
#### PROMPT MANAGEMENT HOOKS ####
|
||||
|
||||
async def async_get_chat_completion_prompt(
|
||||
|
|
|
|||
|
|
@ -10,13 +10,14 @@ from litellm._logging import verbose_logger
|
|||
from litellm._uuid import uuid
|
||||
from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE
|
||||
from litellm.integrations.additional_logging_utils import AdditionalLoggingUtils
|
||||
from litellm.integrations.batch_line_item_objects import batch_line_items_ndjson, batch_line_items_object_name
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
|
||||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
sanitize_cloud_object_component,
|
||||
)
|
||||
from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
|
||||
from litellm.types.integrations.gcs_bucket import *
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
from litellm.types.utils import StandardBatchLineItemsPayload, StandardLoggingPayload
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
|
|
@ -95,6 +96,20 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
except Exception as e:
|
||||
verbose_logger.exception("GCS Bucket logging error: %s", e)
|
||||
|
||||
async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None:
|
||||
"""Write one part of a batch's lines as an NDJSON object under batch_line_items/ in the default bucket."""
|
||||
gcs_logging_config: Final = await self.get_gcs_logging_config()
|
||||
headers: Final = await self.construct_request_headers(
|
||||
vertex_instance=gcs_logging_config["vertex_instance"],
|
||||
service_account_json=gcs_logging_config["path_service_account"],
|
||||
)
|
||||
await self._log_json_data_on_gcs(
|
||||
headers=headers,
|
||||
bucket_name=gcs_logging_config["bucket_name"],
|
||||
object_name=batch_line_items_object_name(payload),
|
||||
logging_payload=await asyncio.to_thread(batch_line_items_ndjson, payload),
|
||||
)
|
||||
|
||||
def _drain_queue_batch(self) -> list[GCSLogQueueItem]:
|
||||
"""
|
||||
Drain items from the queue (non-blocking), respecting batch_size limit.
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ from litellm.constants import (
|
|||
)
|
||||
from litellm.exceptions import Timeout as LiteLLMTimeout
|
||||
from litellm.integrations.adaptive_concurrency import AdaptiveConcurrencyLimiter, PutSample
|
||||
from litellm.integrations.batch_line_item_objects import batch_line_items_ndjson, batch_line_items_object_name
|
||||
from litellm.integrations.s3 import (
|
||||
get_s3_object_download_filename,
|
||||
get_s3_object_key,
|
||||
|
|
@ -60,7 +61,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.integrations.s3_v2 import S3PartitionGranularity, s3BatchLoggingElement
|
||||
from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
|
||||
from litellm.types.utils import StandardAuditLogPayload, StandardBatchLineItemsPayload, StandardLoggingPayload
|
||||
|
||||
from .custom_batch_logger import CustomBatchLogger
|
||||
|
||||
|
|
@ -548,6 +549,32 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
except Exception as e:
|
||||
verbose_logger.exception("S3 audit log error: %s", e)
|
||||
|
||||
async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None:
|
||||
"""Upload one part of a batch's lines as an NDJSON object under <s3_path>/batch_line_items/."""
|
||||
object_name: Final = batch_line_items_object_name(payload)
|
||||
s3_path: Final = cast(str | None, self.s3_path) # cast-ok: s3_path is set from untyped callback params
|
||||
s3_object_key: Final = (f"{s3_path.rstrip('/')}/" if s3_path else "") + object_name
|
||||
element: Final = s3BatchLoggingElement(
|
||||
payload={},
|
||||
body=await asyncio.to_thread(batch_line_items_ndjson, payload),
|
||||
content_type="application/x-ndjson",
|
||||
s3_object_key=s3_object_key,
|
||||
s3_object_download_filename="-".join(object_name.rsplit("/", 2)[-2:]),
|
||||
)
|
||||
# hashing a large object is CPU-bound too, so it is also prepared off the event loop
|
||||
prepared_put: Final = await asyncio.to_thread(self._prepare_put, element)
|
||||
delivered: Final = await self.async_upload_data_to_s3(element, prepared_put)
|
||||
if not delivered:
|
||||
# only a log-queue flush reads and clears the drop marks, and this upload is not part of one
|
||||
self._flush_dropped.pop(id(element), None)
|
||||
verbose_logger.error(
|
||||
"s3 logging: the %d line items of part %d of batch %s were not stored at %s",
|
||||
len(payload["lines"]),
|
||||
payload["part"],
|
||||
payload["batch_id"],
|
||||
s3_object_key,
|
||||
)
|
||||
|
||||
async def _async_log_event_base(self, kwargs, response_obj, start_time, end_time):
|
||||
try:
|
||||
verbose_logger.debug("s3 Logging - Enters logging function for model %s", kwargs)
|
||||
|
|
@ -592,6 +619,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
async def async_upload_data_to_s3(
|
||||
self,
|
||||
batch_logging_element: s3BatchLoggingElement,
|
||||
prepared_put: _PreparedPut | None = None,
|
||||
) -> bool:
|
||||
try:
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
|
|
@ -627,7 +655,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
except _TRANSIENT_ERRORS as error:
|
||||
return error
|
||||
|
||||
prepared: Final = self._prepare_put(batch_logging_element)
|
||||
prepared: Final = prepared_put if prepared_put is not None else self._prepare_put(batch_logging_element)
|
||||
await self._put_with_retries(partial(signed_put, prepared), batch_logging_element.s3_object_key)
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error uploading to s3: %s", e)
|
||||
|
|
|
|||
|
|
@ -5311,6 +5311,9 @@ def init_custom_logger_compatible_class_at_startup(
|
|||
)
|
||||
|
||||
|
||||
init_custom_logger_compatible_class: Final = _init_custom_logger_compatible_class
|
||||
|
||||
|
||||
def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[CustomLogger]) -> "OpenTelemetryV2 | None":
|
||||
"""If ``LITELLM_OTEL_V2`` is on, build (or reuse) a single ``OpenTelemetryV2``
|
||||
instance configured via the preset for ``callback_name``.
|
||||
|
|
|
|||
|
|
@ -424,6 +424,15 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, object]) -> tuple[st
|
|||
return unquote(match["custom_id"]), int(match["index"]), int(match["total"])
|
||||
|
||||
|
||||
def vertex_batch_row_custom_id(vertex_row: Mapping[str, object]) -> str | None:
|
||||
"""The OpenAI `custom_id` a Vertex batch input or output row was written for, or None when it carries none."""
|
||||
request_data: Final = vertex_row.get("request")
|
||||
has_labels: Final = isinstance(request_data, Mapping) and "labels" in request_data
|
||||
if vertex_row.get(_VERTEX_BATCH_KEY_FIELD) is None and not has_labels:
|
||||
return None
|
||||
return _split_vertex_batch_key(vertex_row)[0]
|
||||
|
||||
|
||||
def _vertex_embeddings_rows_to_openai_batch_output_row(
|
||||
custom_id: str,
|
||||
vertex_output_rows: tuple[_VertexEmbeddingBatchRow, ...],
|
||||
|
|
|
|||
|
|
@ -3669,6 +3669,46 @@ class StandardAuditLogPayload(TypedDict):
|
|||
updated_values: str | None
|
||||
|
||||
|
||||
class StandardBatchLineItemUsage(TypedDict):
|
||||
prompt_tokens: ReadOnly[int]
|
||||
completion_tokens: ReadOnly[int]
|
||||
total_tokens: ReadOnly[int]
|
||||
|
||||
|
||||
class StandardBatchLineItemMetadata(TypedDict):
|
||||
user_api_key_hash: ReadOnly[str | None]
|
||||
user_api_key_alias: ReadOnly[str | None]
|
||||
user_api_key_team_id: ReadOnly[str | None]
|
||||
user_api_key_team_alias: ReadOnly[str | None]
|
||||
user_api_key_user_id: ReadOnly[str | None]
|
||||
user_api_key_org_id: ReadOnly[str | None]
|
||||
|
||||
|
||||
class StandardBatchLineItem(TypedDict):
|
||||
"""One line of a finished batch, the request line next to the output or error line the provider wrote for it."""
|
||||
|
||||
batch_id: ReadOnly[str]
|
||||
custom_id: ReadOnly[str | None]
|
||||
status: ReadOnly[Literal["success", "failure"]]
|
||||
request: ReadOnly[Mapping[str, object] | None]
|
||||
response: ReadOnly[Mapping[str, object] | None]
|
||||
usage: ReadOnly[StandardBatchLineItemUsage | None]
|
||||
response_cost: ReadOnly[float]
|
||||
model: ReadOnly[str]
|
||||
custom_llm_provider: ReadOnly[str | None]
|
||||
metadata: ReadOnly[StandardBatchLineItemMetadata]
|
||||
|
||||
|
||||
class StandardBatchLineItemsPayload(TypedDict):
|
||||
"""One part of a finished batch's lines, dispatched to batch_line_item_callbacks. A batch's lines arrive in parts
|
||||
numbered from 0, each holding up to BATCH_LINE_ITEMS_MAX_PART_BYTES of raw request and result lines."""
|
||||
|
||||
batch_id: ReadOnly[str]
|
||||
finished_at: ReadOnly[int]
|
||||
part: ReadOnly[int]
|
||||
lines: ReadOnly[tuple[StandardBatchLineItem, ...]]
|
||||
|
||||
|
||||
class ClassifierAudit(TypedDict, total=False):
|
||||
classifier_input: ReadOnly[Mapping[str, JsonValue]]
|
||||
originating_request_masked: ReadOnly[Mapping[str, JsonValue]]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,902 @@
|
|||
"""Finished batches write their lines to cold storage.
|
||||
|
||||
An owned two-worker proxy lists ``s3_v2`` and ``gcs_bucket`` under ``batch_line_item_callbacks``. Once the batch-cost
|
||||
poller logs a finished batch, each sink gets the batch's lines as NDJSON part objects holding one record per line, a
|
||||
single part for these small batches. The proxy runs on its own database so the shared rig's pollers never claim these
|
||||
batches first, and accepts JWTs signed by a local key next to virtual keys. Four more owned proxies list an unsupported
|
||||
name next to ``s3_v2``, list a null list, check OAuth2 tokens against a local token info endpoint, and close a part
|
||||
after every line
|
||||
"""
|
||||
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qs, unquote, urlsplit
|
||||
|
||||
import jwt
|
||||
import pytest
|
||||
import yaml
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
from integration._support.client import JSON_OBJECT, Gateway, eventually, gateway_from_environment, held, string_value
|
||||
from integration._support.database import read_rows, scratch_database
|
||||
from integration._support.forward_proxy import tunnelling_forward_proxy
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.tls import server_context, write_self_signed_cert
|
||||
from integration._support.upstream import CONTROL_URL, ScenarioHandle, delete_scenario, register_scenario
|
||||
from integration._support.vertex import service_account_json
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse
|
||||
from jwt.algorithms import RSAAlgorithm
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(420)
|
||||
|
||||
S3_BUCKET: Final = "line-items-bucket"
|
||||
S3_PREFIX: Final = "cold"
|
||||
GCS_BUCKET: Final = "line-items-gcs"
|
||||
COMPLETED_AT: Final = 1791494130
|
||||
LINE_DATE: Final = datetime.fromtimestamp(COMPLETED_AT, tz=timezone.utc).strftime("%Y-%m-%d")
|
||||
PROMPT_TOKENS: Final = 100
|
||||
COMPLETION_TOKENS: Final = 50
|
||||
POLLER_SECONDS: Final = 150
|
||||
JWT_KEY_ID: Final = "batch-line-items-signing-key"
|
||||
TWO_ACCOUNT_MODEL: Final = "line-items-two-accounts"
|
||||
FIRST_ACCOUNT: Final = f"batch-line-items-first-account-{uuid.uuid4().hex[:12]}"
|
||||
SECOND_ACCOUNT: Final = f"batch-line-items-second-account-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ObjectSink:
|
||||
"""Stores every accepted object body by name and answers 503 to any upload whose body holds ``reject``. An upload
|
||||
whose body holds ``hang`` gets no answer until ``release``, then a 503."""
|
||||
|
||||
reject: bytes = b"\x00never"
|
||||
hang: bytes = b"\x00never"
|
||||
released: threading.Event = field(default_factory=threading.Event)
|
||||
lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
held: set[str] = field(default_factory=set) # mutable-ok: written by the sink's request threads
|
||||
stored: dict[str, bytes] = field(default_factory=dict) # mutable-ok: written by the sink's request threads
|
||||
attempts: dict[str, int] = field(default_factory=dict) # mutable-ok: written by the sink's request threads
|
||||
|
||||
def accept(self, name: str, body: bytes) -> bool:
|
||||
with self.lock:
|
||||
self.attempts[name] = self.attempts.get(name, 0) + 1
|
||||
hung: Final = self.hang in body
|
||||
if hung:
|
||||
self.held.add(name)
|
||||
if hung:
|
||||
self.released.wait(timeout=300)
|
||||
return False
|
||||
with self.lock:
|
||||
if self.reject in body:
|
||||
return False
|
||||
self.stored[name] = body
|
||||
return True
|
||||
|
||||
def holding(self) -> frozenset[str]:
|
||||
with self.lock:
|
||||
return frozenset(self.held)
|
||||
|
||||
def release(self) -> None:
|
||||
self.released.set()
|
||||
|
||||
def objects(self) -> Mapping[str, bytes]:
|
||||
with self.lock:
|
||||
return MappingProxyType(dict(self.stored))
|
||||
|
||||
def tried(self, name: str) -> int:
|
||||
with self.lock:
|
||||
return self.attempts.get(name, 0)
|
||||
|
||||
|
||||
def s3_peer(sink: ObjectSink) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method != "PUT":
|
||||
return Reply(status=404)
|
||||
name: Final = unquote(urlsplit(request.target).path).removeprefix(f"/{S3_BUCKET}/")
|
||||
if sink.accept(name, request.body):
|
||||
return Reply()
|
||||
return Reply(status=503, body=b"<Error><Code>SlowDown</Code></Error>", content_type="application/xml")
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def gcs_peer(sink: ObjectSink) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
path: Final = urlsplit(request.target).path
|
||||
if request.method == "POST" and path == "/_oauth/token":
|
||||
return Reply(body=b'{"access_token":"scripted-token","expires_in":3600,"token_type":"Bearer"}')
|
||||
if request.method == "POST" and path == f"/upload/storage/v1/b/{GCS_BUCKET}/o":
|
||||
name: Final = parse_qs(urlsplit(request.target).query)["name"][0]
|
||||
if sink.accept(name, request.body):
|
||||
return Reply(body=json.dumps({"name": name, "bucket": GCS_BUCKET}).encode())
|
||||
return Reply(status=503, body=b'{"error":"unavailable"}')
|
||||
return Reply(status=404, body=b'{"error":"not scripted"}')
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
gateway: Gateway
|
||||
database_url: str
|
||||
s3: ObjectSink
|
||||
gcs: ObjectSink
|
||||
signing_key: rsa.RSAPrivateKey
|
||||
|
||||
|
||||
def _jwks_peer(signing_key: rsa.RSAPrivateKey) -> Callable[[Request], Reply]:
|
||||
public_jwk: Final = json.loads(RSAAlgorithm.to_jwk(signing_key.public_key()))
|
||||
jwks: Final = json.dumps({"keys": [{**public_jwk, "kid": JWT_KEY_ID}]}).encode()
|
||||
|
||||
def respond(_request: Request) -> Reply:
|
||||
return Reply(body=jwks)
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _signed_token(signing_key: rsa.RSAPrivateKey, subject: str) -> str:
|
||||
now: Final = int(time.time())
|
||||
return jwt.encode(
|
||||
{"sub": subject, "iat": now, "exp": now + 300, "jti": uuid.uuid4().hex},
|
||||
signing_key,
|
||||
algorithm="RS256",
|
||||
headers={"kid": JWT_KEY_ID},
|
||||
)
|
||||
|
||||
|
||||
def _oauth2_token_info_peer() -> Callable[[Request], Reply]:
|
||||
"""Answers for a bearer token ``oauth_<team id>_<user id>`` with that user and team, and 401 for anything else."""
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
scheme, _, token = request.headers.get("authorization", "").partition(" ")
|
||||
parts: Final = token.split("_")
|
||||
if scheme != "Bearer" or len(parts) != 3 or parts[0] != "oauth":
|
||||
return Reply(status=401, body=b'{"error":"invalid_token"}')
|
||||
return Reply(body=json.dumps({"sub": parts[2], "team_id": parts[1]}).encode())
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _oauth2_token(team_id: str, user_id: str) -> str:
|
||||
return f"oauth_{team_id}_{user_id}"
|
||||
|
||||
|
||||
def _config(directory: Path, s3_url: str, line_item_callbacks: JsonValue, *, oauth2: bool) -> Path:
|
||||
shipped: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
configuration: Final = {
|
||||
**shipped,
|
||||
"model_list": [
|
||||
*shipped["model_list"],
|
||||
*(
|
||||
{
|
||||
"model_name": TWO_ACCOUNT_MODEL,
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "sk-scripted-provider",
|
||||
"api_base": f"{CONTROL_URL}/{account}",
|
||||
"input_cost_per_token": 0.001,
|
||||
"output_cost_per_token": 0.002,
|
||||
"weight": weight,
|
||||
},
|
||||
}
|
||||
for account, weight in ((FIRST_ACCOUNT, 0), (SECOND_ACCOUNT, 1))
|
||||
),
|
||||
],
|
||||
"general_settings": {
|
||||
**shipped["general_settings"],
|
||||
"enable_jwt_auth": True,
|
||||
"litellm_jwtauth": {"user_id_jwt_field": "sub", "user_id_upsert": True},
|
||||
"enable_oauth2_auth": oauth2,
|
||||
},
|
||||
"litellm_settings": {
|
||||
**shipped["litellm_settings"],
|
||||
"callbacks": ["s3_v2"],
|
||||
"batch_line_item_callbacks": line_item_callbacks,
|
||||
"s3_callback_params": {
|
||||
"s3_bucket_name": S3_BUCKET,
|
||||
"s3_region_name": "us-east-1",
|
||||
"s3_endpoint_url": s3_url,
|
||||
"s3_path": S3_PREFIX,
|
||||
"s3_aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
|
||||
"s3_aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
||||
"s3_flush_interval": 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
path: Final = directory / "batch-line-items.yaml"
|
||||
path.write_text(yaml.safe_dump(configuration))
|
||||
return path
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _owned_rig(
|
||||
tmp_path_factory: pytest.TempPathFactory,
|
||||
name: str,
|
||||
line_item_callbacks: JsonValue,
|
||||
*,
|
||||
oauth2: bool = False,
|
||||
max_part_bytes: int | None = None,
|
||||
) -> Iterator[Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp(name)
|
||||
certificate, key = write_self_signed_cert(directory, names=("storage.googleapis.com",))
|
||||
s3: Final = ObjectSink(reject=b"sink-down-", hang=b"sink-hang-")
|
||||
gcs: Final = ObjectSink()
|
||||
signing_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
with (
|
||||
wire_server(_jwks_peer(signing_key)) as jwks_wire,
|
||||
wire_server(_oauth2_token_info_peer()) as token_info_wire,
|
||||
gateway_from_environment() as environment,
|
||||
scratch_database(name.replace("-", "_")) as database_url,
|
||||
wire_server(s3_peer(s3)) as s3_wire,
|
||||
wire_server(gcs_peer(gcs)) as token_wire,
|
||||
wire_server(gcs_peer(gcs), tls=server_context(certificate, key)) as gcs_wire,
|
||||
tunnelling_forward_proxy(int(urlsplit(gcs_wire.url).port or 0)) as tunnel,
|
||||
):
|
||||
service_account: Final = directory / "service-account.json"
|
||||
service_account.write_text(service_account_json("line-items-project", token_wire.url))
|
||||
with owned_proxy(
|
||||
environment,
|
||||
directory,
|
||||
{
|
||||
"DATABASE_URL": database_url,
|
||||
"GCS_BUCKET_NAME": GCS_BUCKET,
|
||||
"GCS_PATH_SERVICE_ACCOUNT": str(service_account),
|
||||
"GCS_FLUSH_INTERVAL": "1",
|
||||
"JWT_PUBLIC_KEY_URL": jwks_wire.url,
|
||||
"OAUTH_TOKEN_INFO_ENDPOINT": token_info_wire.url,
|
||||
"HTTPS_PROXY": tunnel.url,
|
||||
"https_proxy": tunnel.url,
|
||||
"SSL_CERT_FILE": str(certificate),
|
||||
"NO_PROXY": "127.0.0.1,localhost",
|
||||
"no_proxy": "127.0.0.1,localhost",
|
||||
**({"BATCH_LINE_ITEMS_MAX_PART_BYTES": str(max_part_bytes)} if max_part_bytes else {}),
|
||||
},
|
||||
config=_config(directory, s3_wire.url, line_item_callbacks, oauth2=oauth2),
|
||||
remove_environment=("DATABASE_URL_READ_REPLICA",),
|
||||
workers=2,
|
||||
) as candidate:
|
||||
yield Rig(candidate, database_url, s3, gcs, signing_key)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _owned_rig(tmp_path_factory, "batch-line-items", ["s3_v2", "gcs_bucket"]) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig_with_an_unsupported_name(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _owned_rig(tmp_path_factory, "batch-line-items-unsupported", ["s3", "s3_v2"]) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig_with_a_null_list(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _owned_rig(tmp_path_factory, "batch-line-items-null", None) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig_with_oauth2(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _owned_rig(tmp_path_factory, "batch-line-items-oauth2", ["s3_v2"], oauth2=True) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig_with_one_line_parts(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _owned_rig(tmp_path_factory, "batch-line-items-parts", ["s3_v2", "gcs_bucket"], max_part_bytes=1) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
def _input_line(custom_id: str, text: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"custom_id": custom_id,
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": text}]},
|
||||
}
|
||||
|
||||
|
||||
def _success_line(custom_id: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": f"batch_req_{custom_id}",
|
||||
"custom_id": custom_id,
|
||||
"response": {
|
||||
"status_code": 200,
|
||||
"request_id": f"req-{custom_id}",
|
||||
"body": {
|
||||
"id": f"chatcmpl-{custom_id}",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {
|
||||
"prompt_tokens": PROMPT_TOKENS,
|
||||
"completion_tokens": COMPLETION_TOKENS,
|
||||
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
|
||||
},
|
||||
},
|
||||
},
|
||||
"error": None,
|
||||
}
|
||||
|
||||
|
||||
def _error_line(custom_id: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": f"batch_req_{custom_id}",
|
||||
"custom_id": custom_id,
|
||||
"response": {
|
||||
"status_code": 400,
|
||||
"request_id": f"req-{custom_id}",
|
||||
"body": {"error": {"message": "max_tokens is too large", "type": "invalid_request_error"}},
|
||||
},
|
||||
"error": None,
|
||||
}
|
||||
|
||||
|
||||
def _jsonl(lines: tuple[dict[str, JsonValue], ...]) -> str:
|
||||
return "".join(json.dumps(line, separators=(",", ":")) + "\n" for line in lines)
|
||||
|
||||
|
||||
def _file(identity: str, purpose: str) -> JsonResponse:
|
||||
return JsonResponse(
|
||||
content_type="application/json",
|
||||
body={
|
||||
"id": identity,
|
||||
"object": "file",
|
||||
"purpose": purpose,
|
||||
"bytes": 100,
|
||||
"created_at": 1,
|
||||
"filename": f"{identity}.jsonl",
|
||||
"status": "processed",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Lines:
|
||||
inputs: tuple[dict[str, JsonValue], ...]
|
||||
outputs: tuple[dict[str, JsonValue], ...]
|
||||
errors: tuple[dict[str, JsonValue], ...]
|
||||
|
||||
|
||||
def _routes(lines: Lines, *, finished: bool = True) -> RoutedResponse:
|
||||
has_output: Final = bool(lines.outputs)
|
||||
completed: Final[dict[str, JsonValue]] = {
|
||||
"id": "batch-$REQUEST_ID",
|
||||
"object": "batch",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"errors": None,
|
||||
"input_file_id": "file-in-$REQUEST_ID",
|
||||
"completion_window": "24h",
|
||||
"status": "completed",
|
||||
"output_file_id": "file-out-$REQUEST_ID" if has_output else None,
|
||||
"error_file_id": "file-err-$REQUEST_ID" if lines.errors else None,
|
||||
"created_at": COMPLETED_AT - 60,
|
||||
"in_progress_at": COMPLETED_AT - 30,
|
||||
"completed_at": COMPLETED_AT,
|
||||
"expires_at": COMPLETED_AT + 86400,
|
||||
"request_counts": {
|
||||
"total": len(lines.inputs),
|
||||
"completed": len(lines.outputs),
|
||||
"failed": len(lines.errors),
|
||||
},
|
||||
"metadata": None,
|
||||
}
|
||||
return RoutedResponse(
|
||||
content_type="application/x-routed",
|
||||
routes={
|
||||
"POST /files": _file("file-in-$REQUEST_ID", "batch"),
|
||||
"POST /batches": JsonResponse(
|
||||
content_type="application/json",
|
||||
body={**completed, "status": "validating", "output_file_id": None, "error_file_id": None},
|
||||
),
|
||||
"GET /batches/batch-$REQUEST_ID": JsonResponse(
|
||||
content_type="application/json",
|
||||
body=completed
|
||||
if finished
|
||||
else {**completed, "status": "in_progress", "output_file_id": None, "error_file_id": None},
|
||||
),
|
||||
"GET /files/file-in-$REQUEST_ID/content": TextResponse(
|
||||
content_type="application/jsonl", body=_jsonl(lines.inputs)
|
||||
),
|
||||
"GET /files/file-out-$REQUEST_ID": _file("file-out-$REQUEST_ID", "batch_output"),
|
||||
"GET /files/file-out-$REQUEST_ID/content": TextResponse(
|
||||
content_type="application/jsonl", body=_jsonl(lines.outputs)
|
||||
),
|
||||
"GET /files/file-err-$REQUEST_ID": _file("file-err-$REQUEST_ID", "batch_output"),
|
||||
"GET /files/file-err-$REQUEST_ID/content": TextResponse(
|
||||
content_type="application/jsonl", body=_jsonl(lines.errors)
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SubmittedBatch:
|
||||
batch_id: str
|
||||
key: str
|
||||
handle: ScenarioHandle
|
||||
|
||||
|
||||
def _submit(
|
||||
rig: Rig, lines: Lines, key_fields: Mapping[str, JsonValue], *, finished: bool = True, token: str | None = None
|
||||
) -> SubmittedBatch:
|
||||
handle: Final = register_scenario(f"batch-line-items-{uuid.uuid4().hex}", _routes(lines, finished=finished))
|
||||
return _create_batch(rig, lines, key_fields, _add_model(rig, handle), handle, token)
|
||||
|
||||
|
||||
def _add_model(rig: Rig, handle: ScenarioHandle) -> str:
|
||||
return string_value(
|
||||
rig.gateway.post(
|
||||
"/model/new",
|
||||
{
|
||||
"model_name": f"line-items-{uuid.uuid4().hex[:12]}",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "sk-scripted-provider",
|
||||
"api_base": handle.api_base(),
|
||||
"input_cost_per_token": 0.001,
|
||||
"output_cost_per_token": 0.002,
|
||||
},
|
||||
},
|
||||
)["model_name"]
|
||||
)
|
||||
|
||||
|
||||
def _create_batch(
|
||||
rig: Rig, lines: Lines, key_fields: Mapping[str, JsonValue], model: str, handle: ScenarioHandle, token: str | None
|
||||
) -> SubmittedBatch:
|
||||
key: Final = (
|
||||
token
|
||||
if token is not None
|
||||
else string_value(rig.gateway.post("/key/generate", {"models": [model], **key_fields})["key"])
|
||||
)
|
||||
uploaded: Final = rig.gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "target_model_names": model},
|
||||
{"file": ("in.jsonl", _jsonl(lines.inputs).encode(), "application/jsonl")},
|
||||
key=key,
|
||||
)
|
||||
assert uploaded.is_success, uploaded.text
|
||||
created: Final = rig.gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{
|
||||
"input_file_id": string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]),
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"model": model,
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert created.is_success, created.text
|
||||
return SubmittedBatch(string_value(JSON_OBJECT.validate_json(created.content)["id"]), key, handle)
|
||||
|
||||
|
||||
def _quiet_logging() -> dict[str, JsonValue]:
|
||||
return {
|
||||
"logging": [
|
||||
{
|
||||
"callback_name": "s3_v2",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {"turn_off_message_logging": "true"},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def _object_name(prefix: str, batch_id: str, part: int = 0) -> str:
|
||||
return f"{prefix}batch_line_items/{LINE_DATE}/{batch_id.replace('/', '_')}/part-{part:05d}.jsonl"
|
||||
|
||||
|
||||
def _records(body: bytes) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(JSON_OBJECT.validate_json(line) for line in body.splitlines() if line.strip())
|
||||
|
||||
|
||||
def _written(sink: ObjectSink, name: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
return _records(eventually(lambda: sink.objects().get(name), lambda body: body is not None, POLLER_SECONDS) or b"")
|
||||
|
||||
|
||||
def _batch_spend(rig: Rig, key: str) -> float:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND call_type = %s',
|
||||
(sha256(key.encode()).hexdigest(), "aretrieve_batch"),
|
||||
database_url=rig.database_url,
|
||||
),
|
||||
lambda found: len(found) == 1,
|
||||
POLLER_SECONDS,
|
||||
)
|
||||
spend: Final = rows[0]["spend"]
|
||||
assert isinstance(spend, float), rows
|
||||
return spend
|
||||
|
||||
|
||||
def _by_custom_id(records: tuple[dict[str, JsonValue], ...]) -> Mapping[JsonValue, dict[str, JsonValue]]:
|
||||
keyed: Final = {record["custom_id"]: record for record in records}
|
||||
assert len(keyed) == len(records), records
|
||||
return MappingProxyType(keyed)
|
||||
|
||||
|
||||
def test_a_finished_batch_writes_every_line_with_its_request_and_result_to_s3_and_gcs(rig: Rig) -> None:
|
||||
lines: Final = Lines(
|
||||
inputs=(_input_line("ok-1", "first question"), _input_line("bad-1", "second question")),
|
||||
outputs=(_success_line("ok-1"),),
|
||||
errors=(_error_line("bad-1"),),
|
||||
)
|
||||
batch: Final = _submit(rig, lines, {})
|
||||
try:
|
||||
s3_records: Final = _written(rig.s3, _object_name(f"{S3_PREFIX}/", batch.batch_id))
|
||||
gcs_body: Final = eventually(
|
||||
lambda: rig.gcs.objects().get(_object_name("", batch.batch_id)), lambda body: body is not None
|
||||
)
|
||||
spend: Final = _batch_spend(rig, batch.key)
|
||||
finally:
|
||||
delete_scenario(batch.handle)
|
||||
|
||||
assert gcs_body is not None
|
||||
assert _records(gcs_body) == s3_records
|
||||
records: Final = _by_custom_id(s3_records)
|
||||
assert set(records) == {"ok-1", "bad-1"}
|
||||
hashed_key: Final = sha256(batch.key.encode()).hexdigest()
|
||||
succeeded: Final = records["ok-1"]
|
||||
assert succeeded["batch_id"] == batch.batch_id
|
||||
assert succeeded["status"] == "success"
|
||||
assert succeeded["request"] == lines.inputs[0]
|
||||
assert succeeded["response"] == lines.outputs[0]
|
||||
assert succeeded["usage"] == {
|
||||
"prompt_tokens": PROMPT_TOKENS,
|
||||
"completion_tokens": COMPLETION_TOKENS,
|
||||
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
|
||||
}
|
||||
assert succeeded["model"] == "gpt-4o-mini"
|
||||
assert succeeded["custom_llm_provider"] == "openai"
|
||||
assert isinstance(succeeded["metadata"], dict)
|
||||
assert succeeded["metadata"]["user_api_key_hash"] == hashed_key
|
||||
failed: Final = records["bad-1"]
|
||||
assert failed["status"] == "failure"
|
||||
assert failed["request"] == lines.inputs[1]
|
||||
assert failed["response"] == lines.errors[0]
|
||||
assert failed["usage"] is None
|
||||
assert failed["response_cost"] == 0
|
||||
assert isinstance(succeeded["response_cost"], float) and succeeded["response_cost"] > 0
|
||||
assert succeeded["response_cost"] == pytest.approx(spend)
|
||||
assert rig.s3.tried(_object_name(f"{S3_PREFIX}/", batch.batch_id)) == 1
|
||||
|
||||
|
||||
def test_a_key_that_turned_off_message_logging_gets_its_lines_without_request_or_response(rig: Rig) -> None:
|
||||
lines: Final = Lines(
|
||||
inputs=(_input_line("quiet-1", "a private question"),), outputs=(_success_line("quiet-1"),), errors=()
|
||||
)
|
||||
batch: Final = _submit(rig, lines, {"metadata": _quiet_logging()})
|
||||
name: Final = _object_name(f"{S3_PREFIX}/", batch.batch_id)
|
||||
try:
|
||||
records: Final = _written(rig.s3, name)
|
||||
finally:
|
||||
delete_scenario(batch.handle)
|
||||
|
||||
assert len(records) == 1
|
||||
assert records[0]["custom_id"] == "quiet-1"
|
||||
assert records[0]["status"] == "success"
|
||||
assert records[0]["request"] is None
|
||||
assert records[0]["response"] is None
|
||||
assert records[0]["usage"] == {
|
||||
"prompt_tokens": PROMPT_TOKENS,
|
||||
"completion_tokens": COMPLETION_TOKENS,
|
||||
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
|
||||
}
|
||||
assert b"a private question" not in rig.s3.objects()[name]
|
||||
|
||||
|
||||
def test_a_team_that_turned_off_message_logging_gets_its_members_lines_without_request_or_response(rig: Rig) -> None:
|
||||
team_id: Final = string_value(
|
||||
rig.gateway.post("/team/new", {"team_alias": f"quiet-{uuid.uuid4().hex[:8]}", "metadata": _quiet_logging()})[
|
||||
"team_id"
|
||||
]
|
||||
)
|
||||
lines: Final = Lines(
|
||||
inputs=(_input_line("team-1", "a team secret"),), outputs=(_success_line("team-1"),), errors=()
|
||||
)
|
||||
batch: Final = _submit(rig, lines, {"team_id": team_id})
|
||||
name: Final = _object_name(f"{S3_PREFIX}/", batch.batch_id)
|
||||
try:
|
||||
records: Final = _written(rig.s3, name)
|
||||
finally:
|
||||
delete_scenario(batch.handle)
|
||||
|
||||
assert [(record["custom_id"], record["request"], record["response"]) for record in records] == [
|
||||
("team-1", None, None)
|
||||
]
|
||||
assert isinstance(records[0]["metadata"], dict)
|
||||
assert records[0]["metadata"]["user_api_key_team_id"] == team_id
|
||||
assert b"a team secret" not in rig.s3.objects()[name]
|
||||
|
||||
|
||||
def test_a_batch_whose_creating_key_was_deleted_gets_its_lines_without_request_or_response(rig: Rig) -> None:
|
||||
lines: Final = Lines(
|
||||
inputs=(_input_line("orphan-1", "asked before the key was deleted"),),
|
||||
outputs=(_success_line("orphan-1"),),
|
||||
errors=(),
|
||||
)
|
||||
batch: Final = _submit(rig, lines, {}, finished=False)
|
||||
name: Final = _object_name(f"{S3_PREFIX}/", batch.batch_id)
|
||||
try:
|
||||
rig.gateway.post("/key/delete", {"keys": [batch.key]})
|
||||
register_scenario(batch.handle.scenario_id, _routes(lines))
|
||||
records: Final = _written(rig.s3, name)
|
||||
spend: Final = _batch_spend(rig, batch.key)
|
||||
finally:
|
||||
delete_scenario(batch.handle)
|
||||
|
||||
assert [(record["custom_id"], record["status"], record["request"], record["response"]) for record in records] == [
|
||||
("orphan-1", "success", None, None)
|
||||
]
|
||||
assert records[0]["response_cost"] == pytest.approx(spend)
|
||||
assert b"asked before the key was deleted" not in rig.s3.objects()[name]
|
||||
|
||||
|
||||
def test_a_batch_created_with_a_jwt_and_no_key_gets_its_lines_with_request_and_response(rig: Rig) -> None:
|
||||
subject: Final = f"jwt-{uuid.uuid4().hex[:12]}"
|
||||
lines: Final = Lines(
|
||||
inputs=(_input_line("jwt-1", "asked with a signed token"),), outputs=(_success_line("jwt-1"),), errors=()
|
||||
)
|
||||
batch: Final = _submit(rig, lines, {}, token=_signed_token(rig.signing_key, subject))
|
||||
try:
|
||||
records: Final = _written(rig.s3, _object_name(f"{S3_PREFIX}/", batch.batch_id))
|
||||
finally:
|
||||
delete_scenario(batch.handle)
|
||||
|
||||
assert read_rows(
|
||||
'SELECT api_key, created_by FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s',
|
||||
(batch.batch_id,),
|
||||
database_url=rig.database_url,
|
||||
) == [{"api_key": None, "created_by": subject}]
|
||||
assert [(record["custom_id"], record["request"], record["response"]) for record in records] == [
|
||||
("jwt-1", lines.inputs[0], lines.outputs[0])
|
||||
]
|
||||
|
||||
|
||||
def test_a_batch_created_with_an_oauth2_token_follows_its_teams_logging_setting_and_records_only_its_hash(
|
||||
rig_with_oauth2: Rig,
|
||||
) -> None:
|
||||
rig: Final = rig_with_oauth2
|
||||
open_team: Final = string_value(
|
||||
rig.gateway.post("/team/new", {"team_alias": f"open-{uuid.uuid4().hex[:8]}"})["team_id"]
|
||||
)
|
||||
quiet_team: Final = string_value(
|
||||
rig.gateway.post("/team/new", {"team_alias": f"quiet-{uuid.uuid4().hex[:8]}", "metadata": _quiet_logging()})[
|
||||
"team_id"
|
||||
]
|
||||
)
|
||||
user_id: Final = f"oauth-user-{uuid.uuid4().hex[:12]}"
|
||||
open_lines: Final = Lines(
|
||||
inputs=(_input_line("oauth-open", "asked with an oauth token"),),
|
||||
outputs=(_success_line("oauth-open"),),
|
||||
errors=(),
|
||||
)
|
||||
quiet_lines: Final = Lines(
|
||||
inputs=(_input_line("oauth-quiet", "a quiet team's oauth secret"),),
|
||||
outputs=(_success_line("oauth-quiet"),),
|
||||
errors=(),
|
||||
)
|
||||
opened: Final = _submit(rig, open_lines, {}, token=_oauth2_token(open_team, user_id))
|
||||
quieted: Final = _submit(rig, quiet_lines, {}, token=_oauth2_token(quiet_team, user_id))
|
||||
quiet_name: Final = _object_name(f"{S3_PREFIX}/", quieted.batch_id)
|
||||
try:
|
||||
open_records: Final = _written(rig.s3, _object_name(f"{S3_PREFIX}/", opened.batch_id))
|
||||
quiet_records: Final = _written(rig.s3, quiet_name)
|
||||
finally:
|
||||
delete_scenario(opened.handle)
|
||||
delete_scenario(quieted.handle)
|
||||
|
||||
assert read_rows(
|
||||
'SELECT api_key, created_by, team_id FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s',
|
||||
(opened.batch_id,),
|
||||
database_url=rig.database_url,
|
||||
) == [{"api_key": _oauth2_token(open_team, user_id), "created_by": user_id, "team_id": open_team}]
|
||||
assert [(record["custom_id"], record["request"], record["response"]) for record in open_records] == [
|
||||
("oauth-open", open_lines.inputs[0], open_lines.outputs[0])
|
||||
]
|
||||
assert [(record["custom_id"], record["request"], record["response"]) for record in quiet_records] == [
|
||||
("oauth-quiet", None, None)
|
||||
]
|
||||
assert b"a quiet team's oauth secret" not in rig.s3.objects()[quiet_name]
|
||||
open_token: Final = _oauth2_token(open_team, user_id)
|
||||
assert [record["metadata"] for record in open_records] == [
|
||||
{**open_records[0]["metadata"], "user_api_key_hash": sha256(open_token.encode()).hexdigest()}
|
||||
]
|
||||
assert open_token.encode() not in rig.s3.objects()[_object_name(f"{S3_PREFIX}/", opened.batch_id)]
|
||||
|
||||
|
||||
def test_a_batch_bigger_than_one_part_is_written_as_numbered_parts_holding_each_line_once(
|
||||
rig_with_one_line_parts: Rig,
|
||||
) -> None:
|
||||
rig: Final = rig_with_one_line_parts
|
||||
lines: Final = Lines(
|
||||
inputs=(_input_line("part-a", "first"), _input_line("part-b", "second"), _input_line("part-c", "third")),
|
||||
outputs=(_success_line("part-a"), _success_line("part-b")),
|
||||
errors=(_error_line("part-c"),),
|
||||
)
|
||||
batch: Final = _submit(rig, lines, {})
|
||||
try:
|
||||
s3_parts: Final = tuple(
|
||||
_written(rig.s3, _object_name(f"{S3_PREFIX}/", batch.batch_id, part)) for part in range(3)
|
||||
)
|
||||
gcs_parts: Final = tuple(_written(rig.gcs, _object_name("", batch.batch_id, part)) for part in range(3))
|
||||
finally:
|
||||
delete_scenario(batch.handle)
|
||||
|
||||
assert [[record["custom_id"] for record in part] for part in s3_parts] == [["part-a"], ["part-b"], ["part-c"]]
|
||||
assert [record["request"] for part in s3_parts for record in part] == list(lines.inputs)
|
||||
assert gcs_parts == s3_parts
|
||||
assert _object_name(f"{S3_PREFIX}/", batch.batch_id, 3) not in rig.s3.objects()
|
||||
|
||||
|
||||
def test_a_batch_where_every_line_failed_still_writes_its_error_lines(rig: Rig) -> None:
|
||||
lines: Final = Lines(
|
||||
inputs=(_input_line("bad-a", "first"), _input_line("bad-b", "second")),
|
||||
outputs=(),
|
||||
errors=(_error_line("bad-a"), _error_line("bad-b")),
|
||||
)
|
||||
batch: Final = _submit(rig, lines, {})
|
||||
try:
|
||||
records: Final = _by_custom_id(_written(rig.s3, _object_name(f"{S3_PREFIX}/", batch.batch_id)))
|
||||
finally:
|
||||
delete_scenario(batch.handle)
|
||||
|
||||
assert set(records) == {"bad-a", "bad-b"}
|
||||
assert records["bad-a"]["status"] == "failure"
|
||||
assert records["bad-a"]["request"] == lines.inputs[0]
|
||||
assert records["bad-b"]["response"] == lines.errors[1]
|
||||
assert all(record["response_cost"] == 0 for record in records.values())
|
||||
|
||||
|
||||
def test_a_failed_batch_on_the_second_deployment_of_its_model_group_keeps_each_lines_request(rig: Rig) -> None:
|
||||
lines: Final = Lines(
|
||||
inputs=(_input_line("second-a", "first"), _input_line("second-b", "second")),
|
||||
outputs=(),
|
||||
errors=(_error_line("second-a"), _error_line("second-b")),
|
||||
)
|
||||
first: Final = register_scenario(FIRST_ACCOUNT, _routes(lines))
|
||||
second: Final = register_scenario(SECOND_ACCOUNT, _routes(lines))
|
||||
try:
|
||||
batch: Final = _create_batch(rig, lines, {}, TWO_ACCOUNT_MODEL, second, None)
|
||||
records: Final = _by_custom_id(_written(rig.s3, _object_name(f"{S3_PREFIX}/", batch.batch_id)))
|
||||
finally:
|
||||
delete_scenario(first)
|
||||
delete_scenario(second)
|
||||
|
||||
assert read_rows(
|
||||
'SELECT model_object_id FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s',
|
||||
(batch.batch_id,),
|
||||
database_url=rig.database_url,
|
||||
) == [{"model_object_id": f"batch-{SECOND_ACCOUNT}"}]
|
||||
assert {custom_id: record["request"] for custom_id, record in records.items()} == {
|
||||
"second-a": lines.inputs[0],
|
||||
"second-b": lines.inputs[1],
|
||||
}
|
||||
|
||||
|
||||
def test_a_failing_bucket_still_bills_the_batch_and_the_other_sink_keeps_its_lines(rig: Rig) -> None:
|
||||
custom_id: Final = f"sink-down-{uuid.uuid4().hex[:8]}"
|
||||
lines: Final = Lines(
|
||||
inputs=(_input_line(custom_id, "billed anyway"),), outputs=(_success_line(custom_id),), errors=()
|
||||
)
|
||||
batch: Final = _submit(rig, lines, {})
|
||||
try:
|
||||
gcs_body: Final = eventually(
|
||||
lambda: rig.gcs.objects().get(_object_name("", batch.batch_id)),
|
||||
lambda body: body is not None,
|
||||
POLLER_SECONDS,
|
||||
)
|
||||
spend: Final = _batch_spend(rig, batch.key)
|
||||
processed: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT batch_processed FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s',
|
||||
(batch.batch_id,),
|
||||
database_url=rig.database_url,
|
||||
),
|
||||
lambda rows: rows == [{"batch_processed": True}],
|
||||
)
|
||||
finally:
|
||||
delete_scenario(batch.handle)
|
||||
|
||||
assert processed == [{"batch_processed": True}]
|
||||
assert gcs_body is not None
|
||||
assert [record["custom_id"] for record in _records(gcs_body)] == [custom_id]
|
||||
assert spend > 0
|
||||
s3_name: Final = _object_name(f"{S3_PREFIX}/", batch.batch_id)
|
||||
assert rig.s3.tried(s3_name) >= 1
|
||||
assert s3_name not in rig.s3.objects()
|
||||
|
||||
|
||||
def _hanging_lines() -> Lines:
|
||||
custom_id: Final = f"sink-hang-{uuid.uuid4().hex[:8]}"
|
||||
return Lines(inputs=(_input_line(custom_id, "hangs the bucket"),), outputs=(_success_line(custom_id),), errors=())
|
||||
|
||||
|
||||
def test_a_hung_bucket_never_holds_up_billing_the_next_batch(rig: Rig) -> None:
|
||||
"""The first hung write holds one worker's poller if the poll waits on it, and the second batch, billed by the
|
||||
other worker, holds that one too, so only a poll that never waits on a write bills the third batch."""
|
||||
first: Final = _submit(rig, _hanging_lines(), {})
|
||||
second: SubmittedBatch | None = None
|
||||
following: SubmittedBatch | None = None
|
||||
try:
|
||||
eventually(rig.s3.holding, lambda names: _object_name(f"{S3_PREFIX}/", first.batch_id) in names, POLLER_SECONDS)
|
||||
second = _submit(rig, _hanging_lines(), {})
|
||||
_batch_spend(rig, second.key)
|
||||
following = _submit(
|
||||
rig,
|
||||
Lines(inputs=(_input_line("after-1", "asked later"),), outputs=(_success_line("after-1"),), errors=()),
|
||||
{},
|
||||
)
|
||||
spend: Final = _batch_spend(rig, following.key)
|
||||
still_holding: Final = rig.s3.holding()
|
||||
finally:
|
||||
rig.s3.release()
|
||||
for batch in (first, second, following):
|
||||
if batch is not None:
|
||||
delete_scenario(batch.handle)
|
||||
|
||||
assert spend > 0
|
||||
assert _object_name(f"{S3_PREFIX}/", first.batch_id) in still_holding
|
||||
assert _object_name(f"{S3_PREFIX}/", first.batch_id) not in rig.s3.objects()
|
||||
|
||||
|
||||
def _line_item_objects(rig: Rig) -> frozenset[str]:
|
||||
return frozenset(name for sink in (rig.s3, rig.gcs) for name in sink.objects() if "batch_line_items/" in name)
|
||||
|
||||
|
||||
def test_an_unsupported_name_in_the_list_leaves_the_supported_sink_writing(rig_with_an_unsupported_name: Rig) -> None:
|
||||
rig: Final = rig_with_an_unsupported_name
|
||||
lines: Final = Lines(
|
||||
inputs=(_input_line("listed-1", "asked with s3 listed"),), outputs=(_success_line("listed-1"),), errors=()
|
||||
)
|
||||
batch: Final = _submit(rig, lines, {})
|
||||
try:
|
||||
records: Final = _written(rig.s3, _object_name(f"{S3_PREFIX}/", batch.batch_id))
|
||||
spend: Final = _batch_spend(rig, batch.key)
|
||||
finally:
|
||||
delete_scenario(batch.handle)
|
||||
|
||||
assert [(record["custom_id"], record["request"], record["response"]) for record in records] == [
|
||||
("listed-1", lines.inputs[0], lines.outputs[0])
|
||||
]
|
||||
assert records[0]["response_cost"] == pytest.approx(spend)
|
||||
assert _line_item_objects(rig) == {_object_name(f"{S3_PREFIX}/", batch.batch_id)}
|
||||
|
||||
|
||||
def test_a_null_list_bills_the_batch_and_writes_no_line_items(rig_with_a_null_list: Rig) -> None:
|
||||
rig: Final = rig_with_a_null_list
|
||||
lines: Final = Lines(
|
||||
inputs=(_input_line("null-1", "asked with a null list"),), outputs=(_success_line("null-1"),), errors=()
|
||||
)
|
||||
batch: Final = _submit(rig, lines, {})
|
||||
try:
|
||||
spend: Final = _batch_spend(rig, batch.key)
|
||||
eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT batch_processed FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s',
|
||||
(batch.batch_id,),
|
||||
database_url=rig.database_url,
|
||||
),
|
||||
lambda rows: rows == [{"batch_processed": True}],
|
||||
POLLER_SECONDS,
|
||||
)
|
||||
stored: Final = held(lambda: _line_item_objects(rig), lambda names: not names, holding=10)
|
||||
finally:
|
||||
delete_scenario(batch.handle)
|
||||
|
||||
assert spend > 0
|
||||
assert stored == frozenset()
|
||||
434
tests/unit/batches/test_batch_line_items.py
Normal file
434
tests/unit/batches/test_batch_line_items.py
Normal file
|
|
@ -0,0 +1,434 @@
|
|||
import json
|
||||
import logging
|
||||
from hashlib import sha256
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.batches.batch_line_items import (
|
||||
BatchLineItemSource,
|
||||
batch_line_item_parts,
|
||||
dispatch_batch_line_items,
|
||||
line_item_metadata,
|
||||
without_request_and_response,
|
||||
)
|
||||
from litellm.batches.batch_utils import calculate_batch_cost_and_usage, get_file_content_as_dictionary
|
||||
from litellm.constants import BATCH_LINE_ITEMS_MAX_PART_BYTES
|
||||
from litellm.integrations.batch_line_item_objects import batch_line_items_ndjson, batch_line_items_object_name
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig, _OpenAIToVertexBatchUploadStream
|
||||
from litellm.types.utils import ModelInfo, StandardBatchLineItemsPayload
|
||||
|
||||
INPUT_RATE = 0.001
|
||||
OUTPUT_RATE = 0.002
|
||||
FINISHED_AT = 1791494130
|
||||
ALICE_KEY_HASH = sha256(b"sk-alice").hexdigest()
|
||||
|
||||
|
||||
def _jsonl(*lines: dict) -> bytes:
|
||||
return "".join(json.dumps(line) + "\n" for line in lines).encode()
|
||||
|
||||
|
||||
def _request(custom_id: str, text: str) -> dict:
|
||||
return {
|
||||
"custom_id": custom_id,
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": text}]},
|
||||
}
|
||||
|
||||
|
||||
def _success(custom_id: str, prompt_tokens: int, completion_tokens: int) -> dict:
|
||||
return {
|
||||
"id": f"batch_req_{custom_id}",
|
||||
"custom_id": custom_id,
|
||||
"response": {
|
||||
"status_code": 200,
|
||||
"body": {
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}}],
|
||||
"usage": {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": prompt_tokens + completion_tokens,
|
||||
},
|
||||
},
|
||||
},
|
||||
"error": None,
|
||||
}
|
||||
|
||||
|
||||
def _error(custom_id: str) -> dict:
|
||||
return {
|
||||
"id": f"batch_req_{custom_id}",
|
||||
"custom_id": custom_id,
|
||||
"response": {"status_code": 400, "body": {"error": {"message": "max_tokens is too large"}}},
|
||||
"error": None,
|
||||
}
|
||||
|
||||
|
||||
def _source(**overrides) -> BatchLineItemSource:
|
||||
fields = {
|
||||
"batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA",
|
||||
"finished_at": FINISHED_AT,
|
||||
"model": "gpt-4o-mini",
|
||||
"custom_llm_provider": "openai",
|
||||
"cost_provider": "openai",
|
||||
"model_info": ModelInfo(input_cost_per_token_batches=INPUT_RATE, output_cost_per_token_batches=OUTPUT_RATE),
|
||||
"metadata": line_item_metadata(
|
||||
{"user_api_key_hash": ALICE_KEY_HASH, "user_api_key_team_id": "team-alpha", "tags": ["env:prod"]}
|
||||
),
|
||||
}
|
||||
return BatchLineItemSource(**{**fields, **overrides})
|
||||
|
||||
|
||||
def _whole_batch(source: BatchLineItemSource, input_content: bytes, result_contents: tuple[bytes, ...]):
|
||||
"""Every line of a batch small enough for one part."""
|
||||
(payload,) = batch_line_item_parts(source, input_content, result_contents, BATCH_LINE_ITEMS_MAX_PART_BYTES)
|
||||
return payload
|
||||
|
||||
|
||||
def _payload() -> StandardBatchLineItemsPayload:
|
||||
return _whole_batch(
|
||||
_source(),
|
||||
_jsonl(_request("ok-1", "first question"), _request("bad-1", "second question")),
|
||||
(_jsonl(_success("ok-1", 100, 50)), _jsonl(_error("bad-1"))),
|
||||
)
|
||||
|
||||
|
||||
class _RecordingLogger(CustomLogger):
|
||||
def __init__(self, **kwargs) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self.received: list[StandardBatchLineItemsPayload] = []
|
||||
|
||||
async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None:
|
||||
self.received.append(payload)
|
||||
|
||||
|
||||
class _FailingLogger(CustomLogger):
|
||||
async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None:
|
||||
raise RuntimeError("bucket is gone")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_each_result_line_is_paired_with_the_request_that_shares_its_custom_id():
|
||||
payload = _whole_batch(
|
||||
_source(),
|
||||
_jsonl(_request("ok-1", "first"), _request("bad-1", "second"), _request("never-ran", "third")),
|
||||
(_jsonl(_success("ok-1", 100, 50)), _jsonl(_error("bad-1"))),
|
||||
)
|
||||
|
||||
assert payload["batch_id"] == "bGl0ZWxsbV9wcm94eTtiYXRjaA"
|
||||
assert payload["finished_at"] == FINISHED_AT
|
||||
succeeded, failed = payload["lines"]
|
||||
assert succeeded == {
|
||||
"batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA",
|
||||
"custom_id": "ok-1",
|
||||
"status": "success",
|
||||
"request": _request("ok-1", "first"),
|
||||
"response": _success("ok-1", 100, 50),
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
"response_cost": pytest.approx(100 * INPUT_RATE + 50 * OUTPUT_RATE),
|
||||
"model": "gpt-4o-mini",
|
||||
"custom_llm_provider": "openai",
|
||||
"metadata": {
|
||||
"user_api_key_hash": ALICE_KEY_HASH,
|
||||
"user_api_key_alias": None,
|
||||
"user_api_key_team_id": "team-alpha",
|
||||
"user_api_key_team_alias": None,
|
||||
"user_api_key_user_id": None,
|
||||
"user_api_key_org_id": None,
|
||||
},
|
||||
}
|
||||
assert failed["custom_id"] == "bad-1"
|
||||
assert failed["status"] == "failure"
|
||||
assert failed["request"] == _request("bad-1", "second")
|
||||
assert failed["response"] == _error("bad-1")
|
||||
assert failed["usage"] is None
|
||||
assert failed["response_cost"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_line_costs_add_up_to_the_batch_cost_the_spend_row_records():
|
||||
outputs = _jsonl(_success("a", 100, 50), _success("b", 7, 3), _error("c"))
|
||||
|
||||
payload = _whole_batch(_source(), b"", (outputs,))
|
||||
aggregate = await calculate_batch_cost_and_usage(
|
||||
get_file_content_as_dictionary(outputs), "openai", model_name="gpt-4o-mini", model_info=_source().model_info
|
||||
)
|
||||
|
||||
assert [line["response_cost"] for line in payload["lines"]] == pytest.approx(
|
||||
[100 * INPUT_RATE + 50 * OUTPUT_RATE, 7 * INPUT_RATE + 3 * OUTPUT_RATE, 0]
|
||||
)
|
||||
assert sum(line["response_cost"] for line in payload["lines"]) == pytest.approx(aggregate.cost)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_results_without_an_input_file_keep_their_status_and_usage():
|
||||
succeeded = {
|
||||
"custom_id": "a-1",
|
||||
"result": {
|
||||
"type": "succeeded",
|
||||
"message": {"model": "claude-sonnet-5-5", "usage": {"input_tokens": 10, "output_tokens": 4}},
|
||||
},
|
||||
}
|
||||
errored = {"custom_id": "a-2", "result": {"type": "errored", "error": {"type": "invalid_request_error"}}}
|
||||
|
||||
payload = _whole_batch(
|
||||
_source(model="claude-sonnet-5-5", custom_llm_provider="anthropic", cost_provider="anthropic"),
|
||||
b"",
|
||||
(_jsonl(succeeded, errored),),
|
||||
)
|
||||
|
||||
first, second = payload["lines"]
|
||||
assert (first["custom_id"], first["status"], first["request"]) == ("a-1", "success", None)
|
||||
assert first["usage"] == {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}
|
||||
assert (second["custom_id"], second["status"], second["usage"]) == ("a-2", "failure", None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_lines_pair_by_record_id():
|
||||
request = {"recordId": "r-1", "modelInput": {"messages": [{"role": "user", "content": "hi"}]}}
|
||||
result = {**request, "modelOutput": {"usage": {"input_tokens": 3, "output_tokens": 2}}}
|
||||
|
||||
payload = _whole_batch(
|
||||
_source(custom_llm_provider="bedrock", cost_provider="bedrock"), _jsonl(request), (_jsonl(result),)
|
||||
)
|
||||
|
||||
(line,) = payload["lines"]
|
||||
assert (line["custom_id"], line["status"], line["request"]) == ("r-1", "success", request)
|
||||
|
||||
|
||||
def test_vertex_lines_pair_with_the_input_row_that_carries_their_custom_id():
|
||||
vertex_input = b"\n".join(
|
||||
_OpenAIToVertexBatchUploadStream(
|
||||
_jsonl(_request("vx-0", "first"), _request("vx-1", "second")),
|
||||
VertexAIFilesConfig()._map_openai_to_vertex_params,
|
||||
).iter_bytes()
|
||||
)
|
||||
vertex_rows = get_file_content_as_dictionary(vertex_input)
|
||||
assert all("custom_id" not in row for row in vertex_rows)
|
||||
|
||||
payload = _whole_batch(
|
||||
_source(model="gemini-2.5-flash", custom_llm_provider="vertex_ai", cost_provider="vertex_ai"),
|
||||
vertex_input,
|
||||
(_jsonl(_success("vx-1", 4, 2), _success("vx-0", 3, 1)),),
|
||||
)
|
||||
|
||||
assert [(line["custom_id"], line["request"]) for line in payload["lines"]] == [
|
||||
("vx-1", vertex_rows[1]),
|
||||
("vx-0", vertex_rows[0]),
|
||||
]
|
||||
|
||||
|
||||
def test_a_vertex_embeddings_entry_split_into_several_rows_pairs_with_none_of_them():
|
||||
embeddings = {
|
||||
"custom_id": "emb-1",
|
||||
"method": "POST",
|
||||
"url": "/v1/embeddings",
|
||||
"body": {"model": "gemini-embedding-001", "input": ["first", "second"]},
|
||||
}
|
||||
vertex_input = b"\n".join(
|
||||
_OpenAIToVertexBatchUploadStream(
|
||||
_jsonl(embeddings, _request("vx-0", "chat")), VertexAIFilesConfig()._map_openai_to_vertex_params
|
||||
).iter_bytes()
|
||||
)
|
||||
vertex_rows = get_file_content_as_dictionary(vertex_input)
|
||||
assert [row.get("key") for row in vertex_rows] == ["emb-1#0/2", "emb-1#1/2", None]
|
||||
|
||||
payload = _whole_batch(
|
||||
_source(model="gemini-2.5-flash", custom_llm_provider="vertex_ai", cost_provider="vertex_ai"),
|
||||
vertex_input,
|
||||
(_jsonl(_success("emb-1", 2, 0), _success("vx-0", 3, 1)),),
|
||||
)
|
||||
|
||||
assert [(line["custom_id"], line["request"]) for line in payload["lines"]] == [
|
||||
("emb-1", None),
|
||||
("vx-0", vertex_rows[2]),
|
||||
]
|
||||
|
||||
|
||||
def test_a_vertex_input_row_without_a_custom_id_pairs_with_nothing():
|
||||
unlabeled = {"request": {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}}
|
||||
|
||||
payload = _whole_batch(
|
||||
_source(model="gemini-2.5-flash", custom_llm_provider="vertex_ai", cost_provider="vertex_ai"),
|
||||
_jsonl(unlabeled),
|
||||
(_jsonl({**_success("x", 3, 1), "custom_id": None}),),
|
||||
)
|
||||
|
||||
(line,) = payload["lines"]
|
||||
assert (line["custom_id"], line["request"]) == (None, None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"vertex_input",
|
||||
[
|
||||
pytest.param(
|
||||
b"\n".join(
|
||||
_OpenAIToVertexBatchUploadStream(
|
||||
_jsonl(_request("unknown", "first"), {**_request("vx-1", "second"), "custom_id": None}),
|
||||
VertexAIFilesConfig()._map_openai_to_vertex_params,
|
||||
).iter_bytes()
|
||||
),
|
||||
id="custom-id-unknown-next-to-a-null-one",
|
||||
),
|
||||
pytest.param(
|
||||
_jsonl(
|
||||
{"request": {"contents": [{"role": "user", "parts": [{"text": "first"}]}], "labels": {"team": "a"}}},
|
||||
{"request": {"contents": [{"role": "user", "parts": [{"text": "second"}]}]}},
|
||||
),
|
||||
id="native-rows-with-other-labels-next-to-unlabeled-ones",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_a_vertex_line_read_as_unknown_pairs_with_no_input_row(vertex_input):
|
||||
payload = _whole_batch(
|
||||
_source(model="gemini-2.5-flash", custom_llm_provider="vertex_ai", cost_provider="vertex_ai"),
|
||||
vertex_input,
|
||||
(_jsonl(_success("unknown", 3, 1), _success("unknown", 4, 2)),),
|
||||
)
|
||||
|
||||
assert [(line["custom_id"], line["request"]) for line in payload["lines"]] == [("unknown", None)] * 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redaction_drops_only_the_request_and_response():
|
||||
payload = _payload()
|
||||
|
||||
redacted = without_request_and_response(payload)
|
||||
|
||||
assert redacted["batch_id"] == payload["batch_id"]
|
||||
assert redacted["finished_at"] == payload["finished_at"]
|
||||
assert redacted["part"] == payload["part"]
|
||||
assert [{**line, "request": None, "response": None} for line in payload["lines"]] == list(redacted["lines"])
|
||||
assert all(line["request"] is not None for line in payload["lines"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_object_holds_one_json_line_per_batch_line_under_its_finished_day():
|
||||
payload = _payload()
|
||||
|
||||
ndjson = batch_line_items_ndjson(payload)
|
||||
|
||||
assert [json.loads(line) for line in ndjson.splitlines()] == list(payload["lines"])
|
||||
assert ndjson.endswith("}\n")
|
||||
assert (
|
||||
batch_line_items_object_name(payload)
|
||||
== "batch_line_items/2026-10-08/bGl0ZWxsbV9wcm94eTtiYXRjaA/part-00000.jsonl"
|
||||
)
|
||||
assert (
|
||||
batch_line_items_object_name(
|
||||
{**payload, "batch_id": "projects/p/locations/l/batchPredictionJobs/9", "part": 12}
|
||||
)
|
||||
== "batch_line_items/2026-10-08/projects_p_locations_l_batchPredictionJobs_9/part-00012.jsonl"
|
||||
)
|
||||
|
||||
|
||||
def _raw_size(request: dict, result: dict) -> int:
|
||||
return len(json.dumps(request)) + len(json.dumps(result))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"extra_bytes, parts",
|
||||
[(0, [["a", "b"], ["c"]]), (1, [["a", "b", "c"]])],
|
||||
ids=["closes once its raw lines reach the size", "stays open below the size"],
|
||||
)
|
||||
def test_a_part_closes_once_its_raw_request_and_result_lines_reach_the_part_size(extra_bytes, parts):
|
||||
requests = [_request(custom_id, f"question {custom_id}") for custom_id in "abc"]
|
||||
results = [_success(custom_id, 10, 5) for custom_id in "abc"]
|
||||
max_part_bytes = _raw_size(requests[0], results[0]) + _raw_size(requests[1], results[1]) + extra_bytes
|
||||
|
||||
payloads = list(batch_line_item_parts(_source(), _jsonl(*requests), (_jsonl(*results),), max_part_bytes))
|
||||
|
||||
assert [payload["part"] for payload in payloads] == list(range(len(parts)))
|
||||
assert [[line["custom_id"] for line in payload["lines"]] for payload in payloads] == parts
|
||||
assert [line["request"] for payload in payloads for line in payload["lines"]] == requests
|
||||
|
||||
|
||||
def test_blank_and_unreadable_input_lines_leave_the_other_lines_paired():
|
||||
input_content = b"\n".join(
|
||||
[json.dumps(_request("a", "first")).encode(), b"", b"not json", json.dumps(_request("b", "second")).encode()]
|
||||
)
|
||||
|
||||
payload = _whole_batch(_source(), input_content, (_jsonl(_success("a", 1, 1), _success("b", 1, 1)),))
|
||||
|
||||
assert [line["request"] for line in payload["lines"]] == [_request("a", "first"), _request("b", "second")]
|
||||
|
||||
|
||||
def test_files_with_no_readable_line_make_no_part():
|
||||
assert list(batch_line_item_parts(_source(), b"", (b"not json\n", b""), BATCH_LINE_ITEMS_MAX_PART_BYTES)) == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"saved, recorded",
|
||||
[
|
||||
(ALICE_KEY_HASH, ALICE_KEY_HASH),
|
||||
("ya29.opaque-oauth2-access-token", sha256(b"ya29.opaque-oauth2-access-token").hexdigest()),
|
||||
("litellm_proxy_master_key", sha256(b"litellm_proxy_master_key").hexdigest()),
|
||||
(None, None),
|
||||
],
|
||||
ids=["virtual key hash", "oauth2 token", "master key alias", "nothing saved"],
|
||||
)
|
||||
def test_the_creator_credential_is_recorded_only_as_a_hash(saved, recorded):
|
||||
assert line_item_metadata({"user_api_key_hash": saved})["user_api_key_hash"] == recorded
|
||||
|
||||
|
||||
def test_a_line_holding_a_lone_surrogate_is_escaped_so_the_object_still_encodes_as_utf8():
|
||||
payload = _payload()
|
||||
(line, *_) = payload["lines"]
|
||||
surrogate_line = {**line, "custom_id": "emoji-\ud83d", "response": {"body": {"content": "caf\u00e9 \ud83d"}}}
|
||||
readable_line = {**line, "custom_id": "caf\u00e9"}
|
||||
|
||||
ndjson = batch_line_items_ndjson({**payload, "lines": (surrogate_line, readable_line)})
|
||||
|
||||
surrogate_text, readable_text = ndjson.encode("utf-8").decode("utf-8").splitlines()
|
||||
assert json.loads(surrogate_text) == surrogate_line
|
||||
assert "\\ud83d" in surrogate_text
|
||||
assert json.loads(readable_text) == readable_line
|
||||
assert '"custom_id": "caf\u00e9"' in readable_text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_logger_gets_the_request_and_response_unless_message_logging_is_off(monkeypatch):
|
||||
payload = _payload()
|
||||
open_logger, quiet_logger, legacy_logger = (
|
||||
_RecordingLogger(),
|
||||
_RecordingLogger(turn_off_message_logging=True),
|
||||
_RecordingLogger(message_logging=False),
|
||||
)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (open_logger, quiet_logger, legacy_logger))
|
||||
|
||||
await dispatch_batch_line_items(payload, redact_bodies=False)
|
||||
|
||||
assert open_logger.received == [payload]
|
||||
assert quiet_logger.received == [without_request_and_response(payload)]
|
||||
assert legacy_logger.received == [without_request_and_response(payload)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redacted_dispatch_never_hands_a_logger_the_bodies(monkeypatch):
|
||||
payload = _payload()
|
||||
logger = _RecordingLogger()
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
|
||||
await dispatch_batch_line_items(payload, redact_bodies=True)
|
||||
|
||||
assert logger.received == [without_request_and_response(payload)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failing_or_unknown_sink_is_logged_and_the_other_sinks_still_get_the_lines(monkeypatch, caplog):
|
||||
payload = _payload()
|
||||
logger = _RecordingLogger()
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (_FailingLogger(), "langfuse", logger))
|
||||
|
||||
with caplog.at_level(logging.ERROR, logger="LiteLLM"):
|
||||
await dispatch_batch_line_items(payload, redact_bodies=False)
|
||||
|
||||
assert logger.received == [payload]
|
||||
failures = [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR]
|
||||
assert len(failures) == 2
|
||||
assert all(payload["batch_id"] in message for message in failures)
|
||||
assert any("bucket is gone" in message for message in failures)
|
||||
assert any("'langfuse'" in message for message in failures)
|
||||
|
|
@ -330,3 +330,67 @@ async def test_unbatched_gcs_logs_upload_individual_success_and_failure_objects(
|
|||
flush_task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await flush_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_line_items_are_one_ndjson_object_in_the_default_bucket(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
logger, uploads = _gcs_logger_with_upload_status(monkeypatch, 200)
|
||||
|
||||
def stored(request: httpx.Request) -> httpx.Response:
|
||||
uploads(request=request)
|
||||
return httpx.Response(200, json={"name": request.url.params["name"]})
|
||||
|
||||
logger.async_httpx_client = AsyncHTTPHandler(transport=httpx.MockTransport(stored))
|
||||
lines: Final = (
|
||||
{"batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", "custom_id": "ok-1", "status": "success", "response_cost": 0.5},
|
||||
{"batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", "custom_id": "bad-1", "status": "failure", "response_cost": 0.0},
|
||||
)
|
||||
|
||||
await logger.async_log_batch_line_items(
|
||||
{"batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", "finished_at": 1791494130, "part": 3, "lines": lines}
|
||||
)
|
||||
|
||||
request: Final = uploads.call_args.kwargs["request"]
|
||||
assert uploads.call_count == 1
|
||||
assert request.url.path == "/upload/storage/v1/b/test-bucket/o"
|
||||
assert request.url.params["name"] == "batch_line_items/2026-10-08/bGl0ZWxsbV9wcm94eTtiYXRjaA/part-00003.jsonl"
|
||||
assert request.headers["Authorization"] == "Bearer test-access-token"
|
||||
assert [json.loads(line) for line in request.content.decode().splitlines()] == list(lines)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_rejected_batch_line_items_write_raises_for_the_dispatcher(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
logger, uploads = _gcs_logger_with_upload_status(monkeypatch, 503)
|
||||
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
await logger.async_log_batch_line_items(
|
||||
{"batch_id": "b", "finished_at": 1791494130, "part": 0, "lines": ({"custom_id": "ok-1"},)}
|
||||
)
|
||||
|
||||
assert uploads.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_line_items_are_serialized_off_the_event_loop(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import threading
|
||||
|
||||
import litellm.integrations.gcs_bucket.gcs_bucket as gcs_bucket_module
|
||||
from litellm.types.utils import StandardBatchLineItemsPayload
|
||||
|
||||
logger, _ = _gcs_logger_with_upload_status(monkeypatch, 200)
|
||||
logger.async_httpx_client = AsyncHTTPHandler(
|
||||
transport=httpx.MockTransport(lambda request: httpx.Response(200, json={"name": request.url.params["name"]}))
|
||||
)
|
||||
serialize: Final = gcs_bucket_module.batch_line_items_ndjson
|
||||
threads: Final[list[threading.Thread]] = []
|
||||
|
||||
def _serialize(payload: StandardBatchLineItemsPayload) -> str:
|
||||
threads.append(threading.current_thread())
|
||||
return serialize(payload)
|
||||
|
||||
monkeypatch.setattr(gcs_bucket_module, "batch_line_items_ndjson", _serialize)
|
||||
|
||||
await logger.async_log_batch_line_items({"batch_id": "b", "finished_at": 1791494130, "part": 0, "lines": ()})
|
||||
|
||||
(thread,) = threads
|
||||
assert thread is not threading.main_thread()
|
||||
|
|
|
|||
|
|
@ -5848,3 +5848,126 @@ class TestS3Logger(S3Logger):
|
|||
self.recorded_requests[response_obj["id"]] = start_time
|
||||
self.logged_standard_logging_payload = kwargs["standard_logging_object"]
|
||||
return await super().async_log_success_event(kwargs, response_obj, start_time, end_time)
|
||||
|
||||
|
||||
def _batch_line_items_payload() -> dict:
|
||||
line = {
|
||||
"batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA",
|
||||
"custom_id": "ok-1",
|
||||
"status": "success",
|
||||
"request": {"custom_id": "ok-1", "body": {"messages": [{"role": "user", "content": "hi"}]}},
|
||||
"response": {"custom_id": "ok-1", "response": {"status_code": 200, "body": {"id": "chatcmpl-1"}}},
|
||||
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
|
||||
"response_cost": 0.005,
|
||||
"model": "gpt-4o-mini",
|
||||
"custom_llm_provider": "openai",
|
||||
"metadata": {"user_api_key_hash": "hash-alice"},
|
||||
}
|
||||
return {
|
||||
"batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA",
|
||||
"finished_at": 1791494130,
|
||||
"part": 0,
|
||||
"lines": (line, {**line, "custom_id": "ok-2"}),
|
||||
}
|
||||
|
||||
|
||||
def _s3_logger_answering(
|
||||
status: int, s3_path: str | None, content: bytes = b""
|
||||
) -> tuple[S3Logger, list[httpx.Request]]:
|
||||
requests: list[httpx.Request] = []
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
return httpx.Response(status, content=content, request=request)
|
||||
|
||||
logger = S3Logger(
|
||||
s3_bucket_name="test-bucket",
|
||||
s3_region_name="us-east-1",
|
||||
s3_path=s3_path,
|
||||
s3_aws_access_key_id="AKIAIOSFODNN7EXAMPLE",
|
||||
s3_aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
||||
s3_flush_interval=3600,
|
||||
)
|
||||
logger.async_httpx_client = AsyncHTTPHandler(transport=httpx.MockTransport(respond))
|
||||
return logger, requests
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_line_items_are_one_ndjson_object_under_the_configured_path():
|
||||
logger, requests = _s3_logger_answering(200, "cold/")
|
||||
payload = _batch_line_items_payload()
|
||||
|
||||
await logger.async_log_batch_line_items(payload)
|
||||
|
||||
(request,) = requests
|
||||
assert request.method == "PUT"
|
||||
assert request.url.path == "/cold/batch_line_items/2026-10-08/bGl0ZWxsbV9wcm94eTtiYXRjaA/part-00000.jsonl"
|
||||
assert request.headers["Content-Type"] == "application/x-ndjson"
|
||||
assert request.headers["Content-Disposition"] == 'inline; filename="bGl0ZWxsbV9wcm94eTtiYXRjaA-part-00000.jsonl"'
|
||||
assert [json.loads(line) for line in request.content.decode().splitlines()] == list(payload["lines"])
|
||||
assert logger.log_queue == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_line_items_without_an_s3_path_sit_at_the_bucket_root():
|
||||
logger, requests = _s3_logger_answering(200, None)
|
||||
|
||||
await logger.async_log_batch_line_items(_batch_line_items_payload())
|
||||
|
||||
assert [request.url.path for request in requests] == [
|
||||
"/batch_line_items/2026-10-08/bGl0ZWxsbV9wcm94eTtiYXRjaA/part-00000.jsonl"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_rejected_batch_line_items_upload_is_logged_with_the_batch_id(caplog):
|
||||
logger, requests = _s3_logger_answering(400, "cold")
|
||||
|
||||
with caplog.at_level(logging.ERROR, logger="LiteLLM"):
|
||||
await logger.async_log_batch_line_items(_batch_line_items_payload())
|
||||
|
||||
assert len(requests) == 1
|
||||
assert any(
|
||||
"bGl0ZWxsbV9wcm94eTtiYXRjaA" in record.getMessage() and "were not stored" in record.getMessage()
|
||||
for record in caplog.records
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_batch_line_items_object_the_bucket_refuses_for_good_is_not_kept_in_memory():
|
||||
logger, requests = _s3_logger_answering(400, "cold", b"<Error><Code>InvalidArgument</Code></Error>")
|
||||
|
||||
await logger.async_log_batch_line_items(_batch_line_items_payload())
|
||||
await logger.flush_queue()
|
||||
|
||||
assert len(requests) == 1
|
||||
assert logger._flush_dropped == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_line_items_are_serialized_and_hashed_off_the_event_loop(monkeypatch):
|
||||
import threading
|
||||
|
||||
import litellm.integrations.s3_v2 as s3_v2_module
|
||||
|
||||
logger, requests = _s3_logger_answering(200, "cold/")
|
||||
serialize = s3_v2_module.batch_line_items_ndjson
|
||||
prepare = logger._prepare_put
|
||||
threads = {}
|
||||
|
||||
def _serialize(payload):
|
||||
threads["serialize"] = threading.current_thread()
|
||||
return serialize(payload)
|
||||
|
||||
def _prepare(element):
|
||||
threads["hash"] = threading.current_thread()
|
||||
return prepare(element)
|
||||
|
||||
monkeypatch.setattr(s3_v2_module, "batch_line_items_ndjson", _serialize)
|
||||
monkeypatch.setattr(logger, "_prepare_put", _prepare)
|
||||
|
||||
await logger.async_log_batch_line_items(_batch_line_items_payload())
|
||||
|
||||
assert len(requests) == 1
|
||||
assert set(threads) == {"serialize", "hash"}
|
||||
assert threading.main_thread() not in threads.values()
|
||||
|
|
|
|||
|
|
@ -8,7 +8,9 @@ ARN unified_object_id) batches with no managed unified id.
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from contextlib import contextmanager
|
||||
from hashlib import sha256
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -3242,3 +3244,703 @@ class TestMultiPodBatchCostClaim:
|
|||
assert self._claim_calls(prisma) == []
|
||||
assert journal == ["fetch", "bill", "finalize"]
|
||||
logging_obj.async_success_handler.assert_awaited_once()
|
||||
|
||||
|
||||
class TestBatchLineItemsFromThePoller:
|
||||
"""When batch_line_item_callbacks is set, the poller hands every line of a finished batch to those loggers
|
||||
once the batch is billed and retired, in the background, so a slow, hung or failing sink never unbills, re-polls
|
||||
or holds up a batch."""
|
||||
|
||||
_INPUT = b'{"custom_id":"ok-1","body":{"messages":[{"role":"user","content":"hi"}]}}\n'
|
||||
_OUTPUT = (
|
||||
b'{"custom_id":"ok-1","response":{"status_code":200,"body":{"model":"gpt-4",'
|
||||
b'"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}}}\n'
|
||||
)
|
||||
_ERROR = b'{"custom_id":"ok-1","response":{"status_code":400,"body":{"error":{"message":"expired"}}}}\n'
|
||||
|
||||
@staticmethod
|
||||
def _recording_logger(journal: list, error: Exception | None = None):
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
class _Recorder(CustomLogger):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.payloads = []
|
||||
|
||||
async def async_log_batch_line_items(self, payload):
|
||||
journal.append("lines")
|
||||
self.payloads.append(payload)
|
||||
if error is not None:
|
||||
raise error
|
||||
|
||||
return _Recorder()
|
||||
|
||||
@staticmethod
|
||||
def _setup(journal: list, *, status="completed", output_file_id=_CLAIM_OUTPUT_FILE_ID, error_file_id=None):
|
||||
row = _FakeManagedObjectRow()
|
||||
prisma = TestMultiPodBatchCostClaim._prisma(row, journal)
|
||||
router = TestMultiPodBatchCostClaim._router()
|
||||
response = router.aretrieve_batch.return_value
|
||||
response.status = status
|
||||
response.output_file_id = output_file_id
|
||||
response.error_file_id = error_file_id
|
||||
response.input_file_id = "file-input-1"
|
||||
return row, TestMultiPodBatchCostClaim._instance(prisma, router)
|
||||
|
||||
@staticmethod
|
||||
async def _stored(instance) -> None:
|
||||
"""Wait for the line items the poll handed to the background."""
|
||||
await asyncio.gather(*instance._line_item_stores)
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
def _provider(cls, journal: list):
|
||||
"""The provider's batch files by id, and a spend-log writer that journals each write."""
|
||||
files = {"file-input-1": cls._INPUT, _CLAIM_OUTPUT_FILE_ID: cls._OUTPUT, "file-err-1": cls._ERROR}
|
||||
|
||||
async def _afile_content(**kwargs):
|
||||
journal.append(f"read {kwargs['file_id']}")
|
||||
if kwargs["file_id"] not in files:
|
||||
from litellm.exceptions import NotFoundError
|
||||
|
||||
raise NotFoundError(
|
||||
message=f"No file found with id '{kwargs['file_id']}'.", model="gpt-4", llm_provider="openai"
|
||||
)
|
||||
content = MagicMock()
|
||||
content.content = files[kwargs["file_id"]]
|
||||
return content
|
||||
|
||||
def _is_b64(file_id):
|
||||
if file_id == _CLAIM_UNIFIED_BATCH_ID:
|
||||
return "llm_model_id,model-123;llm_batch_id,batch-456;"
|
||||
return False
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.async_success_handler = AsyncMock(side_effect=lambda **kwargs: journal.append("bill"))
|
||||
with (
|
||||
patch(_IS_B64, side_effect=_is_b64),
|
||||
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",
|
||||
),
|
||||
patch("litellm.files.main.afile_content", new=AsyncMock(side_effect=_afile_content)),
|
||||
patch("litellm.litellm_core_utils.litellm_logging.Logging", return_value=logging_obj),
|
||||
):
|
||||
yield
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lines_reach_the_callbacks_after_the_claimed_batch_is_billed_and_retired(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal)
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal == [
|
||||
f"read {_CLAIM_OUTPUT_FILE_ID}",
|
||||
"claim",
|
||||
"bill",
|
||||
"finalize",
|
||||
f"read {_CLAIM_OUTPUT_FILE_ID}",
|
||||
"read file-input-1",
|
||||
"lines",
|
||||
]
|
||||
(payload,) = logger.payloads
|
||||
assert payload["batch_id"] == _CLAIM_UNIFIED_BATCH_ID
|
||||
assert payload["finished_at"] == 2
|
||||
assert payload["part"] == 0
|
||||
(line,) = payload["lines"]
|
||||
assert (line["custom_id"], line["status"]) == ("ok-1", "success")
|
||||
assert line["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
|
||||
assert line["metadata"]["user_api_key_user_id"] == "user-1"
|
||||
assert row.batch_processed is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_lines_are_priced_off_the_event_loop(self, monkeypatch):
|
||||
import threading
|
||||
|
||||
import litellm
|
||||
from litellm.batches import batch_line_items
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (self._recording_logger(journal),))
|
||||
price = batch_line_items._line_item
|
||||
threads = []
|
||||
|
||||
def _price(*args):
|
||||
threads.append(threading.current_thread())
|
||||
return price(*args)
|
||||
|
||||
monkeypatch.setattr(batch_line_items, "_line_item", _price)
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal[-1] == "lines"
|
||||
(thread,) = threads
|
||||
assert thread is not threading.main_thread()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_creator_that_cannot_be_looked_up_gets_lines_without_bodies(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal)
|
||||
row.api_key = sha256(b"sk-deleted-key").hexdigest()
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
(payload,) = logger.payloads
|
||||
assert [(line["request"], line["response"]) for line in payload["lines"]] == [(None, None)]
|
||||
assert "read file-input-1" not in journal
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_creator_with_logging_on_gets_the_request_and_response(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal)
|
||||
row.api_key = "cli-session-user-1"
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
monkeypatch.setattr(litellm, "turn_off_message_logging", False)
|
||||
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal[-4:] == ["finalize", f"read {_CLAIM_OUTPUT_FILE_ID}", "read file-input-1", "lines"]
|
||||
((line,),) = [payload["lines"] for payload in logger.payloads]
|
||||
assert line["request"] == json.loads(self._INPUT)
|
||||
assert line["response"] == json.loads(self._OUTPUT)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failing_sink_leaves_the_batch_billed_and_retired(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal)
|
||||
monkeypatch.setattr(
|
||||
litellm, "batch_line_item_callbacks", (self._recording_logger(journal, RuntimeError("bucket gone")),)
|
||||
)
|
||||
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert [step for step in journal if not step.startswith("read")] == ["claim", "bill", "finalize", "lines"]
|
||||
assert row.batch_processed is True
|
||||
assert row.status == "complete"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_without_callbacks_only_the_output_file_is_read(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
_, instance = self._setup(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", ())
|
||||
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_batch_where_every_line_failed_sends_its_error_lines_after_it_is_retired(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal, status="expired", output_file_id=None, error_file_id="file-err-1")
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal == ["finalize", "read file-err-1", "read file-input-1", "lines"]
|
||||
(payload,) = logger.payloads
|
||||
assert [(line["custom_id"], line["status"], line["response_cost"]) for line in payload["lines"]] == [
|
||||
("ok-1", "failure", 0.0)
|
||||
]
|
||||
assert row.batch_processed is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_batch_retired_with_nothing_billable_reads_its_input_file_by_the_id_the_provider_reported(
|
||||
self, monkeypatch
|
||||
):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
_, instance = self._setup(journal, status="expired", output_file_id=None, error_file_id="file-err-1")
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
|
||||
async def _persist_managed_ids(*, response, **_):
|
||||
response.input_file_id = "managed-id-of-the-copy-on-another-deployment"
|
||||
|
||||
with (
|
||||
self._provider(journal),
|
||||
patch(
|
||||
"litellm.proxy.openai_files_endpoints.common_utils.ensure_batch_response_managed_file_ids",
|
||||
new=AsyncMock(side_effect=_persist_managed_ids),
|
||||
),
|
||||
):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal == ["finalize", "read file-err-1", "read file-input-1", "lines"]
|
||||
((line,),) = [payload["lines"] for payload in logger.payloads]
|
||||
assert line["request"] == json.loads(self._INPUT)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_batch_whose_files_hold_no_readable_line_sends_no_lines(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
_, instance = self._setup(journal, status="expired", output_file_id=None, error_file_id="file-err-1")
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
monkeypatch.setattr(type(self), "_ERROR", b"not json\n")
|
||||
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal == ["finalize", "read file-err-1", "read file-input-1"]
|
||||
assert logger.payloads == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_batch_that_could_not_be_retired_sends_no_lines_until_it_is(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal, status="expired", output_file_id=None, error_file_id="file-err-1")
|
||||
instance.prisma_client.db.litellm_managedobjecttable.update = AsyncMock(side_effect=RuntimeError("db down"))
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal == []
|
||||
assert logger.payloads == []
|
||||
assert row.batch_processed is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("has_batch_processed_column, stored", [(True, 1), (False, 0)])
|
||||
async def test_a_billed_batch_that_could_not_be_marked_complete_sends_lines_only_if_its_claim_retired_it(
|
||||
self, monkeypatch, has_batch_processed_column, stored
|
||||
):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal)
|
||||
instance._has_batch_processed_column = has_batch_processed_column
|
||||
instance.prisma_client.db.litellm_managedobjecttable.update = AsyncMock(side_effect=RuntimeError("db down"))
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal.count("bill") == 1
|
||||
assert len(logger.payloads) == stored
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_batch_whose_output_file_is_gone_still_sends_its_error_lines_after_it_is_retired(
|
||||
self, monkeypatch
|
||||
):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal, output_file_id="file-out-deleted", error_file_id="file-err-1")
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal == ["read file-out-deleted", "finalize", "read file-err-1", "read file-input-1", "lines"]
|
||||
(payload,) = logger.payloads
|
||||
assert [(line["custom_id"], line["status"]) for line in payload["lines"]] == [("ok-1", "failure")]
|
||||
assert row.batch_processed is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_terminal_batch_with_no_lines_sends_nothing(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
_, instance = self._setup(journal, status="failed", output_file_id=None, error_file_id=None)
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal == ["finalize"]
|
||||
assert logger.payloads == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_hung_sink_never_holds_up_billing_and_retiring_the_batch(self, monkeypatch):
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal)
|
||||
released = asyncio.Event()
|
||||
|
||||
class _HungSink(CustomLogger):
|
||||
async def async_log_batch_line_items(self, payload):
|
||||
journal.append("lines started")
|
||||
await released.wait()
|
||||
journal.append("lines")
|
||||
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (_HungSink(),))
|
||||
|
||||
async def _sink_started():
|
||||
while "lines started" not in journal:
|
||||
await asyncio.sleep(0)
|
||||
|
||||
with self._provider(journal):
|
||||
await asyncio.wait_for(instance.check_batch_cost(), timeout=5)
|
||||
await asyncio.wait_for(_sink_started(), timeout=5)
|
||||
billed_while_hung = (row.batch_processed, row.status, list(journal))
|
||||
released.set()
|
||||
await self._stored(instance)
|
||||
|
||||
assert billed_while_hung == (
|
||||
True,
|
||||
"complete",
|
||||
[
|
||||
f"read {_CLAIM_OUTPUT_FILE_ID}",
|
||||
"claim",
|
||||
"bill",
|
||||
"finalize",
|
||||
f"read {_CLAIM_OUTPUT_FILE_ID}",
|
||||
"read file-input-1",
|
||||
"lines started",
|
||||
],
|
||||
)
|
||||
assert journal[-1] == "lines"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lines_are_dropped_with_an_error_while_a_poll_cycle_of_batches_is_still_waiting(
|
||||
self, monkeypatch, caplog
|
||||
):
|
||||
import litellm
|
||||
from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal)
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
waiting = [asyncio.create_task(asyncio.sleep(3600)) for _ in range(MAX_OBJECTS_PER_POLL_CYCLE)]
|
||||
instance._line_item_stores.update(waiting)
|
||||
|
||||
try:
|
||||
with self._provider(journal), caplog.at_level(logging.ERROR):
|
||||
await instance.check_batch_cost()
|
||||
finally:
|
||||
for task in waiting:
|
||||
task.cancel()
|
||||
|
||||
assert logger.payloads == []
|
||||
assert row.batch_processed is True
|
||||
assert f"so the line items of batch batch-456 (job {row.id}) were not stored" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batches_store_their_lines_one_at_a_time_in_the_order_they_were_retired(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
_, instance = self._setup(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (self._recording_logger(journal),))
|
||||
running: list = []
|
||||
|
||||
async def _store(job, response, model_id, batch_id, *, billed):
|
||||
running.append(batch_id)
|
||||
journal.append(f"start {batch_id} with {len(running)} running")
|
||||
await asyncio.sleep(0.01)
|
||||
running.remove(batch_id)
|
||||
|
||||
monkeypatch.setattr(instance, "_store_batch_line_items", _store)
|
||||
for batch_id in ("batch-a", "batch-b", "batch-c"):
|
||||
instance._store_batch_line_items_in_background(MagicMock(), MagicMock(), "model-123", batch_id, billed=True)
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal == [f"start batch-{name} with 1 running" for name in "abc"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lines_are_stored_only_once_the_poll_has_billed_and_retired_every_batch_it_found(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
_, instance = self._setup(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (self._recording_logger(journal),))
|
||||
|
||||
async def _an_earlier_batch_waits_to_store_its_lines():
|
||||
instance._store_batch_line_items_in_background(
|
||||
MagicMock(), MagicMock(), "model-123", "batch-earlier", billed=True
|
||||
)
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
async def _store(job, response, model_id, batch_id, *, billed):
|
||||
journal.append(f"store {batch_id}")
|
||||
|
||||
monkeypatch.setattr(instance, "_cleanup_stale_managed_objects", _an_earlier_batch_waits_to_store_its_lines)
|
||||
monkeypatch.setattr(instance, "_store_batch_line_items", _store)
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal == [
|
||||
f"read {_CLAIM_OUTPUT_FILE_ID}",
|
||||
"claim",
|
||||
"bill",
|
||||
"finalize",
|
||||
"store batch-earlier",
|
||||
"store batch-456",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_output_file_that_cannot_be_read_again_leaves_the_batch_billed_and_stores_nothing(
|
||||
self, monkeypatch, caplog
|
||||
):
|
||||
import litellm
|
||||
from litellm.batches import batch_utils
|
||||
from litellm.exceptions import NotFoundError
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal)
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
fetch = batch_utils.fetch_batch_managed_file_content
|
||||
|
||||
async def _output_gone_after_billing(file_id, **kwargs):
|
||||
if file_id == _CLAIM_OUTPUT_FILE_ID:
|
||||
raise NotFoundError(message="No file found.", model="gpt-4", llm_provider="openai")
|
||||
return await fetch(file_id, **kwargs)
|
||||
|
||||
monkeypatch.setattr(batch_utils, "fetch_batch_managed_file_content", _output_gone_after_billing)
|
||||
with self._provider(journal), caplog.at_level(logging.ERROR):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"]
|
||||
assert logger.payloads == []
|
||||
assert (row.batch_processed, row.status) == (True, "complete")
|
||||
assert f"the line items of batch batch-456 (job {row.id}) were not stored" in caplog.text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_deployment_removed_before_the_lines_are_stored_leaves_the_batch_billed_and_logs_an_error(
|
||||
self, monkeypatch, caplog
|
||||
):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal)
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
deployment = instance.llm_router.get_deployment.return_value
|
||||
instance.llm_router.get_deployment = MagicMock(
|
||||
side_effect=lambda **_: None if instance._poll_idle.is_set() else deployment
|
||||
)
|
||||
|
||||
with self._provider(journal), caplog.at_level(logging.ERROR):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"]
|
||||
assert logger.payloads == []
|
||||
assert (row.batch_processed, row.status) == (True, "complete")
|
||||
assert f"the line items of batch batch-456 (job {row.id}) were not stored" in caplog.text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_batch_bigger_than_one_part_reaches_the_callbacks_as_numbered_parts(self, monkeypatch):
|
||||
import litellm
|
||||
from litellm_enterprise.proxy.common_utils import check_batch_cost
|
||||
|
||||
journal: list = []
|
||||
_, instance = self._setup(journal)
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
monkeypatch.setattr(litellm, "turn_off_message_logging", False)
|
||||
monkeypatch.setattr(check_batch_cost, "BATCH_LINE_ITEMS_MAX_PART_BYTES", 1)
|
||||
monkeypatch.setattr(type(self), "_INPUT", self._INPUT + self._INPUT.replace(b"ok-1", b"ok-2"))
|
||||
monkeypatch.setattr(type(self), "_OUTPUT", self._OUTPUT + self._OUTPUT.replace(b"ok-1", b"ok-2"))
|
||||
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
assert [(payload["part"], [line["custom_id"] for line in payload["lines"]]) for payload in logger.payloads] == [
|
||||
(0, ["ok-1"]),
|
||||
(1, ["ok-2"]),
|
||||
]
|
||||
assert [payload["lines"][0]["request"]["custom_id"] for payload in logger.payloads] == ["ok-1", "ok-2"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_oauth2_creators_token_reaches_the_callbacks_only_as_its_hash(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
journal: list = []
|
||||
row, instance = self._setup(journal)
|
||||
row.api_key = "ya29.opaque-oauth2-access-token"
|
||||
logger = self._recording_logger(journal)
|
||||
monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
|
||||
monkeypatch.setattr(litellm, "turn_off_message_logging", False)
|
||||
|
||||
with self._provider(journal):
|
||||
await instance.check_batch_cost()
|
||||
await self._stored(instance)
|
||||
|
||||
(payload,) = logger.payloads
|
||||
(line,) = payload["lines"]
|
||||
assert line["metadata"]["user_api_key_hash"] == sha256(b"ya29.opaque-oauth2-access-token").hexdigest()
|
||||
assert line["request"] == json.loads(self._INPUT)
|
||||
assert "ya29.opaque-oauth2-access-token" not in json.dumps(payload)
|
||||
|
||||
|
||||
class TestCreatorTurnedOffMessageLogging:
|
||||
"""Batch line items drop the request and response when the creating key or team turned off message logging,
|
||||
when the proxy did, or when the creator can no longer be read."""
|
||||
|
||||
@staticmethod
|
||||
def _instance(key_row=None, team_row=None, error: Exception | None = None):
|
||||
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
|
||||
|
||||
prisma = MagicMock()
|
||||
prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key_row, side_effect=error)
|
||||
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
|
||||
return CheckBatchCost(proxy_logging_obj=MagicMock(), prisma_client=prisma, llm_router=MagicMock())
|
||||
|
||||
@staticmethod
|
||||
def _job(api_key=sha256(b"sk-alice").hexdigest(), team_id=None, created_by="alice"):
|
||||
from types import SimpleNamespace
|
||||
|
||||
return SimpleNamespace(id="job-1", api_key=api_key, team_id=team_id, created_by=created_by)
|
||||
|
||||
@staticmethod
|
||||
def _logging(value: str) -> dict:
|
||||
return {
|
||||
"logging": [
|
||||
{
|
||||
"callback_name": "s3_v2",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {"turn_off_message_logging": value},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("key_value", "proxy_wide", "expected"),
|
||||
[("true", False, True), ("false", True, False), (None, True, True), (None, False, False)],
|
||||
)
|
||||
async def test_the_key_setting_wins_and_the_proxy_setting_fills_in(
|
||||
self, monkeypatch, key_value, proxy_wide, expected
|
||||
):
|
||||
from types import SimpleNamespace
|
||||
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "turn_off_message_logging", proxy_wide)
|
||||
metadata = self._logging(key_value) if key_value is not None else {}
|
||||
instance = self._instance(key_row=SimpleNamespace(metadata=metadata))
|
||||
|
||||
assert await instance._creator_turned_off_message_logging(self._job()) is expected
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_team_that_turned_off_message_logging_covers_its_keys(self, monkeypatch):
|
||||
from types import SimpleNamespace
|
||||
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "turn_off_message_logging", False)
|
||||
instance = self._instance(
|
||||
key_row=SimpleNamespace(metadata={}), team_row=SimpleNamespace(metadata=self._logging("true"))
|
||||
)
|
||||
|
||||
assert await instance._creator_turned_off_message_logging(self._job(team_id="team-alpha")) is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("job_fields", "key_row", "team_row"),
|
||||
[
|
||||
({}, None, None),
|
||||
({"team_id": "team-gone"}, "present", None),
|
||||
({"api_key": None, "team_id": "team-gone"}, None, None),
|
||||
],
|
||||
ids=["deleted-key", "deleted-team", "deleted-team-of-a-batch-created-without-a-key"],
|
||||
)
|
||||
async def test_a_creator_that_is_gone_counts_as_turned_off(self, monkeypatch, job_fields, key_row, team_row):
|
||||
from types import SimpleNamespace
|
||||
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "turn_off_message_logging", False)
|
||||
instance = self._instance(key_row=SimpleNamespace(metadata={}) if key_row else None, team_row=team_row)
|
||||
|
||||
assert await instance._creator_turned_off_message_logging(self._job(**job_fields)) is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("api_key", [None, "opaque-oauth2-access-token"], ids=["jwt", "oauth2"])
|
||||
@pytest.mark.parametrize(
|
||||
("team_value", "proxy_wide", "expected"),
|
||||
[(None, False, False), (None, True, True), ("true", False, True), ("false", True, False)],
|
||||
)
|
||||
async def test_a_batch_created_without_a_virtual_key_follows_its_team_and_the_proxy_setting(
|
||||
self, monkeypatch, api_key, team_value, proxy_wide, expected
|
||||
):
|
||||
from types import SimpleNamespace
|
||||
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "turn_off_message_logging", proxy_wide)
|
||||
team_row = SimpleNamespace(metadata=self._logging(team_value)) if team_value is not None else None
|
||||
instance = self._instance(team_row=team_row)
|
||||
job = self._job(api_key=api_key, team_id="team-alpha" if team_row is not None else None, created_by="alice")
|
||||
|
||||
assert await instance._creator_turned_off_message_logging(job) is expected
|
||||
instance.prisma_client.db.litellm_verificationtoken.find_unique.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_lookup_counts_as_turned_off(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "turn_off_message_logging", False)
|
||||
instance = self._instance(error=Exception("db down"))
|
||||
|
||||
assert await instance._creator_turned_off_message_logging(self._job()) is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("api_key", ["litellm_proxy_master_key", "cli-session-alice"])
|
||||
async def test_the_master_key_and_cli_sessions_have_no_key_row_to_read(self, monkeypatch, api_key):
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "turn_off_message_logging", False)
|
||||
instance = self._instance()
|
||||
|
||||
assert await instance._creator_turned_off_message_logging(self._job(api_key=api_key)) is False
|
||||
instance.prisma_client.db.litellm_verificationtoken.find_unique.assert_not_called()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue