From 5008d18f77a421d669c0ea8f87b04749f16eb690 Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 19 Sep 2026 20:48:36 +0000 Subject: [PATCH] fix(zerobus): hold the queue cap while an insert is in flight Rows arriving during a slow insert are dropped once the queue is at max_queue_size, since trimming the head would corrupt the in-flight batch. Test fakes are typed and record calls as frozen dataclasses; the litellm_logging init and reuse branches are covered. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/zerobus/logger.py | 22 ++-- .../zerobus/test_zerobus_client.py | 90 +++++++++------ .../zerobus/test_zerobus_logger.py | 104 ++++++++++++++---- 3 files changed, 149 insertions(+), 67 deletions(-) diff --git a/litellm/integrations/zerobus/logger.py b/litellm/integrations/zerobus/logger.py index b43610518dc..e2007218c8e 100644 --- a/litellm/integrations/zerobus/logger.py +++ b/litellm/integrations/zerobus/logger.py @@ -1,9 +1,4 @@ -""" -Databricks Zerobus logging integration. - -Buffers one ``TRACE_TABLE_COLUMNS`` row per request and writes each flush to a Unity -Catalog Delta table through the Zerobus Ingest REST API. -""" +"""Databricks Zerobus logging integration.""" import asyncio from collections.abc import Mapping @@ -90,8 +85,6 @@ def connection_for(params: ZerobusInitParams) -> ZerobusConnection: class ZerobusLogger(CustomBatchLogger): - """Batching callback that writes LiteLLM request logs to a Databricks Delta table.""" - preserve_events_added_during_flush = True def __init__( @@ -182,6 +175,10 @@ class ZerobusLogger(CustomBatchLogger): verbose_logger.debug("zerobus: event carried no standard_logging_object, skipping") return + if self._flushing and len(self.log_queue) >= self.max_queue_size: + verbose_logger.warning("zerobus: queue at %s rows during a flush, dropped a row", self.max_queue_size) + return + self.log_queue.append(trace_row(payload)) self._drop_overflow() if len(self.log_queue) >= self.batch_size: @@ -202,6 +199,7 @@ class ZerobusLogger(CustomBatchLogger): return payload def _drop_overflow(self) -> None: + """Trim the oldest rows, except mid flush when the in-flight batch is the head of the queue.""" if self._flushing: return overflow: Final = len(self.log_queue) - self.max_queue_size @@ -220,13 +218,7 @@ class ZerobusLogger(CustomBatchLogger): self._flushing = False async def async_send_batch(self) -> None: - """ - Write everything queued as one insert. - - A retryable failure propagates so ``CustomBatchLogger`` keeps the rows for the next - flush. A failure Zerobus would repeat, a schema mismatch for one, drops the batch, - since holding it would block every row queued behind it. - """ + """A retryable failure propagates so the rows are kept; a permanent one drops them so the queue moves on.""" rows: Final = tuple(self.log_queue) if not rows: return diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_client.py b/tests/test_litellm/integrations/zerobus/test_zerobus_client.py index fe21f70329e..9ed683d1df2 100644 --- a/tests/test_litellm/integrations/zerobus/test_zerobus_client.py +++ b/tests/test_litellm/integrations/zerobus/test_zerobus_client.py @@ -1,6 +1,8 @@ import base64 import json -from collections.abc import Sequence +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain, repeat import httpx import pytest @@ -27,33 +29,53 @@ def _accepted() -> httpx.Response: return httpx.Response(200, text="{}") -class FakeHTTPClient: - """ - Stands in for AsyncHTTPHandler, including its habit of raising on error statuses. +@dataclass(frozen=True, slots=True) +class TokenCall: + url: str + data: Mapping[str, str] + headers: Mapping[str, str] - Results are consumed in order, and the last one repeats. - """ + +@dataclass(frozen=True, slots=True) +class InsertCall: + url: str + content: bytes + headers: Mapping[str, str] + + +def _results(results: Sequence[httpx.Response | Exception]) -> Iterator[httpx.Response | Exception]: + """Results are served in order, and the last one repeats.""" + return chain(results[:-1], repeat(results[-1])) + + +class FakeHTTPClient: + """Stands in for AsyncHTTPHandler, including its habit of raising on error statuses.""" def __init__( self, - token: Sequence[httpx.Response | Exception] | None = None, - insert: Sequence[httpx.Response | Exception] | None = None, + token: Sequence[httpx.Response | Exception] = (), + insert: Sequence[httpx.Response | Exception] = (), ) -> None: - self.token_results = list(token) if token else [_token()] # mutable-ok: results are consumed by popping - self.insert_results = list(insert) if insert else [_accepted()] # mutable-ok: results are consumed by popping - self.token_calls: list[dict] = [] - self.insert_calls: list[dict] = [] + self.token_results = _results(token or (_token(),)) + self.insert_results = _results(insert or (_accepted(),)) + self.token_calls: tuple[TokenCall, ...] = () + self.insert_calls: tuple[InsertCall, ...] = () - async def post(self, url, data=None, content=None, headers=None, **_): + async def post( + self, + url: str, + data: Mapping[str, str] | None = None, + content: bytes | None = None, + headers: Mapping[str, str] | None = None, + ) -> httpx.Response: if url.endswith("/oidc/v1/token"): - self.token_calls.append({"url": url, "data": data, "headers": headers or {}}) - return _next_result(self.token_results, url) - self.insert_calls.append({"url": url, "content": content, "headers": headers or {}}) - return _next_result(self.insert_results, url) + self.token_calls = (*self.token_calls, TokenCall(url, data or {}, headers or {})) + return _raise_like_the_handler(next(self.token_results), url) + self.insert_calls = (*self.insert_calls, InsertCall(url, content or b"", headers or {})) + return _raise_like_the_handler(next(self.insert_results), url) -def _next_result(results: list, url: str) -> httpx.Response: - result = results.pop(0) if len(results) > 1 else results[0] +def _raise_like_the_handler(result: httpx.Response | Exception, url: str) -> httpx.Response: if isinstance(result, Exception): raise result if result.status_code >= 300: @@ -85,12 +107,14 @@ async def test_rows_are_posted_as_one_json_list_to_the_table_insert_endpoint(): assert outcome is None (call,) = http_client.insert_calls - assert call["url"] == ( + # Insert endpoint per the Zerobus Ingest docs, read 2026-09-19: + # https://docs.databricks.com/aws/en/ingestion/lakeflow-connect/zerobus-ingest + assert call.url == ( "https://1234567890123456.zerobus.us-west-2.cloud.databricks.com/zerobus/v1/tables/main.litellm.traces/insert" ) - assert json.loads(call["content"]) == [{"id": "a", "model": "gpt-4o"}, {"id": "b", "model": "gpt-4o"}] - assert call["headers"]["Content-Type"] == "application/json" - assert call["headers"]["Authorization"] == "Bearer tok-1" + assert json.loads(call.content) == [{"id": "a", "model": "gpt-4o"}, {"id": "b", "model": "gpt-4o"}] + assert call.headers["Content-Type"] == "application/json" + assert call.headers["Authorization"] == "Bearer tok-1" @pytest.mark.asyncio @@ -101,11 +125,13 @@ async def test_the_token_is_minted_for_the_zerobus_resource_with_the_table_privi await _client(http_client).insert(ROWS) (call,) = http_client.token_calls - assert call["url"] == "https://dbc-a1b2c3d4-e5f6.cloud.databricks.com/oidc/v1/token" - assert call["data"]["grant_type"] == "client_credentials" - assert call["data"]["scope"] == "all-apis" - assert call["data"]["resource"] == "api://databricks/workspaces/1234567890123456/zerobusDirectWriteApi" - details = json.loads(call["data"]["authorization_details"]) + # Token form per the Zerobus Ingest docs (REST API authentication), read 2026-09-19: + # https://docs.databricks.com/aws/en/ingestion/lakeflow-connect/zerobus-ingest + assert call.url == "https://dbc-a1b2c3d4-e5f6.cloud.databricks.com/oidc/v1/token" + assert call.data["grant_type"] == "client_credentials" + assert call.data["scope"] == "all-apis" + assert call.data["resource"] == "api://databricks/workspaces/1234567890123456/zerobusDirectWriteApi" + details = json.loads(call.data["authorization_details"]) assert [(d["object_type"], d["object_full_path"], d["privileges"]) for d in details] == [ ("CATALOG", "main", ["USE CATALOG"]), ("SCHEMA", "main.litellm", ["USE SCHEMA"]), @@ -120,7 +146,7 @@ async def test_the_service_principal_authenticates_with_http_basic(): await _client(http_client).insert(ROWS) - scheme, credentials = http_client.token_calls[0]["headers"]["Authorization"].split(" ") + scheme, credentials = http_client.token_calls[0].headers["Authorization"].split(" ") assert scheme == "Basic" assert base64.b64decode(credentials).decode() == "sp-client-id:sp-client-secret" @@ -138,7 +164,7 @@ async def test_the_token_is_reused_across_inserts_until_it_nears_expiry(): await client.insert(ROWS) assert len(http_client.token_calls) == 2 - assert [call["headers"]["Authorization"] for call in http_client.insert_calls] == [ + assert [call.headers["Authorization"] for call in http_client.insert_calls] == [ "Bearer tok-1", "Bearer tok-1", "Bearer tok-2", @@ -158,7 +184,7 @@ async def test_a_401_discards_the_token_so_the_next_insert_mints_a_fresh_one(): assert first == ZerobusIngestFailure(detail="insert returned 401, token discarded", retryable=True) assert second is None - assert http_client.insert_calls[1]["headers"]["Authorization"] == "Bearer tok-2" + assert http_client.insert_calls[1].headers["Authorization"] == "Bearer tok-2" @pytest.mark.asyncio @@ -199,7 +225,7 @@ async def test_bad_credentials_fail_the_insert_without_posting_rows(): outcome = await _client(http_client).insert(ROWS) assert outcome == ZerobusIngestFailure(detail="token request returned 401: invalid_client", retryable=False) - assert http_client.insert_calls == [] + assert http_client.insert_calls == () @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_logger.py b/tests/test_litellm/integrations/zerobus/test_zerobus_logger.py index 0e4888cd614..a85a9e2e6a0 100644 --- a/tests/test_litellm/integrations/zerobus/test_zerobus_logger.py +++ b/tests/test_litellm/integrations/zerobus/test_zerobus_logger.py @@ -1,5 +1,6 @@ import asyncio -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence +from itertools import chain, repeat import pytest @@ -12,29 +13,36 @@ WORKSPACE_URL = "https://dbc-a1b2c3d4-e5f6.cloud.databricks.com" SERVER_ENDPOINT = "https://1234567890123456.zerobus.us-west-2.cloud.databricks.com" +Row = Mapping[str, object] + + class FakeIngestClient: - """Records the rows each flush would have written.""" + """Records the rows each flush would have written; outcomes are served in order and the last one repeats.""" - def __init__(self, outcomes: Sequence[ZerobusIngestFailure | None] = (None,)) -> None: - self.outcomes = list(outcomes) # mutable-ok: outcomes are consumed by popping - self.batches: list[tuple[Mapping[str, object], ...]] = [] - self.on_insert: Callable[[], None] | None = None + def __init__( + self, + outcomes: Sequence[ZerobusIngestFailure | None] = (None,), + on_insert: Callable[[], None] | None = None, + ) -> None: + self.outcomes: Iterator[ZerobusIngestFailure | None] = chain(outcomes[:-1], repeat(outcomes[-1])) + self.on_insert = on_insert + self.batches: tuple[tuple[Row, ...], ...] = () - async def insert(self, rows: Sequence[Mapping[str, object]]) -> ZerobusIngestFailure | None: + async def insert(self, rows: Sequence[Row]) -> ZerobusIngestFailure | None: if self.on_insert is not None: self.on_insert() - self.batches.append(tuple(rows)) - return self.outcomes.pop(0) if len(self.outcomes) > 1 else self.outcomes[0] + self.batches = (*self.batches, tuple(rows)) + return next(self.outcomes) - def ids(self) -> list[object]: - return [row["id"] for batch in self.batches for row in batch] + def ids(self) -> tuple[object, ...]: + return tuple(row["id"] for batch in self.batches for row in batch) -def _logger(client: FakeIngestClient, **params) -> ZerobusLogger: - return ZerobusLogger(params=ZerobusInitParams(**params), client=client) +def _logger(client: FakeIngestClient, **params: object) -> ZerobusLogger: + return ZerobusLogger(params=ZerobusInitParams.model_validate(params), client=client) -def _event(request_id: str, **payload) -> dict: +def _event(request_id: str, **payload: object) -> dict[str, object]: return { "standard_logging_object": { "id": request_id, @@ -64,7 +72,7 @@ async def test_a_full_batch_is_written_as_one_insert_of_table_rows(): await _settle(logger) assert len(client.batches) == 1 - assert client.ids() == ["a", "b", "c"] + assert client.ids() == ("a", "b", "c") assert client.batches[0][0]["model"] == "gpt-4o" assert logger.log_queue == [] @@ -76,7 +84,7 @@ async def test_rows_are_held_until_the_batch_is_full(): await logger.async_log_success_event(_event("a"), None, None, None) - assert client.batches == [] + assert client.batches == () assert len(logger.log_queue) == 1 @@ -88,7 +96,7 @@ async def test_failed_requests_are_written_too(): await logger.async_log_failure_event(_event("failed", status="failure", error_str="boom"), None, None, None) await _settle(logger) - assert client.ids() == ["failed"] + assert client.ids() == ("failed",) assert client.batches[0][0]["status"] == "failure" assert client.batches[0][0]["error_str"] == "boom" @@ -100,7 +108,7 @@ async def test_an_event_without_a_standard_payload_is_skipped(): await logger.async_log_success_event({"kwargs": "but no payload"}, None, None, None) - assert client.batches == [] + assert client.batches == () assert logger.log_queue == [] @@ -147,14 +155,44 @@ async def test_a_row_that_arrives_mid_flush_is_kept_for_the_next_one(): await logger.async_log_success_event(_event("first"), None, None, None) await _settle(logger) - assert client.ids() == ["first"] + assert client.ids() == ("first",) assert [row["id"] for row in logger.log_queue] == ["late"] +@pytest.mark.asyncio +async def test_the_queue_cap_holds_while_an_insert_is_in_flight(): + """A slow insert must not let the queue grow past max_queue_size, nor disturb the in-flight head.""" + insert_started = asyncio.Event() + finish_insert = asyncio.Event() + + class SlowClient: + batches: tuple[tuple[Row, ...], ...] = () + + async def insert(self, rows: Sequence[Row]) -> None: + insert_started.set() + await finish_insert.wait() + self.batches = (*self.batches, tuple(rows)) + + client = SlowClient() + logger = ZerobusLogger(params=ZerobusInitParams(batch_size=2), client=client) + logger.max_queue_size = 3 + + for request_id in ("a", "b"): + await logger.async_log_success_event(_event(request_id), None, None, None) + await insert_started.wait() + for request_id in ("c", "d", "e"): + await logger.async_log_success_event(_event(request_id), None, None, None) + finish_insert.set() + await _settle(logger) + + assert [[row["id"] for row in batch] for batch in client.batches] == [["a", "b"]] + assert [row["id"] for row in logger.log_queue] == ["c"] + + @pytest.mark.asyncio async def test_a_client_error_does_not_break_the_request_path(): class ExplodingClient: - async def insert(self, rows): + async def insert(self, rows: Sequence[Row]) -> None: raise RuntimeError("bug") logger = ZerobusLogger(params=ZerobusInitParams(batch_size=1), client=ExplodingClient()) @@ -326,3 +364,29 @@ def test_the_client_is_kept_while_the_connection_is_unchanged_and_rebuilt_when_i assert unchanged is first assert rebuilt is not first assert rebuilt.connection.table_name == "main.litellm.traces_v2" + + +def test_callbacks_zerobus_builds_one_logger_and_reuses_it(monkeypatch): + """`litellm_settings.callbacks: ["zerobus"]` goes through litellm_logging, which must hand back one instance.""" + from litellm.litellm_core_utils import litellm_logging as logging_module + + monkeypatch.setenv("ZEROBUS_WORKSPACE_URL", WORKSPACE_URL) + monkeypatch.setenv("ZEROBUS_SERVER_ENDPOINT", SERVER_ENDPOINT) + monkeypatch.setenv("ZEROBUS_CLIENT_ID", "sp-id") + monkeypatch.setenv("ZEROBUS_CLIENT_SECRET", "sp-secret") + monkeypatch.setenv("ZEROBUS_TABLE_NAME", "main.litellm.traces") + monkeypatch.setattr(litellm, "zerobus_params", None) + monkeypatch.setattr(logging_module, "_in_memory_loggers", []) + + assert logging_module.get_custom_logger_compatible_class("zerobus") is None + + first = logging_module._init_custom_logger_compatible_class( + logging_integration="zerobus", internal_usage_cache=None, llm_router=None, custom_logger_init_args={} + ) + second = logging_module._init_custom_logger_compatible_class( + logging_integration="zerobus", internal_usage_cache=None, llm_router=None, custom_logger_init_args={} + ) + + assert isinstance(first, ZerobusLogger) + assert second is first + assert logging_module.get_custom_logger_compatible_class("zerobus") is first