You ever look up at the stars and marvel at being here as the current version of you rather than the near infinite alternate versions that could have been? Yeah me neither. And for good reason; it turns out that we’re so predictable that any other versions would probably be doing similar stuff. On a related note, if we look at the aggregate of questions we pass to the standard AI assistant interface, the distribution of those questions appears astoundingly similar to a rank-frequency log normal distribution. That means if you send each one of those questions to an LLM, you are leaving money on the table.

Subscribe now

In 2024 Databricks published a semantic caching accelerator that stored cached questions in a Vector Search index. This works great, but it was made before we had Lakebase, and a cache is naturally an OLTP workload: we want to write a row on every miss, delete rows on expiry, and filter by tenant and/or model before the ANN search. These features make pgvector on Lakebase a more compelling cache. Let’s see it in action.

Imagine we have a customer-facing AI assistant in our client application. In the cacheless paradigm, the question goes into the application, maybe gets processed, and then gets handed to the LLM to return an answer to the customer. You pay for the FM API call every time. Now imagine we add a semantic cache. The question gets embedded and looked up in the pgvector cache, and if a high quality match is found, the cached answer is returned without invoking the LLM. If the threshold is not met, the question goes to the LLM and its answer is written back to the cache for future questions. Our cache is self-populating.

The semantic match threshold serves as our main tuning parameter between high cost and latency savings on the one side and risking higher incorrect answers on the other. We have a few possible error paths here:

• The threshold is met, but the wrong answer is returned from the cache

  • i.e. it should have picked a different cached answer

• The threshold is met, but the correct answer wasn’t in the cache

  • i.e. it should have called the LLM instead

• The threshold is met, but the answer saved was factually inaccurate

  • i.e. we need to better curate the cache

• The threshold is not met, and the LLM returned the incorrect answer

  • i.e. we should have had more guardrails, better system context, etc.

  • It’s possible the correct answer was in the cache, but not triggered

From my experience in building this, the threshold sweet spot is somewhere between 0.7 and 0.8. There is an inflection point in this range where the repeats caught drop dramatically, but so do the percent of wrong answers from the cache, regardless of the specific cache failure path. Luckily for us, we don’t need the repeat questions caught to reach a very high percentage of total asked for this to be economically advantageous, because the embedding model and Lakebase instance are both so cheap that intercepting even low single digit percentages of questions makes it worth it at high request volumes.

A big question for enterprises considering a semantic cache is how to know what the tradeoff numbers look like before going to prod. That’s where a labeled evaluation set and an LLM judge come in. The judge takes a sample of hits the cache actually served, generates a fresh answer to each question, and decides whether the cached answer is correct. Before trusting it, check it against your labels. The verdicts land in a Delta table, and a Spark Declarative Pipeline turns them into a false-hit monitor compared against a limit you set. It reports OK, WATCH or ALERT and recommends an action, such as raising the threshold or not caching the categories causing wrong hits. You can run the judge from the app, or you could put it on a job and schedule it.

Usability is also important for real world applications, so we also create a Genie Space over these tables that we can ask questions over in plain text, in realtime.

FAQ

• What if they ask user specific questions like account balance?

  • Partition with scope keys and send questions involving live user data around the cache before lookup

  • Include probes like these in your eval set so you’ll see this behavior in testing

• Why pgvector? Why not Lakebase Search?

  • Our project’s pgvector backend uses HNSW, Lakebase Search uses IVF

  • If the cache is large and loaded in bulk, then you can afford to reindex on a schedule, but IVF clusters cannot be incrementally updated

• What if the answer to a question changes, won’t I have to invalidate the cache to fix it?

  • The system prompt’s hash is part of every cache key, so editing the fact sheet means old answers are never served again, with nothing to invalidate

  • You can also set a TTL for cached items

• How fast are hits?

  • About 10x faster than misses; factor in about 150-200ms for the embedding call and maybe 10ms for the Lakebase lookup

• What does cache maintenance look like?

  • This is one of the most important tunable features here, because more cache entries means more neighbors, and probably higher rates of the wrong answer being returned

  • You could use an async, batch process to leverage an LLM judge to periodically cull closely related semantic items and keep only the best, most popular ones

• What other tuning params are there?

  • Threshold is the most important by far

  • hnsw.ef_search is probably a good one to explore when sweeping your eval set during configuration

  • The embedding model itself is also a swappable component

• Does this come with any guarantees or SLAs?

  • No, this is purely field originated project, so for now it comes with no performance, functionality, or version guarantees

  • If you make a PR or raise an issue in the repo I will get to it on a best effort basis

Repo here. Cheers and happy coding,

-Austin

LoRA adapters make fine-tuning lighter and faster, and multi-LoRA serving makes them inexpensive to run. This post explains when multi-LoRA is the right choice, and walks through serving a base model and several adapters from a single Custom LLM Serving endpoint with vLLM.

Summary

• If you have several fine-tuned variants of the same base model, you can usually serve them all from one GPU endpoint. Multi-LoRA serving keeps one copy of the base model in GPU memory and applies the requested adapter to each request, so you pay for one endpoint instead of one per model.

• On Databricks, you package the base model, the adapters and the vLLM/SGLang command as one MLflow model, register it to Unity Catalog and deploy it with Custom LLM Serving. Each request then picks an adapter, or the base model, by name.

• Multi-LoRA fits best when you have many adapters with moderate traffic each. For a single fine-tuned model with high traffic, where latency matters most, serving the merged weights is usually the better option.

What LoRA is and why it helps

Fine-tuning every weight of a large language model is slow and expensive, and it produces a full copy of the model for every task. LoRA (Low-Rank Adaptation) takes a lighter approach. It freezes the base model and trains a small pair of low-rank matrices next to selected layers, such as the attention and MLP projections. Only those matrices, called the adapter, are saved.

An adapter typically holds about 0.1% to 2% of the base model’s parameters, depending on its rank and the layers it targets. On disk, that is often tens to hundreds of megabytes for a model whose weights take many gigabytes. This brings three practical advantages. Training needs less GPU memory and time, so a single GPU is often enough. Adapters are cheap to store and version, so you can keep one per task, per customer or per experiment. And in many tasks, a well-trained adapter reaches quality close to full fine-tuning, although you should confirm that on your own evaluation data.

Multi-LoRA serving: one base model, many adapters

Because each adapter only adds to a shared base model, a serving engine can load the base model once and keep several adapters next to it. Multi-LoRA serving does exactly that. Every request names the adapter it wants, the engine applies that adapter’s small matrices on the fly, and requests for different adapters can share the same batch on the GPU.

In the example used throughout this post, one endpoint serves Qwen3.5-4B, a small open model, with three adapters: finance, customer_support and text_to_sql. A request with model="finance" gets the finance adapter, and a request with the base model’s name gets the base model, all from the same GPU.

This design brings four main benefits:

• Lower cost. Several fine-tuned models share one GPU instead of each needing its own endpoint, which matters most when each model has moderate or uneven traffic.

• Fewer GPUs to secure. GPUs are in high demand, and every extra endpoint needs its own. Serving many models from one GPU means you need capacity for one endpoint instead of one per model.

• No cold start per model. All adapters are loaded when the endpoint starts, so switching between them does not wait for a model to load.

• Simpler operations. You manage, monitor and secure one endpoint for many specializations.

Good fits include domain adapters on one base model (for example legal, medical, SQL generation and others), one adapter per customer or business unit, and A/B tests that compare adapter versions, with each request naming the version it uses.

When not to use it

Multi-LoRA is not the best choice for every case:

• One high-traffic model where latency matters most. Applying an adapter adds some computation per request, which grows with the adapter rank and the number of different adapters in a batch. Merged weights, with the adapter folded into the base model, avoid that cost and work with engine optimizations that may not support adapters, such as some speculative decoding methods.

• Adapters on different base models. All adapters on one endpoint must share the same base model.

• Adapters that need separate access. Access is controlled per endpoint, so anyone who can query the endpoint can call every adapter on it. Keep such adapters, for example some per-customer ones, on separate endpoints.

• Adapter sets that change many times a day. Adding new adapters to a running endpoint dynamically is not currently supported on Databricks, so every adapter must be included when you log the model. To add one later, you register a new model version and update the endpoint.

• Endpoints that split traffic. To select an adapter, the endpoint must serve a single model with 100% of its traffic. Endpoints that split traffic between models or versions can’t select an adapter.

Custom LLM Serving on Databricks

Custom LLM Serving runs your own LLM on a Model Serving GPU endpoint, with an OpenAI-compatible API. You choose the inference engine and its settings, test the exact command in an AI Runtime notebook, and deploy the same command and environment to production.

Custom LLM Serving is the right choice when you want to serve an open model, or a fine-tuned version of one, that the Databricks Foundation Model APIs offering does not include, when you want to modify the model’s architecture, or when you need control over the engine and its configuration. If Foundation Model APIs already offers the model you need, it is the recommended option, because it is managed, optimized and simple to use.

Serving several LoRA adapters is a natural fit for Custom LLM Serving, because the engine’s own multi-LoRA support does the work. The documentation covers it in Serve LoRA adapters. Custom LLM Serving is in Public Preview, so check the documentation for its current status.

Serving multiple LoRA adapters with vLLM

The workflow has five steps. You put the base model and adapters in one folder, write the vLLM command that serves them, test that command in the notebook, log and register the model to Unity Catalog, and deploy the endpoint. Run these steps in an AI Runtime notebook with a GPU, because the endpoint reuses the notebook’s Python environment. The AI environment comes with vLLM preinstalled, while a Standard environment lets you install the libraries and versions you choose.

Step 1: Prepare the base model and adapters

Everything the endpoint needs lives under a single folder, which is logged as one artifact. Each adapter is a standard PEFT folder with an adapter_config.json and an adapter_model.safetensors file, and the folder names become the names that requests use.

weights/
  base/        # base model weights and tokenizer
  finance/     # adapter_config.json, adapter_model.safetensors
  customer_support/
  text_to_sql/

The adapters can come from your own fine-tuning runs on AI Runtime, saved to a Unity Catalog volume, or from Hugging Face. For example, to copy an adapter from a volume:

import shutil
shutil.copytree("/Volumes/<catalog>/<schema>/<volume>/finance/adapter", "weights/finance")

Before you deploy, check that every adapter’s adapter_config.json names the same base model that you serve. If you use an open base model or public adapters, also check that their licenses allow commercial use by your organization.

Note for Qwen3.5 and Qwen3.8: vLLM and SGLang load these models as vision-language models, so adapters must use model.language_model.layers.* key names. Adapters saved with model.layers.* load without error but behave like the base model, so rename those keys before you log the model.

Step 2: Build the vLLM entrypoint

The entrypoint is the command the endpoint runs to start vLLM. It must listen on port 8080, and its paths are relative to the logged artifacts, so the same command works in the notebook and on the endpoint.

ADAPTERS = ["finance", "customer_support", "text_to_sql"]
ENDPOINT_NAME = "multi-lora-vllm"

def vllm_command(port):
    return " ".join([
        "VLLM_USE_FLASHINFER_SAMPLER=0",
        "python -m vllm.entrypoints.openai.api_server",
        "--model weights/base",
        f"--served-model-name base {ENDPOINT_NAME}",  # names that select the base model
        "--enable-lora",
        "--lora-modules " + " ".join(f"{a}=weights/{a}" for a in ADAPTERS),
        f"--max-loras {len(ADAPTERS)}",  # adapters on the GPU at once
        "--max-lora-rank 16",            # at least the largest adapter rank
        "--max-model-len 2048",
        """--default-chat-template-kwargs '{"enable_thinking": false}'""",  # Qwen3.5 thinks by default
        f"--host 0.0.0.0 --port {port}",
    ])

The LoRA flags carry the multi-LoRA setup.

• --enable-lora turns on adapter support

• --lora-modules registers each adapter under its name

• --max-loras sets how many adapters can be active in one batch

• --max-lora-rank must be at least the largest rank among your adapters, rounded up to a value vLLM accepts. Set it from your adapters’ configs, because higher values reserve more GPU memory.

If you serve more adapters than --max-loras, raise --max-cpu-loras to cover all of them. vLLM then keeps the extra adapters in CPU memory and moves them onto the GPU as requests arrive, so test with your expected mix of traffic to choose --max-loras.

Qwen3.5 thinks before it answers by default, so --default-chat-template-kwargs turns thinking off and the adapters answer directly. A request can still turn it on.

Step 3: Test in the notebook

Before you log anything, start the same command in the notebook and send a request to each adapter. This catches most configuration problems in minutes instead of after a deployment. Notebooks can listen on ports 3000 to 3999, so the test uses port 3080.

import os, signal, subprocess, requests

server = subprocess.Popen(vllm_command(3080), shell=True, start_new_session=True)

# Wait until http://localhost:3080/health responds, then:
print(requests.get("http://localhost:3080/v1/models").json())  # base, the endpoint name and every adapter

for name in ["base", *ADAPTERS]:
    reply = requests.post("http://localhost:3080/v1/chat/completions", json={
        "model": name, "messages": [{"role": "user", "content": "What is EBITDA?"}], "max_tokens": 64,
    })
    print(name, reply.json()["choices"][0]["message"]["content"])

os.killpg(server.pid, signal.SIGTERM)  # stop vLLM and free the GPU

Compare each answer with what the adapter was trained to do. Stopping the whole process group matters, because server.terminate() would stop only the shell and leave vLLM running on the GPU.

Step 4: Log and register the model to Unity Catalog

The endpoint runs the entrypoint directly and never calls the model’s predict method. MLflow still needs a model class to log, so a small placeholder class is enough. The metadata carries the two settings the endpoint reads: task, which sets the API the endpoint serves, and entrypoint, the vLLM command. The entrypoint is logged with port 8080, because Model Serving sends requests to the server on that port, while the notebook test used port 3080.

import mlflow
from mlflow.pyfunc.model import ChatCompletionResponse, ChatModel

class Placeholder(ChatModel):  # never called; the endpoint runs the entrypoint
    def predict(self, context, messages, params):
        return ChatCompletionResponse.from_dict({"choices": []})

mlflow.set_registry_uri("databricks-uc")
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        name="multi_lora_vllm",
        python_model=Placeholder(),
        artifacts={"weights": "weights"},
        metadata={"task": "llm/v1/chat", "entrypoint": vllm_command(8080)},
    )

version = mlflow.register_model(
    model_info.model_uri, "main.models.multi_lora_vllm", env_pack="databricks_model_serving"
).version

Setting env_pack="databricks_model_serving" packages the notebook’s Python environment with the model, so the endpoint runs the same vLLM version that you tested and does not need to rebuild it.

Step 5: Deploy the endpoint

Create the endpoint with the Databricks SDK, choosing a GPU workload type that fits the base model, plus the adapters and the KV cache.

import mlflow
from mlflow.pyfunc.model import ChatCompletionResponse, ChatModel

class Placeholder(ChatModel):  # never called; the endpoint runs the entrypoint
    def predict(self, context, messages, params):
        return ChatCompletionResponse.from_dict({"choices": []})

mlflow.set_registry_uri("databricks-uc")
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        name="multi_lora_vllm",
        python_model=Placeholder(),
        artifacts={"weights": "weights"},
        metadata={"task": "llm/v1/chat", "entrypoint": vllm_command(8080)},
    )

version = mlflow.register_model(
    model_info.model_uri, "main.models.multi_lora_vllm", env_pack="databricks_model_serving"
).version

One A10G fits the 4B model in this example. For a larger base model, such as Qwen3.8-27B, use a 1x H100 workload type (GPU_XLARGE) or RTX Pro 6000 Blackwell, and serve the adapters the same way. Size the endpoint for your peak traffic, and for production, consider turning off scale to zero to avoid the startup wait after idle periods. The documentation lists the available GPU types and the model sizes each one suits.

Calling an adapter or the base model

The recommended way to call the endpoint is the OpenAI client on the chat route, with the endpoint named in the Databricks-Model-Serving-Endpoint header. The request body reaches vLLM unchanged, so the model field selects the adapter, and the base model’s served name selects the base model. The documentation shows the same calls in Query a LoRA adapter.

from openai import OpenAI

client = OpenAI(
    api_key=DATABRICKS_TOKEN,
    base_url=f"{DATABRICKS_HOST}/serving-endpoints",
    default_headers={"Databricks-Model-Serving-Endpoint": ENDPOINT_NAME},
)
messages = [{"role": "user", "content": "What is EBITDA?"}]

finance = client.chat.completions.create(model="finance", messages=messages)  # finance adapter
base = client.chat.completions.create(model="base", messages=messages)        # base model

If you are not using Python, send the same request over REST to the same route. The header names the endpoint, and the model field in the body names the adapter.

curl "$DATABRICKS_HOST/serving-endpoints/chat/completions" \
  -H "Authorization: Bearer $DATABRICKS_TOKEN" \
  -H "Databricks-Model-Serving-Endpoint: multi-lora-vllm" \
  -H "Content-Type: application/json" \
  -d '{"model": "text_to_sql", "messages": [{"role": "user", "content": "Count orders by status."}]}'

A name that was not loaded returns an error, so a typo does not silently fall back to the base model.

Use the chat route whenever a request needs an adapter. Other routes, such as /serving-endpoints/<endpoint>/invocations (used by the Databricks SDK’s query() method), serve the base model with vLLM.

Serving adapters with SGLang

SGLang also supports multi-LoRA serving, and the Databricks workflow is the same once you install SGLang in the notebook environment. Only the command changes: use --enable-lora and --lora-paths to register the adapters, and set --max-loras-per-batch to the number of adapters plus one, because base-model requests take a slot too.

To call an adapter, use the OpenAI client with the same base_url but without the endpoint header, set model to the endpoint name, and pass the adapter name with extra_body={"lora_path": "finance"}. Without lora_path, the request goes to the base model. Note that model="finance", which selects the adapter in vLLM, serves the base model in SGLang.

Get started

Multi-LoRA serving lets one GPU endpoint do the work of many, which lowers cost and keeps operations simple when you have several fine-tuned variants of one base model. To try it, fine-tune few adapters on AI Runtime, deploy them together with Custom LLM Serving, and compare the cost of that one endpoint with the cost of one endpoint per model.

For a ready-made start, the documentation includes a Qwen3.5-4B and LoRA adapters starter notebook. Evaluate each adapter on your own data, and benchmark your own workload on each candidate GPU before you choose a configuration for production.

TL;DR

• Cursor and Claude Code call https://<your-app-url>/mcp/ directly; Genie One goes through a Unity Catalog MCP Service and HTTP connection.

• Either way, the app gets a token its workspace issued to the human user, and the Databricks Apps proxy forwards that user’s identity.

• The only account-admin step: register one public OAuth client with each AI client’s redirect URI. The rest is client config and grants.

How it works: two paths, one identity

Direct path. Cursor and Claude Code connect to the app’s Streamable HTTP MCP endpoint, https://<your-app-url>/mcp/.

Governed path. Genie One calls a Unity Catalog MCP Service, which calls the same app URL through a Unity Catalog HTTP connection.

Identity. On both paths, the app receives a token issued by the app’s workspace for the human user. The Databricks Apps proxy validates it and forwards the user’s identity (X-Forwarded-Email, X-Forwarded-User). The app resolves the user from those headers, so data is owned by the person.

Discovery. Databricks Apps publish OAuth protected-resource metadata (RFC 9728) at /.well-known/oauth-protected-resource/mcp/, pointing at the workspace’s OIDC server. Databricks OIDC supports PKCE and public clients (no secret; PKCE protects the sign-in instead), but not dynamic client registration. So AI clients can’t register themselves; they need a pre-registered public OAuth client, and one can serve all three.

That last point explains most of the setup, and most of the errors.

The one account-admin step: register a public OAuth client

An account admin of the account that owns the app’s workspace creates a public (no secret) custom OAuth app connection: Account Console → Settings → App connections, or databricks account custom-app-integration create.

• Scopes: all-apis offline_access (optionally openid email profile).

• Redirect URIs:

  • Cursor: http://localhost:8787/callback, the callback current versions use

  • Claude Code: http://localhost:<callbackPort>/callback, for example port 8080

  • Genie One: https://<genie-workspace-host>/login/oauth/http.html, for every hostname that workspace uses, vanity ones included

Updates replace the whole redirect list, so always include the existing URIs. Then grant CAN_USE on the app to the users and groups who should use it.

Connect Cursor to the MCP endpoint

Add the server to ~/.cursor/mcp.json. Current Cursor versions read auth.CLIENT_ID and ignore the Claude-style oauth.clientId.

{
  "mcpServers": {
    "my-app": {
      "url": "https://<your-app-url>/mcp/",
      "auth": { "CLIENT_ID": "<client-id>", "scopes": ["all-apis", "offline_access"] }
    }
  }
}

Then open Settings → Tools & MCP, click Connect, and sign in once in the browser.

Connect Claude Code to the MCP endpoint

Add the server with the CLI, then add scopes so its mcpServers entry in ~/.claude.json matches the JSON below.

claude mcp add --transport http --client-id <client-id> --callback-port 8080 -s user my-app https://<your-app-url>/mcp/

"my-app": {
  "type": "http",
  "url": "https://<your-app-url>/mcp/",
  "oauth": { "clientId": "<client-id>", "callbackPort": 8080, "scopes": "all-apis offline_access" }
}

Run /mcp, select the server, and choose Authenticate. Both clients store and refresh tokens themselves. No proxy scripts, no token commands.

Connect Genie One through Unity Catalog

In the Genie One workspace, create a Unity Catalog HTTP connection:

• Type: OAuth U2M (user-to-machine), per user

• Host: https://<your-app-url>, with base path /mcp/

• Authorization endpoint: https://<app-workspace-host>/oidc/v1/authorize

• Token endpoint: https://<app-workspace-host>/oidc/v1/token

• Client ID and scope: your public client ID and all-apis offline_access

Then create a Unity Catalog MCP Service on top of it, such as <catalog>.<schema>.<service>. Each user signs in once (from Genie One’s connections or the service’s sign-in page), attaches the service in a chat (+ → More connections), and asks Genie to use the tools.

How Unity Catalog governs Genie One access

• The MCP Service is the access check. It decides who may call: USE CATALOG and USE SCHEMA on its catalog and schema, plus EXECUTE on the service.

• The HTTP connection is token management. It stores each user’s token from their one-time sign-in, refreshes it, and attaches it to calls.

• Callers need no privileges on the connection itself.

• Tokens are per user. Every Genie One call carries the caller’s own token, never a shared one.

• Use a catalog and schema your audience can reach. A personal catalog only works for you.

Genie One in a different workspace? Same setup

Genie One can live in another workspace, even on another cloud: the connection signs users in against the app’s workspace, which issues the token. Register the Genie workspace’s callback on the public client, and if the app workspace uses IP access lists, allow the Genie workspace’s gateway egress.

Five traps we hit, and their fixes

The first two stem from false assumptions: that you can borrow the app’s own OAuth client, and that clients register themselves.

• invalid_client at token exchange. You reused the app’s built-in OAuth client. It’s confidential; only the Apps platform has its secret. Use a separate public client.

• does not support dynamic client registration. The client isn’t reading your client ID (wrong config key), or a running session still has old config. Fix the key, then reconnect.

• redirect_uri ... not registered. The OAuth client is missing that AI client’s callback, or the Genie workspace’s vanity hostname. Add it.

• Default scopes. Unless you pin them, clients request the scopes the app advertises, often iam.*, which may not be on your client. Pin all-apis offline_access.

