[mason] tracing: managed MLflow tracing, on by default, experiment_id-centric by jamesbxwu · Pull Request #554 · databricks/databricks-ai-bridge

jamesbxwu · 2026-09-08T06:18:01Z

Summary

Make tracing default on for mason dev / mason deploy and writing to a default mlflow experiment derived from the current agent
Introduced mason tracing configure to bind to an existing experiment
Remove support for UC based tracing due to its setup complexity (will address in a follow up)

What changed

mason tracing: configure [--experiment <id>] (on / rebind, id-only override), disable (off), list, get. Dropped setup --catalog/--schema/--warehouse-id, the UC set_experiment_trace_location linking, and instrument.
On by default: mason dev and mason deploy send traces to a per-app experiment (/Users/<user>/mason-traces/<app>), auto-created on first run. No --with-traces flags.
One identifier: the runtime binding is just MLFLOW_TRACKING_URI=databricks + MLFLOW_EXPERIMENT_ID=<id>.
agent.toml: a single [tracing] table (experiment_id / disabled).
Deploy grant: the experiment is declared as an AppResourceExperiment(experiment_id, CAN_EDIT) app resource so the app's service principal can write traces with no manual grant.
Provisioning is best-effort: dev/deploy never block on tracing (no auth, no mlflow, or offline all degrade to running without traces).

Managed-only for now (UC-backed rejected)

mason manages managed (non-UC) experiments. If configure --experiment <id> is pointed at a UC-backed experiment, it errors clearly (UC-backed MLflow tracing is not supported by mason.) rather than silently wiring a config that would fail at read/deploy - reading UC traces needs a SQL warehouse and a deployed app's SP needs UC grants mason doesn't set up. Supporting UC-backed experiments (require a warehouse + grant the app SP USE_CATALOG/USE_SCHEMA/SELECT/MODIFY on the schema) is deferred to a follow-up.

Testing done

Unit: 401 passed, ty clean, ruff clean.

Live (e2-dogfood, fresh env - installed the CLI, scaffolded, dev, deploy):

mason dev: the per-app experiment was auto-created and the two MLflow env vars wired into app.yaml; invoking the local agent produced a trace that mason tracing list shows.
mason deploy: the deploy reported Trace access - granted to app service principal (the experiment app-resource), and invoking the deployed app via OAuth produced a second trace in the same experiment - written by the app's service principal with no manual grant. Both traces (dev + deployed) appear in mason tracing list.
UC-backed guard: created a UC-linked experiment and confirmed configure --experiment <it> errors cleanly; a managed experiment still configures fine.

Isaac Review found three MAJOR bugs that unit tests and the initial live run missed - dev/deploy keying the experiment on different names, a non-AgentCliError from provisioning aborting the deploy, and an offline mason dev regression - each now fixed and covered by a regression test.

Note

This supersedes the UC-only tracing PR #545 (different branch); that approach is abandoned in favor of this simpler managed model.

This pull request and its description were written by Isaac.

…-centric Replace the UC-schema tracing surface with managed MLflow tracing, keyed entirely on the experiment id. Tracing is on by default: dev/deploy send an agent's traces to a per-app experiment (/Users/<user>/mason-traces/<app>), auto-created on first run. - tracing.py: commands are `configure [--experiment <id>]` (on/rebind, id-only override), `disable` (off), `list`, `get`. Drops UC schema/warehouse, `set_experiment_trace_location` linking, and `instrument`. `ensure_experiment` creates the experiment (with the parent workspace-folder fix) and returns its id. - Runtime binding is just MLFLOW_TRACKING_URI=databricks + MLFLOW_EXPERIMENT_ID=<id>. - agent_project.py: a single `[tracing]` table (`experiment_id` / `disabled`) with `configure_tracing` / `disable_tracing`. - deploy.py/dev.py: resolve/create the experiment via `resolve_trace_experiment` (best-effort; never blocks a run). Deploy grants the app's SP write access by declaring the experiment as an app resource (CAN_EDIT) — no manual SQL grant. - store_access.py: `apply_experiment_resource` (the platform-managed grant). - Rewrote tracing/deploy/dev tests for the managed model. Co-authored-by: Isaac <no-reply@databricks.com>

Isaac Review + live dev/deploy verification against e2-dogfood surfaced three real bugs (all now fixed, with regression tests): - dev/deploy keyed the experiment differently: deploy used the mason-prefixed deployment name, dev the source-dir name, so a normal agent traced to two experiments. deploy now keys on source_dir.name too, matching dev. - tracing provisioning is genuinely best-effort now: dev/deploy catch any Exception (not just AgentCliError) from experiment creation, so an MLflow/network/permission error no longer aborts the run. - `mason dev` no longer builds the workspace client eagerly outside the tracing try, so a purely-local run with no auth/stores still runs (tracing degrades to off). - list: use search_traces(locations=...) instead of the deprecated experiment_ids= (was leaking a FutureWarning into CLI output). Live-verified on e2-dogfood: dev traces land in the auto-created per-app experiment; the deployed app's SP writes traces via the experiment app-resource grant with no manual grant. Co-authored-by: Isaac <no-reply@databricks.com>

