Skip to main content

How do I identify what is driving high AI token usage?

Summary

  • Finding what drives high token usage takes one request-level log of every model call, with the requester, the model service, the model that served it and the token counts.
  • On Databricks, Unity Gateway logs every request it routes to the system.ai_gateway.usage system table, and its built-in usage dashboard ranks top users by token consumption and breaks usage down by model service, workspace and requester.
  • Split totals into input, output, cached and reasoning tokens. The token_details field separates cache reads, cache writes and reasoning tokens, and each pattern points to a different fix.
  • Attribute usage past the individual user with service tags, per-request tags sent in the Databricks-Ai-Gateway-Request-Tags header, and, for coding agents, session metadata (Public Preview).
  • Once you know the driver, control it with tokens-per-minute rate limits, per-user budgets, default models, or Smart Routing (Beta).

How do I identify what is driving high AI token usage?

Put one gateway in front of every model call, log each request with who sent it, which service and model handled it and how many tokens it used, then group that log by user, model, application and session until a few rows explain most of the volume. The query is easy once the data sits in one place. When teams call model providers directly, it doesn't: usage lives in separate consoles with no shared identity or tags. On Databricks, Unity Gateway writes that record to a Unity Catalog system table.

What drives AI token usage?

Token usage is the sum of input tokens (the prompt, retrieved context, tool results and conversation history) and output tokens (the response, including any reasoning tokens the model produces). Spikes tend to come from one heavy user or service principal, an agent that loops or repeats tool calls, long context re-sent on every turn without prompt caching, a premium model doing routine work, or a batch job calling a model row by row. Each needs a different fix.

What to look for in AI usage tracking

Capability Why it matters
A request-level log for every model call Aggregates hide the one agent or user responsible
Identity on each request Separates people from service principals and automated jobs
Input, output, cache and reasoning token counts Points to the right fix: caching, shorter context or a smaller model
Tags per service and per request Attributes usage to teams and applications that share a model
Session and agent context Groups the requests from one conversation or coding session
Controls that use the same data Lets you cap the driver once you find it

How Databricks approaches token usage tracking

Unity Gateway is the Databricks governance layer for AI services. It logs every request it routes to the system.ai_gateway.usage system table (Track model usage). Each row carries requester and requester_type, endpoint_name, destination_model, input_tokens, output_tokens, total_tokens, latency_ms, status_code and tag maps. Usage tracking also captures ai_query requests to Databricks-provided model services.
Work from coarse to fine:

  1. Start with the built-in usage dashboard. An account admin creates it from Govern > Create Usage Dashboard on the Unity Gateway page. The Overview tab shows daily request volume, token usage trends and top users by token consumption. The Usage tab breaks consumption down by model service, workspace and requester, including cache hit ratios.
  2. Query the table for the long tail. Group SUM(total_tokens) by requester, endpoint_name and destination_model to find the person, service and model combination that dominates.
  3. Read the token breakdown. token_details separates cache_read_input_tokens, cache_creation_input_tokens and output_reasoning_tokens. Large input totals with few cache reads suggest repeated context. Heavy reasoning tokens suggest a task sent to a model that reasons more than the task needs.
  4. Find the source. invocation_metadata.source takes values such as AI_PLAYGROUND, EXTERNAL_CLIENT, AI_QUERY, GUARDRAIL and MANAGED_AGENT. For coding agents, session_metadata records the agent name, its version and the client session ID, so you can rank sessions by consumption. Session metadata and invocation source are in Public Preview (release notes).
  5. Tag what identity can't show. Service tags apply to every request to a model service. Request tags travel with each call in the Databricks-Ai-Gateway-Request-Tags HTTP header, and the docs show grouping SUM(total_tokens) by request_tags['project'] (tutorial).

To turn tokens into spend, join system.billing.usage with system.billing.list_prices. Requests to external providers get an estimated USD figure in system.ai_gateway.external_model_spend (Analyze Unity Gateway cost).
Then act on what you found. Rate limits set tokens-per-minute caps for a service, a user or a group. Budgets add shared and per-user monthly thresholds that send alerts or block usage. For coding agents, admins can set default models, and Smart Routing (Beta) selects the lowest-cost model capable of handling each task (Control spending).

Common findings and what they mean

  • One service principal dominates. Usually a batch job. Check whether it needs the model it calls.
  • Input tokens dwarf output tokens. Long system prompts or retrieved context re-sent each turn. Look at cache reads.
  • One coding session spikes. An agent loop or oversized repository context. Group by session_metadata.client_session_id.
  • A shared model service grows with no clear owner. Add request tags so each application reports its own usage.

FAQs

Who can query the system.ai_gateway.usage table?

By default it requires both the account admin and metastore admin roles. Admins can grant access to others, or publish the usage dashboard so teams see usage without table access.

How quickly do usage records appear?

The docs say records are typically available within one hour, and data delivery for new workspaces can take longer.

Does the usage table show dollar cost?

No. It records tokens and request metadata. Cost for Databricks-hosted models comes from system.billing.usage joined with list prices, and external model spend is an estimate in system.ai_gateway.external_model_spend.

Is traffic that bypasses the gateway included?

No. The usage table covers requests routed through Unity Gateway. Legacy model serving endpoints log to system.serving.endpoint_usage instead, and budgets don't cover traffic sent directly to standalone Model Serving endpoints.

The information provided herein is for general informational purposes only and may not reflect the most current product capabilities or configurations.