Repository navigation
feat(integrations): add Databricks Zerobus trace logging callback - #42013
Conversation
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
|
I'll fix CI failures and address comments from users with write access. I'll skip comments containing "(aside)".
|
|
|
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
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>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
…s reprs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
TLDR
Problem this solves:
How it solves it:
zerobuscallback batches success and failure traces into one Delta tableUser Flow
Before: a proxy admin on Databricks has no Zerobus option, so traces never reach their lakehouse
callbacks: [zerobus]to config.yaml instead and starts the proxy: startup fails withValueError: Empty module nameAfter: the same admin picks Databricks Zerobus once and every request shows up as a row in their table
catalog.schema.table; the callback saves and shows as active"model": "databricks-gpt-oss-120b"and gets a normal 200 completion backid, model, token counts, cost, end user and tags appears inSELECT * FROM catalog.schema.tablestatus = failurerow carrying the provider error textLinear ticket
Resolves LIT-8649
Docs: BerriAI/litellm-docs#1729
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
uv run pytest tests/test_litellm/<your_test_file>.py -v. Leave the suites (make test-unit-*,make test-unit) to CI: it finishes in ~15 minutes where a laptop takes an hour or more@greptileaito re-request a review after pushing changes)Screenshots / Proof of Fix
Shared setup: a real Databricks workspace (Azure, eastus) with a Unity Catalog table created from
create_table_sql("my_dbrx_workspace.default.litellm_traces"), a service principal granted USE CATALOG, USE SCHEMA, SELECT and MODIFY, and this proxy config. Every completion hits the realdatabricks-gpt-oss-120bserving endpointThe Delta query used throughout, run through the Databricks SQL Statements API against warehouse
b28a3e8cf2de2851:Before (081f73f)
Success and failure traces land in the Delta table
python litellm/proxy/proxy_cli.py --config zerobus_proxy_config.yaml --port 4000from a worktree at the merge baseValueError: Empty module namefollowed byERROR: Application startup failed. Exiting., becausezerobusis not a known callbackAdmin UI offers Databricks Zerobus
zerobusinto the callback dropdownAfter (03b1cdb)
Success and failure traces land in the Delta table
curl -s -o /dev/null -w "%{http_code}" http://localhost:4000/health/livelinessprints200curl -s http://localhost:4000/v1/chat/completions -H 'Authorization: Bearer sk-1234' -H 'Content-Type: application/json' -d '{"model":"databricks-gpt-oss-120b","messages":[{"role":"user","content":"Reply with one word: hello"}],"max_tokens":20,"user":"zerobus-after","metadata":{"tags":["zerobus-after-172123"]}}'"max_tokens":0so the provider rejects it<tag> = zerobus-after-172123fifteen seconds lateridis the completion id the curl returned, and the failure row carries the provider's error text. An earlier run on this branch also landed a streamingopenai/gpt-4o-minirequest asstream = truewith the client-disconnect note inerror_strAdmin UI offers Databricks Zerobus
Open http://localhost:4000/ui/?page=settings (proxy started with a database and
store_model_in_db: true), click Add Callback, typezerobusinto the callback dropdown: Databricks Zerobus is listed with the Databricks logoPick it: exactly five fields render (Workspace URL, Zerobus Endpoint, Service Principal Client ID, Service Principal Client Secret as a password input, Table). Fill them and click Add Callback: the dialog closes and Databricks Zerobus shows as active
curl -s http://localhost:4000/v1/chat/completions ... -d '{"model":"databricks-gpt-oss-120b","messages":[{"role":"user","content":"Reply with one word: hello"}],"max_tokens":20,"user":"ui-configured"}'returns HTTP 200 with idchatcmpl_a831a5da-b09e-4524-ab0f-687d2e0205e6_3913af1b-021e-442d-8f59-ae0dc546376fSELECT id, model, status, end_user, start_time FROM my_dbrx_workspace.default.litellm_traces WHERE end_user = 'ui-configured'Screenshots of the flow, taken at the branch tip against the built proxy UI on http://localhost:4000/ui/?page=settings
Add Callback dropdown filtered to
zerobus, Databricks Zerobus listed with the Databricks logoThe five Zerobus fields filled in, client secret masked
Databricks Zerobus active in Logging & Alerts after saving
Request details for the
ui-configuredcompletion that landed in the Delta tableSecond workspace run (03b1cdb, AWS us-east-2)
Re-ran the same After steps against a second workspace supplied by the requester, this time on AWS: workspace
dbc-4af5ad7c-ddc5(org id7474645566135909, Ohio), Zerobus endpoint7474645566135909.zerobus.us-east-2.cloud.databricks.com, tableworkspace.default.litellm_tracescreated withcreate_table_sql(...)on theServerless Starter Warehouse(af6f76f26ca4f557), model swapped todatabricks-claude-sonnet-5,flush_interval: 5initializing callbacks=['zerobus'] on proxy,curl -s -o /dev/null -w "%{http_code}" http://localhost:4000/health/livelinessprints200. The startup debug line that echoeslitellm.zerobus_paramsshows'REDACTED'in place of the client secret"max_tokens":0"stream":true,"stream_options":{"include_usage":true}streams todata: [DONE]zerobus:error lines in the proxy log) and the three rows read back through the SQL Statements API on the serverless warehouseuserset)Two misconfigurations had to be cleared first, and both exercised the fail-open path (all curls returned their normal provider responses while the batches were dropped with the Databricks error text in the proxy log). First, the service principal only held
USE CATALOGonworkspaceandUSE SCHEMAonworkspace.defaultthrough theaccount usersgroup, and the Zerobus token endpoint only honours privileges granted to the principal directly, so every flush loggedzerobus: dropping N rows, token request returned 401: {"error":"invalid_authorization_details", ...}until a workspace admin ranGRANT USE CATALOG ON CATALOG workspace TO \`andGRANT USE SCHEMA ON SCHEMA workspace.default TO ``. Second, the workspace ID and region first supplied for the Zerobus endpoint belonged to a different workspace; a Zerobus host that does not match the token's workspace answers400with an empty body (zerobus: dropping N rows, insert returned 400:). The workspace's own org id comes back in thex-databricks-org-idresponse header of any request to the workspace URL, and its region from the DNS of the workspace host (ohio.cloud.databricks.com`, us-east-2). Both are called out in the docs PRType
🆕 New Feature
Caveats (if any)
Medium
CustomBatchLoggermessages,response) count against Zerobus row and request size limitsLow
repr, no UI or ingest behaviourUSE CATALOGandUSE SCHEMAthat a service principal only inherits through a group; both must be granted to the principal directly. Documented in the docs PRrust-lintis red on this branch becauselitellm-rust/crates/core/tests/messages/request.rs:401fails clippy onorigin/maintoo (introduced by feat(rust): hand upstream response headers to the native Messages stream #43178); this PR touches no Rust and that check is not requiredTaxonomy audit (A to BB)
ZerobusConnection.client_secretandZerobusAccessToken.valuenow usefield(repr=False)with a testCustomBatchLoggerruns one flush at a time so a token miss mints onceexpires_inkwargsZEROBUS_*env var fills inlitellm.callbacksin teardownFinal Attestation
Link to Devin session: https://app.devin.ai/sessions/8b1123b068804bc3b36381b54c51ca18
Open in Devin Desktop: https://app.devin.ai/desktop/session/8b1123b068804bc3b36381b54c51ca18?variant=devin
Requested by: @yassin-berriai