scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

fix(zerobus): hold the queue cap while an insert is in flight

Browse source 
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>
This commit is contained in:
yassin 2026-09-19 20:48:36 +00:00
parent cf6d04175d
commit 5008d18f77
3 changed files with 149 additions and 67 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
litellm/integrations/zerobus/logger.py
View file

@ -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

90
tests/test_litellm/integrations/zerobus/test_zerobus_client.py
View file

@ -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

104
tests/test_litellm/integrations/zerobus/test_zerobus_logger.py
View file

@ -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