Skip to content

feat(integrations): add Databricks Zerobus trace logging callback - #42013

Merged
yassin-berriai merged 8 commits into
mainfrom
litellm_zerobus_logging_integration
Sep 25, 2026
Merged

yassin-berriai merged 8 commits into
mainfrom
litellm_zerobus_logging_integration

Conversation

@devin-ai-integration

@devin-ai-integration devin-ai-integration Bot commented Sep 19, 2026 •

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • No first party way to land LiteLLM traces in a Databricks Delta table
  • Databricks users had to write a custom callback and their own Zerobus auth

How it solves it:

  • New zerobus callback batches success and failure traces into one Delta table
  • Rows go through the Zerobus Ingest REST API with a service principal OAuth token
  • Retries transient failures, drops permanent ones, never breaks the LLM request
  • Configurable from config.yaml, env vars, or the Admin UI Logging & Alerts page

User Flow

Before: a proxy admin on Databricks has no Zerobus option, so traces never reach their lakehouse

  1. Admin opens http://localhost:4000/ui/?page=settings, clicks Add Callback and searches the dropdown for Zerobus: nothing is listed
  2. Admin adds callbacks: [zerobus] to config.yaml instead and starts the proxy: startup fails with ValueError: Empty module name
  3. No request ever reaches the proxy, so the Delta table stays empty

After: the same admin picks Databricks Zerobus once and every request shows up as a row in their table

  1. Admin opens http://localhost:4000/ui/?page=settings, clicks Add Callback, picks Databricks Zerobus (Databricks logo) and fills workspace URL, Zerobus endpoint, service principal client ID and secret (masked), and catalog.schema.table; the callback saves and shows as active
  2. A developer sends POST http://localhost:4000/v1/chat/completions with "model": "databricks-gpt-oss-120b" and gets a normal 200 completion back
  3. Within the flush interval a row with that completion's id, model, token counts, cost, end user and tags appears in SELECT * FROM catalog.schema.table
  4. A request the provider rejects (400) lands too, as a status = failure row carrying the provider error text

Linear 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

  • I have added meaningful tests
  • The handful of test files covering my change pass locally, e.g. 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
  • My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to 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 real databricks-gpt-oss-120b serving endpoint

model_list:
  - model_name: databricks-gpt-oss-120b
    litellm_params:
      model: databricks/databricks-gpt-oss-120b
      api_key: os.environ/DATABRICKS_API_KEY
      api_base: os.environ/DATABRICKS_API_BASE

general_settings:
  master_key: sk-1234

litellm_settings:
  callbacks:
    - zerobus
  zerobus_params:
    workspace_url: os.environ/ZEROBUS_WORKSPACE_URL
    server_endpoint: os.environ/ZEROBUS_SERVER_ENDPOINT
    client_id: os.environ/ZEROBUS_CLIENT_ID
    client_secret: os.environ/ZEROBUS_CLIENT_SECRET
    table_name: os.environ/ZEROBUS_TABLE_NAME
    flush_interval: 5

The Delta query used throughout, run through the Databricks SQL Statements API against warehouse b28a3e8cf2de2851:

SELECT id, status, model, prompt_tokens, completion_tokens, response_cost, substr(error_str, 1, 70) AS error_str, start_time
FROM my_dbrx_workspace.default.litellm_traces
WHERE end_user = 'zerobus-after' AND request_tags::string LIKE '%<tag>%'
ORDER BY start_time

Before (081f73f)

Success and failure traces land in the Delta table

  1. python litellm/proxy/proxy_cli.py --config zerobus_proxy_config.yaml --port 4000 from a worktree at the merge base
  2. Startup aborts before serving traffic: ValueError: Empty module name followed by ERROR: Application startup failed. Exiting., because zerobus is not a known callback
  3. No curl is possible and the Delta query returns zero rows

Admin UI offers Databricks Zerobus

  1. Open http://localhost:4000/ui/?page=settings, click Add Callback, type zerobus into the callback dropdown
  2. The dropdown shows no match