• Not Acceptable: Client must accept text/event-stream. You opened /mcp/ in a browser. That’s normal.

Why it matters

One app, one sign-in model, every AI client: real-user identity on every call, Unity Catalog governance for Genie One, and no shared token or proxy script to babysit. For more, see the Databricks docs on MCP, custom MCP servers, and Apps authorization, which covers both identity models.

FAQ

Q: What is an MCP server on Databricks Apps? A: An MCP server you build yourself and run on . Apps provides the compute and an HTTPS URL that AI clients connect to over Streamable HTTP. You can start from a template, wrap an existing REST API, or deploy an MCP server you already have. See and the MCPs overview.

Q: How is this different from Databricks-provided or external MCP servers? A: Databricks-provided MCPs (Genie One, Databricks SQL, Slack, GitHub and others) live in system.ai and need no hosting. External MCP servers are third-party or self-hosted servers you register in Unity Catalog. A custom server on Databricks Apps is your own code. This post registers it as an external MCP Service so Genie One can use it.

Q: Do I need a separate MCP server for each AI client? A: No. Cursor, Claude Code and Genie One all use the same https://<your-app-url>/mcp/ endpoint and the same public OAuth client.

Q: What’s the difference between app authorization and user authorization? A: With app authorization, everyone shares the app’s service principal. With user authorization (on-behalf-of-user), Databricks forwards the user’s identity and token, so Unity Catalog row filters and column masks apply to that person. See and Authentication for agents.

Q: Can I use a personal access token instead of OAuth? A: Not for servers on Databricks Apps; those require OAuth. PATs only work for local testing against Databricks-provided or registered MCPs.

Q: Why can’t Cursor or Claude Code register their own OAuth client? A: Databricks MCP endpoints (RFC 7591). So an account admin pre-registers one public client that uses PKCE. Discovery follows the MCP authorization spec and RFC 9728 protected-resource metadata.

Q: How do I create the public OAuth client? A: In the account console, go to Settings → App connections → Add connection and clear Generate a client secret. Or run databricks account custom-app-integration create. See . all-apis covers non-SQL APIs, and offline_access lets clients get refresh tokens.

Q: What scopes should I request? A: For an app-hosted server, include the app’s user-authorization scopes plus offline_access, and grant CAN USE on the app. See .

Q: How do I connect Cursor and Claude Code? A: Start with the Databricks guides for Cursor and Claude Code. Each client’s own MCP docs cover its config details: Cursor MCP docs and Claude Code MCP docs. Using Codex instead? See Connect Codex.

Q: What is a Unity Catalog HTTP connection, and do Genie One users need privileges on it? A: An is a Unity Catalog securable that stores endpoint and credential details and handles OAuth token refresh and per-user sign-in. When you go through an MCP Service, users need USE CATALOG, USE SCHEMA and EXECUTE on the service, not privileges on the connection. See Unity Catalog connections.

Q: Can I also use Genie One from Claude or ChatGPT? A: Yes. The Genie One MCP server works in the other direction: it exposes Genie to outside MCP clients.

Q: Does this work when Genie One is in a different workspace or cloud? A: Yes, as long as the connection signs users in against the app’s workspace. If that workspace uses IP access lists, allow the Genie workspace’s egress IPs.

Q: I changed my OAuth client and nothing happened. Why? A: Changes to an OAuth app connection . Wait, then reconnect the client.

• Domains are a Unity Catalog organizing and discovery layer, built on governed tags — not a new securable type, and not an access-control mechanism.

• They sit on top of your existing catalog/schema layout, letting you overlay business structure (Finance, Customer, Risk) without moving data.

• Design top-level domains around business capabilities — Function, Region, Brand, or Business Unit — and spend your one subdomain level deliberately.

• Keep domains, meaning (Pages), measures (Metric Views), and access (grants/ABAC) as separate concerns; collapsing them is the most common design mistake.

• Access always stays with Unity Catalog privileges — a domain’s Owner and Contact fields are metadata only.

Where domains fit in Genie Ontology

Genie Ontology is the AI-ready business semantic layer that grounds Genie and every business-facing search — it’s how the platform knows what your data means, not just where it sits. It combines a Modelled Ontology you curate and govern with inferred context Genie learns automatically from existing tables, queries, dashboards, and notebooks. Each piece of context gets an authority score based on source, usage, and freshness; Genie ranks the relevant pieces, resolves conflicts, and answers using only sources each user is permitted to see.

The Modelled Ontology has three components, grouped as Unity Catalog Semantics: Domains, Pages, and Metric Views. The most common design mistake is collapsing them — treating a domain as if it also carried meaning, metric logic, or access rights. It doesn’t. Domains organize assets by business purpose and assign an owner; Pages supply meaning; Metric Views supply measures. Each Page belongs to exactly one domain or subdomain, and a domain typically holds many terms and metrics, so domains are the scaffold the other two hang from.

Figure 1. Where domains sit: above the physical layout, granting no access.

That scoping matters as the semantic layer grows. Once an estate holds hundreds of Metric Views across dozens of business areas, the hard part stops being reaching a definition and becomes choosing the right one — and domain, owner, and certification are exactly the filters that narrow that choice. The same curation pays off twice: it makes the estate easier for people to browse and gives Genie a trusted, business-scoped signal for retrieval.

Why domains, and why now

Most Unity Catalog estates are organized the way they were built — by source system, pipeline, or environment — which rarely matches how the business thinks about its data (”Marketing,” “Risk,” “Customer 360”). A domain is the top-level unit of data ownership: a business function, region, brand, or unit with a real, accountable team. Subdomains add one level of nesting (e.g., Finance → Tax).

This separates physical data ownership (who runs the pipeline) from business ownership (who defines what the data means) — one of the main reasons domains exist. In a small team, the people who build a pipeline also know what the data means and who should use it. As organizations grow, these jobs come apart: pipeline teams own how data arrives, while business stewards own what it means. Domains give the second group a place to stand, independent of where the bytes live.

It helps to name the two areas of responsibility explicitly, so problems land with the right owner. A pipeline owner is accountable for freshness and SLAs, job failures, schema changes, and cost — “the table didn’t refresh this morning” gets fixed in the pipeline or job. A domain owner or data steward is accountable for business logic, metric definitions, certification, and being the person to ask — “revenue here excludes refunds, and it shouldn’t” gets fixed in the domain’s Pages, Metric Views, and asset metadata.

Figure 2. Two kinds of ownership, and where each kind of problem gets fixed.

When everything is owned by the pipeline team by default, business-logic questions end up in engineering backlogs and nobody signs off on what a number means. Domains make the business side of ownership visible and assignable.

The idea isn’t new — it’s common across enterprise catalogs and Data Mesh/Data Fabric architectures. What’s different is that Databricks makes domains native to the platform, built on Unity Catalog governed tags, rather than bolted on via a separate product.

As of this writing, Domains, subdomains, and Discover are in Public Preview on AWS, Azure, and GCP; Pages is in Beta; Genie Ontology is in Public Preview; Metric Views are GA. Check the docs for current status.

How domains actually work

A few mechanics drive almost every design decision:

• A domain is a governed tag designated as a domain — not a new securable. A subdomain is a child tag ({parentDomain}/{subdomain}), nested exactly one level deep.

• Assets join a domain by applying the backing tag, via the “Add to Domain” picker, or indirectly through a subdomain. Membership is many-to-many — any taggable securable or workspace asset can belong to multiple domains at once.

• Each domain is its own tag key, not a value — that’s what makes multi-domain membership possible, since tag values are mutually exclusive but keys aren’t.

• Tagging a catalog or schema cascades to everything inside it — the fastest way to fill a domain when a schema maps cleanly, and a blunt instrument when it doesn’t.

• Browsing a parent domain rolls up its subdomains, but only shows what the viewer is already permitted to see.

• Domains are account-level (governed tags replicate globally) while Discover is enabled per workspace. A single domain can even span multiple metastores — it isn’t tied to one. Two consequences: never put sensitive text in domain names or descriptions, and expect the same domain to look different across workspaces.

• There’s a ceiling: 300 domains/subdomains per account, and each also counts against the account’s 1,000 governed-tag limit (at the time of writing).

• Domains confer no access. Owner and Contact are metadata only.

Figure 3. Why each domain is its own tag key.

The tradeoff of one key per domain. Every domain and subdomain uses up a governed tag (see the limit above), and governed tags are account-wide, so adopt a naming convention that keeps domain tags distinct from your other governed tags. Questions like “which domains is this asset in?” mean checking several keys rather than reading one value. And if you ever write ABAC policies against domain tags, you need a condition per domain — one more reason to keep access on its own tags (see “Domains and access control” below).

Managing domains needs curator permissions: account-level MANAGE DISCOVER, or MANAGE on a specific domain or subdomain. Day-to-day curation needs ASSIGN TAG on the domain’s tag plus APPLY TAG on the asset.

Design principles for your domain taxonomy

Model the business, not the plumbing

Base top-level domains on business capabilities, never source systems or pipelines. Source-system groupings feel natural to engineers and useless to everyone else: if your top-level domains read like a list of your ingestion connectors, start over.

So what should the top level look like? The four most common patterns we see with customers are below. Treat them as illustrative examples, not a fixed taxonomy — most real organizations blend them.

• Function — grouped by type of work (Sales, Finance, HR). Owned by a cost-center lead with no standalone P&L. Common in single-product or single-market orgs with one central data team.

• Region — grouped by geography (North America, EMEA, APAC). Owned by regional leadership, often with a regional P&L, and usually driven by regulation or data residency. Common in banks, insurers, telcos, and multinationals.

• Brand — grouped by market-facing identity, such as the separate brands of a consumer-goods or automotive group. Owned by a brand team, usually with its own P&L. Common in CPG, retail groups, hospitality, and automotive.

• Business Unit — a self-contained business (Retail, Commercial, Wealth, Markets) with a full P&L, often its own ERP and data team. Common in conglomerates and diversified groups.

Figure 4. Four common ways to draw the top level. Most organizations blend them.

Rather than a decision tree, a few signals tell you which pattern dominates: separate P&Ls lean Business Unit; hard data-residency lines lean Region; brand-owned P&Ls lean Brand; one business on one platform leans Function. Whichever dominates, add a shared Corporate domain for consolidated data no single domain owns — group finance, enterprise reference data, cross-business KPIs.

You have exactly two levels — spend them deliberately

Keep 8–15 coarse top-level domains and add subdomains only where a domain is genuinely large or has distinct sub-owners. Where you do subdivide, find the greatest common denominator of how your domains naturally break down — the subdivision that holds up across most of them. That’s a modelling exercise, not a mechanical rule: the dimensions you didn’t choose for the top level don’t automatically become your subdomains. If you’re wishing for a third level, a top-level domain is really two.

Signs you’ve mis-sized: subdomains with one or two assets (that’s a tag, not a subdomain); a single “Analytics” domain that owns half the estate (too coarse); or a top level so broad that subdomains are doing all the real work.

One clear owner per domain and subdomain

Every domain needs a named, accountable owner, not a committee. Owner/Contact are metadata, not permissions — their job is answering “who do I ask?”, turning a Slack scavenger hunt into a one-click answer. If two teams both claim a domain, that’s a signal to split it into subdomains with distinct owners. If an asset sits in several domains, own it where it’s mastered; membership elsewhere is a reference that helps people find it, not co-ownership.

Keep domains decoupled from physical layout

A domain can span multiple catalogs and schemas with no data movement — design the taxonomy you want, then map assets into it. That decoupling is the feature’s superpower: it lets you present a clean business view without a risky migration.

Design for discovery first

Optimize for a person (or agent) browsing for data they don’t know exists: business-recognizable names, no jargon, shallow trees, and certification signals so people can tell the golden asset from the abandoned one. A good test: could a new analyst find “customer churn” data in under three clicks without asking anyone?

Stable at the top, flexible below

Top-level domains should map to durable capabilities; absorb reorgs in subdomains, which are cheaper to reshape. A volatile top level erodes trust and breaks bookmarks, dashboards, and mental models. Get names right from day one — a domain’s name is its tag key and can’t be changed later. You can edit a domain’s subtitle and description, but changing its name means creating a new domain, re-tagging assets, and deleting the old one.

Mapping domains onto catalogs and schemas

Three common layouts: schema-per-domain (a clean greenfield 1:1 mapping), catalog-per-business-unit (a top-level domain gathers schemas within one catalog, subdomains map to individual schemas), and domains spanning multiple catalogs (a “Customer” concept scattered across raw, curated, and data-science catalogs can be pulled into one business view without touching the physical layout).

Rule of thumb: map domains to the gold/consumption layer first, then decide case-by-case whether raw/silver assets belong in the same domain or a separate “raw/engineering” domain. Tag a schema when it maps cleanly to one domain; tag individual assets when a schema mixes domains or holds scratch tables you don’t want surfaced.

Figure 5. One estate, two views: domains pick consumption-ready assets from any catalog, with no data movement.

A domain is an overlay, not a remedy — confirm resolved identity, documented grain, and readable metadata on your gold layer before laying domains over it.

Domains as the enabler of a Data Mesh on Databricks

Domains are the Unity Catalog primitive that makes a domain-driven Data Mesh practical: each business domain gets an owner, a home for its data products, and a place in a shared discovery experience — without splitting the platform into separate stacks. Domain ownership, data-as-a-product, and federated governance map directly onto UC domains, grants, and ABAC, while Discover serves as the self-serve browsing index. The key insight: you get decentralized ownership without decentralized infrastructure — one metastore, one governance model, one discovery surface, with a business-shaped view on top.

Domains and access control: what they do not do

This is the most common misconception, so it’s worth stating plainly: domains are a curation and discovery layer, not an authorization mechanism. Assigning an asset to a domain does not grant or revoke access — users still need BROWSE/VIEW to see it listed and the underlying Unity Catalog privileges to query it. You also can’t make a domain private: there’s no setting that restricts a published domain to specific users or groups, and the only role you can assign on a domain is the curator role that lets someone manage it. (Before publishing, a draft domain is visible only to curators.) What each person sees inside a domain follows their permissions on the assets, not the domain — so don’t reach for domains to hide a business area. Owner/Contact remain metadata only.

Figure 6. Domains help people find data; Unity Catalog decides who gets it.

Governed tags, including domain tags, are valid ABAC attributes. Row filter and column mask policies use them to protect data a user can already access; they don’t grant access to tables (yet). ABAC GRANT policies can grant privileges by tag, but today only on models, model and MCP services, and (in Beta) agent services and skills. So for tables and views, domain membership doesn’t provision access.

Subscribe now

If your access policies need a business-area attribute, model it as its own tag — a single key with controlled values — rather than reusing your discovery domain tags. The two shapes look similar but serve opposite purposes: one decides what a user can see, the other helps them find it. Keep access in grants and ABAC policies; use domains only to organize and surface what people can (or can request to) access.

Figure 7. Keep discovery tags and policy tags separate.

Operating model, ownership, and anti-patterns to avoid

Stand up one high-value domain first, prove the pattern, then expand — don’t attempt an enterprise-wide rollout at once. Assign a named owner per domain/subdomain, govern top-level domain creation centrally, and delegate curation below. Link Pages and Metric Views to domain assets rather than duplicating descriptions, certify production-ready assets, and route feedback to the layer that actually owns the fault — a definition belongs in a Page, a measure in a Metric View, organization in the domain. Manage the taxonomy itself (domains, subdomains, owners) as code via the REST API or Terraform, and review it quarterly.

Common anti-patterns: mirroring the org chart instead of modeling capabilities; expecting domains to enforce access; reusing one tag for both discovery and ABAC policy; wishing for a third nesting level; source-system domains at the top (the exception: a pipeline team can keep a clearly labeled, source-aligned “Raw / Engineering” domain for the assets it owns, separate from the business taxonomy); one-asset subdomains; co-owning shared assets instead of picking a home; treating a domain as the semantic model itself; and placeholder names you’ll regret, since domains can’t be renamed.

Getting started checklist

1. Confirm the foundation: resolved identity, documented grain, readable metadata on gold tables.

2. Inventory your gold/consumption assets.

3. Pick the dominant pattern — Function, Region, Brand, or Business Unit — plus a shared Corporate domain.

4. Draft 8–15 durable top-level domains with names you’ll keep.

5. Name an owner for each domain.

6. Add subdomains only where scale or ownership demands it.

7. Map assets into domains — tag schemas where they map cleanly, assets individually otherwise — and certify production-ready ones.

8. Operationalize one high-value domain end-to-end first, then expand.

9. Link Pages and Metric Views to domain assets.

10. Keep access in grants/ABAC, on its own tags.

11. Set a quarterly review cadence.

FAQ

Q: Is a Unity Catalog domain a new type of securable object?

A: No. A domain is a governed tag that has been designated as a domain. It doesn’t change where data lives or create a new access boundary.

Q: Can one asset belong to more than one domain?

A: Yes. Because each domain is modelled as its own tag key rather than a value, membership is many-to-many — an asset like “Customer” can sit in both “Marketing” and “Risk” at once.

Q: Do domains control who can see or query data?

A: No. Domain membership and the Owner/Contact fields are metadata only. Access is governed entirely by Unity Catalog privileges and, where relevant, ABAC policies.

Q: How many levels of subdomains can I create?

A: Exactly one. A domain can have subdomains, but subdomains cannot have their own sub-subdomains. If you need a third level, your top-level domain probably needs to be split.

Q: Can I rename a domain later?

A: No. A domain’s name is its governed tag’s key, which is immutable. Renaming requires creating a new domain, re-tagging assets, and deleting the old one — so choose names you’ll keep.

Q: Do domains require restructuring my catalogs and schemas?

A: No. Domains are a logical overlay built on tags, so a single domain can span multiple catalogs and schemas without moving any data.

Documentation

Official Databricks docs for every feature mentioned in this post:

• Genie Ontology & Unity Catalog Semantics: Genie Ontology · Unity Catalog Semantics · Genie

• Domains and subdomains: Domains and subdomains · Domains REST API (Beta) · Terraform databricks_domain

• Pages: Pages

• Metric Views: Metric views

• Governed tags: Governed tags · Manage governed tags · Governed tag permissions · Tag propagation

• Discover: Discover data · Discover page

• Certification: Certify and deprecate data

• Access control: Unity Catalog privileges · ABAC · Row filters and column masks · ABAC GRANT policies

• MCP services, agent services, and skills: Govern MCP services · Register agents (agent services) · Govern skills

• Limits: Resource limits

Further reading

• Operationalizing Genie Ontology across your data stack — the Databricks blog on where domains fit among the other layers of Genie Ontology, and how to roll them out.

• Making UC Metric Views Discoverable for Agents — the Metric Views side of the semantic layer.

• Migrating Existing Dashboards to Databricks AI/BI, Part 3: User Filters and Row-Level Security with Unity Catalog — the access-control half that domains deliberately leave out.

• Why Your information_schema Queries Are Killing the Metastore — querying Unity Catalog metadata at scale.

• Stop Guessing AI Agent Integrations: The Databricks Decision Tree — patterns for connecting agents to governed data.

• Multi-Genie Slack Bot: Production-Grade Conversational Analytics in Slack — putting Genie in front of business users.

“Great, but can I build a managed service on top of it? And how do I do transformations before the data hits the table?”

This post answers both questions in depth. If you’re a platform engineer, principal architect, or anyone designing a multi-tenant ingestion layer on Zerobus Ingest — this is for you.

The Two Questions Engineers Ask

When I sit down with customers who are evaluating Zerobus Ingest for production, the conversation almost always lands on two things:

1. “Can we do transformations on the fly?” — They want to scrub PII, normalize schemas, filter noise, or enrich records before they land in Delta.

2. “Can we build a managed service on top of Zerobus?” — They’re building a platform for internal teams or external tenants, and they want Zerobus as the ingestion backbone.

The short answer to both: yes, but not the way you might expect. Zerobus Ingest is intentionally simple — it does one thing extremely well (durable, high-throughput 12+ Gb/Second writes to Delta), and it does not try to be a transformation engine. That’s a feature, not a limitation.

Why Zerobus Doesn’t Do Transformations (And Why That’s Good)

Let me be direct: Zerobus Ingest does not offer native transformations.

This is a deliberate design choice. Zerobus is built in Rust for raw speed — WAL-durable ACKs in ~150ms, data queryable in Delta within ~5 seconds, 100 MB/s per stream, 10+ GB/s per table. Adding a transformation layer inside the ingestion path would compromise that performance and add complexity to a service whose entire value proposition is simplicity.

But here’s the thing: you don’t need Zerobus to do transformations. You need a tiered transformation strategy — and if you’re coming from a Kafka architecture, you probably already have the pieces in place.

The Tiered Transformation Architecture

This is the pattern I recommend to customers building on Zerobus.

The Four Tiers

TIER 1: PRODUCERS Mobile apps, web apps, IoT devices, APIs. Raw events: unvalidated, may contain PII, mixed formats.

TIER 2: YOUR INGESTION API SERVICE FastAPI / Flask / custom Go or Rust micro-service.

This is where in-flight transforms happen:

• ✅ PII scrubbing / anonymization

• ✅ IP → Geo mapping (database lookups)

• ✅ Compliance flags (privacy service calls)

• ✅ Format conversion (XML→JSON, CSV→JSON)

• ✅ Tenant identification & tagging

Latency budget: milliseconds per record. Then calls: stream.ingest_record_offset({cleaned_record})

TIER 3: ZEROBUS INGEST Databricks-managed, serverless.

• Schema validation against Delta table

• Durable ACK in ~150ms

• Queryable in ~5 seconds

• Auto-compaction via Predictive Optimization

TIER 4: LAKEHOUSE (Downstream)

• Bronze → Silver: Deduplication, type casting

• Silver → Gold: Multi-stage enrichment, joins

• Complex aggregations, ML feature engineering

• SCD Type 2, slowly changing dimensions

• Tools: Lakeflow Declarative Pipelines / DLT / dbt

Why This Works

The key insight is that Tier 2 already exists in most architectures. If you’re currently using Kafka, you have an ingestion API service that receives events from producers and forwards them to Kafka. That same service can call the Zerobus SDK instead of the Kafka producer — and you can add your lightweight transforms there.

• In-Flight (Tier 2): Record-level transforms via API calls to rules-based services and database lookups (IP-to-Geo mapping, privacy flags, anonymization)

• Compliance & Privacy: Anonymization and pseudonymization happen before storage — critical for GDPR/CCPA

• Downstream (Tier 4): Deduplication from Bronze to Silver, multi-stage enrichment from Silver to Gold

The beauty of this pattern is that your transformation logic is yours — you own it, you version it, you test it. Zerobus stays fast and simple.

Reference Implementation: Zerobus Ingest Station

Databricks has published a reference implementation called Zerobus Ingest Station — a FastAPI-based container service that wraps the Zerobus SDK to create customizable ingestion endpoints. This is exactly the Tier 2 pattern:

• Define validation logic and transformations before pushing records to Zerobus

• Ideal for external partners, IoT devices, or REST integrations

• Deploy on Databricks Apps, Kubernetes, or any container runtime

Check it out: Zerobus Ingest Station on GitHub

Building a Multi-Tenant Managed Service

Now let’s talk about the second question: building a managed service. This is the pattern I’ve been recommending to healthcare, manufacturing, and SaaS platform customers.

The Architecture

Your managed service platform sits in front of Databricks:

Zerobus Supports Both Row-Level AND Bulk Ingestion

One misconception I keep hearing: “Zerobus is only for row-by-row ingestion.” Not true. Here’s the full picture:

