From b51762556d2bcfb489d7c03275d3710b77727879 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 30 Sep 2026 01:40:34 +0000 Subject: [PATCH] fix(s3_v2): reuse the cold storage key only when s3_v2 owns cold storage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/s3_v2.py | 6 ++- .../test_s3_v2_partition_granularity.py | 41 ++++++++++++++++++- tests/unit/integrations/test_s3_v2.py | 23 +++++++++++ 3 files changed, 67 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 6119cbdc87b..d470705c91b 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -840,7 +840,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): partition_granularity=partition_granularity, ) - cold_storage_object_key: Final = standard_logging_payload.get("metadata", {}).get("cold_storage_object_key") + cold_storage_object_key: Final = ( + standard_logging_payload.get("metadata", {}).get("cold_storage_object_key") + if litellm.cold_storage_custom_logger == "s3_v2" + else None + ) s3_object_key: Final = ( cold_storage_object_key if cold_storage_object_key is not None diff --git a/tests/integration/observability/test_s3_v2_partition_granularity.py b/tests/integration/observability/test_s3_v2_partition_granularity.py index 7730618241d..ee9a2623116 100644 --- a/tests/integration/observability/test_s3_v2_partition_granularity.py +++ b/tests/integration/observability/test_s3_v2_partition_granularity.py @@ -468,6 +468,43 @@ def test_s3_v2_hour_cold_storage_key_names_the_uploaded_object_and_reads_back(ga assert prompts[1] not in json.dumps(missing.json()["response"]), missing.text +def test_s3_v2_hour_layout_holds_when_another_logger_owns_cold_storage(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hgcs" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + lock: Final = threading.Lock() + puts: Final[dict[str, bytes]] = {} # mutable-ok: filled per PUT by the bucket thread under lock + + def bucket_reply(request: Request) -> Reply: + assert request.method == "PUT", request.method + with lock: + puts[unquote(request.target)] = request.body + return Reply(status=200) + + def uploaded() -> Mapping[str, bytes]: + with lock: + return dict(puts) + + with ( + wire_server(upstream.respond) as provider, + wire_server(bucket_reply) as bucket, + _s3_proxy( + gateway, tmp_path, bucket.url, {**HOUR, "s3_path": ""}, {"cold_storage_custom_logger": "gcs_bucket"} + ) 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) + objects: Final = eventually(uploaded, lambda values: len(values) >= len(prompts), seconds=60) + cold_keys: Final = tuple(_cold_storage_key(prompt) for prompt in prompts) + hour_object: Final = re.compile(rf"/{BUCKET}/\d{{4}}-\d{{2}}-\d{{2}}/(\d{{2}})/time-(\d{{2}})-[^/]+\.json") + matches: Final = tuple(hour_object.fullmatch(target) for target in objects) + assert returned == prompts + assert sorted(str(object_value(json.loads(body))["id"]) for body in objects.values()) == sorted(prompts) + assert all(re.fullmatch(r"\d{4}-\d{2}-\d{2}/time-[^/]+\.json", cold_key) for cold_key in cold_keys), cold_keys + assert all(match is not None and match.group(1) == match.group(2) for match in matches), sorted(objects) + + def test_s3_v2_hour_cold_storage_rebuilds_previous_response_id_history_from_the_hour_object( gateway: Gateway, tmp_path: Path ) -> None: @@ -822,7 +859,6 @@ def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway 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, @@ -833,7 +869,8 @@ def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway 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 frozenset(landed) <= frozenset((*sent, *later)), "only ids this test sent may land" + assert len(results) == len(sent), results assert len(landed) == len(set(landed)), "no id may land twice" assert _outside_layout(objects, "hour") == () diff --git a/tests/unit/integrations/test_s3_v2.py b/tests/unit/integrations/test_s3_v2.py index a50ca31d367..d8d7527e7d2 100644 --- a/tests/unit/integrations/test_s3_v2.py +++ b/tests/unit/integrations/test_s3_v2.py @@ -2739,6 +2739,29 @@ def test_cold_storage_key_matches_upload_when_env_var_changes_mid_request(monkey assert cold_key == uploaded.s3_object_key == f"coldlogs/2026-09-29/time-14-05-09-123456_{_PARTITION_ID}.json" +def test_hour_upload_ignores_a_cold_storage_key_owned_by_another_logger(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False) + monkeypatch.setattr( + litellm, "s3_callback_params", {"s3_bucket_name": "test-bucket", "s3_partition_granularity": "hour"} + ) + monkeypatch.setattr(litellm, "cold_storage_custom_logger", "gcs_bucket") + logger = S3Logger() + cold_key = StandardLoggingPayloadSetup._generate_cold_storage_object_key( + start_time=_PARTITION_START, response_id=_PARTITION_ID + ) + uploaded = logger.create_s3_batch_logging_element( + _PARTITION_START, + StandardLoggingPayload(id=_PARTITION_ID, metadata={"cold_storage_object_key": cold_key}, messages=[]), + ) + + assert cold_key == f"2026-09-29/time-14-05-09-123456_{_PARTITION_ID}.json" + assert uploaded is not None + assert uploaded.s3_object_key == f"2026-09-29/14/time-14-05-09-123456_{_PARTITION_ID}.json" + + @pytest.mark.parametrize("callback_name", ["s3", "s3_v2"]) def test_partition_granularity_is_exposed_to_admin_ui(callback_name: str) -> None: from litellm.integrations.custom_logger import CustomLogger