After (03b1cdb)

Success and failure traces land in the Delta table

  1. Same command, the proxy comes up: curl -s -o /dev/null -w "%{http_code}" http://localhost:4000/health/liveliness prints 200
  2. curl -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"]}}'
    {"id":"chatcmpl_fb204cd6-3012-41eb-84c0-966ff17f2c0d_4458869d-7351-4e72-9f04-5dd7facec8e3","created":1790356884,"model":"databricks-gpt-oss-120b","object":"chat.completion","choices":[{"finish_reason":"length","index":0,"message":{"content":"","role":"assistant","reasoning_content":"The user: \"Reply with one word: hello\". So they want a single word", ...}}],"usage":{"completion_tokens":20,"prompt_tokens":107,"total_tokens":127}}
    
  3. Same curl with "max_tokens":0 so the provider rejects it
    {"error":{"message":"litellm.BadRequestError: DatabricksException - {\"error_code\":\"BAD_REQUEST\",\"message\":\"Bad request: parameter \\\"max_tokens\\\" must be at least 1\\n\"}\n\nLiteLLM: model group 'databricks-gpt-oss-120b' failed with the error above. No fallback was attempted.","type":"invalid_request_error","param":null,"code":"400"}}
    
  4. Delta query with <tag> = zerobus-after-172123 fifteen seconds later
    SUCCEEDED
    id | status | model | prompt_tokens | completion_tokens | response_cost | error_str | start_time
    chatcmpl_fb204cd6-3012-41eb-84c0-966ff17f2c0d_4458869d-7351-4e72-9f04-5dd7facec8e3 | success | databricks/databricks-gpt-oss-120b | 107 | 20 | 2.8050469999999996E-5 |  | 2026-09-25T17:21:23.721Z
    630129e5-b580-432f-9334-ca418054e4b0 | failure | databricks/databricks-gpt-oss-120b | 0 | 0 | 0.0 | litellm.BadRequestError: DatabricksException - {"error_code":"BAD_REQU | 2026-09-25T17:21:24.602Z
    
    The success row's id is 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 streaming openai/gpt-4o-mini request as stream = true with the client-disconnect note in error_str

Admin UI offers Databricks Zerobus

  1. Open http://localhost:4000/ui/?page=settings (proxy started with a database and store_model_in_db: true), click Add Callback, type zerobus into the callback dropdown: Databricks Zerobus is listed with the Databricks logo

  2. Pick 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

  3. 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 id chatcmpl_a831a5da-b09e-4524-ab0f-687d2e0205e6_3913af1b-021e-442d-8f59-ae0dc546376f

  4. SELECT id, model, status, end_user, start_time FROM my_dbrx_workspace.default.litellm_traces WHERE end_user = 'ui-configured'

    SUCCEEDED
    id | model | status | end_user | start_time
    chatcmpl_a831a5da-b09e-4524-ab0f-687d2e0205e6_3913af1b-021e-442d-8f59-ae0dc546376f | databricks/databricks-gpt-oss-120b | success | ui-configured | 2026-09-25T17:00:21.238Z
    

    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 logo

    Add Callback dropdown showing Databricks Zerobus

    The five Zerobus fields filled in, client secret masked

    Databricks Zerobus form with the five fields filled

    Databricks Zerobus active in Logging & Alerts after saving

    Databricks Zerobus active in Logging and Alerts

    Request details for the ui-configured completion that landed in the Delta table

    Request details page for the UI configured completion

Second 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 id 7474645566135909, Ohio), Zerobus endpoint 7474645566135909.zerobus.us-east-2.cloud.databricks.com, table workspace.default.litellm_traces created with create_table_sql(...) on the Serverless Starter Warehouse (af6f76f26ca4f557), model swapped to databricks-claude-sonnet-5, flush_interval: 5

  1. Proxy boots with the callback: initializing callbacks=['zerobus'] on proxy, curl -s -o /dev/null -w "%{http_code}" http://localhost:4000/health/liveliness prints 200. The startup debug line that echoes litellm.zerobus_params shows 'REDACTED' in place of the client secret
  2. Success request
    $ curl -s http://localhost:4000/v1/chat/completions -H 'Authorization: Bearer <master-key>' -H 'Content-Type: application/json' -d '{"model":"databricks-claude-sonnet-5","messages":[{"role":"user","content":"Reply with one word: zerobus"}],"max_tokens":20,"user":"zerobus-aws","metadata":{"tags":["zerobus-aws-204005"]}}'
    {"id":"msg_bdrk_a5iyu4hbtoz6nkajso3mfsqiuiwtyqrvkbwhctprjf4r5oa5hpza","created":1790368807,"model":"databricks-claude-sonnet-5","object":"chat.completion","choices":[{"finish_reason":"stop","index":0,"message":{"content":"zerobus","role":"assistant"}}],"usage":{"completion_tokens":6,"prompt_tokens":17,"total_tokens":23, ...}}
    http=200
    
  3. Failure request, same curl with "max_tokens":0
    {"error":{"message":"litellm.BadRequestError: DatabricksException - {\"error_code\":\"INVALID_PARAMETER_VALUE\",\"message\":\"INVALID_PARAMETER_VALUE: Invalid value for the parameter max_tokens. Please provide a value greater than 0.\"}\n\nLiteLLM: model group 'databricks-claude-sonnet-5' failed with the error above. No fallback was attempted.","type":"invalid_request_error","param":null,"code":"400"}}
    http=400
    
  4. Streaming request with "stream":true,"stream_options":{"include_usage":true} streams to data: [DONE]
  5. Both flushes were accepted (no zerobus: error lines in the proxy log) and the three rows read back through the SQL Statements API on the serverless warehouse
    SELECT id, status, model, call_type, custom_llm_provider, end_user, request_tags, prompt_tokens, completion_tokens, total_tokens, response_cost, error_str, from_unixtime(unix_timestamp(start_time)) AS start_time FROM workspace.default.litellm_traces ORDER BY start_time DESC LIMIT 10
    SUCCEEDED
    id | status | model | call_type | custom_llm_provider | end_user | request_tags | prompt_tokens | completion_tokens | total_tokens | response_cost | error_str | start_time
    5ec232bb-1d0f-4d29-b006-fdd266d91a0c | failure | databricks/databricks-claude-sonnet-5 | acompletion | databricks |  | ["User-Agent: curl","User-Agent: curl/7.81.0"] | 0 | 0 | 0 | 0.0 | litellm.BadRequestError: DatabricksException - {"error_code":"INVALID_PARAMETER_VALUE","message":"INVALID_PARAMETER_VALUE: Invalid value for the parameter max_tokens. Please provide a value greater than 0."} | 2026-09-25 20:40:07
    msg_bdrk_sel2ncrow6ycir6p2vqfuklcqi3jryabbd4rmgu26yjv46kclyua | success | databricks/databricks-claude-sonnet-5 | acompletion | databricks |  | ["User-Agent: curl","User-Agent: curl/7.81.0"] | 14 | 20 | 34 | 3.4200026E-4 |  | 2026-09-25 20:40:07
    msg_bdrk_a5iyu4hbtoz6nkajso3mfsqiuiwtyqrvkbwhctprjf4r5oa5hpza | success | databricks/databricks-claude-sonnet-5 | acompletion | databricks | zerobus-aws | ["zerobus-aws-204005","User-Agent: curl","User-Agent: curl/7.81.0"] | 17 | 6 | 23 | 1.4099995E-4 |  | 2026-09-25 20:40:05
    
    The middle row is the streaming request (14 prompt, 20 completion tokens, no user set)

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 CATALOG on workspace and USE SCHEMA on workspace.default through the account users group, and the Zerobus token endpoint only honours privileges granted to the principal directly, so every flush logged zerobus: dropping N rows, token request returned 401: {"error":"invalid_authorization_details", ...} until a workspace admin ran GRANT 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 answers 400 with an empty body (zerobus: dropping N rows, insert returned 400: ). The workspace's own org id comes back in the x-databricks-org-id response 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 PR

