mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
402fa63366
* feat(proxy): cap batch file records, daily batch uploads, and per-file downloads
Adds three opt-in limits for batch jobs, each settable in general_settings as a
per-key default and overridable in key or team metadata by a proxy admin:
max_batch_file_records rejects a purpose=batch upload with more request lines
than allowed with a 413 before it reaches the provider.
max_batch_file_uploads_per_day counts accepted batch uploads per key and per
team in a UTC day and returns 429 with Retry-After once the count is used.
max_file_downloads_per_minute counts GET /v1/files/{id}/content per key and per
team for each file in a one-minute window and returns 429 with Retry-After.
The existing admin-only guard for batch_enqueued_token_limit now covers all four
metadata keys.
* fix(proxy): keep file usage counters in their own store and gate batch limits on user creation
* fix(proxy): count keyless JWT callers per user, require positive file caps, and take the upload slot after request validation
* refactor(proxy): end file usage cap describers with an explicit return after the match
* fix(proxy): keep file usage counters when more than 200 are live without Redis
The file usage counter store used a default in-memory cache, which holds 200
entries and evicts the one that expires soonest. Without Redis, a caller got a
fresh per-file download allowance after touching about 200 other file ids in
the same minute, and a key got a fresh daily upload allowance once about 200
other keys had uploaded that day. The store now tracks up to 20,000 live
counters per worker, the same bound the login throttle uses
* test(proxy): move the file usage cap tests into the directory the proxy shard runs
Main's shard coverage check found tests/unit/proxy/openai_files_endpoints
claimed by no shard, so its tests would not run in CI. The file moves next to
the other files endpoint tests in tests/unit/proxy/openai_files_endpoint, which
the proxy-endpoints shard already runs
* fix(proxy): declare the file usage counters as rate limit calls
Main's redis producer gate requires every module that writes a shared cache to name its key family, and the file usage counters wrote theirs without one.
* test(files): audit batch file usage caps across processes, Redis outages, and config reloads
* test(files): guard the chaos cells against minute boundaries and open Redis breakers
Two chaos cells each failed once in the audit run. The restart check ran
three sequential downloads with no guard against straddling a UTC minute,
and the exact-cap probe after a Redis outage ran while both workers' Redis
circuit breakers were still open (60 s default recovery), so it counted in
per-process memory and the two workers split the cap
Every burst now carries a window guard, the chaos fixture lowers the
breaker recovery to 2 s, and the post-outage check drives a fresh key to its
cap through a one-worker sibling proxy and then expects the two-worker
candidate to refuse the whole burst, which only the shared Redis count can
produce, polled until the breakers close
* test(files): release the held uploads when the killed-worker cell fails early
---------
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
85 lines
3.2 KiB
Python
85 lines
3.2 KiB
Python
import uuid
|
|
from collections.abc import Callable, Mapping
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import yaml
|
|
from integration._support.client import Gateway, eventually
|
|
from integration._support.otlp_sink import Span, SpanSinks, recorded_spans, spans_for_trace
|
|
from integration._support.process import owned_proxy
|
|
from integration._support.wire import wire_server
|
|
from integration.management._batch_file_caps import (
|
|
UPLOADS,
|
|
batch_file,
|
|
caps_config,
|
|
marker,
|
|
provider,
|
|
provider_environment,
|
|
uploads_seen,
|
|
)
|
|
from pydantic import JsonValue
|
|
|
|
COUNTER_SPAN: Final = "redis.incr rate_limits"
|
|
SERVER: Final = 2
|
|
CAP: Final = 3
|
|
|
|
|
|
def _config(directory: Path, otel: Path, provider_url: str) -> Path:
|
|
tracing: Final = yaml.safe_load(otel.read_text())
|
|
caps: Final = yaml.safe_load(caps_config(directory, provider_url, {UPLOADS: CAP}).read_text())
|
|
path: Final = directory / "otel-file-usage-caps.yaml"
|
|
path.write_text(
|
|
yaml.safe_dump(
|
|
{
|
|
**tracing,
|
|
"model_list": caps["model_list"],
|
|
"files_settings": caps["files_settings"],
|
|
"general_settings": {**tracing["general_settings"], UPLOADS: CAP},
|
|
}
|
|
)
|
|
)
|
|
return path
|
|
|
|
|
|
def _names(spans: tuple[Span, ...]) -> list[str]:
|
|
return sorted(span["name"] for span in spans)
|
|
|
|
|
|
def _traced_upload(candidate: Gateway, content: bytes, trace_id: str) -> int:
|
|
response: Final = candidate.client.post(
|
|
"/v1/files",
|
|
data={"purpose": "batch"},
|
|
files={"file": ("batch.jsonl", content, "application/jsonl")},
|
|
headers={
|
|
"Authorization": f"Bearer {candidate.key}",
|
|
"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01",
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
return response.status_code
|
|
|
|
|
|
def test_a_counted_upload_exports_its_counter_increment_as_a_rate_limits_span(
|
|
gateway: Gateway,
|
|
tmp_path: Path,
|
|
audit_sinks: SpanSinks,
|
|
otel_audit_config: Callable[[Path, Mapping[str, JsonValue]], Path],
|
|
) -> None:
|
|
with wire_server(provider) as files:
|
|
config: Final = _config(tmp_path, otel_audit_config(tmp_path, {}), files.url)
|
|
environment: Final = {**provider_environment(files.url), "LITELLM_OTEL_V2": "1"}
|
|
with owned_proxy(gateway, tmp_path, environment, config=config) as candidate:
|
|
since, _ = recorded_spans(audit_sinks.operator)
|
|
mark: Final = marker()
|
|
trace_id: Final = uuid.uuid4().hex
|
|
assert _traced_upload(candidate, batch_file(mark, 1), trace_id) == 200
|
|
trace: Final = eventually(
|
|
lambda: spans_for_trace(recorded_spans(audit_sinks.operator, since)[1], trace_id),
|
|
lambda spans: COUNTER_SPAN in _names(spans),
|
|
seconds=40,
|
|
)
|
|
assert _names(trace).count(COUNTER_SPAN) == 1, _names(trace)
|
|
assert sum(1 for span in trace if span["kind"] == SERVER) == 1, _names(trace)
|
|
(counter,) = (span for span in trace if span["name"] == COUNTER_SPAN)
|
|
assert counter["attributes"].get("litellm.service.target") == "rate_limits", counter["attributes"]
|
|
assert len(uploads_seen(files.drain(), mark)) == 1
|