Arrow Flight is the key for bulk scenarios:

• Sends entire RecordBatch (thousands of rows) per call

• Splits large batches into smaller transport messages automatically

• Best when your producer already works with columnar data (Polars, DataFusion, pyarrow)

• Supports ZSTD compression for improved throughput

If your managed service needs to handle both real-time events and periodic bulk loads, Zerobus has you covered with a single endpoint.

Regional Endpoints: There Is No Multi-Region Endpoint

This catches people off guard, so let me be clear: Zerobus endpoints are regional. There is no global or multi-region endpoint.

Your producers call the endpoint in the region where your workspace (and Delta tables) reside. Cross-region calls work but incur cloud provider egress charges. For best throughput, keep producer and endpoint in the same region.

What You Need to Call Zerobus

For your managed service’s Tier 2 layer, here are the parameters:

Critical permission gotcha: The service principal needs explicit grants — USE CATALOG, USE SCHEMA, SELECT, and MODIFY on the target table. ALL_PRIVILEGES does NOT work and will give you an invalid_authorization_details error.

The Honest Part: When NOT to Use This Pattern

Staying true to the spirit of my original post — here’s when this architecture doesn’t fit:

Customer Proof Points

This isn’t theoretical. Here are real customers stories:

Beusa Energy (Oil & Gas IoT): >99% cost reduction, ~3 seconds end-to-end latency, 22M+ rows/day from 6,000+ devices. “Zerobus Ingest saved us a lot of money, simplified our stack, and put us on a managed path forward.”

Joby Aviation (eVTOL Manufacturing): Streams gigabytes per minute from manufacturing sites. Reduced telemetry resolution latency from days to minutes. Custom on-premises forwarding agents push directly to Zerobus.

Toyota Motor Corporation (Factory IoT): Detects overheating factory conditions in minutes instead of hours. Uses Soracom Beam for global IoT connectivity → Zerobus → Delta.

Getting Started

If you want to build this pattern yourself:

1. Start with the docs: Zerobus Ingest Overview

2. Try the Zerobus Ingest Station: GitHub — a ready-made Tier 2 service

3. Read the deep dive: Deep Dive on Zerobus Ingest (Now GA)

4. See it at scale: Ingesting the Milky Way: Petabyte-Scale

5. Build an end-to-end app: Near Real-Time App with Zerobus + Lakebase

Conclusion

Building a managed service on Zerobus Ingest is not only possible — it’s the pattern Databricks designed for. Your service owns the API gateway, authentication, rate limiting, and lightweight transformations. Zerobus owns durable, high-throughput writes to governed Delta tables. The Lakehouse owns deep enrichment downstream.

The architecture is simple because Zerobus is simple. And sometimes the best architecture isn’t the most sophisticated — it’s the one that solves your problem.

Thanks for reading!

Jitesh Soni is a Specialist Solutions Architect at Databricks and the author of canadiandataguy.com. He writes about streaming, data engineering, and the things that break at 3 AM.

It overlaps the Plan, Execute, and Commit phases across multiple batches simultaneously, shrinking the gap between batches to near zero and delivering meaningfully fresher data. It’s on by default. It’s serverless only. And it’s the reason your continuous pipeline on serverless Lakeflow get more from the same hardware.

This post breaks down what it does, how it works, and why it matters for data freshness.

How Streaming Normally Works: One Batch at a Time

Every Spark Structured Streaming pipeline processes data in microbatches. Each microbatch goes through three phases:

1. Plan — Determine the next slice of data to process (read offsets from the source, discover new files, define batch boundaries)

2. Execute — Run the Spark SQL transformations on the executors

3. Commit — Write the results to Delta and persist the checkpoint

On classic compute, these phases run strictly sequentially. Batch 1 must complete all three phases — Plan, Execute, Commit — before Batch 2 even begins planning.

If each batch takes ~1.45 seconds end-to-end, processing 4 batches takes ~5.8 seconds. During the Plan and Commit phases, your executors are largely idle — waiting for the driver to finish its work before they get the next batch to process.

What Microbatch Pipelining Changes

On serverless Lakeflow pipelines, Microbatch Pipelining is enabled by default. Instead of waiting for Batch 1 to finish entirely, the engine determines the data boundaries for Batch 2 while Batch 1 is still executing. While Batch 2 executes, Batch 3 is already planning. Up to 4 microbatches can be in flight at the same time.

The key insight is that the phases of different batches overlap — the planning of the next batch runs concurrently with the execution of the current one, and commits are staggered so they don’t conflict.

How the Overlap Actually Works

Each streaming query gets its own thread pool (prefixed microbatch-runner) that manages up to 4 concurrent batches. The overlap isn’t random — there are synchronization points between batches to maintain correctness:

• Batch N+1 can begin planning as soon as Batch N’s offsets are committed to the write-ahead log (after the Plan phase), even while Batch N is still executing.

• Batch N+1 can begin executing once its own planning is complete, even if Batch N hasn’t committed yet.

• Commits are serialized — Batch N+1 waits for Batch N’s commit to complete before committing its own results. This preserves exactly-once semantics and checkpoint ordering.

This means the expensive Execute phase — where the actual data transformations happen on the executors — runs in parallel across multiple batches, while the lightweight Plan and the ordering-sensitive Commit phases are carefully coordinated.

What This Means for Data Freshness

The practical impact is significant. With sequential execution, your data freshness is bounded by the full batch cycle time — if each batch takes 10 seconds, new data waits up to 10 seconds before the next batch even starts processing it.

With pipelining, the next batch starts planning almost immediately after the current batch begins executing. Data that arrives during Batch 1’s execution is picked up by Batch 2’s planning phase — which is already running. The effective freshness improves because the gap between batches shrinks to near zero.

In a benchmark on serverless Lakeflow pipelines, the Silver query showed:

• Average batch duration: 2,242 ms (against a 1-second trigger interval)

• Batch overlap observed in 99.7% of batches

• Up to 3 microbatches in flight simultaneously

Nearly every batch was pipelined with its neighbors. The pipeline was continuously processing overlapping batches rather than waiting in a sequential queue.

Key Details

• Serverless only. Microbatch Pipelining is enabled by default on serverless Lakeflow pipelines. It is not available on classic compute and there are no plans to enable it there.

• Up to 4 concurrent batches. The default maximum is 4 in-flight microbatches per streaming query, controlled by spark.databricks.streaming.execution.pipelined.maxInflightBatches.

• Concurrency compounds. For a pipeline with N streaming tables, each table can have up to 4 concurrent batches. A pipeline with 30 streaming tables could have up to 120 concurrent batches across all queries — all sharing the same driver.

  You can turn it off. If your workload doesn’t benefit from pipelining (e.g., the pipeline is already keeping up with the source rate), you can disable it:

  spark.databricks.streaming.forceDisablePipelinedExecution = true

  This reverts to sequential execution.

When Pipelining Helps Most

Microbatch Pipelining delivers the biggest freshness improvement when:

• Batch duration exceeds the trigger interval. If your trigger interval is 1 second but each batch takes 3 seconds, sequential execution falls behind. Pipelining keeps up by overlapping batches.

• The pipeline has many flows. More streaming tables means more opportunities for parallel batch execution across queries.

• The workload is stateless. Stateless transformations (appends, filters, maps) pipeline cleanly. Stateful operations (aggregations, deduplication) have additional synchronization requirements.

It helps less when:

• The pipeline is already idle between batches. If each batch finishes well within the trigger interval, there’s nothing to overlap — pipelining has no work to pipeline.

• The driver is memory-constrained. Each in-flight batch holds execution state on the driver. With many streaming tables and 4 concurrent batches each, driver memory pressure can increase. For pipelines with 30+ flows, monitor driver memory if pipelining is enabled.

FAQ

Q: Is Continuous mode always better than Triggered mode?
A: No. Triggered mode is often preferable for periodic or bursty data when minutes or hours of latency are acceptable. Continuous mode is more appropriate for steady arrivals and low-latency requirements.

Q: What is stream pipelining?
A: Stream pipelining allows eligible queries in Serverless Spark Declarative Pipelines to overlap work from successive micro-batches rather than requiring strictly serial end-to-end completion.

Q: Is micro-batch pipelining the same as Spark task parallelism?
A: No. Task parallelism occurs within one micro-batch. Stream pipelining allows work associated with different micro-batches to overlap.

Q: Does pipelining eligibility guarantee overlapping batches?
A: No. A query may be eligible but show no overlap when its batch duration is shorter than its trigger interval, as occurred with the Gold query.

Q: How should I select a trigger interval?
A: Choose an interval that meets the workload’s freshness SLO while avoiding sustained backlog or instability. Do not reduce the interval solely to produce concurrent batches.

References

1. https://community.databricks.com/t5/technical-blog/triggered-vs-continuous-mode-a-deep-dive-into-serverless/ba-p/164327 Triggered vs. Continuous Mode: A Deep Dive into Serverless Lakeflow Spark Declarative Pipelines

2. Configure a serverless pipeline — this is the official doc that mentions Stream Pipelining by name https://learn.microsoft.com/en-us/azure/databricks/ldp/serverless/ “Stream pipelining: To improve utilization, throughput, and latency for streaming data workloads such as data ingestion, microbatches are pipelined... enabled by default in serverless pipelines.”

3. Streaming on serverless compute — confirms stream pipelining is on by default, has the use-case table https://learn.microsoft.com/en-us/azure/databricks/compute/serverless/streaming/ “Stream pipelining is enabled by default in serverless Lakeflow pipelines. Microbatches run concurrently rather than sequentially.”

4. Pipeline event log — for readers who want to measure freshness using flow_progress and stream_progress events https://learn.microsoft.com/en-us/azure/databricks/ldp/monitor-event-logs/

5. Best practices for Lakeflow pipelines — triggered vs continuous mode guidance https://learn.microsoft.com/en-us/azure/databricks/ldp/best-practices/index/

6. Configure Structured Streaming trigger intervals — explains the trigger model that MBP builds on https://learn.microsoft.com/en-us/azure/databricks/structured-streaming/triggers/

7. Monitor pipelines — overview of all monitoring options (UI, event log, system tables) https://learn.microsoft.com/en-us/azure/databricks/ldp/observability/

• Structured Streaming checkpoints traditionally identify sources by position, so changing the source list can require a fresh checkpoint and discarded state.

• Source evolution tracks offsets by stable .name() values instead.

• On-demand state repartitioning (Public Preview, RocksDB) can resize state partitions on restart.

• This Ethereum archive-to-live test preserved one checkpoint and query id through source addition, source removal, and state repartitioning without rereading archive rows or losing keys.

• Name every source when creating the checkpoint, and plan for tombstoned names and retained state.

Problem: We hope for permanent sources, and life give us transient ones

Think about what a streaming job is actually for. It maintains a metric: orders per region per minute, active sessions per tenant, fraud score per card, blocks per proposer per window. That metric has consumers who expect it to be continuous. It has a state store carrying every open window and every in-flight key. It is the thing the business depends on.

Now think about what feeds it. A Kafka cluster that will be migrated next quarter. A vendor feed on a contract that ends in March. A region launching in six months. An S3 archive that exists only until the live feed catches up. A tenant who churns. None of these are permanent, and everyone knows it on the day the job is written.

Yet the checkpoint treats the source list as permanent. Structured Streaming identifies sources by position—0, 1, 2—and the offset log is positional too. Add, remove, or reorder a source and the query refuses to start; for a stateful job, the only fix was a fresh checkpoint—discarding state and replaying history while consumers see a gap or a double count.

Teams worked around that in three ways: replay inside a maintenance window (duplicates and downtime), one job per source (the metric assembled downstream from N partial metrics), or a bus in the middle (an extra system /table whose only job is making the source list look permanent).

Each workaround compensates for the engine; source evolution removes the need—the job and its state stay put, and sources become what they always were: replaceable.

What changed: named sources and resizable state

With source evolution enabled, every streaming source carries a stable, user-defined name via .name(), and the checkpoint tracks offsets by that name instead of by position. Three operations become safe across restarts on the same checkpoint:

• Reorder sources. Each one resumes from its own last committed offset.

• Add a source. The new source starts from its beginning; existing sources continue where they stopped.

• Remove a source. Its name is tombstoned in the checkpoint and cannot be reused, which protects you from silently rewinding a feed.

There is a second thing a checkpoint used to freeze: the number of state partitions, set by spark.sql.shuffle.partitions on the day the checkpoint was born. A backfill wants a wide cluster and many partitions; a live feed of a few records per second wants a small node and a handful. On-demand state repartitioning (Public Preview, RocksDB) unfreezes that too. Together, the two features cover the whole lifecycle: change what feeds the query, then right-size what runs it, without ever starting over.

The proof: one checkpoint through a full ingestion lifecycle

We wanted a workload where losing state would hurt, on data anyone can reproduce: AWS publishes the entire Ethereum chain as date-partitioned Parquet in a public S3 bucket, and public JSON-RPC endpoints serve the live tip with no signup. The query counts blocks per 5-minute window, per source and proposer (miner), with a 2-day watermark, update mode, a foreachBatch MERGE into Delta, and RocksDB state—run through five steps on one checkpoint, as tasks of one Databricks job, scored from the run's own StreamingQueryProgress

Step 1: name the sources

For a newly created compatible query, the following two configs plus a stable, unique .name() on every source enable supported source-list evolution. This cannot be enabled retroactively on an existing checkpoint; removed names are permanently tombstoned; and support requires Databricks Runtime 18.2 or later.

spark.conf.set("spark.sql.streaming.queryEvolution.enableSourceEvolution", "true")
spark.conf.set("spark.sql.streaming.offsetLog.formatVersion", "2")

archive = (spark.readStream.format("parquet").schema(BLOCK_SCHEMA)
           .option("maxFilesPerTrigger", 365)
           .name("s3_backfill")
           .load(landing))

Naming requires a fresh checkpoint, and it is one-way. You cannot enable source evolution on a checkpoint created without it, and you cannot turn it off later. The practical consequence: name every source on day one, including the ones you are certain will never change.

Step 2: backfill the archive on a cluster that turns itself off

Backfill is the expensive part: We used 8 workers, Spark’s default 200 shuffle partitions, and trigger(availableNow=True) to read everything and stop.

(events.writeStream
       .option("checkpointLocation", checkpoint)
       .trigger(availableNow=True)  # drain, then stop paying
       .foreachBatch(upsert_into_delta)
       .start())

The query processed 26,086,569 Ethereum blocks from 4,080 files in 13 micro-batches, leaving 3,154 keys in state. The checkpoint now retained 200 state partitions.. That number is now baked into the checkpoint. We come back to it in Step 5.

Step 3: add the live feed, and watch the archive replay nothing

This is the cutover that used to cost a fresh checkpoint. The live feed is a Python data source over a public Ethereum RPC endpoint, configured to start one block after the archive’s last. It runs on a single small node, on the same checkpoint, with the archive source still present.

The important thing is what does not change. Both sources project to the same seven columns, so the union, the watermark, the aggregation and the sink are exactly the code that ran in Step 2. The only edit is one more entry in the list of sources:

archive = (spark.readStream.format("parquet").schema(BLOCK_SCHEMA)
           .option("maxFilesPerTrigger", 365)
           .name("s3_backfill")                         # unchanged
           .load(landing))

live = (spark.readStream.format("eth_rpc")             # Python DataSource
        .option("start_block", archive_max_block + 1)  # one past the archive
        .name("eth_live")                              # NEW named source
        .load())

# Same projection for every source, so the union is by name.
events = (archive.select(*BLOCK_COLUMNS).withColumn("source", lit("s3_backfill"))
          .unionByName(
          live.select(*BLOCK_COLUMNS).withColumn("source", lit("eth_live"))))

# The metric. Unchanged from Step 2.
counts = (events
          .withWatermark("block_ts", "2 days")
          .groupBy(window("block_ts", "5 minutes"), "source", "miner")
          .count())

(counts.writeStream
       .outputMode("update")
       .option("checkpointLocation", checkpoint)        # SAME checkpoint as Step 2
       .foreachBatch(upsert_into_delta)                 # MERGE on (window, source, miner)
       .start())

Spark reads the offset log by name, finds s3_backfill at its last committed offset, finds no entry for eth_live, and starts the new source from its beginning. The RocksDB state for every open window carries straight over, because the stateful operator did not change.

To see what happened on restart, we flattened the query’s own StreamingQueryProgress events (query.recentProgress, or a StreamingQueryListener) into one row per source per microbatch. Each event carries a sources[] array with a SourceProgress per input: its description, numInputRows, startOffset and endOffset.

Three things to read off this table:

• The batch ids start at 13, not 0. Step 2 committed batches 0 through 12. A fresh checkpoint would have restarted the count; this one continued it.

• The FileStreamSource rows (srcIdx 0) read 0 input rows in every batch, with startOffset and endOffset both held at {"logOffset": 11}. The archive resumed exactly where it stopped and re-read nothing.

• The new source (srcIdx 1) starts from null in batch 13, then walks forward 20 blocks per batch from 26,086,590. It had no committed offset, so it began at its configured start block.

The new source started from its beginning; the old one resumed exactly where it stopped, and no history was re-read. Nothing about the aggregation state was touched.

Why is the name column empty when every source has a .name()?

.name() is a checkpoint identity, not a runtime label: on Databricks Runtime 18.3, SourceProgress describes sources by class and description. Names live in the checkpoint’s offsets/ directory, in V2 entries keyed by source name (Step 4 reads one); verify by srcIdx, offsets, and row counts.

Step 4: Retire the source or keep it running, depends on your use-case

Once the live feed has been running for a while, the archive has nothing left to contribute. Retiring it is not a migration; it is a restart with a shorter list. Stop the query, edit the same pipeline code and delete s3_backfill code, and start it against the same checkpoint. Nothing else is edited:

query.stop()

live = (spark.readStream.format("eth_rpc")
        .option("start_block", archive_max_block + 1)  # ignored: offset comes from the checkpoint
        .name("eth_live")                              # unchanged
        .load())

# No union any more. Same projection, same metric, same sink.
events = live.select(*BLOCK_COLUMNS).withColumn("source", lit("eth_live"))

counts = (events
          .withWatermark("block_ts", "2 days")
          .groupBy(window("block_ts", "5 minutes"), "source", "miner")
          .count())

(counts.writeStream
       .outputMode("update")
       .option("checkpointLocation", checkpoint)        # SAME checkpoint
       .foreachBatch(upsert_into_delta)
       .start())

On restart Spark compares the named sources in the plan with the named sources in the newest offset-log entry. eth_live is in both, so it resumes from its committed offset (the start_block option is only honoured the first time a name is seen). s3_backfill is in the log but not in the plan, so Spark records its departure. The offset log keeps the name, marked with a tombstone, and the live source carries on:

Two behaviours are worth noticing. First, removing a source does not remove its state. At removal, RocksDB still held 3,060 archive keys, exactly one per archive row in Delta whose window the watermark had not yet closed. Those windows will close on their own schedule, as they should. Second, the name is gone for good. Re-attach s3_backfill and the query refuses to start:

[STREAMING_QUERY_EVOLUTION_ERROR.TOMBSTONE_SOURCE_NAME_REUSE] Cannot reuse tombstoned source names: s3_backfill. These source names were previously used and then removed from the streaming query at checkpoint location ... Reusing tombstoned source names can lead to data correctness issues. Please use different source names.

You want that error. A reusable name would let a source silently rewind to an offset from two restarts ago and double-count every window it re-read.

Step 5: shrink the state to fit the live workload

With the archive gone, the query is processing 20 blocks per batch on one node, but its state is still spread across the backfill’s 200 partitions. Changing spark.sql.shuffle.partitions in the session does nothing; the checkpoint remembers 200.

# Resizes state on the next restart. shuffle.partitions cannot.
spark.conf.set("spark.sql.streaming.stateStore.partitions", "8")

The restart’s first progress event reports 66.8 seconds of controlBatch.REPARTITION, with no keys lost: 3,155 before, 3,157 after, 0 removed. From then on the same 20-block batches took a median 15.0 seconds of addBatch instead of 35.5, and state memory fell from 1,323 MB to 1,067 MB. Same checkpoint, same query id, a third of the partitions’ overhead gone.

Proof: two clusters, four phases, one query id

Spark derives a streaming query’s id from checkpoint metadata, so an unchanged id across restarts is direct evidence that all four phases were the same logical query resuming, not four queries that happened to write to the same table. Both clusters’ Spark UIs show it:

Nothing fell through the seams. Delta holds 26,087,209 blocks: the archive’s 26,086,569 plus 640 live blocks starting at 26,086,570, exactly one past the archive’s last.

What this means for how you design streams

The demo is a blockchain, but the principal is universal: a long-lived metric fed by sources which will come and go. Once the job no longer has to be rebuilt when its inputs change, several patterns become routine restarts instead of migration projects:

And four rules to carry into every new stream:

1. Assume every source is temporary, and name it on day one. Naming needs a fresh checkpoint, so the cheapest moment is before the first batch, and the sources you are most certain about are the ones you will be most surprised by.

2. Treat names as permanent identifiers. A removed name is tombstoned; a rename is a remove plus an add that starts from the beginning.

3. Removing a source keeps its state until the watermark closes it. Plan for that in memory sizing.

4. Remember to alter spark.sql.streaming.stateStore.partitions: when you’re backfilling set this number high so you can backfill fast, and then remember to bring down that number when you go in live state.

Common misconceptions about streaming source lifecycle management

The lifecycle test exposes several common misconceptions; confusing identity, offsets, and retained state causes replay, double counting, or failed restarts.

• Misconception: a source list is fixed for the life of its checkpoint. Correction: with source evolution from a fresh checkpoint and stable names, sources can be added, removed, or reordered across restarts without discarding state.

• Misconception: adding a source causes existing sources to replay. Correction: existing named sources resume from their committed offsets; only the newly introduced name starts from its configured beginning.

• Misconception: removing a source immediately deletes its state. Correction: removal tombstones the name; existing state remains until normal watermark eviction closes it.

• Misconception: a removed source name can later be reused. Correction: tombstoned names cannot be reused, because doing so could rewind offsets and double-count data.

• Misconception: source names appear as labels in StreamingQueryProgress. Correction: on the tested runtime, .name() is a checkpoint identity; verify by source index, offsets, row counts, or V2 offset-log entries.

• Misconception: changing spark.sql.shuffle.partitions resizes an existing state store. Correction: the checkpoint keeps its count; use spark.sql.streaming.stateStore.partitions with on-demand repartitioning.

Scope of this run

• Runtime and compute. Databricks Runtime 18.3 on classic compute. Source evolution requires Databricks Runtime 18.2 or above; on-demand state repartitioning is in Public Preview and requires the RocksDB state store.

• Structured Streaming jobs. This run used plain Structured Streaming with foreachBatch. It did not exercise Lakeflow Spark Declarative Pipelines or real-time mode.

• Not a benchmark. One workload, one run. The partition-count timings are illustrative of the mechanism, not a performance claim. Things can be done in milli-seconds end to end as well.

Get started

Two configs and a .name() on each source is all it takes to decouple the job from its inputs:

spark.conf.set("spark.sql.streaming.queryEvolution.enableSourceEvolution", "true")
spark.conf.set("spark.sql.streaming.offsetLog.formatVersion", "2")

orders_us = spark.readStream.name("orders_us").table("catalog.schema.orders_us")
orders_eu = spark.readStream.name("orders_eu").table("catalog.schema.orders_eu")
all_orders = orders_us.union(orders_eu)