Type

🆕 New Feature

Caveats (if any)

Medium

  • Rows still queued at shutdown are lost, same as every other CustomBatchLogger
  • Queue caps at 50,000 rows; when a flush is in flight, extra rows are dropped with a warning
  • Wide rows (messages, response) count against Zerobus row and request size limits

Low

  • The UI leg of the After proof was captured at cb5acb4; the two later commits only touch a test and the dataclass repr, no UI or ingest behaviour
  • The Zerobus token endpoint rejects USE CATALOG and USE SCHEMA that a service principal only inherits through a group; both must be granted to the principal directly. Documented in the docs PR
  • rust-lint is red on this branch because litellm-rust/crates/core/tests/messages/request.rs:401 fails clippy on origin/main too (introduced by feat(rust): hand upstream response headers to the native Messages stream #43178); this PR touches no Rust and that check is not required

Taxonomy audit (A to BB)

  • M1: found and fixed, ZerobusConnection.client_secret and ZerobusAccessToken.value now use field(repr=False) with a test
  • B3: not applicable, CustomBatchLogger runs one flush at a time so a token miss mints once
  • J4: not applicable, queue bounded at 50,000 rows with oldest-first drop
  • M6 and M7: not applicable, token cached per client bound to one connection, refreshed 60s before expires_in
  • U5: not applicable, every connection field is required and validated before the URL is built
  • V1: not applicable, the destination is operator config, never taken from the request
  • W1: not applicable, the callback only reads kwargs
  • X1: intended, an empty config value means unset so the ZEROBUS_* env var fills in
  • T5: not applicable, tests inject a fake client and clock and restore litellm.callbacks in teardown

Final Attestation

  • The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

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

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

I'll fix CI failures and address comments from users with write access. I'll skip comments containing "(aside)".

  • Disable automatic comment, CI, and merge conflict monitoring

@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@codspeed

codspeed Bot commented Sep 19, 2026 •

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_zerobus_logging_integration (03b1cdb) with main (081f73f)

Open in CodSpeed

@greptile-apps

greptile-apps Bot commented Sep 19, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerConfidence Score: 5/5

[Medium risk] Adds a new logging integration for Databricks Zerobus.

The PR appears safe to merge based on the findings reviewed

Summary

This PR adds a Databricks Zerobus callback for batched success and failure traces, with configuration through YAML, environment variables, and the Admin UI.

  • The changes since the previous review strengthen the row serialization test and keep credentials and access tokens out of dataclass representations.
  • The previous findings are resolved or withdrawn; no new actionable issue was identified.

Reviews (3) · Last reviewed commit: "fix(zerobus): keep the client secret and..."

Comment thread litellm/integrations/zerobus/logger.py
Comment thread litellm/integrations/zerobus/logger.py
Comment thread tests/test_litellm/integrations/zerobus/test_zerobus_client.py Outdated
Comment thread tests/test_litellm/integrations/zerobus/test_zerobus_client.py Outdated
Comment thread tests/test_litellm/integrations/zerobus/test_zerobus_client.py Outdated
Comment thread litellm/integrations/zerobus/logger.py Outdated
@codecov

codecov Bot commented Sep 19, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 95.84665% with 13 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
litellm/integrations/zerobus/logger.py 92.30% 11 Missing ⚠️
litellm/integrations/zerobus/client.py 97.29% 2 Missing ⚠️

📢 Thoughts on this report? Let us know!

yassin-berriai and others added 2 commits September 19, 2026 20:31
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>
@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

@greptileai

Comment thread tests/test_litellm/integrations/zerobus/test_zerobus_row.py Outdated
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>
@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

@greptileai

@yassin-berriai
yassin-berriai merged commit 191305e into main Sep 25, 2026
104 of 106 checks passed

This branch is waiting to be deployed

1 waiting deployment
e2e-changed — 03b1cdb4 Waiting Sep 25, 2026 by devin-ai-integration[bot] via oauth #1119
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants