Skip to content

[mason] tracing: UC-only and opt-in via mason tracing setup - #545

Open
jamesbxwu wants to merge 8 commits into
mainfrom
mason-tracing-dev-uc
Open

[mason] tracing: UC-only and opt-in via mason tracing setup#545
jamesbxwu wants to merge 8 commits into
mainfrom
mason-tracing-dev-uc

Conversation

@jamesbxwu

@jamesbxwu jamesbxwu commented Sep 4, 2026

Copy link
Copy Markdown

UX Change

  • Off until configured. mason tracing setup --trace-location <catalog.schema> records a UC schema in agent.toml. Tracing is disabled until then.
  • Never blocks. mason deploy and mason dev displays a message to remind users to set up tracing using mason tracing setup

Key changes

  • dev.py / deploy.py: share provision_trace_experiment (returns None when unconfigured; no gate).
  • pyproject.toml: add opentelemetry-exporter-otlp-proto-http to the runtime/runtime-openai extras - UC trace export uses HTTP OTLP.
  • runtime/tracing.py: reverted to the shipped gate (has_destination and has_experiment); we always wire both, so no runtime change is required.
  • Removed the managed [dev] experiment machinery (dev_experiment_name, MasonClient.ensure_experiment, the [dev] fallback in tracing list).

Manual testing

Full CUJ against a real workspace, from a fresh venv with the CLI installed from this branch, for both LangGraph and OpenAI. Sequencing: no tracing first (confirm nothing is logged), then mason tracing setup, then confirm traces land - in both mason dev and a real mason deploy.

Stage LangGraph OpenAI
mason init scaffolds, tracing off scaffolds, tracing off
mason dev, no tracing invoke OK, 0 traces, off-hint, no MLFLOW_* env invoke OK, 0 traces, off-hint, no MLFLOW_* env
mason tracing setup --trace-location james_wu.mason recorded in agent.toml recorded in agent.toml
mason dev, tracing on invoke OK, trace in james_wu.mason invoke OK, trace in james_wu.mason
mason deploy, tracing on real rollout, UC experiment created, env wired
live deployed app SP granted, app invoked, its trace in james_wu.mason

Warning message when tracing is off

mason dev with no mason tracing setup run - no MLFLOW_* env is wired, nothing is exported, and the startup panel nudges (never blocks):

✓ Starting agent
  Chat UI   http://localhost:8020
  Next steps   ($ = run in your terminal)
  • Open http://localhost:8020 to chat with your agent
  $ mason tools add mcp <service>      Give the agent a tool
  $ mason dev -m <store> -s <store>    Attach a memory / session store
  $ mason deploy cuj-langgraph         Deploy it to Databricks
  • Tracing is off. Run `mason tracing setup --trace-location <catalog.schema>` to enable it.

mason deploy on an unconfigured project shows the same one-line hint as a next step instead of gating, so the deploy still succeeds.

After mason tracing setup (tracing on)

mason dev now names the destination in its panel:

✓ Starting agent
  Chat UI   http://localhost:8020
  Traces    james_wu.mason
  ...

mason deploy wires it and creates the UC-linked experiment:

✓ Deployed agent 'mason-cuj-openai'
  URL              https://mason-cuj-openai-….staging.aws.databricksapps.com
  Traces           /Users/james.wu@databricks.com/mason-traces/cuj-openai (james_wu.mason)
  Next steps
  • Grant the app's service principal USE CATALOG + USE SCHEMA + MODIFY/SELECT on james_wu.mason so it can write traces.

Traces actually land

mason tracing list after dev invokes (both frameworks) and a live invoke of the deployed OpenAI app - all three in the configured UC schema:

Agent Traces · james_wu.mason
TRACE ID                                  STATUS   LATENCY (MS)   CREATED
trace:/james_wu.mason/cbeca5bccee60cf0…   ● Ok     1833           just now         ← deployed OpenAI app
trace:/james_wu.mason/b792fba9600c702c…   ● Ok     1695           12 minutes ago   ← dev (OpenAI)
trace:/james_wu.mason/31b138b84d6d2436…   ● Ok     1353           20 minutes ago   ← dev (LangGraph)
3 items

Notes

  • Reading traces back (mason tracing list) uses the workspace default SQL warehouse; pass --warehouse-id (or mason tracing setup --warehouse-id) where no default warehouse resolves in the target workspace.
  • The HTTP OTLP exporter dep ships in the extras here; to validate before this merges/publishes it was injected into the agent venvs, which is what the merged extras provide automatically.

Rework `mason tracing` into one experiment-centric model with two backends,
so local iteration has zero setup friction while production traces are governed.

- `mason dev`: tracing on by default to a managed "[dev] <app>" experiment
  under the user's workspace home (no Unity Catalog access needed). The runtime
  auto-creates it on first trace, so dev makes no MLflow call.
- `mason tracing setup --trace-location <catalog.schema> [--warehouse-id]`:
  records the UC schema (+ optional warehouse) in agent.toml.
- `mason deploy`: gated on UC tracing. When configured it creates a UC-linked
  experiment and wires MLFLOW_TRACKING_URI/EXPERIMENT_NAME into app.yaml;
  otherwise it stops and points at `mason tracing setup`. A UC experiment can
  only be linked while empty, so an experiment that already holds non-UC traces
  raises a clear error pointing at the trace-to-UC migration docs.