• Change streaming sources with source evolution

• On-demand state repartitioning for stateful streaming queries

• Configure Structured Streaming trigger intervals

• AWS Public Blockchain Data

FAQ

Q: What is source evolution in Structured Streaming?
A: Source evolution lets a new checkpoint track each source by a stable .name() rather than query-plan position. With offset-log format V2 and an otherwise compatible plan, named sources can be added, removed, or reordered across restarts.

Q: Do I need a new checkpoint to enable source evolution?
A: Yes. Source evolution cannot be enabled retroactively, and it cannot be disabled after the checkpoint is created. It requires Databricks Runtime 18.2 or later.

Q: Does adding a source replay existing sources?
A: No. Existing named sources resume from their committed offsets. Only the new source name starts from its configured beginning.

Q: What happens to state when a source is removed?
A: The source name is tombstoned, but its existing aggregation state remains until normal watermark eviction closes the associated windows.

Q: Can I change the number of state partitions without losing state?
A: With supported on-demand state repartitioning and RocksDB, set spark.sql.streaming.stateStore.partitions and restart. Changing spark.sql.shuffle.partitions does not resize state already recorded in the checkpoint.

Q: Can a tombstoned source name be reused?
A: No. Spark blocks reuse because an old committed offset could rewind the feed and double-count data. A returning feed requires a different name and therefore starts from that new name’s configured beginning.

• Use asynchronous checkpointing to UC Volumes with UCVolumeWriter and checkpoint every 30 to 60 minutes to balance resilience against I/O overhead.

• Connect your local IDE (VS Code, Cursor) or terminal directly to AI Runtime GPUs via SSH tunnel for interactive development and debugging.

• Use Genie Code for AI Runtime work. It natively understands your cluster, environment, and AI Runtime APIs, and more than doubled the success rate of leading coding agents on real data science tasks.

• Productionize training with Declarative Automation Bundles (DABs) for version-controlled, CI/CD-friendly GPU workload deployments.

• Run Ray Core, Ray Data, Ray Train, and Ray Tune natively on AI Runtime serverless GPU compute.

• Monitor and attribute GPU spend using system.billing.usage with Serverless Usage Policies for team-level cost tracking.

• Migrate deep learning workloads from classic DBR ML to AI Runtime for on-demand GPUs, zero idle cost, and a simpler infrastructure experience.

This is a continuation of our blog that covers practical lessons we’ve learned from working with different customers shipping real GPU workloads on AI Runtime. Think of it as the “things we wish someone told us on day one” guide.

In Part 1, we covered hardware options, environments, entry points (notebooks, CLI, IDE), MLflow observability, and data loading.

In Part 2, we cover checkpointing, remote development workflows, Genie Code for AI Runtime, productionizing training runs, Ray on AI Runtime, migrating from classic DBR ML, and cost governance.

For the full Part 1, see Best Practices for GPU Workloads on Databricks AI Runtime — Part 1

6. Checkpointing: Make AI workloads resilient and resumable

Model Checkpointing

Use asynchronous checkpointing to UC volumes with UCVolumeWriter and PyTorch DCP to avoid blocking training. These backends stage data on fast local NVMe (/tmp) before uploading, which is faster than writing directly to FUSE and ensures metadata is only committed after all shards are uploaded.

Checkpoint frequently enough to limit lost work, but not so often that I/O slows training ,typically every 30–60 minutes. For async saves, initialize the process group with a CPU backend (e.g., cpu:gloo, cuda:nccl).

import torch.distributed.checkpoint as dcp

from serverless_gpu.data import UCVolumeWriter

checkpoint_path = “/Volumes/my_catalog/my_schema/model/checkpoints”

writer = UCVolumeWriter(checkpoint_path)

future = dcp.async_save(state_dict, storage_writer=writer)

# ...continue training...

future.result()  # blocks until the upload lands on the UC volume

# Load a checkpoint with UCVolumeReader

from serverless_gpu.data import UCVolumeReader

reader = UCVolumeReader(checkpoint_path)

dcp.load(state_dict, storage_reader=reader)

Data pipeline checkpointing

A model checkpoint captures model and optimizer state, but not the position of your data pipeline within the dataset, so a resumed run cannot fast-forward to the exact sample where it stopped. Account for this in how you resume: restart from an epoch boundary, or track processed samples or shards in your own training state so you can skip them on resumption.

7. Remote Development for AI Runtime: A local-to-Databricks workflow

You can connect your local IDE or terminal directly to AI Runtime serverless GPU compute using an SSH tunnel. This lets you interactively run and debug code from VS Code, Cursor, or the CLI while keeping a consistent environment between your local setup and the Databricks workspace.

Tools like Cursor and VS Code (with the Remote SSH extension) work out of the box. Other coding agents like Claude Code can also be used once the SSH tunnel connection is active.

Setup is straightforward. Establish the connection via a single command in the Databricks CLI:

databricks ssh connect --accelerator <GPU_type>

where <GPU_type> can be set to GPU_1xA10 or GPU_8xH100

To launch directly into your IDE, add the --ide vscode or --ide cursor parameter. You can also launch the tunnel through the latest release of the IDE extension.

Note that the SSH tunnel gives you an interactive session on a single node. For long-running training jobs or multi-node distributed training, use the AIR CLI instead (covered in Part 1).

Additional capabilities to simplify using your preferred IDE as your primary development environment:

• Handling project dependencies: Use the --base-environment option to start the SSH connection with pre-configured Python libraries from a workspace base environment.

• Tracking costs and resource utilization: Apply a serverless usage policy using the --usage-policy-id parameter to attribute SSH connection costs across teams or projects.

• Exploring data assets via Unity Catalog: Inspect catalogs, schemas, and resources directly within the IDE without leaving your development workflow.

• File persistence: Files under /Workspace, /Volumes, and /dbfs persist across restarts. Anything in /home or /root is ephemeral.

8. Genie Code: Build AI Runtime notebooks and jobs faster

Genie Code is the AI coding and data assistant built into the Databricks workspace. It has a deep, native understanding of the Databricks platform including your data assets, lineage, governance policies, compute environment, and runtime context. It’s equipped with the latest Databricks-specific skills and APIs that general-purpose coding agents don’t have access to, and its agentic harness is purpose-built for the full spectrum of Databricks workloads, from data engineering and analytics to ML training and production operations.

The result: on real-world data science tasks, Genie Code more than doubled the success rate of a leading coding agent, even when that agent was connected to Databricks via MCP.

When to Use Genie Code vs. Other Assistants
If the work touches Databricks, start with Genie Code. For everything else, use whatever you’re already productive with, whether that’s Claude Code, Cursor, or Copilot. They’re great for general software engineering, building apps, and working across repos. These tools aren’t competitors, they’re complementary.

If Genie Code can’t be used, connect your preferred coding agent to Databricks using MCP and the open-source Databricks AI Tools. It won’t match Genie Code’s native depth, especially on AI Runtime, but it brings meaningful Databricks awareness to tools like Claude Code, Cursor, or Copilot.

Leveraging Genie Code for AI Runtime Development
Genie Code can help you develop and troubleshoot deep learning workloads in notebooks connected to AI Runtime. It can generate distributed training code, resolve environment and dependency issues, and investigate GPU workload failures.

Steps to use Genie Code with AI Runtime:

1. Connect the notebook to AI Runtime and open the Genie Code pane (Genie lamp on top right)

2. Describe the model training workload that you want to develop or the problem that you want to resolve.

Here are some example Genie code prompts for the following use cases to get you started

• Develop distributed training workloads:

  • “Adapt this training workload for AI Runtime and explain the important changes.”

• Resolve environment and dependency issues:

  • “Diagnose this package compatibility error and recommend the smallest appropriate change.”

• Debug GPU and distributed workloads:

  • “Analyze this failed distributed training run using the available output from every rank and identify the root cause.”

For more details refer to the Use Genie Code with AI Runtime documentation.

9. Productionize: take your training to production

Once your training code is working in a notebook, the next step is making it reproducible, scheduled, and deployable across environments. Databricks supports this through Declarative Automation Bundles (DABs), which let you define your AI Runtime training workload as code. In Part 1, we covered the AIR CLI for submitting training jobs. DABs are the next step when you need version-controlled, CI/CD-friendly deployments.

The ai_runtime_task is a job task type purpose-built for this. You point it at your training code, specify the GPU type and count, and declare your dependencies, all in a YAML bundle definition that lives in source control. From there, you can deploy to dev or production with a single CLI command, schedule it on a cadence, and compose it with other tasks like data preparation steps that run on CPU.

Key things to know:

• Your training code can be packaged as a tgz artifact and shipped with the bundle, or referenced directly from a workspace or volume path.

• Multi-node runs are supported. Set accelerator_count higher than the per-node GPU count and AI Runtime handles the distributed environment variables (WORLD_SIZE, MASTER_ADDR, etc.) for you.

• GPU and CPU tasks can be combined in multi-task jobs using depends_on, so you only use GPU time for actual training. Note that ai_runtime_task does not support job task values for passing data between steps, so use a shared location like a UC volume instead.

• Dev to production promotion works the same way as any other Databricks bundle. Schedules are paused by default, so deploying won’t accidentally trigger runs.

For the full configuration reference, task fields, and a complete working example, see the Productionize training workloads documentation.

10. Ray on AI Runtime: Practical scaling with Ray Core, Ray Data, Ray Train and Ray Tune

Ray is a popular framework for scaling Python workloads. Databricks AI Runtime supports Ray natively, so you can create Ray clusters and run Ray applications on serverless GPU compute without managing infrastructure.

AI Runtime supports Ray’s main libraries, including Ray Core, Ray Data, Ray Train, and Ray Tune. The Databricks AI environment includes Ray out of the box. With the Standard environment, add ray or the relevant extra (ray[data], ray[train], ray[tune]) to your workload dependencies.

Use Ray in notebooks

For single node, use ray_init() from the serverless_gpu package to start Ray on the notebook attached to AI Runtime compute:

from serverless_gpu import ray_init
ray_init()

Ray Dashboard: ray_init() automatically configures the Ray dashboard and prints the dashboard URL in the notebook output. Use the dashboard to inspect Ray jobs, tasks, and resource usage while your code runs. This is in addition to the MLflow System Metrics and Model Metrics monitoring you get for AI Runtime runs.

Use Ray with the AI Runtime CLI

The AI Runtime CLI supports Ray in single-node and multi-node configurations. Include a bootstrap script under the workload’s command, which starts the Ray cluster and coordinates head and worker nodes when the workload launches.

For Ray examples with CLI, refer to the AI Runtime Ray documentation.

11. Migrate GPU workloads from Classic DBR ML to Databricks AI Runtime

Databricks AI Runtime is the recommended compute for GPU workloads.

If you’re still on classic DBR ML, here are the key steps to migrate deep learning workloads to AIR:

• Replace cluster-dependent code. Remove any references to Spark-based distributed training (for example, TorchDistributor) and replace them with the @distributed decorator from serverless_gpu. Or you can go for the air cli approach.

• Reinstall dependencies. Do not rely on Databricks Runtime ML pre-installed libraries. Add explicit %pip install commands for all required packages with exact versions.

• Update data loading.

  • Replace direct DBFS paths with Unity Catalog volumes paths.

  • Replace local Spark DataFrame operations with Spark Connect.

  • For streaming file-based data from volumes, use UCVolumeDataset from serverless_gpu.data.

• Update checkpoint paths.

  • Move checkpoints from DBFS or local storage to Unity Catalog volumes.

  • For distributed checkpointing, use UCVolumeWriter and UCVolumeReader from serverless_gpu.data, which stage I/O through local NVMe.

• Update MLflow configuration. Ensure experiment names use absolute paths and configure run names so they can be easily restarted.

• Test interactively first. Validate your workload in an interactive notebook before scheduling it as a job.

Several of these steps, such as environments, data loading, checkpointing, and MLflow configuration, are covered in detail earlier in the blog. Refer to those sections for deeper guidance.

Recommendation: Use Genie Code to accelerate the migration. It understands both the classic DBR ML patterns you’re moving from and the AI Runtime APIs you’re moving to, so it can help rewrite your training code, resolve dependency conflicts, and test the result in your live environment.

12. Cost and Usage: Measure, govern, and optimize GPU compute

AI Runtime will appear on your bill as DBUs on the Model Training SKU. The DBUs will be based on the number of GPU hours and the GPU type selected.

You can view on-demand pricing for the available GPU instance types in your cloud provider and region on the Databricks AI Runtime pricing page.

You can monitor your AI Runtime GPU spend by querying the billable usage system table (system.billing.usage).

The following query returns total usage for serverless GPU workloads:

SELECT SUM(usage_quantity)

FROM system.billing.usage

WHERE product_features.serverless_gpu IS NOT NULL

Spend by team and model (via custom tags):

SELECT custom_tags[’team’]  AS team,

       custom_tags[’model’] AS model,

       SUM(usage_quantity)  AS gpu_hours

FROM   system.billing.usage

WHERE  product_features.serverless_gpu IS NOT NULL

  AND  usage_date >= current_date() - INTERVAL 7 DAYS

GROUP  BY team, model

ORDER  BY gpu_hours DESC;

Tips

• Create Serverless Usage Policies with tags for team, project, or model so GPU spend is attributable from day one. These tags propagate automatically to system.billing.usage in the custom_tags column.

• For notebooks, select the usage policy from the compute dropdown. For AIR CLI jobs, set usage_policy_name in your workload YAML. For DABs jobs, set it in the job details.

• Build a Lakeview dashboard on top of these queries for weekly cost review. Databricks also offers a pre-built cost monitoring dashboard that account admins can import.

That covers it. Across Part 1 and Part 2, we’ve covered the full lifecycle of GPU workloads on AI Runtime. From choosing the right entry point, hardware, and environment, to data loading, distributed training, and observability in Part 1, through to checkpointing, developer experience, productionizing, and cost governance in Part 2. Everything here comes from design patterns we’ve seen work across real teams shipping real workloads.

We hope it gives you a practical understanding of how to make the best use of AI Runtime for your use cases. If you’re just getting started, pick the section most relevant to you and build from there. And as AI Runtime continues to evolve, we’ll keep updating these guides with what we learn.

Thanks for reading. Feel free to share any feedback on topics you’d like to see next in the comments section.

Databricks Genie is great at answering questions about your data. But you can’t listen to the details on your walk to standup, while driving to the office, or during your commute. So let’s build a Databricks App that asks Genie your questions and reads the answers back to you — as a short, spoken podcast you can play on your phone- anytime, anywhere!

Features

• 🗣️ Narrated answers — Each Genie summary is spoken aloud by an in-app voice model

• 🧠 No second LLM — Genie itself writes the spoken summary

• 📊 Details on tap — Exact figures, generated SQL, Genie’s reasoning, and result rows stay on screen

• 🔐 Runs as you — Every Genie query runs on-behalf-of the signed-in user, with their permissions

• 💾 Durable — Episodes live in Lakebase, audio in a UC Volume

• 📱 Mobile-first — GenieOne App enables your Databricks App to be available on mobile devices

Architecture

Prequisites

• Databricks workspace with Databricks Apps enabled

• Genie space you can query

• Lakebase Instance

• Unity Catalog Volume

• Databricks App

• Databricks CLI (v1.10+) with an auth profile

Setup Instructions

Download the application code and make the necessary updates to the config values in app.yaml. Download all files from repo:

https://github.com/adgitdemo/ad_databricks/tree/main/my-genie-podcaster

1. Provision storage. Create a Lakebase database (note its branch + database identifiers). Create a UC Volume in an existing catalog/schema (note /Volumes/<catalog>/<schema>/<volume>).

2. Configure the bundle. In databricks.yml, set the workspace host, the postgres resource (branch, database, CAN_CONNECT_AND_CREATE), two uc_securable VOLUME resources (READ_VOLUME + WRITE_VOLUME), and user_api_scopes: [dashboards.genie]. In app.yaml, set LAKEBASE_ENDPOINT (valueFrom: postgres) and DATABRICKS_VOLUME_AUDIO.

3. Stage the model. npm install && npm run setup:model to download the Kokoro weights into .hf-cache/…. Then databricks fs mkdir dbfs:/Volumes/<…>/_models/onnx-community/Kokoro-82M-v1.0-ONNX/onnx and databricks fs cp the four files (config.json, tokenizer.json, tokenizer_config.json, onnx/model_quantized.onnx) to the matching Volume paths. (On first boot the app copies them to local disk — no Hugging Face egress.)

4. Build, deploy, run. npm run build, then databricks bundle deploy --profile <p>, then databricks bundle run app --profile <p>. Open the printed app URL.

Note: users need access to their Genie space and its data, or Genie returns a permission error.

Usage

• Tap New.

  • Enter a Genie space (by name or ID) and one or more questions.

  • Hit Create podcast. The app asks Genie, captures each SUMMARY and DETAILS, and pre-synthesizes the narration in the background.

• Open the episode and press Play. Switch narrators from the voice picker. Expand “Detailed analysis” on any answer to see the figures, SQL, Genie’s reasoning, and a preview of the rows.

Demo

Open the “My Genie Podcaster” app from the One mobile App and use it!

This blog explores the art of the possible with Databricks Genie!

Nested JSON never arrives flat. Open telemetry metrics are the textbook case: one payload carries an array of resources, each resource carries scopes, each scope carries metrics, each metric carries data points. On Spark, you land the whole thing in a single VARIANT column and unroll it later. That unroll is what most raw-to-refined jobs actually spend their time doing.

The obvious way is one explode_outer per level. It works. It is also the part of the job that quietly costs you the most. The rest of this is the mechanical rewrite from N nested explode_outer calls into a single explode_outer of a flatten/transform tree, plus the one subtlety that makes the naive version silently wrong. General Spark technique for nested arrays in VARIANT or JSON.

The working example below uses a public OpenTelemetry payload.

Why chaining explode_outer calls hurts cost and correctness

Cost. Each explode_outer compiles to a Generate operator. Four nested arrays mean four Generate hops, with parent columns — including the raw payload — carried between them.

Correctness. explode drops rows for empty arrays; explode_outer preserves the parent and emits a null child for null or empty input. Any rewrite must reproduce that behavior exactly.

Mapping the nested array structure of an OTLP payload

payload (VARIANT column `raw`)
└── $.resourceMetrics[]                          rm
    └── $.scopeMetrics[]                         sm
        └── $.metrics[]                          metric
            └── $.sum|gauge|summary|histogram|
                exponentialHistogram .dataPoints[]   dp   <- leaf

Chain form:  raw --Generate--> rm --Generate--> sm --Generate--> metric --Generate--> dp
               (4 Generate hops, fat parent row carried through each)
Single form: raw --transform/flatten (no Generate)--> array<struct<rm,sm,metric,dp>>
               --Generate--> leaf
               (1 Generate hop, narrow leaf struct)

Before: chaining four explode_outer calls

Assume a table otlp_metrics_raw(payload_id STRING, raw VARIANT) with one payload per row:

from pyspark.sql import functions as F
df = (
    spark.table("otlp_metrics_raw")
    .withColumn("rm", F.explode_outer(
        F.variant_get(F.col("raw"), "$.resourceMetrics", "array<variant>")))
    .withColumn("sm", F.explode_outer(
        F.variant_get(F.col("rm"), "$.scopeMetrics", "array<variant>")))
    .withColumn("metric", F.explode_outer(
        F.variant_get(F.col("sm"), "$.metrics", "array<variant>")))
    .withColumn("dp", F.explode_outer(F.coalesce(
        F.variant_get(F.col("metric"), "$.sum.dataPoints", "array<variant>"),
        F.variant_get(F.col("metric"), "$.gauge.dataPoints", "array<variant>"),
        F.variant_get(F.col("metric"), "$.summary.dataPoints", "array<variant>"),
        F.variant_get(F.col("metric"), "$.histogram.dataPoints", "array<variant>"),
        F.variant_get(F.col("metric"), "$.exponentialHistogram.dataPoints", "array<variant>"),
    )))
)

The rewrite rule: build the leaf struct from the inside out

1. At each parent level, transform the child array into an array of leaf structs, then flatten the result back to one level.

2. The leaf transform builds the struct — struct(rm AS rm, sm AS sm, m AS metric, dp AS dp) — carrying every ancestor handle the projection reads after the explode (need $.schemaUrl off the resource? Then rm must be in the struct).

3. The last step — and only the last step — is one explode_outer of that array.

4. Restore outer semantics on every array — the step people skip, and the whole point of this post.

The hidden risk: flatten and transform lack outer semantics

Swapping a chain of explode_outer calls for one explode_outer over a flatten(transform(...)) tree trades an operator with outer behavior for two without it. explode_outer on a null or empty array keeps the parent row and emits one NULL-child row; flatten and transform make no such promise:

• Empty inner array → the parent disappears. transform of an empty array yields an empty array, flatten drops it, and the parent never reaches the explode — explode_outer would have emitted one row.

• NULL inner array → the whole row collapses. flatten returns NULL if any element is NULL [4]. One resource missing scopeMetrics nulls the entire leaf array and takes every sibling leaf with it. The damage isn’t local to the bad element.

def null_to_singleton(array_expr: str) -> str:
    """explode_outer parity: NULL/empty child array -> one NULL placeholder."""
    return (
        f"CASE WHEN ({array_expr}) IS NULL OR size({array_expr}) = 0 "
        f"THEN array(CAST(NULL AS VARIANT)) ELSE ({array_expr}) END"
    )

Wrap every array in the chain — leaf included — and the rewrite is row-for-row identical to the chain it replaces.

After: one explode_outer call replaces the chain

METRIC_KINDS = ("sum", "gauge", "summary", "histogram", "exponentialHistogram")

def null_to_singleton(array_expr: str) -> str:
    return (
        f"CASE WHEN ({array_expr}) IS NULL OR size({array_expr}) = 0 "
        f"THEN array(CAST(NULL AS VARIANT)) ELSE ({array_expr}) END"
    )

def dp_array(metric: str) -> str:
    kinds = ",\n            ".join(
        f"variant_get({metric}, '$.{k}.dataPoints', 'array<variant>')" for k in METRIC_KINDS
    )
    return f"coalesce(\n            {kinds}\n          )"

LEAF_ARRAY_EXPR = """explode_outer(
    flatten(transform({rm_arr}, rm ->
      flatten(transform({sm_arr}, sm ->
        flatten(transform({m_arr}, m ->
          transform({dp_arr}, dp -> struct(rm AS rm, sm AS sm, m AS metric, dp AS dp))
        ))
      ))
    ))
  )""".format(
    rm_arr=null_to_singleton("variant_get(raw, '$.resourceMetrics', 'array<variant>')"),
    sm_arr=null_to_singleton("variant_get(rm, '$.scopeMetrics', 'array<variant>')"),
    m_arr=null_to_singleton("variant_get(sm, '$.metrics', 'array<variant>')"),
    dp_arr=null_to_singleton(dp_array("m")),
)

FLATTEN_SQL = """
SELECT
  payload_id,
  rm,
  sm,
  metric,
  dp,
  variant_get(metric, '$.name', 'string')  AS metric_name,
  variant_get(metric, '$.unit', 'string')  AS metric_unit,
  CAST(variant_get(dp, '$.startTimeUnixNano', 'long') / 1e9 AS TIMESTAMP)
                                           AS start_timestamp,
  CAST(variant_get(dp, '$.timeUnixNano', 'long') / 1e9 AS TIMESTAMP)
                                           AS event_timestamp,
  variant_get(dp, '$.asInt', 'long')       AS value_as_int,
  variant_get(dp, '$.asDouble', 'double')  AS value_as_double
FROM (
  SELECT
    payload_id,
    leaf.rm     AS rm,
    leaf.sm     AS sm,
    leaf.metric AS metric,
    leaf.dp     AS dp
  FROM (
    SELECT payload_id, {leaf_array_expr} AS leaf
    FROM otlp_metrics_raw
  )
)
""".format(leaf_array_expr=LEAF_ARRAY_EXPR)

flat = spark.sql(FLATTEN_SQL)

rm, sm, metric, and dp remain VARIANT handles in the projection; resource attributes, scope name, bucket counts, exemplars — anything you need later is a variant_get away, without touching the unroll.

General checklist for flattening any nested JSON array

1. List the arrays from root to leaf.

2. Define a leaf struct containing the leaf and every required ancestor.

3. Replace each null or empty array with a typed singleton-NULL array.

4. Build nested transform and flatten expressions from the inside out.

5. Call explode_outer exactly once, at the top.

6. Verify Generate counts and compare rows against the original chain.

Rule: N nested explode_outer calls become one explode_outer over a guarded flatten(transform(...)) tree that pre-builds the leaf structs.

FAQ

Q: What if my input data is not balanced, it is never balanced ?

A: Explicitly run a repartition command so the input gets balanced.

Q: What if my input data is not balanced, it is never balanced ?

A: Set maxFilesPerTrigger or maxBytesPerTrigger, or similar options to control the size of the MicroBatch

Q: How to check spark behavior before executing ?

A: Run explain() on both forms and count the Generate nodes. Four becomes one — the entire structural claim, verifiable in one command:

spark.sql(FLATTEN_SQL).explain()   # count `Generate` nodes: expect exactly 1
r.

Q: Does this technique require the VARIANT type?

A: No. transform, flatten, and explode_outer are standard Spark SQL functions. For parsed structs, replace variant_get with field access and cast each NULL placeholder to the array’s element type.

Q: Does the same rule apply to OTLP logs?

A: Yes. OTLP logs use three array levels: resourceLogs[] → scopeLogs[] → logRecords[] [3]. Remove one flatten(transform(...)) level and the corresponding leaf-struct field..

Q: Is one explode_outer always faster than a chain?

A: Almost always. It removes Generate operators but may widen the row before the remaining Generate and adds generated-SQL complexity. Benchmark both implementations on representative data.

Q: What happens if one array is left unguarded?

A: You can lose rows silently. An empty array removes its parent leaf, while a NULL inner array can null the entire flattened array and remove otherwise-valid sibling rows.

References

[1] opentelemetry-proto/examples/metrics.json — Apache-2.0. Raw: https://raw.githubusercontent.com/open-telemetry/opentelemetry-proto/main/examples/metrics.json

[2] opentelemetry/proto/metrics/v1/metrics.proto

[3] opentelemetry-proto/examples/logs.json

Spark SQL references

[4] Spark SQL, Built-in Functions — flatten — “Returns NULL if any element of the input array is NULL.”

[5] Spark SQL, Built-in Functions — transform — higher-order function; returns an empty array for empty input.

[6] Spark SQL, Built-in Functions — explode_outer — “Returns null for entries that are null or empty.”

Every company depends on a handful of KPIs: revenue, active users, churn, margin, and each one usually gets defined in a dozen places at once: a SQL query here, a spreadsheet formula there, a dashboard filter somewhere else, each subtly different, until two teams walk into a meeting with two revenue figures and no one can say which is right. A semantic layer fixes that by providing a single governing definition that every tool reads, and Unity Catalog Metric Views are Databricks’ version of it.

That consistency has always mattered, but AI raises the stakes. When an agent answers a business question at machine speed, a wrong or ambiguous definition no longer just misleads one analyst; it gets baked into automated decisions and customer-facing answers. The semantic layer is where trust in your data is enforced, and your agents are only as good as the definitions behind them. This post is about the newest consumer of that layer, the AI agent, and a lightweight pattern that lets it reach the semantic layer using only existing Databricks services.

Agents can’t find the right metric definition

A person finds the right metric view without thinking about it. An agent starts from nothing. When a custom agent receives a question like “what was our average order value this quarter?”, it first has to find the right metric view from among hundreds across dozens of domains, and then understand its measures, dimensions, and query syntax.

Today, there’s no first-class, public way to search across them. An agent can already reach the definitions: it can enumerate catalog objects and read a metric view’s definition through the Unity Catalog REST API or SQL, and run one with the MEASURE() function. The trouble is choosing. Enumerating and parsing raw definitions on every call doesn’t scale past a small catalog. It re-scans each time, burns tokens, and surfaces overlapping or conflicting definitions with no way to rank them.

The semantic layer is a valuable, governed asset, but there’s no discovery surface that lets an agent ask, “Given this question, which metric view should I use, and what does it expose?” the way it would for any other object it can reach. This post builds that surface.

Expose the semantic layer as MCP

The idea is a small, customer-owned Databricks App that acts as an MCP server with a single job: answer “which metric views match this ask, and how do I query them?”

At startup, the app builds an in-memory index of trimmed metric-view metadata: each view’s name, description, measures, and dimensions. It exposes one MCP tool, search_metric_views. An agent calls it with a natural-language ask and gets back the matching metric views, their measures and dimensions, and an example query so the agent knows how to use them without having to learn the metric-view syntax on its own.

Two design choices matter:

• It’s discovery, not execution. The app returns which metric views fit and an example of how to query them.

• You control what’s exposed. The app can index only certified metric views, so external agents see governed, production-ready definitions and nothing that’s still in development or QA.

Architecture

Multiple agents, inside or outside Databricks, connect to one Databricks App, which acts as the MCP server and looks up the semantic layer in Unity Catalog. The pattern improves discoverability: the app returns matching metric definitions and an example query. The agent can take necessary actions to generate SQL or manage the semantic layer to avoid duplication, etc., based on the use case.

Where does the metadata live?

The index can live in a few places, and each earns its place in a different situation:

• In-memory, inside the app: the simplest option, and the one this post uses. No extra infrastructure, and reads are fast because they never leave the process. The catch is that the index rebuilds on every restart; each running copy of the app maintains its own copy; definitions can go stale between refreshes; and search is at the keyword level. Enough for a bounded catalog and a single team.

• Metadata table via a SQL endpoint: fits a bounded catalog that changes rarely. Every app instance reads one shared, governed copy, refreshed by a scheduled job.

• Lakebase: fits when many app instances and agents hit the index at once, and you want durable, low-latency lookups with structured filters such as domain, owner, or certified status, with room to keep usage or audit state alongside it.

• Managed AI Search (vector) index: fits when search quality is the point, whether a large catalog or plain-language questions, where ranking by meaning beats keyword matching, at the cost of more infrastructure and an index to keep in sync.

The right choice comes down to the size of the catalog, the number of requests per minute you expect, and the amount of latency you can tolerate.

Governance

The app runs under a service principal, and you govern what it can reach with Unity Catalog grants (for example, restricting it to certified metric views). This fits multi-agent and cross-environment scenarios where the calling agents don’t carry a Databricks user identity.

If you need user-level access control, on-behalf-of (OBO) user authorization is supported. The agent acts as the calling user, and row and column governance is preserved end-to-end. It adds identity-passing complexity and isn’t required for the service-principal pattern this post uses.

Prerequisites

Before you build the app, get a few things in place:

1. Databricks workspace

• Active Databricks workspace

• One or more Unity Catalog metric views

2. An MCP-capable agent or client: anything that speaks the Model Context Protocol can connect and call the search tool.

Setup Instructions

Step 1: Get the code

Download the five files from the repo folder:

https://github.com/adgitdemo/ad_databricks/tree/main/mcp-metric-view-search

You’ll have app.py, app.yaml, config.yaml, requirements.txt, and README.md.

Step 2: Point it at your warehouse and your data

Two files need real values before you deploy.

In app.yaml, set your SQL Warehouse ID. You’ll find it under SQL Warehouses > your warehouse > Connection details.

env:

- name: DATABRICKS_WAREHOUSE_ID

value: “your_warehouse_id”

In config.yaml, specify which catalogs and schemas the app should scan for metric views. Scope it down to keep the index focused, or leave it empty to scan everything the app can see.

scope:

- catalog: your_catalog

schemas:

- your_schema

Step 3: Deploy the app

Upload the folder to your workspace and deploy it. Deploying is what creates the app’s service principal, and that service principal is the identity you’ll grant data access to in the next step. So deploy first.

databricks workspace import-dir ./mcp-metric-view-search \

/Workspace/Users/<email id>/mcp-metric-view-search --overwrite

databricks apps create metric-view-search-app

databricks apps deploy metric-view-search-app \

--source-code-path /Workspace/Users/<email id>/mcp-metric-view-search

You can do the same thing from the Apps UI if you prefer.

Step 4: Find the app’s service principal

databricks apps get metric-view-search-app

Look at the service_principal_name field. It’ll be something like xxxx-mcp-metric-view-search. The app runs as this identity, so this is what needs to read your data, not your own user account.

Step 5: Grant the service principal access to your data

The app reads metric views through the warehouse, so the service principal needs three Unity Catalog grants plus warehouse access. Run each GRANT on its own. Databricks won’t accept several of them in a single statement.

GRANT USE CATALOG ON CATALOG your_catalog TO `xxxx-mcp-metric-view-search`;

GRANT USE SCHEMA ON SCHEMA your_catalog.your_schema TO `xxxx-mcp-metric-view-search`;

GRANT SELECT ON SCHEMA your_catalog.your_schema TO `xxxx-mcp-metric-view-search`;

Then give it the warehouse:

GRANT CAN USE ON WAREHOUSE `your_warehouse_id` TO `xxxx-mcp-metric-view-search`;

Repeat the schema grants for every schema you listed in config.yaml.

Step 6: Confirm it’s actually indexing

Redeploy so the app re-scans with its new permissions, then check any one of these:

- Open the app URL in a browser. The landing page shows how many metric views it has indexed.

- Hit the /health endpoint. It returns {”status”: “ok”, “indexed_metric_views”: N}.

- Read the app logs. You want to see the Indexed N metric view(s). If you instead see Cannot list catalogs or the scope still shows your_catalog_name, the config or the grants didn’t land.

Step 7: Let your teammates and their agents in

Everything above enables the app to read your data. To let other people connect their agents to it, grant them permission on the app itself. There are two ways to do this.

From the UI: open the app’s page, click Permissions in the top right, and add users, groups, or service principals. CAN USE is enough to connect and call the tools. CAN MANAGE also lets them redeploy and change settings.

Without the UI: use the permissions API. A PATCH adds people without disturbing anyone who already has access (a PUT would replace the entire list, so use PATCH here).

databricks api patch /api/2.0/permissions/apps/metric-view-search-app \

--json ‘{

“access_control_list”: [

{”user_name”: “teammate@yourcompany.com”, “permission_level”: “CAN_USE”}

]

}’

You can grant a group or another service principal the same way, and add several entries at once:

databricks api patch /api/2.0/permissions/apps/metric-view-search-app \

--json ‘{

“access_control_list”: [

{”user_name”: “teammate@yourcompany.com”, “permission_level”: “CAN_MANAGE”},

{”group_name”: “data-team”, “permission_level”: “CAN_USE”},

{”service_principal_name”: “<client-id>”, “permission_level”: “CAN_USE”}

]

}’

To confirm who has access, read the permissions back:

databricks api get /api/2.0/permissions/apps/metric-view-search-app

This is a different permission from Step 5. Step 5 controls what data the app can read. Step 7 controls who can use the app.

Step 8: Connect an agent

Register the app’s /mcp endpoint as an MCP server in your agent, give the connection a name like metric-view-search, and authenticate with a Databricks token.

For a connection that won’t expire after an hour, use a non-expiring Personal Access Token instead of a short-lived OAuth token. Then:

In the Claude Code CLI, that’s one command:

claude mcp add metric-view-search https://<your-app-url>/mcp \

--header “Authorization: Bearer <databricks-pat>”

In a coded agent using the Claude Agent SDK, you configure the same three things (name, URL, and the auth header) in the mcpServers option, and add the tools to allowedTools so the agent can call them. See steps here for wiring the app using the Claude Agent SDK:

options = ClaudeAgentOptions(

mcp_servers={

“metric-view-search”: {

“type”: “http”,

“url”: “https://<your-app-url>/mcp”,

“headers”: {”Authorization”: “Bearer <databricks-pat>”},

}

},

allowed_tools=[”mcp__metric-view-search__*”],

)

Either way, metric-view-search is just the connection name, and <databricks-pat> is the token. Once connected, the agent can search your metric views, read their definitions, and generate valid SQL. Other agents follow the same pattern.

▎ Note: a PAT is a standing personal credential that acts on your behalf until revoked. Fine for a demo; for shared or production use, prefer a service principal.

Other agents follow the same shape; they just differ in where you put the URL and the header. These steps may change based on your agent/tools.

Demo

With a few metric views registered, an agent connected to the app can explore the semantic layer directly:

• Ask “what metric views do we have around revenue?” and the app returns the matching views with their measures and dimensions. Where definitions overlap (e.g., a fee-only revenue view and a blended-revenue view), both are surfaced with their descriptions so the agent or user can tell them apart.

• For each match, the app returns an example query showing how to use the metric view.

When do I need it?

A few situations where searching the semantic layer directly is useful:

• Use a semantic layer from your own agents. You’ve standardized KPIs as metric views and want custom agents to use them without having to stand up additional consumer surfaces.

• Partners and BI tools built on Databricks that want to explore available metrics programmatically to build more intelligent experiences.

• Governing a large semantic layer. Data stewards find overlapping or duplicate metric views by keyword and owner across thousands of definitions, then audit and consolidate them.

• Standardizing across a multicloud, multi-agent estate. Metric views are being open-sourced, so the same pattern can front a semantic layer wherever it runs.

• Choose the right entry point: use notebooks with the @distributed decorator for rapid prototyping, and switch to the AIR CLI with YAML-based configs for long-running, multi-node production training.

• Pick the right accelerator for the job: attach a low-cost A10 or single H100 for orchestration, and reserve multi-GPU or multi-node H100 configurations for actual training workloads.

• Use UCVolumeDataset and serverless_gpu.data.DataLoader for efficient data loading with local caching, automatic distributed partitioning, and optimized prefetching keep GPUs fed without Spark overhead.

• Enable MLflow autologging and monitor the System Metrics tab for GPU utilization and memory; set a custom run name and MLFLOW_RUN_ID to resume previous training runs.

This post covers the practical lessons we’ve learned from working with teams in Databricks and with customers shipping real GPU training workloads. Think of it as the “things we wish someone had told us on day one” guide.

If you’re training or fine-tuning deep learning models on Databricks, AI Runtime is how you get access to GPUs without the usual infrastructure headaches. But like any powerful tool, knowing how to use it well makes all the difference.

Update: You can find Part 2 of this blog here.

What is AI Runtime? How Can It Help?

AI Runtime is a serverless model training product at Databricks built specifically for deep learning workloads. It brings GPU support to Databricks Serverless, so you can train and fine-tune models using your favorite frameworks (PyTorch, DeepSpeed, vLLM, Ray, etc.) without managing any infrastructure.

Before AI Runtime, running GPU workloads meant configuring clusters, picking the right drivers, managing autoscaling, and often negotiating GPU capacity directly with your cloud provider (AWS, Azure, or GCP). AI Runtime takes all of that off your plate.

Key benefits:

• Fully managed GPU infrastructure: You request the GPUs you need and start working. Pay only for what you use. On-demand by nature.

• Seamless scaling. Go from a single GPU experiment to multi-node distributed training without reconfiguring anything.

• A runtime dedicated for deep learning: Choose either a minimal base environment for full control over your dependencies, or a full-featured AI environment pre-loaded with popular ML frameworks. Either way, CUDA, NCCL, and drivers are already set up. For multi-node H100s, RDMA networking is configured out of the box.

• Natively integrated across notebooks, jobs, Unity Catalog, and MLflow. Your development workflow, data access, experiment tracking, and governance all work together.

Here are the 5 AI Runtime best practices we will cover in this Part 1 blog. Feel free to jump directly to the section that’s interesting for you:

• What are my hardware accelerator options?

• What are environments in AI Runtime?

• What are the 3 ways of using AI Runtime- notebooks, CLI and IDE?

• How can I use MLflow for tracing and observability?

• What are the different options for data loading for training?

1. Hardware Accelerator Options

The following accelerators are currently supported. All accelerators provision a single node.

• 1xA10- 24 GB-Small to medium ML and deep learning tasks

• 1xH100- 80 GB- Workloads that need more GPU memory than 1xA10 provides but do not require multi-GPU distributed training.

• 8xH100- 80 GB per GPU- Large-scale AI workloads that need distributed training or multiple GPUs

For multi-GPU workloads, choose a single-node 8xH100 instance (one node with 8 H100 GPUs). Partial multi-GPU configurations like 2xH100 or 4xH100 are not supported though.

For multi-node workloads, scale out using multiple 8xH100 nodes; for example, 4 nodes with 8 H100 GPUs each give you a total of 32 H100 GPUs. Or you can use multiple nodes of 1xA10.

Besides the options listed above, be sure to consult the latest documentation, as additional accelerator types may become available over time.

2. Environments

AI Runtime offers two managed Python environments, the default base environment and the Databricks AI environment.

Default base environment

• Minimal, includes only required packages: torch, cuda, and torchvision

• You want full control over your dependency stack and prefer to install only what you need

Databricks AI environment

• Pre-loaded with popular ML frameworks (PyTorch, Transformers, and more)

• You want a complete environment for training, fine-tuning, and experimentation without manual dependency management

You can also use a workspace base environment that a workspace admin has built for serverless GPU compute.

Environment Caching: Environments are cached across sessions to speed up startup times. When you reconnect to AI Runtime with the same environment configuration, previously installed packages may be available from cache, reducing setup time. However, this behavior is not guaranteed, so always ensure your notebook includes the necessary %pip install commands for reproducibility..

Import custom modules: Import custom modules by placing them under /Workspace/shared and adding the path to sys.path. Or You can also upload module files as Workspace files and import them directly.

import sys

sys.path.append(”/Workspace/Shared/my-project/src”)

from my_module import my_function

Some differences as compared to Classic Compute:

• Spark: AI Runtime is a Python-only environment; Spark is not available as a local runtime. So you cannot use PySpark functions but you can use Spark Connect for data loading.

• MLR: Pre-installed AI Runtime packages do not fully replace Databricks Runtime ML, because some ML libraries available in Databricks Runtime ML are not included.

3. Main Ways to Get Started: Notebooks, CLI and IDE

You have two primary ways for running workloads on AI Runtime: Notebooks and AIR CLI. Each has its strengths, and picking the right one for your use case makes the journey easier.

3a. Notebooks

Notebooks on AI Runtime let you develop interactively, iterate quickly, and visualize results in place. Simply choose your preferred accelerator (A10, 1×H100, or 8×H100), attach it to your notebook, and your code will execute directly on that GPU-backed node.

Multi-GPU (Single Node)

To run jobs across multiple GPUs on a single node, use the @distributed decorator. Wrap your training workflow inside a function and apply the decorator to it. This function becomes the entry point for distributed execution, so all the components - model initialization, data loading, and training logic - should be defined within it.
The @distributed API integrates with major distributed training libraries: PyTorch, DDP, FSDP and DeepSpeed

@distributed(gpus=8, gpu_type=’H100’)

def run_train(num_epochs, batch_size):

    # training code here


run_train.distributed(num_epochs=3, batch_size=1)

Ray

For a single node (A10 or 8xH100), you can just use ray_init from serverless_gpu

from serverless_gpu import ray_init; 

ray_init()

For multi-node, use air CLI.

Best practices:

• Pin your library versions: AI Runtime has regular releases. If you need reproducibility, pin the versions that work for your code.

• For distributed training, any environment variable needed by your training code should be set within @distributed. Note: Some environment variables must be set before importing a particular library so that it initializes with the correct configuration and hardware settings. This is specific to the library, not AI runtime.

• Keep experimentation in notebooks, production in jobs: Notebooks are great for prototyping and debugging, but once your code is stable, you can move it to a scheduled job. For scheduling notebook jobs, use the Jobs API or Declarative Automation Bundles (fka Databricks Asset Bundles)

• For operations that do not require GPUs (for example, converting data formats, or exploratory data analysis), attach your notebook to a CPU cluster to preserve GPU resources.

• Note: Connection to your interactive compute auto-terminates after 60 minutes of inactivity.

3b. AI Runtime CLI

air CLI is a command-line tool for submitting and managing training jobs on Databricks Serverless GPU compute directly from your terminal, instead of using notebooks.The CLI uses YAML-based job configuration, and supports workspace-based and git-based code workflows. It’s aimed for users who prefer repos, IDEs, git, and CI/CD-style workflows over notebooks.

air CLI follows a simple submit → schedule → execute → observe loop:

1. Define your workload in a train.yaml file (compute shape, environment, launch command).

2. Submit via air run -f workload.yaml

3. AIR automatically provisions the requested GPU nodes, installs dependencies from requirements.yaml, and injects coordination environment variables

4. Your training code runs (e.g., via torchrun) across the provisioned nodes.

5. Observe progress through air get logs, air get status, or the MLflow UI for GPU metrics.

Here is a simple train.yaml

experiment_name: simple-training

environment:

  dependencies: requirements.yaml

env_variables:

  NCCL_DEBUG: “INFO”

secrets:

  HF_TOKEN: ‘my_scope/hf_token’

compute:

  num_accelerators: 8

  accelerator_type: GPU_8xH100

code_source:

  type: snapshot

    snapshot:

      root_path: /home/username/repo

      git:

        branch: main

command: torchrun --nproc_per_node=8 train.py

max_retries: 2

timeout_minutes: 90

usage_policy_id: abcd123-25b8-3e87-9a2c-f86eb19d101c

Here is the doc for AI CLI for more information.

Best practices for AIR CLI:

• Use CLI for multi-node, long-duration jobs: If your training run takes hours or days and spans multiple nodes, AIR CLI is the right tool. It handles the orchestration cleanly.

• Treat the workload.yaml as the source of truth. Keep the launch definition in version-controlled YAML rather than relying on notebook/UI state, so runs are easier to reproduce, hand off, and compare later.

• Check workload.yaml and requirements.yaml into your repo along with train.py for version tracking and reproducibility.

Custom Docker Images: You can run your own Docker images with AI Runtime using a simple CLI and YAML-based configuration. This is useful for workloads that need specific libraries, complex dependencies, fully reproducible environments, or standard images defined by your organization.

When to use AIR CLI vs. Notebooks:

Use notebooks (with the @distributed decorator) for rapid iteration and prototyping.

Use the AI Runtime CLI when you need to:

• Submit GPU training jobs directly from your terminal.

• Define jobs in YAML for version control and reproducibility.

• Run long (multi-day) or distributed (multi-node) training that shouldn’t depend on an interactive session.

3c. Connecting from the IDE terminal like Visual Studio or Cursor

You can connect to AI Runtime on serverless GPU compute directly from a terminal in your IDE through an SSH tunnel.

To connect from a terminal within your IDE, run the databricks ssh connect command with the --accelerator option.

databricks ssh connect --accelerator=GPU_1xA10

To connect and start the session in Visual Studio Code or Cursor, use the --ide option. The CLI opens an IDE window pointing to the home workspace folder.

databricks ssh connect --ide=vscode

4. Use MLflow for Tracking and Observability

On AI Runtime, MLflow is the main place to track and inspect training runs. AIR is natively integrated with MLflow for experiment tracking, model logging, metric visualization, and observability.

It’s recommended to use the latest version of mlflow (minimum 3.7 version)

For each run of your experiment:

• The System Metrics tab provides an overview of GPU, CPU, and disk performance, including key metrics such as utilization and memory usage. These are generated automatically for each run. Use these metrics to proactively identify bottlenecks, tune workloads (for example by adjusting batch size or data loading), and ensure resources are right-sized so you maintain high utilization without over-provisioning

• The Model Metrics tab surfaces key performance indicators logged during training, such as accuracy, loss, and other task-specific scores. Use these metrics to validate and compare experiments, ensuring that changes to model architecture, training configuration (for example, batch size, learning rate, and data preprocessing),etc actually improve model quality over time/

• Logs for processes running on each GPU are stored as artifacts and can be accessed from the MLflow UI. For notebook-based runs, the same logs are also visible directly in the notebook cell outputs. Use these logs to monitor training progress, debug training issues, and trace performance regressions across runs so you can quickly identify and fix problems in your training pipeline.

System Metrics

Logs are also available as artifacts in the MLflow run

Best practices:

• Autologging: MLflow supports automatic logging for popular frameworks like PyTorch, TensorFlow etc. You can enable it for your training job by adding just one of code, mlflow.autolog(). It automatically captures:

  • Metrics: standard training and evaluation measures (such as accuracy and F1 score);

  • Parameters: hyper-parameters, such as learning rate and number of estimators

  • Artifacts: output files, such as your trained model weights.

• Manual Logging: you can manually log any parameter, metric or artifact in your training code.

• The Serverless GPU API automatically launches an MLflow experiment with default name /Users/{WORKSPACE_USER}/{get_notebook_name()}. To overwrite it, set

import os

os.environ[”MLFLOW_EXPERIMENT_NAME”] = “/Users/<username>/my-experiment”

• It’s recommended to customize the MLflow run name. It enables you to restart from a previous run.You can customize the run name using the run_name parameter in mlflow.start_run(run_name=”your-custom-name) or in third-party libraries that support MLflow (for example, Hugging Face Transformers). Otherwise, the default run name is jobTaskRun-xxxxx.

• Resume previous training by setting the MLFLOW_RUN_ID from the earlier run: mlflow.start_run(run_id=”<previous-run-id>”)

• Set the `step` parameter in MLFlowLogger to reasonable batch numbers. MLflow has a limit of 10 million metric steps, so logging every single batch on large training runs can hit this limit

Monitor GPU resources:

If you are using Notebook, you can use GPU resources in the right side pane to monitor GPU health and utilization while your code runs on AI Runtime. The pane supports both single-node and multi-node workloads. The pane displays the following metrics for each GPU: GPU utilization percentage, GPU memory usage, Temperature.

GPU metrics

If you are using CLI, you can utilize System Metrics to view GPU metrics.

5. Data Loading

Tabular Data:

• Small Dataset: For small Delta tables, use Spark Connect to load the data. Then convert the Spark DataFrame to a pandas DataFrame with toPandas(). If needed, you can further convert it to a NumPy array using to_numpy(). This works well for Single Node

# Load Delta table

df = spark.read.format(”delta”).table(”your_catalog.your_schema.your_table”)

# Convert Spark DataFrame to pandas

pdf = df.toPandas()

# Optional: convert pandas DataFrame to NumPy array

arr = pdf.to_numpy()

• Large dataset: For large Delta tables that are too big to convert with toPandas(), export the data to a Unity Catalog volume and load it directly using PyTorch or Hugging Face. This approach avoids Spark overhead during training and works well for both single-GPU and distributed training workflows.

# Step 1: Export the Delta table to Parquet files in a UC volume

output_path = “/Volumes/catalog/schema/my_volume/training_data”

spark.table(”catalog.schema.my_table”).write.mode(”overwrite”).parquet(output_path)

# Step 2: Load the exported data directly using Hugging Face datasets

from datasets import load_dataset