…-only for now) `mason tracing configure --experiment <id>` now detects a UC-backed experiment (via its `mlflow.experiment.databricksTraceDestinationPath` tag) and errors with "UC-backed MLflow tracing is not supported by mason." rather than silently pinning a config that later fails (reads need a SQL warehouse; a deployed app's SP needs UC grants mason doesn't set up yet). Managed experiments are unaffected. Supporting UC-backed experiments (require a warehouse + grant the app SP UC access on the schema) is deferred to a follow-up. Co-authored-by: Isaac <no-reply@databricks.com>

# Conflicts: # integrations/mason/src/databricks_mason/agent_project.py # integrations/mason/src/databricks_mason/deploy.py

jamesbxwu · 2026-09-08T20:54:10Z

+        trace_experiment_id: str | None = None,
+        trace_disabled: bool = False,


is it possible to derive trace_disabled from trace_experiment_id == None?

Not cleanly - trace_experiment_id and disabled are orthogonal, so folding one into the other loses a state:

experiment_id answers which experiment. None = "use the per-app default", not "off".

disabled answers on/off. Tracing is on by default, so a fresh/absent [tracing] (experiment_id == None) must mean on-with-default. If None also meant disabled, we could never express on-by-default. And clearing experiment_id to turn tracing off would not stick - the next dev/deploy just re-derives the default and turns it back on.

That said, your underlying point (the special-cased "recompute the default each run" was weird) is fixed in 107ec39: resolve_trace_experiment now pins the resolved default id into agent.toml on the first dev/deploy, so later runs reuse it by id. After first run there is no separate "default" state - just a pinned id or disabled = true.

can we create the toml with the trace_experiment_id set to the default experiment path always? and dev and deploy never overwrite trace_experiment_id in agent.toml? then we can remove disabled, if trace_experiment_id is unset then the user must have disabled it

the question here is when to create the default experiment and get back the id to write to the toml file. While we can derive the experiment name (path) from the mason project name, we cannot get the id until we actually create the experiment which is done in mason dev / deploy and not mason init.

ah I see... makes sense. What about setting the experiment name in toml file? as the agent author it would be nice to see the experiment name rather than the id in my configuration file.

also i think similar motivation to your first comment, having both trace_disabled and trace_experiment_id is a little confusing to me.

What about setting the experiment name in toml file

I chose id since it's the canonical identifier as mlflow experiment names can be changed.

having both trace_disabled and trace_experiment_id is a little confusing to me

Good call out. I went back and forth on this and open to ideas. The reason I kept both is I wanted to provide a knob to disable tracing even if it was bound to an experiment previously.

- deploy: `trace_env` -> `mlflow_tracing_config` returning a typed `MlflowTracingConfig` dataclass (`.env()` renders the two MLflow env vars) instead of a bare dict. - tracing: `_trace_json` -> `_trace_to_json`; reword the configure hints to "use the mason default". - dev: reword the on-by-default docstring; "Running without traces" -> "Proceeding without tracing". Co-authored-by: Isaac <no-reply@databricks.com>

…irst run `resolve_trace_experiment` now writes the auto-created per-app experiment id back into `agent.toml [tracing] experiment_id` on the first `mason dev` / `mason deploy`. Later runs read the pinned id instead of re-deriving the default from the app name, so there is no special-cased "default" state once tracing has run once. The `disabled` flag is retained: it remains the only representation of tracing-off (a fresh/absent config still means on-by-default, which cannot be encoded by a cleared experiment id). Co-authored-by: Isaac <no-reply@databricks.com>

…app" The default tracing experiment is keyed on the source directory's basename (the Mason project name), not the deployed app name (which is `mason-<name>`). Rename the parameter (`resolve_trace_experiment`'s `app` -> `project_name`, `default_experiment`'s `app` -> `project`) and align all user-facing help, docs, and comments on "per-project" so the terminology matches what the value actually is. No behavior change: the experiment path is still `/Users/<you>/mason-traces/<project>`. Co-authored-by: Isaac <no-reply@databricks.com>

elainewang-db · 2026-09-08T23:33:11Z

