diff --git a/tests/integration/observability/test_s3_v2_partition_granularity.py b/tests/integration/observability/test_s3_v2_partition_granularity.py new file mode 100644 index 00000000000..b8f52951296 --- /dev/null +++ b/tests/integration/observability/test_s3_v2_partition_granularity.py @@ -0,0 +1,855 @@ +import json +import re +import threading +import uuid +from collections.abc import Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass, field +from datetime import datetime, timedelta +from pathlib import Path +from typing import Final +from urllib.parse import quote, unquote + +import httpx +import openai +import psutil +import pytest +from _s3_v2_support import ( + BUCKET, + PREFIX, + SURFACES, + RecordingS3Sink, + call_surface, + collect_payloads, + matched_ids, + mixed_burst, + s3_config, + surface_reply, +) +from integration._support.client import Gateway, JsonValue, Scenario, eventually, object_value +from integration._support.database import read_rows, scratch_database +from integration._support.process import OwnedProxy, group_members, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server + +FLUSH: Final = {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "1"} +HOUR: Final = {"s3_partition_granularity": "hour"} +ANTHROPIC_MODEL: Final = "anthropic/claude-sonnet-4-5-20250929" +WARNING: Final = "s3 logging: s3_partition_granularity=" +SINK_CREDENTIALS: Final = { + "s3_bucket_name": BUCKET, + "s3_region_name": "us-east-1", + "s3_path": PREFIX, + "s3_aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", + "s3_aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", +} + + +@dataclass(slots=True) +class CountingUpstream: + """Scripted provider that answers every surface and fails any prompt ending in -fail with a 401.""" + + lock: threading.Lock = field(default_factory=threading.Lock) + prompts: list[str] = field(default_factory=list) # mutable-ok: appended per upstream request under lock + + def respond(self, request: Request) -> Reply: + if request.method != "POST" or not request.body: + return Reply(status=404) + body: Final = json.loads(request.body) + prompt: Final = str(body["input"] if "input" in body else body["messages"][0]["content"]) + with self.lock: + self.prompts.append(prompt) + if prompt.endswith("-fail"): + return Reply(status=401, body=b'{"error": {"message": "synthetic upstream rejection", "code": "401"}}') + return surface_reply(request) + + def received(self) -> tuple[str, ...]: + with self.lock: + return tuple(self.prompts) + + +def _prompt(payload: Mapping[str, JsonValue]) -> str: + messages: Final = payload["messages"] + if isinstance(messages, str): + return messages + assert isinstance(messages, list) and len(messages) == 1, payload + first: Final = messages[0] + return first if isinstance(first, str) else str(object_value(first)["content"]) + + +def _start(payload: Mapping[str, JsonValue]) -> datetime: + return datetime.fromtimestamp(float(str(payload["startTime"]))) + + +def _folder(payload: Mapping[str, JsonValue], granularity: str, prefix: str = "") -> str: + start: Final = _start(payload) + hour: Final = f"{start:%H}/" if granularity == "hour" else "" + return f"/{BUCKET}/{PREFIX}/{prefix}{start:%Y-%m-%d}/{hour}" + + +def _object_pattern(payload: Mapping[str, JsonValue], granularity: str, prefix: str = "") -> re.Pattern[str]: + return re.compile( + re.escape(_folder(payload, granularity, prefix)) + rf"time-{_start(payload):%H-%M-%S}-\d{{6}}_[^/]+\.json" + ) + + +def _outside_layout(objects: Mapping[str, bytes], granularity: str, prefix: str = "") -> tuple[str, ...]: + return tuple( + target + for target, body in objects.items() + if not _object_pattern(object_value(json.loads(body)), granularity, prefix).fullmatch(unquote(target)) + ) + + +def _batches_outside_layout(objects: Mapping[str, bytes], granularity: str) -> tuple[str, ...]: + def folders(body: bytes) -> frozenset[str]: + return frozenset(_folder(object_value(json.loads(line)), granularity) for line in body.splitlines()) + + return tuple( + target + for target, body in objects.items() + if len(folders(body)) != 1 + or not re.fullmatch( + re.escape(next(iter(folders(body)))) + r"batch_\d{2}-\d{2}-\d{2}_[0-9a-f]{32}\.jsonl", unquote(target) + ) + ) + + +@contextmanager +def _s3_proxy( + gateway: Gateway, + tmp_path: Path, + sink_url: str, + extra: Mapping[str, JsonValue], + settings: Mapping[str, JsonValue] | None = None, + environment: Mapping[str, str] | None = None, + workers: int = 2, +) -> Iterator[OwnedProxy]: + config: Final = s3_config(tmp_path, sink_url, extra, settings) + with owned_proxy_process( + gateway, tmp_path, {**FLUSH, **(environment or {})}, config=config, workers=workers + ) as owned: + yield owned + + +def _models(scenario: Scenario, provider_url: str, **key_fields: JsonValue) -> tuple[str, str, str]: + openai_model: Final = scenario.model(api_base=provider_url + "/v1", api_key="synthetic-provider-key") + anthropic_model: Final = scenario.model( + model=ANTHROPIC_MODEL, api_base=provider_url, api_key="synthetic-provider-key" + ) + return openai_model, anthropic_model, scenario.key(models=[openai_model, anthropic_model], **key_fields) + + +def _sdk_chats(candidate: Gateway, model: str, key: str, prompts: tuple[str, ...]) -> tuple[str, ...]: + client: Final = openai.OpenAI(base_url=f"{str(candidate.client.base_url).rstrip('/')}/v1", api_key=key) + + def send(prompt: str) -> str: + reply: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": prompt}], extra_body={"cache": {"no-cache": True}} + ) + assert reply.choices[0].finish_reason == "stop", reply.model_dump_json() + return reply.id + + with ThreadPoolExecutor(max_workers=16) as pool: + return tuple(pool.map(send, prompts)) + + +def _surface_prompts(marker: str, per_surface: int) -> frozenset[str]: + return frozenset(f"{marker}-{surface}-{index}" for surface in SURFACES for index in range(per_surface)) + + +def _cold_storage_key(request_id: str, database_url: str | None = None) -> str: + rows: Final = eventually( + lambda: read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,), database_url=database_url + ), + lambda values: len(values) == 1, + seconds=60, + ) + metadata: Final = rows[0]["metadata"] + return str(object_value(json.loads(metadata) if isinstance(metadata, str) else metadata)["cold_storage_object_key"]) + + +def _update_environment(candidate: Gateway, values: Mapping[str, JsonValue]) -> None: + candidate.post( + "/config/update", + {"environment_variables": dict(values), "litellm_settings": {"success_callback": ["s3_v2"]}}, + ) + + +def test_s3_v2_hour_granularity_files_every_surface_under_its_hour_folder(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hour" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, anthropic_model, key = _models(scenario, provider.url) + answered: Final = mixed_burst(owned.gateway, openai_model, anthropic_model, key, marker, per_surface=2) + payloads: Final = collect_payloads(sink, len(answered)) + objects: Final = sink.objects() + log: Final = owned.log.read_text() + sent: Final = _surface_prompts(marker, 2) + assert len(answered) == len(sent) and len(payloads) == len(sent), payloads + assert matched_ids(payloads, answered) == frozenset(str(payload["id"]) for payload in payloads) + assert sorted(upstream.received()) == sorted(sent) + assert len(objects) == len(sent) + assert sorted(_prompt(payload) for payload in payloads) == sorted(sent) + assert all(payload["status"] == "success" for payload in payloads), payloads + assert _outside_layout(objects, "hour") == (), "every object must sit in YYYY-MM-DD/HH/ of its start time" + assert WARNING not in log + + +@pytest.mark.parametrize( + "extra", + [ + pytest.param({}, id="missing"), + pytest.param({"s3_partition_granularity": "day"}, id="day"), + pytest.param({"s3_partition_granularity": ""}, id="empty"), + pytest.param({"s3_partition_granularity": None}, id="null"), + ], +) +def test_s3_v2_missing_day_empty_or_null_granularity_keeps_the_daily_layout( + gateway: Gateway, tmp_path: Path, extra: Mapping[str, JsonValue] +) -> None: + marker: Final = "s3day" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, extra) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, anthropic_model, key = _models(scenario, provider.url) + answered: Final = mixed_burst(owned.gateway, openai_model, anthropic_model, key, marker, per_surface=1) + payloads: Final = collect_payloads(sink, len(answered)) + objects: Final = sink.objects() + log: Final = owned.log.read_text() + sent: Final = _surface_prompts(marker, 1) + assert len(answered) == len(sent) and len(payloads) == len(sent), payloads + assert matched_ids(payloads, answered) == frozenset(str(payload["id"]) for payload in payloads) + assert sorted(upstream.received()) == sorted(sent) + assert sorted(_prompt(payload) for payload in payloads) == sorted(sent) + assert len(objects) == len(sent) + assert _outside_layout(objects, "day") == () + assert WARNING not in log + + +@pytest.mark.parametrize( + ("extra", "environment", "shown"), + [ + pytest.param({"s3_partition_granularity": "hourly"}, {}, "'hourly'", id="unknown_word"), + pytest.param({"s3_partition_granularity": "HOUR"}, {}, "'HOUR'", id="wrong_case"), + pytest.param({"s3_partition_granularity": 1}, {}, "1", id="integer"), + pytest.param({"s3_partition_granularity": ["hour"]}, {}, "['hour']", id="list"), + pytest.param({"s3_partition_granularity": "h" * 5120}, {}, "'[base64_data truncated: 3.8KB]'", id="five_kb"), + pytest.param({}, {"S3_PARTITION_GRANULARITY": "weekly"}, "'weekly'", id="env_unknown_word"), + ], +) +def test_s3_v2_unrecognized_granularity_warns_once_per_worker_and_keeps_the_daily_layout( + gateway: Gateway, tmp_path: Path, extra: Mapping[str, JsonValue], environment: Mapping[str, str], shown: str +) -> None: + marker: Final = "s3bad" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, extra, environment=environment) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, anthropic_model, key = _models(scenario, provider.url) + answered: Final = mixed_burst(owned.gateway, openai_model, anthropic_model, key, marker, per_surface=2) + payloads: Final = collect_payloads(sink, len(answered)) + objects: Final = sink.objects() + warning: Final = f"{WARNING}{shown} is not one of day, hour, using day" + log: Final = eventually(owned.log.read_text, lambda text: warning in text, seconds=15) + sent: Final = _surface_prompts(marker, 2) + assert len(answered) == len(sent) and len(payloads) == len(sent), payloads + assert matched_ids(payloads, answered) == frozenset(str(payload["id"]) for payload in payloads) + assert sorted(upstream.received()) == sorted(sent) + assert sorted(_prompt(payload) for payload in payloads) == sorted(sent) + assert _outside_layout(objects, "day") == () + assert 1 <= log.count(warning) <= 2, "the warning is memoized per distinct value in each of the two workers" + + +@pytest.mark.parametrize( + ("extra", "environment", "granularity"), + [ + pytest.param({}, {"S3_PARTITION_GRANULARITY": "hour"}, "hour", id="env_hour_applies"), + pytest.param({"s3_partition_granularity": "day"}, {"S3_PARTITION_GRANULARITY": "hour"}, "day", id="yaml_wins"), + ], +) +def test_s3_v2_env_granularity_applies_only_when_callback_params_leave_it_unset( + gateway: Gateway, tmp_path: Path, extra: Mapping[str, JsonValue], environment: Mapping[str, str], granularity: str +) -> None: + marker: Final = "s3env" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, extra, environment=environment) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + prompts: Final = tuple(f"{marker}-{index}" for index in range(8)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + payloads: Final = collect_payloads(sink, len(prompts)) + objects: Final = sink.objects() + assert returned == prompts + assert sorted(upstream.received()) == sorted(prompts) + assert frozenset(str(payload["id"]) for payload in payloads) == frozenset(prompts) + assert _outside_layout(objects, granularity) == () + + +def test_s3_v2_hour_batch_files_group_lines_under_the_hour_folder(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hbat" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, {**HOUR, "s3_batch_file_upload": True}) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, anthropic_model, key = _models(scenario, provider.url) + answered: Final = mixed_burst(owned.gateway, openai_model, anthropic_model, key, marker, per_surface=4) + payloads: Final = collect_payloads(sink, len(answered)) + objects: Final = sink.objects() + sent: Final = _surface_prompts(marker, 4) + assert len(answered) == len(sent) and len(payloads) == len(sent), payloads + assert matched_ids(payloads, answered) == frozenset(str(payload["id"]) for payload in payloads) + assert sorted(upstream.received()) == sorted(sent) + assert sorted(_prompt(payload) for payload in payloads) == sorted(sent) + assert _batches_outside_layout(objects, "hour") == () + + +def test_s3_v2_hour_folder_sits_below_the_team_and_key_prefix(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hpre" + uuid.uuid4().hex[:8] + team_alias: Final = f"alpha-{uuid.uuid4().hex[:8]}" + key_alias: Final = f"beta-{uuid.uuid4().hex[:8]}" + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + extra: Final = {**HOUR, "s3_use_team_prefix": True, "s3_use_key_prefix": True} + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, extra) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + team: Final = scenario.team(team_alias=team_alias, models=[openai_model]) + key: Final = scenario.key(team_id=team, key_alias=key_alias, models=[openai_model]) + prompts: Final = tuple(f"{marker}-{index}" for index in range(6)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + payloads: Final = collect_payloads(sink, len(prompts)) + objects: Final = sink.objects() + assert returned == prompts + assert sorted(upstream.received()) == sorted(prompts) + assert frozenset(str(payload["id"]) for payload in payloads) == frozenset(prompts) + assert _outside_layout(objects, "hour", f"{team_alias}/{key_alias}/") == () + + +def test_s3_v2_hour_failure_and_rejected_requests_keep_the_hour_layout(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hfail" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + + def send(prompt: str, model: str = openai_model, caller: str = key) -> httpx.Response: + return owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": prompt}], "cache": {"no-cache": True}}, + key=caller, + ) + + successes: Final = tuple(f"{marker}-{index}" for index in range(4)) + failures: Final = tuple(f"{marker}-{index}-fail" for index in range(3)) + with ThreadPoolExecutor(max_workers=8) as pool: + responses: Final = tuple(pool.map(send, (*successes, *failures))) + ghost: Final = send(f"{marker}-ghost", model=f"ghost-{uuid.uuid4().hex}") + unauthenticated: Final = send(f"{marker}-anon", caller="sk-not-a-real-key") + after: Final = send(f"{marker}-after") + payloads: Final = collect_payloads(sink, len(successes) + len(failures) + 1) + objects: Final = sink.objects() + assert [response.status_code for response in responses[:4]] == [200] * 4, [r.text for r in responses] + assert tuple(response.json()["id"] for response in responses[:4]) == successes + assert all(response.status_code == 401 for response in responses[4:]), [r.text for r in responses[4:]] + assert all("synthetic upstream rejection" in response.text for response in responses[4:]) + assert ghost.status_code == 403 and "key_model_access_denied" in ghost.text, ghost.text + assert unauthenticated.status_code == 401 and "error" in unauthenticated.json(), unauthenticated.text + assert after.status_code == 200 and after.json()["id"] == f"{marker}-after", after.text + assert sorted(upstream.received()) == sorted((*successes, *failures, f"{marker}-after")) + succeeded: Final = frozenset(str(payload["id"]) for payload in payloads if payload["status"] == "success") + assert succeeded == frozenset((*successes, f"{marker}-after")) + failed: Final = tuple(payload for payload in payloads if payload["status"] == "failure") + assert frozenset(str(payload["litellm_call_id"]) for payload in failed) >= frozenset( + response.headers["x-litellm-call-id"] for response in responses[4:] + ) + assert _outside_layout(objects, "hour") == () + + +def test_s3_v2_hour_cache_hit_twins_land_one_object_each_under_the_hour_folder( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "s3hcache" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, anthropic_model, key = _models(scenario, provider.url) + first: Final = tuple( + call_surface(owned.gateway, surface, openai_model, anthropic_model, key, f"{marker}-{surface}", False) + for surface in ("chat", "responses") + ) + eventually(lambda: len(sink.objects()), lambda count: count >= 2, seconds=30) + repeated: Final = tuple( + call_surface(owned.gateway, surface, openai_model, anthropic_model, key, f"{marker}-{surface}", False) + for surface in ("chat", "responses") + ) + payloads: Final = collect_payloads(sink, 4) + objects: Final = sink.objects() + assert first[0][0] == f"{marker}-chat" and repeated[0][0] == first[0][0] + assert matched_ids(payloads, first + repeated) == frozenset(str(payload["id"]) for payload in payloads) + assert sorted(_prompt(payload) for payload in payloads) == sorted((f"{marker}-chat", f"{marker}-responses") * 2) + assert sorted(upstream.received()) == sorted((f"{marker}-chat", f"{marker}-responses")) + assert len(objects) == 4, list(objects) + assert sum(1 for payload in payloads if payload["cache_hit"] is True) == 2 + assert _outside_layout(objects, "hour") == () + + +def test_s3_v2_hour_cold_storage_key_names_the_uploaded_object_and_reads_back(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hcold" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR, {"cold_storage_custom_logger": "s3_v2"}) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + prompts: Final = (f"{marker}-kept", f"{marker}-missing") + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + collect_payloads(sink, len(prompts)) + objects: Final = sink.objects() + keys: Final = {prompt: _cold_storage_key(prompt) for prompt in prompts} + with sink.lock: + sink.store.pop(f"/{BUCKET}/{quote(keys[prompts[1]], safe='/')}") + kept: Final = eventually( + lambda: owned.gateway.request("GET", f"/spend/logs/ui/{prompts[0]}"), + lambda reply: reply.status_code == 200 and bool((reply.json() or {}).get("messages")), + seconds=30, + ) + missing: Final = owned.gateway.request("GET", f"/spend/logs/ui/{prompts[1]}") + assert returned == prompts + assert sorted(upstream.received()) == sorted(prompts) + assert frozenset(f"/{BUCKET}/{quote(key, safe='/')}" for key in keys.values()) == frozenset(objects) + assert _outside_layout(objects, "hour") == () + assert kept.json()["messages"] == [{"role": "user", "content": prompts[0]}], kept.text + assert prompts[0] in json.dumps(kept.json()["response"]), kept.text + assert missing.status_code == 200, missing.text + assert prompts[1] not in json.dumps(missing.json()["response"]), missing.text + + +def test_s3_v2_audit_logs_follow_the_audit_params_granularity_not_the_request_logs( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "s3haudit" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with wire_server(upstream.respond) as provider, wire_server(sink.respond) as bucket: + settings: Final = { + "store_audit_logs": True, + "audit_log_callbacks": ["s3_v2"], + "s3_audit_callback_params": {**SINK_CREDENTIALS, "s3_endpoint_url": bucket.url, **HOUR}, + } + with ( + _s3_proxy(gateway, tmp_path, bucket.url, {}, settings, workers=1) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url, key_alias=marker) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, (marker,)) + audit_prefix: Final = f"/{BUCKET}/{PREFIX}/audit_logs/" + eventually( + lambda: tuple(target for target in sink.objects() if target.startswith(audit_prefix)), + lambda targets: len(targets) >= 1, + seconds=30, + ) + collect_payloads(sink, 2) + objects: Final = sink.objects() + audits: Final = { + target: object_value(json.loads(body)) for target, body in objects.items() if target.startswith(audit_prefix) + } + requests: Final = {target: body for target, body in objects.items() if not target.startswith(audit_prefix)} + assert returned == (marker,) + assert upstream.received() == (marker,) + assert _outside_layout(requests, "day") == () + created: Final = tuple(audit for audit in audits.values() if audit["action"] == "created") + assert "LiteLLM_VerificationToken" in frozenset(str(audit["table_name"]) for audit in created), audits + for target, audit in audits.items(): + located: Final = re.fullmatch( + re.escape(audit_prefix) + + rf"(\d{{4}}-\d{{2}}-\d{{2}})/(\d{{2}})/(\d{{2}})-\d{{2}}-\d{{2}}_{re.escape(str(audit['id']))}\.json", + unquote(target), + ) + assert located and located[2] == located[3], (target, audit["updated_at"]) + folder: Final = datetime.fromisoformat(f"{located[1]}T{located[2]}:00:00+00:00") + updated: Final = datetime.fromisoformat(str(audit["updated_at"])) + assert timedelta(0) < folder + timedelta(hours=1) - updated <= timedelta(hours=1, minutes=1), ( + target, + audit["updated_at"], + ) + + +def test_s3_v2_admin_ui_granularity_update_moves_live_traffic_on_both_workers(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hui" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + scratch_database() as database_url, + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, {}, environment={"DATABASE_URL": database_url}) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + before: Final = _sdk_chats(owned.gateway, openai_model, key, (f"{marker}-before",)) + eventually(lambda: len(sink.objects()), lambda count: count >= 1, seconds=30) + listed: Final = owned.gateway.get("/get/config/callbacks") + _update_environment(owned.gateway, {"callback": "s3_v2", "s3_partition_granularity": "hour"}) + probe_round: Final = iter(range(1000)) + + def probe() -> Mapping[str, bytes]: + round_id: Final = next(probe_round) + prompts: Final = tuple(f"{marker}-probe{round_id}-{index}" for index in range(8)) + _sdk_chats(owned.gateway, openai_model, key, prompts) + eventually( + lambda: frozenset(str(payload["id"]) for payload in sink.payloads()), + lambda landed: frozenset(prompts) <= landed, + seconds=20, + ) + return {target: body for target, body in sink.objects().items() if f"-probe{round_id}-" in target} + + eventually(probe, lambda probed: len(probed) == 8 and _outside_layout(probed, "hour") == (), seconds=60) + prompts: Final = tuple(f"{marker}-after-{index}" for index in range(16)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + eventually( + lambda: frozenset(str(payload["id"]) for payload in sink.payloads()), + lambda landed: frozenset(prompts) <= landed, + seconds=30, + ) + after: Final = {target: body for target, body in sink.objects().items() if f"{marker}-after-" in target} + before_objects: Final = { + target: body for target, body in sink.objects().items() if f"{marker}-before" in target + } + readback: Final = owned.gateway.get("/get/config/callbacks") + s3_rows: Final = tuple(row for row in listed["callbacks"] if object_value(row)["name"] in ("s3", "s3_v2")) + assert s3_rows and all( + "S3_PARTITION_GRANULARITY" in object_value(object_value(row)["variables"]) for row in s3_rows + ), listed + after_rows: Final = tuple(row for row in readback["callbacks"] if object_value(row)["name"] in ("s3", "s3_v2")) + assert all( + object_value(object_value(row)["variables"])["S3_PARTITION_GRANULARITY"] == "hour" for row in after_rows + ), readback + assert before == (f"{marker}-before",) + assert returned == prompts + assert _outside_layout(before_objects, "day") == () + assert len(after) == len(prompts) + assert _outside_layout(after, "hour") == () + + +def test_s3_v2_granularity_toggles_mid_burst_keep_every_cold_storage_key_on_its_object( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "s3htog" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + scratch_database() as database_url, + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy( + gateway, + tmp_path, + bucket.url, + {}, + {"cold_storage_custom_logger": "s3_v2"}, + environment={"DATABASE_URL": database_url}, + ) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + prompts: Final = tuple(f"{marker}-{index}" for index in range(32)) + with ThreadPoolExecutor(max_workers=1) as burst: + pending: Final = burst.submit(_sdk_chats, owned.gateway, openai_model, key, prompts) + for value in ("hour", "day", "hour", "day", "hour", "day"): + _update_environment(owned.gateway, {"s3_partition_granularity": value}) + returned: Final = pending.result() + collect_payloads(sink, len(prompts)) + objects: Final = sink.objects() + keys: Final = {prompt: _cold_storage_key(prompt, database_url) for prompt in prompts} + assert returned == prompts + assert sorted(upstream.received()) == sorted(prompts) + assert len(objects) == len(prompts) + assert frozenset(f"/{BUCKET}/{quote(key, safe='/')}" for key in keys.values()) == frozenset(objects), ( + "every spend log cold_storage_object_key must name the object the logger uploaded" + ) + assert all( + _object_pattern(object_value(json.loads(body)), "hour").fullmatch(unquote(target)) + or _object_pattern(object_value(json.loads(body)), "day").fullmatch(unquote(target)) + for target, body in objects.items() + ) + + +def test_legacy_s3_callback_ignores_hour_granularity(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3v1hour" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR, {"callbacks": [], "success_callback": ["s3"]}) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + prompts: Final = tuple(f"{marker}-{index}" for index in range(3)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + payloads: Final = collect_payloads(sink, len(prompts)) + objects: Final = sink.objects() + assert returned == prompts + assert frozenset(str(payload["id"]) for payload in payloads) == frozenset(prompts) + assert _outside_layout(objects, "day") == (), "legacy s3 keeps the daily layout, the setting is s3_v2 only" + + +def test_s3_v2_hour_sink_outage_mid_mixed_burst_lands_every_id_exactly_once(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hout" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05, fail_until=float("inf"), fail_status=503) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, anthropic_model, key = _models(scenario, provider.url) + answered: Final = mixed_burst(owned.gateway, openai_model, anthropic_model, key, marker, per_surface=6) + eventually(lambda: sink.attempts, lambda attempts: attempts >= 1, seconds=30) + during: Final = owned.gateway.client.get("/health/readiness") + rejected: Final = sink.attempts + sink.fail_until = 0.0 + payloads: Final = collect_payloads(sink, len(answered), seconds=60) + objects: Final = sink.objects() + sent: Final = _surface_prompts(marker, 6) + assert len(answered) == len(sent) and len(payloads) == len(sent), payloads + assert matched_ids(payloads, answered) == frozenset(str(payload["id"]) for payload in payloads) + assert sorted(upstream.received()) == sorted(sent) + assert during.status_code == 200, during.text + assert rejected >= 1 and sink.attempts > len(objects) + assert sorted(_prompt(payload) for payload in payloads) == sorted(sent), "every burst id lands exactly once" + assert len(objects) == len(sent) + assert _outside_layout(objects, "hour") == () + + +def test_s3_v2_hour_coded_403_retries_reuse_the_same_hour_key(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3h403" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05, fail_attempts=10, fail_status=403, fail_code="AccessDenied") + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + prompts: Final = tuple(f"{marker}-{index}" for index in range(16)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + payloads: Final = collect_payloads(sink, len(prompts), seconds=60) + objects: Final = sink.objects() + attempted: Final = dict(sink.attempt_counts) + assert returned == prompts + assert sorted(str(payload["id"]) for payload in payloads) == sorted(prompts) + assert frozenset(attempted) == frozenset(objects), "a retried upload must reuse the key of its first attempt" + assert sum(attempted.values()) == len(objects) + 10 + assert _outside_layout(objects, "hour") == () + + +def test_s3_v2_hour_slow_sink_batches_never_duplicate_an_upload(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hslow" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=1.5) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, {**HOUR, "s3_batch_file_upload": True}) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + prompts: Final = tuple(f"{marker}-{index}" for index in range(32)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + + def delivered() -> int: + readiness: Final = owned.gateway.client.get("/health/readiness") + assert readiness.status_code == 200, readiness.text + return sum(len(body.splitlines()) for body in sink.objects().values()) + + eventually(delivered, lambda total: total >= len(prompts), seconds=60) + payloads: Final = sink.payloads() + objects: Final = sink.objects() + targets: Final = tuple(put.target for put in bucket.drain()) + assert returned == prompts + assert len(set(targets)) == len(targets), "the same batch object was PUT more than once" + assert sorted(str(payload["id"]) for payload in payloads) == sorted(prompts) + assert _batches_outside_layout(objects, "hour") == () + + +def _worker_processes(owned: OwnedProxy) -> tuple[int, ...]: + return tuple( + process.pid + for process in group_members(owned.process.pid) + if process.pid != owned.process.pid and "spawn_main" in " ".join(process.cmdline()) + ) + + +def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hkill" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + workers: Final = _worker_processes(owned) + sent: Final = tuple(f"{marker}-{index}" for index in range(40)) + + def send(prompt: str) -> tuple[str, bool]: + try: + response: Final = owned.gateway.request( + "POST", + "/v1/chat/completions", + { + "model": openai_model, + "messages": [{"role": "user", "content": prompt}], + "cache": {"no-cache": True}, + }, + key=key, + ) + except httpx.HTTPError: + return prompt, False + return prompt, response.status_code == 200 and response.json()["id"] == prompt + + with ThreadPoolExecutor(max_workers=16) as pool: + futures: Final = tuple(pool.submit(send, prompt) for prompt in sent) + eventually(lambda: len(upstream.received()), lambda count: count >= 8, seconds=30) + psutil.Process(workers[0]).kill() + results: Final = tuple(future.result() for future in futures) + later: Final = tuple(f"{marker}-later-{index}" for index in range(8)) + later_results: Final = tuple(send(prompt) for prompt in later) + answered: Final = frozenset(prompt for prompt, ok in (*results, *later_results) if ok) + eventually( + lambda: frozenset(str(payload["id"]) for payload in sink.payloads()), + lambda landed: frozenset(later) <= landed, + seconds=45, + ) + payloads: Final = sink.payloads() + objects: Final = sink.objects() + assert len(workers) == 2, workers + assert all(ok for _, ok in later_results), "the surviving worker must keep serving after the kill" + landed: Final = tuple(str(payload["id"]) for payload in payloads) + assert frozenset(landed) <= answered, "only answered ids may land; the killed worker's unflushed queue is lost" + assert len(landed) == len(set(landed)), "no id may land twice" + assert _outside_layout(objects, "hour") == () + + +def test_s3_v2_hour_proxy_restart_mid_burst_keeps_the_layout_without_duplicates( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "s3hterm" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with wire_server(upstream.respond) as provider, wire_server(sink.respond) as bucket: + model_name: Final = f"integration-{marker}" + + def register(candidate: Gateway) -> str: + return str( + candidate.post( + "/model/new", + { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "synthetic-provider-key", + "api_base": provider.url + "/v1", + }, + "model_info": {}, + }, + )["model_info"]["id"] + ) + + def send(candidate: Gateway, key: str, prompt: str) -> tuple[str, bool]: + try: + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": model_name, + "messages": [{"role": "user", "content": prompt}], + "cache": {"no-cache": True}, + }, + key=key, + ) + except httpx.HTTPError: + return prompt, False + return prompt, response.status_code == 200 + + sent: Final = tuple(f"{marker}-{index}" for index in range(40)) + with _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as first: + model_id: Final = register(first.gateway) + first_key: Final = str(first.gateway.post("/key/generate", {"models": [model_name]})["key"]) + with ThreadPoolExecutor(max_workers=16) as pool: + futures: Final = tuple(pool.submit(send, first.gateway, first_key, prompt) for prompt in sent) + eventually(lambda: len(upstream.received()), lambda count: count >= 8, seconds=30) + first.process.terminate() + results: Final = tuple(future.result() for future in futures) + first.process.wait(timeout=30) + answered: Final = frozenset(prompt for prompt, ok in results if ok) + landed_before_restart: Final = frozenset(str(payload["id"]) for payload in sink.payloads()) + with _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as second: + restarted: Final = tuple(f"{marker}-restart-{index}" for index in range(8)) + second_key: Final = second.gateway.post("/key/generate", {"models": [model_name]})["key"] + restart_results: Final = tuple(send(second.gateway, str(second_key), prompt) for prompt in restarted) + eventually( + lambda: frozenset(str(payload["id"]) for payload in sink.payloads()), + lambda landed: frozenset(restarted) <= landed, + seconds=30, + ) + second.gateway.post("/model/delete", {"id": model_id}) + payloads: Final = sink.payloads() + objects: Final = sink.objects() + assert all(ok for _, ok in restart_results) + assert landed_before_restart <= answered, "a delivered object has no answered request" + landed: Final = tuple(str(payload["id"]) for payload in payloads) + assert len(landed) == len(set(landed)), "no id may land twice across the restart" + assert frozenset(restarted) <= frozenset(landed) + targets: Final = tuple(put.target for put in bucket.drain()) + assert len(set(targets)) == len(targets) + assert _outside_layout(objects, "hour") == ()