dataset = load_dataset(”parquet”, data_files=”/Volumes/catalog/schema/my_volume/training_data/*.parquet”)

Unstructured Data:

For unstructured data such as images, audio, and text files stored in Unity Catalog volumes, our recommendation is to use UCVolumeDataset from the serverless_gpu.data package.

UCVolumeDataset is a PyTorch IterableDataset that copies each file from the volume to a fast local cache on first access and yields the cached local file path. Besides local caching it supports automatic partitioning of files and divided across DataLoader workers, so each (rank, worker) pair receives a non-overlapping slice.

For optimal performance, pair UCVolumeDataset with serverless_gpu.data.DataLoader, rather than the stock PyTorch DataLoader, as it is tuned for serverless GPU I/O and fetches and caches files concurrently while the GPU computes.

from serverless_gpu.data import UCVolumeDataset

from torch.utils.data import IterableDataset

from PIL import Image

import torchvision.transforms.functional as TF

class ImageDataset(IterableDataset):

    “”“Decodes each cached file path from UCVolumeDataset into a tensor.”“”

    def __init__(self, path_dataset: UCVolumeDataset):

        self._path_dataset = path_dataset

    def __iter__(self):

        for local_path in self._path_dataset:

            image = Image.open(local_path).convert(”RGB”)

            yield TF.to_tensor(image)

path_dataset = UCVolumeDataset(”/Volumes/catalog/schema/my_volume/images”)

dataset = ImageDataset(path_dataset)

For more info about data loading, here is the doc.

Best Practices:

• Load data inside the @distributed decorator: For distributed training, move data loading code inside the @distributed decorator. The dataset size can exceed the maximum size allowed by pickle, so it is recommended to generate the dataset inside the decorator.

from serverless_gpu import distributed

# This may cause a pickle error if the dataset is too large

dataset = get_dataset(file_path)

@distributed(gpus=8, gpu_type=’H100’)

def run_train():

   # Load data inside the decorator to avoid pickle serialization issues

   dataset = get_dataset(file_path)

   ...

When you build a UCVolumeDataset in the decorator, it uses torch.distributed rank info at runtime to automatically split files across ranks, so you don’t need a DistributedSampler for file-based volume data.

• Use UCVolumeDataset for multi-epoch training so files are cached locally on first access and reused, instead of copying the entire volume upfront with shutil.copytree when you may only read part of it.

• Cache data locally for multi-epoch training. For large datasets, copy datasets to /tmp for faster access across epochs:

import shutil
shutil.copytree(”/Volumes/catalog/schema/volume/dataset”, “/tmp/dataset”)

• Parallelize data fetching: Use serverless_gpu.data.DataLoader to parallelize fetching. It’s a torch DataLoader subclass optimized for serverless GPUs, with higher defaults for worker count and prefetching so files are loaded and cached while the GPU runs. It also records per-batch load times in MLflow to help identify data bottlenecks.

from serverless_gpu.data import DataLoader

loader = DataLoader(

    dataset,

    batch_size=32,

    pin_memory=True,

    # num_workers=6, by default

    # prefetch_factor=4, by default

    # raise num_workers to increase parallel reads, or prefetch_factor to deepen each worker’s queue.

)

• Adjust batch size: For large datasets, use a larger batch size so the cost of loading data per batch is spread over more samples, reducing file fetches each step. If GPU memory is tight, pair a larger batch with gradient accumulation to keep the same effective batch size.

• Use Streaming datasets: For very large datasets that do not fit in memory, use streaming approaches.

  • UCVolumeDataset from serverless_gpu.data for streaming files from Unity Catalog volumes with local caching and automatic distributed partitioning.

  • PyTorch IterableDataset for custom streaming logic.

  • Hugging Face datasets with streaming for datasets hosted on the Hub or in volumes.

  • RayData for distributed batch data processing

FAQs

Q: My training run is slow and GPU utilization is low. What should I check first?

A: Start by inspecting the System Metrics tab in MLflow, focusing on GPU utilization, GPU memory usage, and data loading throughput. If GPU utilization is low while the GPU is not memory-bound, your data pipeline is likely the bottleneck. We recommend increasing your batch size or switching to serverless_gpu.data.DataLoader, which is specifically tuned for serverless I/O. Additionally, ensure you are using UCVolumeDataset to cache files locally instead of refetching them from Unity Catalog during every step.

Q: Should I use notebooks or the AIR CLI for my training workload?

A: It depends on your workflow phase. Use notebooks with the @distributed decorator for rapid prototyping and interactive debugging. Switch to the AIR CLI when you need production-grade features like version-controlled YAML configurations, multi-day training runs, or a CI/CD-style approach directly from your IDE.

Q: How do I resume a previous MLflow training run?
A: Set a custom run name using mlflow.start_run(run_name=”your-custom-name”) so you can identify the run later. To resume, call mlflow.start_run(run_id=”<previous-run-id>”) with the earlier run’s ID. This lets you continue logging metrics and artifacts into the same run record.

Q: Which environment should I choose — the base environment or the AI environment?
A: Start with the AI environment (recommended) as it has popular ML frameworks (PyTorch, DeepSpeed, vLLM, Ray, etc.) pre-installed. Use the base environment if you need full control over your dependencies and want a minimal setup. You can also use a base environment defined by your workspace admin, built on organization-standard images. Always pin library versions for reproducibility.

Thanks for reading. Feel free to share any feedback on topics you’d like to see next in the comments section.

Shout out to the AI Runtime PM, Tejas Sundaresan for helping us review the content in this blog.

We plan to work on Part 2 for this AI Runtime blog series next. In this post, we will share some more best practices around remote development options, cost and usage monitoring, and using genie code with AI Runtime. Stay tuned!

Leave a comment

Here are five lessons we learned over a week of benchmarking on Spark 4

Rule 1: Upgrade the Silicon (M8/M9, Not M5/M6)

Before tuning shuffle partitions or threading, pick the right generation.

On the same workload — same Databricks Runtime, same node count, same pipeline — we moved from M6gd to M8gd and saw a rare three-way win:

• 2× throughput (more records per second)

• ~38% lower cost per record processed

• Lower batch latency end to end

Newer Graviton generations (M8, and M9 when available) bring better CPU, memory bandwidth, and local NVMe behavior. That matters most when every micro-batch does shuffle, spill, and multi-table writes.

Practice: Don’t lock onto M5/M6 families because they’re familiar. Benchmark newest generation up on the same job before you optimize code. Often the biggest lever is hardware, not another spark.conf tweak.

Rule 2: Multi-Thread Your foreachBatch Sinks

foreachBatch runs once per micro-batch. If you write to more than one table (main, audit, dead-letter), sequential writes leave the cluster idle between commits.

Practice: Do transforms once, split the batch, submit writes in parallel from the driver.

from concurrent.futures import ThreadPoolExecutor

def process_batch(batch_df, batch_id):
    with ThreadPoolExecutor(max_workers=2) as executor:
        futures = [
            executor.submit(write_to_table_a, batch_df, batch_id),
            executor.submit(write_to_table_b, batch_df, batch_id),
        ]
        for f in futures:
            f.result()

Parallel writes mean more Spark jobs per batch — you’re trading wall-clock time for driver coordination. Profile before raising max_workers.

For a fuller production version of this — caching the batch once, fanning out to N tables instead of 2, propagating thread exceptions properly — see the reference notebook. Two companion videos walk through it: foreachBatch fundamentals first, then parallelizing the writes.

Rule 3: Spill to Local Disk, Don’t Recache What’s Already There

Every task writes its output to the disk on the machine it runs on. You can let that spill hit the network over EBS, or you can pick an instance family with disk physically attached to the node. The attached disk wins on performance, every time.

Default is simple: don't cache. When the fork happens after a shuffle — the batch already passed through a join, aggregation, or stateful dedup — the shuffle write already left a durable copy on local disk that Spark reuses automatically for any second write path, retry or independent action alike (look for a "skipped" stage in the Spark UI), so caching here only duplicates what's already there.

The judgment call lives in the other case: a fork before any shuffle, where a shared read and a few light transforms split straight into two writes with no durable file to fall back on. Skip the cache and both paths redo that shared work from scratch, source read included; cache it and you pay a serialize-then-deserialize round trip through the block manager instead. Which one wins depends on the weight of the shared transforms — a wide UDF or an expensive parse earns the cache, a couple of column selects don’t, and recomputing them twice is cheaper than the round trip.

Here’s the number that makes this rule worth enforcing: memory-optimized (R-series) instances carry at least a ~20% price premium over general-purpose compute — on every major cloud, not just one. Cut the unnecessary caching and you cut the memory pressure that pushes you onto R-series in the first place. That 20% doesn’t vanish — you either bank it as savings or spend it on more CPU. Cache when profiling proves it pays off, not by habit.

Rule 4: Optimize for Finishing on Time, Not Just Finishing

In a latency-bound streaming job, “did the job succeed” is the wrong question. The right one is “did it finish inside the SLA.” A batch that completes in 90 seconds against a 60-second budget hasn’t crashed — it’s succeeded and missed the point.

That reframes how you pick a join strategy. Sort-merge joins are stable: they spill to disk gracefully and rarely blow up outright. But they don’t respond to more hardware as easily — you can’t throw machines at a sort-merge join and expect a proportional speedup. (Full mechanics here.)

Shuffle hash joins do respond to more hardware. They’re less forgiving under memory pressure, but they scale with compute in a way sort-merge doesn’t. On an SLA-bound job, that tunability is worth more than the extra stability. (Full breakdown here.)

Practice: If your streaming job has an SLA attached, treat join strategy as a lever you tune on purpose — not whatever Spark defaults to.

Rule 5: Set shuffle partitions to match cluster cores

Highly discouraged, so unless you absolutely need to, you can just let auto do it for you

The only reason we set this is because we were heavily latency-bound and did not want to leave it to the auto algorithm. That means we were overly conservative.

On join-heavy streams, spark.sql.shuffle.partitions drives task size and scheduler load.

• Too low → fat tasks, spill, OOM

• Too high → tiny tasks, overhead, small files
  

Recommended Companion Reads

Now that you’ve upgraded your hardware and parallelized your sinks, here is how to take your high-performance streaming architecture to the next level:

1. Optimize Your Compute Storage: Moving beyond default node disk settings can dramatically reduce streaming latency. Learn how to Deploy and Manage Local SSDs on AWS or upgrade your block storage setup with GP3 Support on Databricks to instantly maximize write IOPS.

2. Master Advanced Join Tactics: Choosing the wrong join strategy can derail your SLA during transaction spikes. Read our step-by-step guides on How Shuffle Hash Joins Scale and why traditional Sort-Merge Joins Might Limit Your Streaming Performance.

3. Download the Implementation Notebook: Skip the boilerplate. Grab our ready-to-import Parallel ForeachBatch Script from GitHub, and check out the walkthrough videos on ForeachBatch Parallelization (Part 1) and ThreadPool Tuning (Part 2).

FAQ

When should I parallelize foreachBatch writes?

When a single micro-batch writes to multiple independent sinks, such as a main table and a dead-letter path. Parallelization helps when sink writes are a meaningful share of batch duration and the driver can coordinate the extra jobs safely.

Is caching inside foreachBatch usually a good idea?

No. In many streaming pipelines, especially after joins or aggregations, Spark has already materialized shuffle output. Adding .cache() can increase memory use and GC overhead without enough reuse to justify it.

Why does local NVMe matter for streaming jobs?

Because shuffle, spill, and write-heavy micro-batches are sensitive to disk performance. Faster local storage can reduce batch-time variability and improve throughput for join-heavy workloads.

Should I always change spark.sql.shuffle.partitions manually?

No, almost never unless you absolutely can live without them. Auto-optimized settings may be sufficient.

What does “success” mean for a latency-bound streaming job?

Success is not just that the query keeps running. Success means the pipeline consistently finishes each micro-batch within the target latency window and prevents backlog from growing during spikes.

I’ve spent the last 5 years designing and running thousands of streaming jobs. Here’s how I actually see the tradeoff.

1. Zero-Downtime Runtime Upgrades

This is one of the most under-appreciated advantages of SDP. With a classic Structured Streaming job on regular compute, a runtime upgrade looks like this:

1. Stop the job (downtime begins)

2. Update the runtime version on the cluster

3. Restart the job

4. Wait for startup and checkpoint recovery

Every upgrade becomes a scheduled maintenance window. Because of this, teams constantly defer them, accumulate technical debt, and eventually face a forced migration under pressure.

With SDP, your pipelines are versionless. Databricks manages the runtime lifecycle for you. New compute spins up in the background, your pipeline migrates over automatically, and there’s no restart sequence to plan around. If you are running a few dozen streaming pipelines, this alone eliminates an entire category of operational toil.

2. Infrastructure You Don’t Babysit

• Serverless SDP: There is no cluster to configure. SDP auto-scales, provisions, and tears down compute automatically. You write the pipeline logic; the platform handles the rest. No sizing debates, no idle costs between micro-batches, and SLA-backed startup times.

• Classic SDP: If you aren’t on serverless, you still get Enhanced Autoscaling. This feature is purpose-built for streaming, right-sizing your clusters based on actual backlog and throughput rather than rough guesses.

3. Rewind and Recover, Instead of Rebuild

Production pipelines break. Bad data ships, an unexpected code change causes corruption, or a bug slips through code review. On raw Structured Streaming, recovery usually means manually resetting checkpoints or reprocessing everything from scratch.

SDP’s Rewind API changes that story. With a single operation, you can restore a pipeline—including table data, source offsets, and operator state together—to a clean point in time, apply your fix, and resume.

• No checkpoint deletion required.

• No full refreshes needed.

• The question shifts from “How long until we’re caught up?” to “Rewind, fix, resume.”

(Note: This is currently available on request—reach out via Databricks to get access.)

4. Quality, Dependencies, and Schema: Native, Not Bolted On

On raw Structured Streaming, you have to build and maintain operational guardrails yourself. On SDP, they come baked into the platform:

• Data Quality: Declare @dp.expect rules directly in your pipeline code. Choose whether to warn, drop, or fail per rule. Violations are tracked and visualized in the pipeline UI rather than being silently dropped.

@dp.expect("valid_order_id", "order_id IS NOT NULL")
@dp.expect_or_drop("positive_amount", "amount > 0")
@dp.expect_or_fail("no_duplicates", "COUNT(*) = COUNT(DISTINCT order_id)")
def silver_orders():
    return dlt.read_stream("bronze_orders")

• Multi-Table Dependencies: SDP automatically resolves the full DAG (Directed Acyclic Graph) on its own. You don’t need an external orchestrator for pipeline-internal chains.

• Schema Evolution: Unexpected or changed fields automatically route to a rescue column instead of crashing your production run.

• CDC (Change Data Capture): The APPLY CHANGES INTO syntax gives you SCD Type 1/2 out of the box—no hand-written, complex MERGE logic required.

Add a visual pipeline graph UI, a queryable event log, and row-level lineage through Unity Catalog.

5. When Structured Streaming Is Still the Right Call

While you should default to SDP, two situations are worth reconsidering:

• Library Compatibility: Your pipeline depends on third-party or custom libraries that don’t play well with serverless compute environments. Classic compute gives you full control over the environment.

• An Actual SDP Ceiling: You’ve hit a hard, documented limitation for your specific use case, rather than a hypothetical one.

For everything else, choose SDP first.

The Bottom Line

SDP doesn’t replace Structured Streaming—it removes the operational tax most teams end up paying anyway by building this layer poorly by hand. You ship faster, operate with less toil, and keep raw Structured Streaming as an escape hatch for the rare cases that genuinely need lower-level engine control.

If you are evaluating SDP for the first time or weighing whether to migrate existing Structured Streaming pipelines, talk to your Databricks account team. They can walk through your specific architecture and help you identify where the operational wins are real versus theoretical.

FAQ

What is SDP?

SDP stands for Spark Declarative Pipelines. It provides a declarative way to build data pipelines with built-in operations, quality enforcement, dependency management, and observability.

How is SDP different from Structured Streaming?

Structured Streaming is the underlying core engine. SDP builds on top of it, adding production-focused capabilities like managed upgrades, pipeline orchestration, quality rules, and native monitoring.

Is SDP only useful for large teams?

No. Smaller teams often benefit even more because SDP removes the heavy operational lifting that smaller engineering teams don’t have the time or staffing to manage manually.

Is SDP only useful for streaming workloads?

No. While it is exceptionally strong for streaming and incremental pipelines, it also works seamlessly for multi-hop data pipelines that combine both batch and streaming patterns.

Does SDP replace orchestration tools?

For dependencies inside a single pipeline, often yes. SDP manages execution order and downstream dependencies automatically, reducing the need for an external orchestrator within that specific pipeline boundary.

❌ The Anti-Pattern: Full Catalog Scans in a Loop

Here is the code that can causes pipelines to hit metastore rate limits or fail entirely:

# This forces a full scan of the entire metastore catalog on every execution
spark.read.table("catalog.information_schema.tables") \
    .filter((col("table_schema") == my_schema) & (col("table_name") == my_table))

The Technical Bottleneck

The information_schema is built for wide-scope catalog introspection—such as auditing, broad governance mapping, or checking metadata across an entire environment at once.

When you chain a .filter() in Spark after loading information_schema.tables, there is zero predicate pushdown to the underlying metastore database. The metastore is forced to enumerate and return every single object inside that catalog to the Spark engine, leaving Spark to do the actual row filtering client-side in memory.

The Solution: Use True Point-Lookup APIs

If you only need to verify the existence of a single table, you must bypass catalog-wide views and target a single metadata record directly.

Option 1: Programmatic Check via tableExists() (Recommended)

For pure automation pipelines, the native Spark catalog API handles this efficiently without compiling a heavy relational query plan. It hits a single metadata record instantly.

# Target a single metadata record directly
exists = spark.catalog.tableExists("catalog.schema.table")

Option 2: SQL Scoped Check via SHOW TABLES

If you are working strictly within raw SQL strings, limit your query scope using an explicit schema filter and a wildcard match rather than querying the entire catalog.

-- Fast, schema-scoped point lookup
SHOW TABLES IN catalog.schema LIKE 'table_name';

Option 3: Structural Validation via DESCRIBE TABLE

Often, you don’t just need to know if a table exists—you need to ensure its current column layout matches what your engine expects. You can leverage structural evaluation enclosed in an exception block.

from pyspark.sql.utils import AnalysisException

# Validates existence and structure simultaneously 
try:
    spark.sql("DESCRIBE TABLE catalog.schema.table")
    exists = True
except AnalysisException:
    # Catch AnalysisException specifically. A bare 'except' turns 
    # permission errors and transient timeouts into false "table missing" signals.
    exists = False

When Do You Actually Need information_schema?

If your objective is broad discovery—like mapping end-to-end data lineage across the entire workspace—using information_schema is completely appropriate. However, you should follow this single operational rule: Read it exactly once at the top of the job, cache it, and filter from memory.

# Execute exactly one metastore scan total
all_tables = spark.read.table("catalog.information_schema.tables").cache()

# Filter from the cached memory layer as many times as needed
bronze_tables = all_tables.filter(col("table_schema") == "bronze")
silver_tables = all_tables.filter(col("table_schema") == "silver")

The Ultimate Fix: Cut the Check Entirely

Pre-flight existence checks usually exist because a pipeline doesn’t know whether to create a new table or append to an existing one. If you are using Delta Lake, Delta already solves this for you.

Using .mode("append") will automatically create the table on the first write, and seamlessly append to it on every run after:

df.write \
  .format("delta") \
  .mode("append") \
  .saveAsTable("catalog.schema.target_table")

SQL pipelines get the exact same guarantee from a single DDL line at setup:

CREATE TABLE IF NOT EXISTS catalog.schema.target_table (
    id INT,
    data STRING,
    event_time TIMESTAMP
) USING DELTA;

INSERT INTO catalog.schema.target_table SELECT * FROM staging_view;

The Bottom Line: When you see an existence check before a write in a code review, ask what it’s actually guarding against. Usually, there’s no good answer—and removing it will save your metastore from a world of hurt.

And before you accuse me of writing clickbait titles that have little to do with the content of the blog (I would never do that), I assure you that Mbappé will be mentioned no fewer than 4 times.

Subscribe now

The landscape of LLMs evolves so rapidly that it’s nearly impossible to keep up with every model, let alone every dependency change that may be required by them. With AI Runtime and Custom GPU Model Serving, it’s trivially easy to serve the latest OSS models, because we can edit the versions at the outset, just as we would any other notebook.

In honor of their World Cup victory, let’s use a Mistral model: Devstral-Small-2-24B-Instruct-2512, a great mid-sized model for SWE tasks that deploys easily on Databricks custom GPU model serving.

Mbappé.

Most of these package versions are shipped natively with AI Runtime v5, but just to underscore my above point, I’ll set them here as well:

%pip install databricks-sdk==0.102.0 vllm==0.13.0 transformers==4.57.6 mistral_common openai==2.17.0 opencv-python-headless==4.12.* mlflow==3.12.0 hf_transfer==0.1.9
%restart_python

We’ll now set the wd to local disk and configure the endpoint. We set SKIP_LOCAL_SERVE_TEST=True unless we want to load the whole thing into RAM here in the serverless notebook. That would work on a single node H100, but not on an A10. We can also set the max model length here if we don’t need to service long prompts and want to conserve VRAM during serving. We’ll also set the location in UC we want to save this to, the endpoint name, and the sizing here as well.

# Set working directory to local disk (/Workspace doesn't support large files).
import os, tempfile
workdir = tempfile.mkdtemp()
os.chdir(workdir)
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
# Hugging Face model to download.
MODEL_REPO_ID = "mistralai/Devstral-Small-2-24B-Instruct-2512"
ARTIFACTS_PATH = "devstral"        # Local dir the weights download to
SERVED_MODEL_NAME = "devstral"     # Name vLLM exposes the model under.

# Devstral 2 Small FP8 (~24 GB) won't fit an A10
SKIP_LOCAL_SERVE_TEST = True
MAX_MODEL_LEN = 24576
GPU_MEMORY_UTILIZATION = 0.90

# Allowlisted ports for Serverless GPU notebooks are 3000-3999. Model Serving requires 8080.
LOCAL_PORT = 3080
SERVING_PORT = 8080

# Unity Catalog destination
UC_MODEL_NAME = "catalog.schema.devstral_small_2_24b_fp8"

# Serving endpoint configuration.
ENDPOINT_NAME = "austin-devstral-24b-h100"
WORKLOAD_SIZE = "Small"
SCALE_TO_ZERO_ENABLED = False   # GPU_XLARGE/H100 does NOT support scale-to-zero in Beta

from databricks.sdk.service.serving import ServingModelWorkloadType
WORKLOAD_TYPE = ServingModelWorkloadType.GPU_XLARGE

Mbappé.

Now we can download the model:

from huggingface_hub import snapshot_download

snapshot_download(
   repo_id=MODEL_REPO_ID,
   local_dir=ARTIFACTS_PATH,
)

And log it to our custom entry point and print the model URI in case we want to reference that later. This will take a few minutes. Note that we do NOT download from HF in the serving container.

import mlflow
from mlflow.pyfunc.model import ChatModel, ChatCompletionResponse

# Required placeholder. Serving runs the entrypoint, not python_model.predict.
class LLMModel(ChatModel):
   def predict(self, context, messages, params):
       return ChatCompletionResponse.from_dict({"choices": []})

mlflow.set_registry_uri("databricks-uc")

model_info = mlflow.pyfunc.log_model(
   name=SERVED_MODEL_NAME,
   python_model=LLMModel(),
   artifacts={
       "model_dir": ARTIFACTS_PATH,
   },
   metadata={
       "task": "llm/v1/chat",
       "entrypoint": entrypoint(SERVING_PORT),
   },
   extra_pip_requirements=[
       "mlflow==3.12.0",
   ],
)
model_info.model_uri

Now we register the model. Note that env_pack="databricks_model_serving" packages the env and the bundled weights for Serverless Optimized Deployments. This step will take like 15 minutes.

model_version = mlflow.register_model(model_info.model_uri, UC_MODEL_NAME, env_pack="databricks_model_serving")
print(f"Registered {UC_MODEL_NAME} version {model_version.version}")

Mbappé.

Our model is now ready for deployment. This cell provisions its own GPU_XLARGE (i.e. H100) serving endpoint independent of this A10 notebook.

from databricks.sdk import WorkspaceClient
from datetime import timedelta
from databricks.sdk.service.serving import EndpointCoreConfigInput, ServedEntityInput

config = EndpointCoreConfigInput(
   name=ENDPOINT_NAME,
   served_entities=[
       ServedEntityInput(
           entity_name=UC_MODEL_NAME,
           entity_version=str(model_version.version),
           workload_type=WORKLOAD_TYPE,
           workload_size=WORKLOAD_SIZE,
           scale_to_zero_enabled=SCALE_TO_ZERO_ENABLED,
       )
   ]
)

w = WorkspaceClient()
w.serving_endpoints.create_and_wait(name=ENDPOINT_NAME, config=config, timeout=timedelta(minutes=45))

And you’ve done it. With that, all that’s left to do is query the endpoint to confirm it’s working.

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ChatMessage, ChatMessageRole

w = WorkspaceClient()
resp = w.serving_endpoints.query(
   name=ENDPOINT_NAME,
   messages=[ChatMessage(role=ChatMessageRole.USER, content="Hi, will Mbappé lead France to another World Cup victory? And by extension make Mistral the world favorite OSS LLM for SWE tasks?")],
)
print(resp.choices[0].message.content)

Happy coding,

-Austin

Business users do most of their work outside of a data platform: sales reps draft proposals in Outlook, operations leads run their day in Teams, finance models everything in Excel.Yet the answers to most of their questions sit in Databricks: pipeline metrics, supply chain status, support SLAs, financial actuals, etc.

That gap between “where data lives” and “where people work” used to be bridged by reports and dashboards. Now the bridge is AI agents: assistants embedded in Teams, Outlook, Excel, Copilot Hub, Microsoft Foundry projects, or an in-house chat app. The user asks a question in their flow of work. The agent calls back to Databricks. The answer returns with proper governance applied.

The pattern is simple in theory. In practice, the integration involves a few interlocking decisions: which Databricks resources to wire up, how authentication flows, whether the agent supports MCP or only REST, whether users and Databricks share a Microsoft Entra tenant, and what governance controls you need. Get them wrong and you lose all the benefits of governance inherent to Databricks.

This post walks you through a decision tree to get to the right architecture, with the exact set of artifacts to share between the Databricks side and the agent side at the end.

Architecture Components

Before walking the tree, here are some quick definitions of the different Databricks functions we will be dealing with.

Question 1: Which Databricks Resources Need to Be Connected?

Databricks provides four managed MCP servers out of the box. Each connects a different resource type with its own URL pattern and OAuth scope.

All four are governed by Unity Catalog permissions and visible in AI Gateway > MCPs.

Note on Vector Search: The managed MCP server works directly with indexes that use Databricks-managed embeddings; you send plain text and get plain text results. If you use self-managed embeddings, wrap the index in a Model Serving endpoint (for example, via VectorSearchRetrieverTool in a deployed agent) or a UC Function, and call that from the agent instead.

Question 2: How Many Databricks Endpoints, and Does the Agent Support Multiple Tool Endpoints?

This determines whether you connect a single URL or several.

Path A: Single resource

Connect one managed MCP URL directly. Simplest path: no orchestration needed.

Path B: Multiple resources, agent supports multiple endpoints

Create one OAuth application with combined scopes and hand the agent team all the MCP URLs. The same OAuth token works across all of them.

OAuth app scopes (combine as needed):

vector-search genie unity-catalog sql offline_access

Or use all-apis for blanket access (less restrictive).

One OAuth flow, one token, multiple MCP servers. The user authenticates once and the token carries the correct scopes for every resource.

Path C: Multiple resources, agent supports only one endpoint

Two options.

Option 1: Supervisor Agent (recommended for most cases)

Register each Databricks resource as a sub-agent of an AgentBricks Multi-Agent Supervisor Agent. The external agent connects to one Model Serving endpoint.

The external agent connects to a single REST endpoint: POST https://<workspace>/serving-endpoints/<supervisor>/invocations. The supervisor’s LLM handles routing: it reads the user’s query and delegates to the right sub-agent based on descriptions you provide.

Option 2: Custom MCP Server (Databricks App)

Deploy a Databricks App that proxies to multiple managed MCP servers behind a single URL: https://<app-url>/mcp.

Decision rule: Minimal code with LLM-driven routing goes to Supervisor Agent. Lightweight deterministic proxying goes to Custom MCP App.

Question 3: What Authentication Model Is Required?

This is the most consequential decision. It determines how identity flows from the external agent to Databricks.

The key question: does the end user’s identity need to flow through to Databricks?

Yes: On-Behalf-Of (OBO) Authentication

The external agent authenticates on behalf of each individual user. Databricks receives the user’s identity, and Unity Catalog permissions are enforced per-user.

OBO-supported resources on Model Serving:

If you need OBO access to resources outside this list (e.g., UC Volumes), deploy on Databricks Apps instead of Model Serving. Apps support additional OAuth scopes.

No: Service Principal Authentication

The external agent authenticates as a single identity (service principal). All queries run under that SP’s permissions.

Authentication Method Comparison

The question to ask the data security team: do different users need different access levels to the data behind these resources, or is a shared service account acceptable?

Question 4: Are Users and Databricks in the Same Azure Tenant?

This question only matters if you chose OBO or token federation in Q3. With a simple PAT or client credentials inside the same org, skip to Q5.

Same tenant

Standard OAuth flows work directly:

• Users already exist in Databricks via SCIM sync from the same Entra ID

• OAuth authorization code flow authenticates users against the same tenant

• No federation policy needed

Different tenants (cross-tenant)

Common in enterprise scenarios: users live in Tenant A (corporate Entra ID), but Databricks is associated with Tenant B.

Path A: Per-user OBO (users in external tenant need individual identity)

1. Create an account-wide federation policy:

   • Issuer: https://login.microsoftonline.com/<external-tenant-id>/v2.0

   • Subject claim: sub or oid (to match users by their Entra ID object ID)

   • Audience: Databricks account ID or custom value

2. Provision users in Databricks. Federation validates identity but does NOT create accounts. Use SCIM sync from the external tenant, or provision manually.

3. Each user’s Entra ID token is exchanged for a Databricks token carrying their identity.

4. Unity Catalog permissions apply per-user.

Path B: Service Principal (automated agent, no per-user identity)

1. Create a Databricks service principal in your Databricks account.

2. Create a service principal federation policy:

   • Issuer: https://login.microsoftonline.com/<external-tenant-id>/v2.0

   • Subject: The Entra ID app/service principal’s sub or oid claim

   • Audience: Databricks account ID

3. The external agent obtains a JWT from its own Entra ID tenant.

4. Exchanges it via:

POST https://<workspace>/oidc/v1/token

grant_type=urn:ietf:params:oauth:grant-type:token-exchange

subject_token=<entra-id-jwt>

subject_token_type=urn:ietf:params:oauth:token-type:jwt

1. Zero Databricks secrets cross the tenant boundary.

Cross-tenant requirements checklist

Question 5: What Governance Controls Are Needed?

AI Gateway is the enterprise control plane that sits between MCP servers and external agents. It provides the operational controls needed in production.

AI Gateway Capabilities for MCP

Find it in: Workspace sidebar > AI Gateway > MCPs.

Recommended governance setup

1. Rate limiting: set conservative limits initially, increase based on observed usage.

2. Usage tracking: enable from day one to establish a baseline.

3. Audit logging: required for any integration where per-user identity flows through (OBO).

4. IP allowlisting: restrict MCP server access to known external agent IPs.

5. Scope-limited OAuth: use granular scopes (vector-search, genie) instead of all-apis for production.

Question 6: What Information Needs to Be Exchanged Between Teams?

Setting up the integration requires a handshake between the Databricks team and the external agent team. Here is exactly what each side provides.

From the Databricks side

From the external agent team

Setup steps on the Databricks side

1. Get redirect URLs and IPs from the external agent team.

2. Create the OAuth app: Account Console > Settings > App Connections > Add connection.

• Name: descriptive (e.g., foundry-mcp-client)

• Redirect URLs: from agent team

• Client type: as specified by agent team

• Scopes: combine as needed

1. Allowlist IPs in the workspace IP access list (if applicable).

2. Grant UC permissions: ensure users or SPs have access to the vector search index, Genie space, UC functions, or tables.

3. Share credentials: send the agent team the client ID, endpoints, and MCP server URLs.

4. Configure AI Gateway: set up rate limits, enable usage tracking.

Automation-friendly OAuth app creation via CLI:

databricks account custom-app-integration create --json ‘{

“name”: “foundry-mcp-client”,

“redirect_urls”: [”https://<redirect-from-agent-team>”],

“confidential”: true,

“scopes”: [”vector-search”, “genie”, “unity-catalog”, “offline_access”]

}’

The Decision Tree at a Glance

Common Scenarios (Quick Paths)

Closing Thoughts

Working through these six questions before writing any code keeps you from over-engineering a Supervisor Agent when a single MCP URL would have done the job, or from committing to OBO before confirming that your agent framework supports authorization code flow.

If you take only one thing from this post: answer the cross-tenant question early. Tenant topology determines whether you need a federation policy, and federation policies require account admin access on the Databricks side. Discovering that requirement after you have built half the integration is painful.

Done well, none of this complexity ever reaches the user. They open Teams, ask “what was our on-time delivery rate in the Northeast last week,” and get back numbers that respect their permissions, their region, their role. They never see the OAuth flow, the supervisor routing, or the federation policy. They just get the answer where they were already working. That is the whole point.

References

• Use Databricks managed MCP servers

• Connect non-Databricks clients to Databricks MCP servers

• Authentication for AI agents (Model Serving)

• OAuth token federation

• Configure a federation policy

• Supervisor Agent (multi-agent system)

• Host custom MCP servers using Databricks Apps

• Build a multi-agent system on Databricks Apps

• Unity AI Gateway

• Connect agents to unstructured data (Vector Search)

Glossary

Term

Definition

MCP

Model Context Protocol, an open standard for connecting AI agents to tools and data sources

OBO

On-behalf-of authentication: the agent acts with the end user’s identity and permissions

Token Federation

Exchanging an external IdP JWT for a Databricks OAuth token without storing Databricks secrets

AI Gateway

Databricks’ central governance layer for LLM endpoints, MCP servers, and coding agents

Supervisor Agent

A multi-agent orchestrator that routes queries to specialized sub-agents behind a single endpoint

Databricks Apps adoption can outpace the visibility into what’s running, who owns each app, and what it’s costing. A workspace can quickly accumulate hundreds of apps across business teams - some delivering ongoing value, many running 24/7 long after anyone last touched them, with the compute bill climbing alongside.

Effective app management requires two complementary signals:

• System tables : for inventory, ownership, lifecycle events, and cost attribution

• OpenTelemetry export to Unity Catalog : for request-level usage and idle detection [At the time this blog was published, App telemetry is in Public Preview.]

Used together, these answer the full set of questions a platform team needs: what do we have, who owns it, what does it cost, and is anyone using it? This post walks through both layers and shows how to combine them into a daily workflow that includes automated idle-app shutdown.

Pillar 1: Governance and Cost Attribution with System Tables

System tables provide the foundation for Apps governance. Four tables, combined, answer most operational questions:

Questions this layer answers

System tables provide the core signals needed to understand the operational footprint of Databricks Apps across a workspace:

• Ownership : Creation events help identify the original app owner, while ACL change events help surface who currently has management access.

• Cost : Usage data from system.billing.usage, combined with pricing data from system.billing.list_prices, makes it possible to estimate app cost over time.

• Orphaned Apps: Apps with no apparent current manager can be flagged for review as potentially orphaned.

• Deletion history : Delete events provide a historical view of app removals, including timing and actor.

• Recreate cycles: Creation and deletion history can reveal create → delete → recreate patterns that may indicate experimentation, churn, or unstable ownership.

System tables answer what do we have and what does it cost. The next dimension - is anyone actually using it - requires a usage signal.

Pillar 2: Usage Signals with OpenTelemetry

Databricks Apps support native OpenTelemetry export to Unity Catalog. Three telemetry tables can be configured for an app, but they differ in what is required to populate them:

• otel_logs - request logs and platform-emitted auth events. Populated automatically once OTel export is configured and the app is redeployed. No code changes required.

• otel_metrics - request counts, latencies, and other application-level metrics. Only populated after custom instrumentation is added to the app code (e.g. using the OpenTelemetry SDK in the application).

• otel_spans - distributed traces. Only populated after custom instrumentation is added to the app code.

For governance use cases - and idle detection in particular - otel_logs is the primary signal, and the one that works out of the box. Each request lands in otel_logs with an app.name attribute and timestamp, providing per-request activity data that is independent of compute state.

If teams add custom instrumentation later, otel_metrics and otel_spans open up deeper analysis (latency SLOs, error budgets, user journey traces), but they are not required for the workflow described in this post.

Enabling OTel export on an app

Enablement is a single asynchronous update call per app, using the :createUpdate endpoint with an explicit update_mask.

‼️Do not use the direct PATCH /api/2.0/apps/{name} endpoint for this purpose because it performs a full replacement of the app object, which will silently wipe resources, env_vars, and permission attachments when those fields aren’t included in the request body. :createUpdate with update_mask: “telemetry_export_destinations” ensures only that field is modified.

def patch_telemetry(app_name):
   """Enables OpenTelemetry export to Unity Catalog for the specified Databricks App."""
   payload = {
       "app": {
           "telemetry_export_destinations": [
               {"unity_catalog": TELEMETRY_TABLES}
           ]
       },
       "update_mask": "telemetry_export_destinations"
   }
   return requests.post(
       f"{workspace_url}/api/2.0/apps/{app_name}:createUpdate",
       headers=headers,
       json=payload,
   )

Looping this over GET /api/2.0/apps lets you enable telemetry across the full app fleet. After each app is redeployed, otel_logs begins populating automatically once the app receives traffic, while otel_metrics and otel_spans remain empty until the application code is instrumented.

What OTel adds beyond billing data

Both layers are needed for confident operational decisions. An app may be running and accruing DBUs while no user has touched it for weeks.

🚨 Caveats and Operational Considerations with OTel

• Cold-start gap after OTel enablement. The workflow trusts otel_logs to mean ‘this app had no requests in the activity window.’ If OTel export was just enabled, no telemetry has had time to accumulate - every app looks idle, and a first-day run will flag every app for shutdown. Don’t act on the workflow’s recommendations until telemetry has been collecting for at least as long as your idle-detection window.

• PII in telemetry. OTel auto-instrumentation captures request URLs, query parameters, headers, and request bodies. If apps pass user data through URLs, it will land in otel_logs. Apply row-level security on telemetry tables from day one.

• Staggered enablement and redeployment. Patching telemetry config and redeploying causes a user-visible app restart. For workspaces with hundreds of apps, schedule the enablement in off-hours batches and stage tier-1 apps separately to minimize disruption.

• Explicit protection for critical apps. Maintain a PROTECTED_APPS list in the shutdown job. Incident-response tools, on-call dashboards, and tier-1 business apps should be hard-coded as exempt regardless of activity level.

• Telemetry table retention.otel_logs grows quickly at scale. Set retention windows and VACUUM schedules on day one, partition by date, and monitor storage cost. Otherwise, telemetry storage will offset the savings from idle-app shutdown.

Governance using System Tables and OTel

With both System tables and OTel metrics in place, identifying and stopping idle apps becomes a 5-step workflow.

Idle Detection and Shutdown in 5 simple steps

Step 1: Find apps with recent activity

Query the OpenTelemetry logs for any app that has had user requests or auth events in the activity window. A 7-day window is a safe production default; tighten only after observing baseline traffic patterns.

WITH activity AS (
SELECT DISTINCT
  COALESCE(
    service_name,
    variant_get(resource.attributes, '$["app.name"]', 'STRING')
  ) AS app_name
FROM my_catalog.my_schema.otel_logs
WHERE `time` >= current_timestamp() - INTERVAL 7 DAYS
)
SELECT app_name FROM activity WHERE app_name IS NOT NULL

Step 2: Compute idle candidates

Combine the activity set with the live app inventory (GET /api/2.0/apps) and apply three safety filters: only ACTIVE apps are candidates, anything in PROTECTED_APPS is exempt regardless of activity, and any app younger than MIN_AGE_DAYS is skipped so newly-deployed apps have room to ramp.

from datetime import datetime, timezone
from dateutil.parser import parse as parse_ts

PROTECTED_APPS = {
    "oncall-dashboard",
    "incident-router",
    "exec-monitoring",
}
MIN_AGE_DAYS = 7   # grace period for newly-deployed apps

active = {
    row["app_name"]
    for row in spark.sql("<Step 1 SQL>").collect()
}

candidates = []
for app in list_all_apps():
    state = (app.get("compute_status") or {}).get("state")
    age_days = (datetime.now(tz=timezone.utc) - parse_ts(app["create_time"])).days
    if (
        state == "ACTIVE"
        and app["name"] not in active
        and app["name"] not in PROTECTED_APPS
        and age_days >= MIN_AGE_DAYS
    ):
        candidates.append(app["name"])

Step 3: Resolve owners and notify

This step combines both pillars: query system.access.audit (Pillar 1) to derive a per-app ownership table, then resolve owners for the candidate set generated from OpenTelemetry logs (Pillar 2), and send a 7-day warning to each candidate’s owner. The notification mechanism is intentionally a placeholder - Slack, email, or a Databricks Apps inbox notification all work. The audit trail produced in Step 4 is what matters.

# Build the per-app ownership table from System table system.access.audit.
# Run this once during setup; refresh weekly via a separate job.

spark.sql("""
  CREATE OR REPLACE TABLE apps_ownership AS
  WITH created_latest AS (
    SELECT
      get_json_object(request_params['app'], '$.name') AS app_name,
      workspace_id,
      user_identity.email AS creator_email,
      ROW_NUMBER() OVER (
        PARTITION BY get_json_object(request_params['app'], '$.name'),
                     workspace_id
        ORDER BY event_time DESC
      ) AS rn
    FROM system.access.audit
    WHERE service_name = 'apps'
      AND action_name = 'createApp'
      AND event_date >= current_date() - INTERVAL 365 DAYS
  )
  SELECT app_name, workspace_id, creator_email FROM created_latest WHERE rn = 1
""")

# Daily run: scope the ownership lookup to today's candidate set
spark.createDataFrame(
    [(name,) for name in candidates],
    ["app_name"],
).createOrReplaceTempView("idle_app_candidates")

owners = (
    spark.sql("SELECT app_name, creator_email FROM apps_ownership WHERE app_name IN (SELECT app_name FROM idle_app_candidates) AND workspace_id = '<your workspace_id>'")
    .toPandas()
    .set_index("app_name")["creator_email"]
    .to_dict()
)

for name in candidates:
    owner = owners.get(name, ADMIN_FALLBACK_EMAIL)
    notify(  # replace with your Slack / email / paging integration
        to=owner,
        subject=f"App '{name}' scheduled for shutdown in 7 days",
        body=(
            "No activity detected in the last 7 days. "
            "Touch the app or add it to PROTECTED_APPS to keep it running."
        ),
    )

Step 4: Reconcile the journal and enforce the warning period

This step does two things to the journal:

• First, it sweeps stale rows - any journal entry whose app showed activity in the current window is deleted, since the warning’s premise no longer holds. This includes rows from previously-stopped apps whose owners have restarted them, ensuring the reactivation is acknowledged and a fresh warning cycle begins if the app later goes idle again.

• Second, it MERGEs today’s idle candidates into the journal. The actual stop call in Step 5 only fires for apps that have been on the journal for at least WARNING_DAYS and are still idle on the next run, so this guarantees no app is stopped without a full warning period elapsing, and an app that bounces between active and idle starts each new idle cycle with a fresh warning_sent_at.

# Sweep stale journal entries first.
spark.createDataFrame(
    [(n,) for n in active],
    ["app_name"],
).createOrReplaceTempView("active_apps_today")

spark.sql("""
  DELETE FROM shutdown_journal
  WHERE app_name IN (SELECT app_name FROM active_apps_today)
""")

# Then MERGE today's warnings (idle_app_candidates was registered in Step 3).
spark.sql("""
  MERGE INTO shutdown_journal target
  USING (
    SELECT app_name, current_timestamp() AS warning_sent_at
    FROM idle_app_candidates
  ) source
  ON target.app_name = source.app_name
  WHEN NOT MATCHED THEN INSERT (app_name, warning_sent_at, stopped_at, status)
    VALUES (source.app_name, source.warning_sent_at, NULL, NULL)
""")

Step 5: Stop expired candidates with an audit record

For each app that’s been in the journal at least WARNING_DAYS and is still idle and still not protected, call the stop endpoint and write the outcome back to the journal.

‼️ Stop, don’t delete - stopping is reversible, deletion is not.

expired = spark.sql("""
  SELECT app_name
  FROM shutdown_journal
  WHERE warning_sent_at <= current_timestamp() - INTERVAL 7 DAYS
    AND stopped_at IS NULL
""").collect()

for row in expired:
    name = row["app_name"]
    if name in active or name in PROTECTED_APPS:
        continue  # last-second safety check
    r = requests.post(
        f"{workspace_url}/api/2.0/apps/{name}/stop",
        headers=headers,
        timeout=60,
    )
    spark.sql(
        "UPDATE shutdown_journal "
        "SET stopped_at = current_timestamp(), status = :status "
        "WHERE app_name = :name",
        args={"status": str(r.status_code), "name": name},
    )

Orchestration

Wrap the five steps in a single Databricks Job:

• Cadence: daily, off-hours (e.g., 02:00 UTC) so restart impact lands outside business hours

• Identity: a dedicated service principal (apps-governance-sp) with the minimum permissions:

  • SELECT on otel_logs and apps_ownership

  • SELECT / INSERT / UPDATE /DELETE on shutdown_journal

  • CAN_MANAGE on the apps you intend to govern (required to call /stop)

• Parameters: DRY_RUN (default true for the first two weeks of rollout), MIN_AGE_DAYS (default 7), WARNING_DAYS (default 7)

• Alerts: route job failures to the platform-team channel. A silent failure of this job means orphaned compute keeps running - the cost of missed runs is real money.

• Materialization: keep apps_ownership as a Delta table (not a view), refreshed weekly by a separate job. Daily shutdown runs read from the pre-materialized object in milliseconds instead of re-scanning audit logs every time. Apps younger than MIN_AGE_DAYS can’t be candidates anyway, so weekly refresh is sufficient.

Apps monitoring dashboard

The daily shutdown job is one side of the governance loop, an automated action. The other side is decision support: Databricks dashboard that surfaces the same signals to platform reviewers, app owners, and executives.

Built on the same two pillars, an account-wide dashboard answers the questions the platform team gets asked every week:

• Inventory KPIs: total apps, active vs stopped vs error states, OTel enablement coverage

• Cost trend: weekly DBU consumption per app, with hover-to-USD via system.billing.list_prices

• Usage signals: apps used in the last 7 days, apps unused in the last 7 days (the shutdown candidate set), broken down by workspace

• Lineage/ownership: owner email, last access, last admin action - all derived from system.access.audit

Admins see a single account-wide view instead of N per-workspace views. The visual interface turns raw signal into context that lets people decide which apps to keep, retire, or invest in - without manually correlating audit logs, OTel data, and billing tables.

Summary

A complete governance posture for Databricks Apps combines three building blocks:

1. System tables for inventory, ownership, cost, lifecycle, and compliance reporting.

2. OpenTelemetry `otel_logs` export for request-level usage and idle detection - available out of the box once OTel is configured. otel_metrics and otel_spans add deeper application-level observability when teams add custom instrumentation.

3. Automated controls - a daily job that identifies idle apps, notifies owners, and stops stale apps.

System tables answer what do we have, and what does it cost. OpenTelemetry tells you whether anyone is using it. Together, the two enable platform teams to support broad Apps adoption while keeping costs predictable and ownership accountable.

FAQ

1. Why aren’t system tables alone enough to find idle apps?

System tables show inventory, ownership, lifecycle events, and billing, but they do not reliably show whether users are actively interacting with an app. For idle detection, you need request-level activity from telemetry.

2. Which OpenTelemetry table is most important for this workflow?

otel_logs is the key table because it captures request and auth activity without requiring application code changes, once telemetry export is configured and the app is redeployed.

3. Why should the workflow stop apps instead of deleting them?

Stopping is reversible and preserves the app for later use. Deletion is permanent and introduces unnecessary operational risk for an automated governance process.

4. What is the biggest rollout risk when enabling telemetry?

If you run idle detection immediately after enabling telemetry, every app can appear idle because no activity history exists yet. Wait at least as long as your idle window before enforcing shutdown actions.

5. How should critical apps be handled?

Maintain an explicit protected-apps allowlist so business-critical or operational apps are never stopped automatically, regardless of recent activity.

6. How often should the idle-app workflow run?

A daily off-hours run is a practical default. It gives enough frequency to control spend without creating unnecessary operational churn.

• Most production agent failures come from three architectural decisions: data serving path, data access abstraction, and observability.

• Start with the simplest serving path that meets latency needs: use SQL Warehouse first, optimize table layout and queries, then add session caching or OLTP only if measurements justify it.

• Put all data access behind one execution interface like execute_sql(...) so you can centralize routing, caching, retries, parameterization, and safety controls.

• Treat observability as a launch requirement: trace tool calls, store traces in queryable tables, and add a synthetic heartbeat to measure true end-to-end freshness and latency.

• The winning pattern is not “most advanced architecture first”; it is measure, optimize the hot path, and preserve the ability to swap components without rewriting the agent.

The Framework

You can stand up an agent in an afternoon. Standing one that holds latency SLAs under real concurrency is a different exercise — one most teams underestimate.

The good news: most production agents need to talk to data and they will eventually run into the same three forks in the road. Pick well on each, and most of the other problems become tractable. Pick poorly, and you’ll be rebuilding under pressure.

This is a practitioner’s framework for the three decisions. They’re independent — you can tackle them in any order, and most teams end up working on all three in parallel. What matters is that you’ve made a deliberate call on each one.

Decision 1: How will the agent fetch data?

The instinct here is to jump to the most complex serving pattern before tuning the simplest one. Resist it.

You have three realistic options:

• Option A — Direct SQL Warehouse Serverless. Your tool calls go straight to a SQL warehouse. No new infrastructure, no reverse ETL. Worth knowing: Serverless keeps a remote result cache that survives warehouse shutdowns, so repeat queries stay fast even after the warehouse scales to zero. The fastest thing to ship.

• Option B — In-session caching. If your agent makes repeated queries against the same dataset within a single session, fetch it once and cache it in the session so subsequent calls don’t re-hit the warehouse. The implementation can be as simple as a dict or a DataFrame in the agent process, or an embedded SQL engine like DuckDB if you need to issue varied SQL against the cached data. The decision is about recognizing intra-session repetition, not about picking a cache library.

• Option C — Lakebase (OLTP). You reverse-ETL the tables the agent reads into Databricks Lakebase, a managed Postgres, and the agent queries Postgres instead of Delta. Sub-50ms point lookups. Highest performance ceiling, highest operational cost. Lakebase also doubles as a place to store online features for sub-second feature lookups and can persist agent session state — if you’re already running it for those reasons, the marginal cost of using it as your serving layer drops.

The mistake I see most: teams pattern-match on “we need low latency” and jump to Option C before they’ve measured Option A. They end up running a reverse-ETL pipeline they didn’t need.

Here’s the order that actually works:

1. Start with Option A. Get the agent end-to-end on SQL Warehouse/ DBSQL Serverless. Size the warehouse sensibly — Medium for dev, Large for prod is a reasonable starting point — X-Large gets expensive fast and is rarely the right place to begin. Start with min 1, max 3–4 clusters, then watch the query history; if queries are getting queued, add a cluster. Keep the warehouse warm if you can’t tolerate cold-start tail latency on the first query after a scale-down.

2. Tune the simple path before adding a complex one. Did you actually optimize the layout of the table the agent reads? Are your queries hitting the right clustering keys? Are you using native parameterized queries, not inline params? In our experience, half of the “we need OLTP” cases evaporate once the table layout is right.

3. Only then evaluate B or C. The decision turns on access pattern, not vibes. If a single session issues several tool calls that hit the same underlying data, Option B’s in-session caching pays for itself — the trigger is repetition within a session, not session length. If your agent makes single point lookups with hard sub-second SLAs across thousands of concurrent sessions, Option C is the answer.

A note on what actually matters for performance: it’s the engine serving the agent and the layout of the specific table that engine reads — nothing else. Medallion labels are organizational naming, not architecture. The agent doesn’t care what tier a table sits in; it cares whether the table it queries is laid out for fast reads. If your agent reads from a Delta table via SQL Warehouse/DBSQL, that table’s clustering and file layout determine your latency. If your agent reads from Lakebase, the Delta layout upstream is mostly irrelevant; what matters is the Postgres indexing and the sync freshness. Either way: optimize whatever the serving engine actually reads at query time, and let Predictive Optimization handle upstream tables that aren’t in the agent’s hot path. Avoid materialized views with non-deterministic refresh behavior anywhere in that hot path.

For Delta tables specifically, use liquid clustering and put real thought into the clustering keys. Sortable types make great keys — integers, dates, and timestamps in particular, because they have clean min/max statistics that let Delta skip files efficiently. String keys can work, but they’re far more sensitive to encoding choices (lexicographic ordering, padding, casing) and tend to underperform. Pick keys that match the actual filter patterns the agent will hit — usually some tenant or entity ID plus a time dimension — and verify the agent’s queries are pruning files the way you expect by reading the query profile.

Quick tip: If you’re merging into a table continuously — say every 5 seconds — schedule explicit liquid clustering at a regular cadence (every 30–60 minutes). Otherwise you’re relying on Predictive Optimization to catch up on its own schedule, and for hot, agent-facing tables that’s usually not fast enough.

Force yourself to write down, before you build:

• The p95 latency budget per agent tool call, end-to-end, not per-query.

• Average and p90 queries per agent session. If it’s two, the cache barely helps. If it’s ten, it’s a no-brainer.

• The size of the largest object the agent will scan. Some OLTP stores have hard per-instance ceilings; if you’re past them, the architecture has to bend.

Decision 2: How will you expose data to the agent?

Some tool in your agent that needs data —they needs a way to talk to the data plane. The instinct is to give each tool its own SQL connector when needed, its own connection management, its own retry logic. Don’t.

Use one execution interface. A single function — call it execute_sql(query, max_staleness:int = 60) — that every data-bearing tool calls. Routing, caching, parameterization, observability, retries: all internal to that one function. Tools pass SQL, get results back, and don’t know or care whether the query went to SQL Warehouse, an in-memory cache, or Lakebase.

Without abstraction each tool manages its own connector, retries, parameters. With execute_sql One place to route, optimize, and enforce safety. Without the abstraction, every tool reinvents connection management. With it, you change one function to swap engines, tune queries, or harden security.

This is the single highest-leverage abstraction you can put in your agent. Three reasons:

• One place to optimize. Tuning one query path improves every tool simultaneously. Tuning nine separate paths means you’ll tune three and forget the rest.

• One place to swap implementations. When you decide to add caching (Decision 1, Option B) or move to Lakebase (Option C), no tool code changes. You rewrite execute_sql and you’re done.

• One place to enforce safety. Parameterization, query limits, allowed-table lists — all enforced once.

On safety specifically, two things to bake in from day one.

Native parameterized queries, always. If you’re on databricks-sql-connector >= 3.0, use paramstyle=”named”. Inline parameters bypass server-side caching, won’t survive a future connector release, and open up injection vectors.

Keep sensitive identifiers out of the model’s parameter space. If your agent operates on a specific tenant, customer, or dataset, the model should not be the thing choosing which tenant to query. Use a factory/closure pattern: bind the tenant ID into the tool at construction time, so the model can only call the tool — it can’t substitute a different tenant via prompt injection. Save the parameterized version for arguments the model genuinely should control, like a column name to inspect.

def make_detect_missing_data(tenant_id: str):
    @tool
    def detect_missing_data(column_name: str) -> str:
        """Detect null values in a column."""
        return _run(tenant_id, column_name)  # tenant_id closed over
    return detect_missing_data

Compare that to the naive pattern where tenant_id is a tool parameter the LLM fills in. The naive version is one cleverly-worded user message away from a tenant-isolation incident.

Let the caller declare freshness tolerance. execute_sql should accept a max_staleness parameter — the maximum age of cached data the caller will tolerate for this query. If max_staleness is generous (say, 60 seconds), the query runner is free to serve from the in-session cache, the SQL Warehouse remote result cache, or a Lakebase replica with known lag. If it’s zero, the runner bypasses caches and goes to the source of truth. This pushes the consistency-vs-latency tradeoff to the call site, where the tool author actually knows whether eventual consistency is acceptable. Without it, you end up either over-caching (and serving stale answers in places that can’t tolerate them) or under-caching (and paying full warehouse latency for queries that didn’t need fresh data).

Decision 3: How will you know the agent is working?

If you can’t see what your agent is doing in production, you don’t have an agent — you have a black box that occasionally answers questions. Most teams get tracing late, after the first incident. Don’t be that team.

Use MLflow tracing, and route it to Unity Catalog. MLflow’s tracing is OpenTelemetry-compatible. Link your experiment to a Unity Catalog trace location with

syntax trace_location=UnityCatalog(
catalog_name=catalog_name,
schema_name=schema_name,
table_prefix=table_prefix)

and your traces land as Delta tables in UC that you can query with SQL. You’ll thank yourself the first time someone asks “show me every slow trace for tenant X in the last 24 hours.” Without UC-backed traces, that’s an afternoon of grepping logs. With them, it’s a SQL query.

Instrument the right boundaries. At minimum: input validation, tool dispatch, each tool’s SQL execution, and the model call itself. If your agent framework supports autolog, lean on it — but verify it’s actually emitting the spans you expect. Auto-instrumentation has bugs and gets disabled by version mismatches. When autolog isn’t trustworthy in your stack, fall back to manual mlflow.start_span() for the boundaries that matter. Don’t trace everything; trace the things you’d want to filter on during an incident — tenant_id, query_pattern, tool_name.

Add a synthetic heartbeat. Inject a known record at the source on a fixed cadence — every 30 or 60 seconds — and measure when it shows up at the agent. This is the one honest measure of end-to-end latency. System tables and per-stage metrics will tell you each component is healthy while the whole pipeline is silently 90 seconds behind. A heartbeat catches that immediately. It’s about a hundred lines of code; build it before you scale.

Set a latency budget per stage and defend it. Source → ingestion → transformation → serving table → agent. Every stage gets a number. Every stage has an SLO. When something breaks, you know which stage to look at without guessing. The budget also forces an honest design conversation: if your end-to-end target is 30 seconds and CDC alone takes 12, you don’t have 18 seconds for everything else — you have to decide what to cut.

Three observability tactics in one picture: a per-stage budget you can defend, trace spans landing in queryable storage, and a heartbeat that measures end-to-end the same way a user would.

What this looks like in practice

Start simple. Ship Option A for serving. Wrap data access in execute_sql. Turn on MLflow tracing to Unity Catalog and add a heartbeat. Resist the urge to add caching, OLTP, or stateful streaming gymnastics until you’ve measured a real problem.

The teams that ship reliable agents they do it boringly. They benchmark before they scale. They write down their latency budget. They abstract their data access. They trace the things that matter.

The teams that ship unreliable agents skip ahead to the interesting parts and find out, in production, that the boring parts were the load-bearing parts.

You don’t need to pick the optimal architecture on day one. You need to pick one you can swap pieces of without rewriting everything. Get the three decisions above right and you will thank yourself later.

FAQ

1. What is the best default serving option for a new production agent?

Start with SQL Warehouse. It is the fastest path to production, requires the least new infrastructure, and is often sufficient once query patterns and table layout are optimized.

2. When should I add in-session caching?

Add it when a single agent session repeatedly queries the same underlying data. If sessions only issue one or two reads, caching usually adds complexity without much benefit.

3. When does an OLTP serving layer make sense?

Use it when you need very fast point lookups under high concurrency and have hard latency targets that the warehouse path cannot meet even after optimization.

4. Why should every tool call the same execute_sql(...) function?

A shared execution layer gives you one place to manage routing, retries, caching, parameterization, observability, and safety policies. It also makes future architecture changes much easier.

5. What is the biggest security mistake in agent data access?

Letting the model choose sensitive identifiers such as tenant or customer IDs. Those should be bound into the tool outside the model’s control to reduce isolation and prompt-injection risk.

6. What should I instrument first?

At minimum, trace input validation, tool dispatch, SQL execution, and model calls. Then add a synthetic heartbeat so you can measure actual end-to-end lag, not just component-level health.

Check out Databricks native Slack Integration before implementing a custom solution!

https://docs.databricks.com/aws/en/genie-one/genie-slack

Connect multiple Databricks Genie spaces to Slack with per-user authentication, intelligent routing, semantic caching, and audit logging — all deployed as a Databricks App. A production-grade Slack bot that brings Databricks Genie directly into Slack, allowing users to ask natural-language questions across multiple data domains without leaving their workspace.

Special Thanks

I extend my sincere thanks to Sunil Patil (https://www.linkedin.com/in/pppsunil/) for his invaluable help during the brainstorming phase. His insights were instrumental in solving the complex challenge of managing multiple genies across multiple Slack channels. This project definitely levelled up thanks to his help!

Features

• Multi-Genie Support: Route questions to the right Genie space based on channel, thread, AI classification

• Per-User Authentication: Use OAuth 2.0 User-to-Machine with PKCE so queries run under each user’s Unity Catalog permissions

• Threaded Conversations: Maintain Genie context across Slack threads

• Semantic Cache: Reuse answers for similar questions using Databricks Vector Search

• Delta-Aware Cache Invalidation: Automatically invalidate cached answers when underlying Delta tables change

• Audit Logging: Track user activity, cache hits, latency, errors, and Genie usage in Delta tables

• Reliable Genie Calls: Handle rate limits with retry, exponential backoff, and adaptive polling

• Native Deployment: Run the entire application as a Databricks App with no external infrastructure

Architecture

The application runs as a Databricks App and uses two services:

• Slack Bolt in Socket Mode to listen for messages, mentions, button clicks, and threaded replies

• Flask to handle OAuth callbacks and health checks

When a user asks a question in Slack, the bot follows this flow:

1. Checks whether the Slack user has a valid Databricks OBO token

2. Routes the question to the correct Genie space

3. Looks for a semantically similar cached response

4. Validates cached answers against Delta table versions

5. Sends the question to Genie on cache miss

6. Formats the response in Slack with SQL, tables, and suggested follow-ups

7. Writes cache metadata and audit logs to Delta

Prerequisites

1. Databricks workspace

• Active Databricks workspace

• One or more Genie spaces

• SQL Warehouse

• Unity Catalog permissions

• Vector Search endpoint and index

2. Slack App

• Slack workspace with admin access

• Slack app with Socket Mode enabled

• Bot token and app-level token

3. OAuth App

• Databricks OAuth application configured for User-to-Machine authentication

• Redirect URI pointing to the Databricks App callback endpoint

4. Download the application code and make the necessary updates to the config values in app.yaml. Download all files from repo: https://github.com/adgitdemo/ad_databricks/tree/main/multi-genie-multi-slack-channel-app

Setup Instructions

1. Slack App Setup

• Create a Slack app and enable Socket Mode.

• Under OAuth & Permissions, add these bot token scopes:

  • app_mentions:read

  • chat:write

  • im:history

  • im:read

  • im:write

  • channels:history

  • groups:history

  • users:read

• Subscribe to these bot events:

  • app_mention

  • message.im

  • message.channels

  • message.groups

• Install the app to your workspace and save:

  • Bot User OAuth Token

  • App-Level Token

  • Signing Secret

2. Databricks OAuth Setup

• Create a Databricks OAuth application for User-to-Machine authentication.

• Configure the redirect URI to point to the Flask callback route exposed by the Databricks App.

The bot uses OAuth 2.0 with PKCE to securely link each Slack user to their Databricks identity. After authentication, the user’s pending question is automatically processed, so they do not need to re-type it.

3. Databricks Workspace Setup

• Prepare the required Databricks resources:

  • Genie spaces

  • SQL Warehouse

  • Cache Delta table

  • Audit Delta table

  • Vector Search endpoint

  • Vector Search index

  • Catalog and schema permissions

• Permission Grants Summary

4. Configure the Application

• Update the application configuration with:

  • Slack bot token

  • Slack app token

  • Slack signing secret

  • Databricks workspace URL

  • OAuth client ID

  • OAuth redirect URI

  • Genie space mappings

  • Default Genie space alias

  • SQL Warehouse ID

  • Cache table name

  • Audit table name

  • Vector Search endpoint and index

  • Cache similarity threshold

  • Optional Service Principal fallback flag

• Example routing configuration:

{

“trips-data”: {

“space_id”: “your-trips-genie-space-id”,

“description”: “Questions about trips, pickup locations, fares, and distance”

},

“finance”: {

“space_id”: “your-finance-genie-space-id”,

“description”: “Questions about revenue, spend, forecasts, and financial KPIs”

}

}

Usage

In Slack

Direct Messages:

Send a DM to your bot

Ask questions like: “Show me sales data for last month”

Channel Mentions:

Invite the bot to a channel: /invite @YourBotName

Mention the bot: @YourBotName what are the top 10 customers?

Threaded Conversations:

Continue asking questions in a thread

The bot maintains conversation context within threads

Demo Time

In the General Slack Channel, the bot will use the default Genie Space to answer questions.

DM Thread, the bot will use the default Genie Space to answer questions.

In a dedicated Trip Slack Channel, the bot will use the Trip Genie Space to answer questions.

In a dedicated weather Slack Channel, the bot will use the Weather Genie Space to answer questions.