-    only checks a bound store still exists (a typo or unbound clone fails here, not at runtime) and
-    returns the trace env.
+def validate_stores(client, *, memory_store: Optional[str], session_store: Optional[str]) -> None:
+    """Validate the agent's bound stores exist. Shared by `mason deploy` and `mason dev`.


follow up: move this and _upsert_manifest_env to some shared util file and resolve_trace_experiment to mason.tracing?

will do as a follow up

elainewang-db · 2026-09-08T23:38:50Z

-        return _DEFAULT_EXPERIMENT
-    return f"/Users/{user}/mason-traces/{app}"

+def default_experiment(user: str, project: Optional[str]) -> str:


can this be used in agent_project.py as the default experiment?

see above comment that we store the experiment_id in the project but this default is the experiment_name (that is derived from the mason project name)

elainewang-db · 2026-09-08T23:42:05Z

+        trace_experiment_id: str | None = None,
+        trace_disabled: bool = False,


can we create the toml with the trace_experiment_id set to the default experiment path always? and dev and deploy never overwrite trace_experiment_id in agent.toml? then we can remove disabled, if trace_experiment_id is unset then the user must have disabled it

jamesbxwu · 2026-09-09T04:54:56Z

-        return _DEFAULT_EXPERIMENT
-    return f"/Users/{user}/mason-traces/{app}"

+def default_experiment(user: str, project: Optional[str]) -> str:


Suggested change

def default_experiment(user: str, project: Optional[str]) -> str:

def default_experiment_name(user: str, project: Optional[str]) -> str:

jamesbxwu · 2026-09-09T04:55:25Z

-    if warehouse_id:
-        os.environ["MLFLOW_TRACING_SQL_WAREHOUSE_ID"] = warehouse_id

+def ensure_experiment(profile: Optional[str], client, name: str) -> str:


Suggested change

def ensure_experiment(profile: Optional[str], client, name: str) -> str:

def create_experiment_idempotent(profile: Optional[str], client, name: str) -> str:

jamesbxwu · 2026-09-09T04:59:15Z

+        trace_experiment_id: str | None = None,
+        trace_disabled: bool = False,


the question here is when to create the default experiment and get back the id to write to the toml file. While we can derive the experiment name (path) from the mason project name, we cannot get the id until we actually create the experiment which is done in mason dev / deploy and not mason init.

jamesbxwu · 2026-09-09T05:05:21Z

-    only checks a bound store still exists (a typo or unbound clone fails here, not at runtime) and
-    returns the trace env.
+def validate_stores(client, *, memory_store: Optional[str], session_store: Optional[str]) -> None:
+    """Validate the agent's bound stores exist. Shared by `mason deploy` and `mason dev`.


will do as a follow up

jamesbxwu · 2026-09-09T05:06:24Z

-        return _DEFAULT_EXPERIMENT
-    return f"/Users/{user}/mason-traces/{app}"

+def default_experiment(user: str, project: Optional[str]) -> str:


see above comment that we store the experiment_id in the project but this default is the experiment_name (that is derived from the mason project name)

Rename tracing helpers so a function's name states whether it deals in an experiment *name* (workspace path) or an *id*, and what it does: - default_experiment -> default_experiment_name (returns the name/path) - ensure_experiment -> create_experiment_idempotent (creates-or-gets, returns id) - resolve_trace_experiment -> resolve_trace_experiment_id (returns the id) - _configure -> _set_tracking_uri (it sets MLflow's tracking uri; the old name collided with the user-facing `tracing configure` command) Pure rename, no behavior change. Co-authored-by: Isaac <no-reply@databricks.com>

Follow-up to b627196, which by mistake committed only dev_test.py. This adds the actual source renames and the other test files, so the branch is consistent: - default_experiment -> default_experiment_name - ensure_experiment -> create_experiment_idempotent - resolve_trace_experiment -> resolve_trace_experiment_id - _configure -> _set_tracking_uri Pure rename, no behavior change. Co-authored-by: Isaac <no-reply@databricks.com>

jamesbxwu and others added 5 commits September 8, 2026 05:03

Merge remote-tracking branch 'origin/main' into mason-tracing-managed

452f085

Merge remote-tracking branch 'origin/main' into mason-tracing-managed

147606f

# Conflicts: # integrations/mason/src/databricks_mason/agent_project.py # integrations/mason/src/databricks_mason/deploy.py

jamesbxwu commented Sep 8, 2026

View reviewed changes

jamesbxwu and others added 2 commits September 8, 2026 21:02

jamesbxwu requested a review from elainewang-db September 8, 2026 23:13

elainewang-db reviewed Sep 8, 2026

View reviewed changes

jamesbxwu commented Sep 9, 2026

View reviewed changes

jamesbxwu and others added 2 commits September 9, 2026 05:24

jamesbxwu requested a review from elainewang-db September 9, 2026 05:39

elainewang-db approved these changes Sep 9, 2026

View reviewed changes

jamesbxwu merged commit 530d9f3 into main Sep 9, 2026
50 checks passed

jamesbxwu deleted the mason-tracing-managed branch September 9, 2026 18:22

jamesbxwu mentioned this pull request Sep 10, 2026

[mason] drop the [tracing] extra and make MLflow a base dependency #570

Merged

		trace_experiment_id: str \| None = None,
		trace_disabled: bool = False,

	def default_experiment(user: str, project: Optional[str]) -> str:
	def default_experiment_name(user: str, project: Optional[str]) -> str:

	def ensure_experiment(profile: Optional[str], client, name: str) -> str:
	def create_experiment_idempotent(profile: Optional[str], client, name: str) -> str:

Conversation

jamesbxwu commented Sep 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What changed

Managed-only for now (UC-backed rejected)

Testing done

Note

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jamesbxwu commented Sep 8, 2026 •

edited

Loading