- `mason tracing list`: resolves --trace-location -> agent.toml -> the [dev]
  experiment, and queries by `locations=` (drops the deprecated experiment_ids
  path that also crashed on current MLflow). `get` stays id-only (self-locating).
- render: wrap field/cell values in Text so Rich doesn't parse "[dev]" as markup
  and silently drop it.

Removes the `--with-traces` / `--traces-experiment` flags on dev/deploy in favor
of the agent.toml binding. Unit tests updated; the deploy-time UC provisioning is
covered with mocked mlflow (verified live separately).

Co-authored-by: Isaac <no-reply@databricks.com>
@jamesbxwu jamesbxwu changed the title [mason] tracing: on-by-default in dev, gate deploy on Unity Catalog [mason] tracing CLI improvements: on-by-default in dev, gate deploy on UC enabled tracing Sep 4, 2026
jamesbxwu and others added 2 commits September 4, 2026 03:26
Listing the managed [dev] experiment from inside a project that has a UC
warehouse configured was needlessly pulling that warehouse in and cold-starting
it. A managed experiment is read from the tracking store and needs no warehouse;
only a UC catalog.schema does. Gate the warehouse on the location being a UC
schema.

Co-authored-by: Isaac <no-reply@databricks.com>
Replace the split dev(non-UC)/deploy(UC) tracing model with one Unity Catalog
path shared by `mason dev` and `mason deploy`.

- Tracing is off until `mason tracing setup --trace-location <catalog.schema>`
  records a UC schema in agent.toml. Both dev and deploy then create/link a
  per-app UC experiment and wire the same MLFLOW_TRACKING_URI +
  MLFLOW_EXPERIMENT_NAME + MLFLOW_TRACING_DESTINATION (+ warehouse).
- Deploy no longer gates on tracing: an unconfigured project deploys with a
  one-line hint instead of being blocked, so developers without catalog/schema
  access are never stuck.
- `mason dev` checks the trace config (a cheap agent.toml read) before touching
  the workspace, so a plain unconfigured or offline dev run makes no auth call.
- Add opentelemetry-exporter-otlp-proto-http to the runtime/runtime-openai
  extras: UC trace export uses HTTP OTLP.
- Revert the runtime configure_tracing gate change; we always wire both the
  experiment and the destination, which the published runtime already accepts.
- Drop the managed [dev] experiment machinery (dev_experiment_name,
  MasonClient.ensure_experiment) and the `[dev]` fallback in `tracing list`.

Verified end to end for LangGraph and OpenAI: dev with no tracing logs nothing;
after setup, dev and a real deploy both land traces in the configured UC schema.
@jamesbxwu jamesbxwu changed the title [mason] tracing CLI improvements: on-by-default in dev, gate deploy on UC enabled tracing [mason] tracing: UC-only and opt-in via mason tracing setup Sep 4, 2026
jamesbxwu and others added 5 commits September 4, 2026 18:04
provision_trace_experiment now returns Optional[tuple]; narrow it with an
`assert result is not None` before unpacking so `ty check` passes.
After tracing is configured, `mason dev` and `mason deploy` now show the MLflow
experiment id and a direct link to its traces page
(`<host>/ml/experiments/<id>/traces`, built from the profile's workspace host)
instead of just the experiment name.

- ensure_uc_experiment now returns the experiment id.
- provision_trace_experiment returns (experiment_id, catalog_schema).
- add experiment_ui_url(host, experiment_id) and render it in both panels.
# Conflicts:
#	integrations/mason/src/databricks_mason/tracing.py
… + messaging

- get/list now require a resolvable SQL warehouse (--warehouse-id, the project's
  configured warehouse, or MLFLOW_TRACING_SQL_WAREHOUSE_ID) and fail fast with a
  clear hint. get previously ignored the configured warehouse entirely.
- provision_trace_experiment now creates the parent /Users/<user>/mason-traces
  workspace folder before create_experiment, so first-time dev/deploy on a
  workspace that lacks it no longer fails (create_experiment does not create
  intermediate folders).
- Refresh tracing help/messages: accurate warehouse help, verb-first success
  title, imperative next-steps, and corrected/removed stale help.py examples.

Co-authored-by: Isaac <no-reply@databricks.com>
Reconcile main's #540 agent.toml store-binding model with the UC-only tracing
redesign so both stores and tracing are configured via agent.toml (no flags):

- Stores: adopt main's model - `mason memory/sessions bind` records bindings in
  agent.toml; dev/deploy read them via store_bindings and validate they exist
  (validate_stores). Dropped --memory/--session/--no-create-stores flags and the
  store-env-into-app.yaml path (stores are read from agent.toml at runtime).
- Tracing: keep UC-only via `mason tracing setup` + provision_trace_experiment.
  Dropped main's --with-traces/--traces-experiment flags and the trace-env branch
  of the old validator; removed the stale `mason tracing instrument` help example.
- agent_project.py: keep both additions - main's memory_store_id and our
  trace_location/trace_warehouse.
- Kept main's #515 --instances deploy feature and #538 UI changes as-is.

Co-authored-by: Isaac <no-reply@databricks.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant