diff --git a/integrations/mason/pyproject.toml b/integrations/mason/pyproject.toml index acc41b3f..94499506 100644 --- a/integrations/mason/pyproject.toml +++ b/integrations/mason/pyproject.toml @@ -35,6 +35,7 @@ runtime = [ "mlflow>=3.10.1", "uuid-utils>=0.10.0", "opentelemetry-exporter-otlp-proto-grpc>=1.25.0", + "opentelemetry-exporter-otlp-proto-http>=1.25.0", "databricks-agents>=1.9.3", ] runtime-openai = [ @@ -44,6 +45,7 @@ runtime-openai = [ "mlflow>=3.10.1", "uuid-utils>=0.10.0", "opentelemetry-exporter-otlp-proto-grpc>=1.25.0", + "opentelemetry-exporter-otlp-proto-http>=1.25.0", "databricks-agents>=1.9.3", ] diff --git a/integrations/mason/src/databricks_mason/agent_project.py b/integrations/mason/src/databricks_mason/agent_project.py index ec08a170..ccffc337 100644 --- a/integrations/mason/src/databricks_mason/agent_project.py +++ b/integrations/mason/src/databricks_mason/agent_project.py @@ -17,6 +17,11 @@ from databricks_mason.errors import AgentCliError from databricks_mason.runtime.tool_manifest import MEMORY_STORE_TABLE, SESSION_STORE_TABLE +# The UC trace destination binding (`mason tracing setup`): a "catalog.schema" (+ optional +# warehouse id). The runtime doesn't read this from agent.toml - `mason deploy` uses it to create a +# UC-linked experiment and wires MLFLOW_TRACING_DESTINATION into app.yaml, which the agent reads. +TRACE_LOCATION_TABLE = "trace_location" + _SCHEMA_VERSION = 1 _SUPPORTED_FRAMEWORKS = {"langgraph", "openai"} _SUPPORTED_SCOPE_KINDS = {"table", "volume", "workspace"} @@ -193,12 +198,16 @@ def _required_string(value: object, description: str) -> str: def _store_name_from_manifest(value: object, table: str) -> str | None: - """Read the ``name`` from a ``[memory_store]`` / ``[session_store]`` table, or None if absent.""" + """Read the ``name`` from a ``[memory_store]`` / ``[session_store]`` table, or None if absent. + + Coerced to a plain ``str``: tomlkit hands back a ``String`` subclass, which yaml's type-keyed + dumper (used to materialize app.yaml) can't represent, so return a bare str. + """ if value is None: return None if not isinstance(value, Mapping): raise AgentCliError(f"agent.toml [{table}] must be a table.") - return _required_string(cast(Mapping[str, Any], value).get("name"), f"[{table}] name") + return str(_required_string(cast(Mapping[str, Any], value).get("name"), f"[{table}] name")) def _store_id_from_manifest(value: object) -> str | None: @@ -296,6 +305,8 @@ def __init__( memory_store: str | None = None, session_store: str | None = None, memory_store_id: str | None = None, + trace_location: str | None = None, + trace_warehouse: str | None = None, ) -> None: self.root = root self.path = root / "agent.toml" @@ -307,6 +318,9 @@ def __init__( self.memory_store = memory_store self.session_store = session_store self.memory_store_id = memory_store_id + # UC trace destination ("catalog.schema") + optional SQL warehouse; None = not configured. + self.trace_location = trace_location + self.trace_warehouse = trace_warehouse @classmethod def load(cls, root: pathlib.Path | str) -> "AgentProject": @@ -347,6 +361,16 @@ def load(cls, root: pathlib.Path | str) -> "AgentProject": session_store = _store_name_from_manifest( document.get(SESSION_STORE_TABLE), SESSION_STORE_TABLE ) + trace_location = _store_name_from_manifest( + document.get(TRACE_LOCATION_TABLE), TRACE_LOCATION_TABLE + ) + trace_table = document.get(TRACE_LOCATION_TABLE) + raw_warehouse = ( + trace_table.get("warehouse_id") if isinstance(trace_table, Mapping) else None + ) + trace_warehouse = ( + str(raw_warehouse) if isinstance(raw_warehouse, str) and raw_warehouse else None + ) return cls( project_root, document, @@ -355,6 +379,8 @@ def load(cls, root: pathlib.Path | str) -> "AgentProject": memory_store, session_store, memory_store_id, + trace_location, + trace_warehouse, ) @classmethod @@ -429,6 +455,33 @@ def unbind_session_store(self) -> bool: """Remove the session store binding from agent.toml. Returns True if it was present.""" return self._clear_store(SESSION_STORE_TABLE) + def bind_trace_location(self, location: str, warehouse_id: str | None = None) -> bool: + """Declare the UC trace destination (``catalog.schema``) + optional SQL warehouse. + + Stored as ``[trace_location] name = ""`` with an optional ``warehouse_id``. + Returns True if anything changed. + """ + location = _required_string(location, f"[{TRACE_LOCATION_TABLE}] name") + if self.trace_location == location and self.trace_warehouse == warehouse_id: + return False + existing = self._document.get(TRACE_LOCATION_TABLE) + table = existing if isinstance(existing, Mapping) else tomlkit.table() + table["name"] = location + if warehouse_id: + table["warehouse_id"] = warehouse_id + elif "warehouse_id" in table: + del table["warehouse_id"] + if not isinstance(existing, Mapping): + self._document.append(TRACE_LOCATION_TABLE, table) + self.trace_location = location + self.trace_warehouse = warehouse_id + return True + + def unbind_trace_location(self) -> bool: + """Remove the trace destination from agent.toml. Returns True if it was present.""" + self.trace_warehouse = None + return self._clear_store(TRACE_LOCATION_TABLE) + def _set_store(self, table: str, name: str, store_id: str | None = None) -> bool: name = _required_string(name, f"[{table}] name") if getattr(self, table) == name and getattr(self, f"{table}_id", None) == store_id: diff --git a/integrations/mason/src/databricks_mason/deploy.py b/integrations/mason/src/databricks_mason/deploy.py index 89d30346..8d928eff 100644 --- a/integrations/mason/src/databricks_mason/deploy.py +++ b/integrations/mason/src/databricks_mason/deploy.py @@ -15,7 +15,7 @@ import json import pathlib import time -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional import click import yaml @@ -24,7 +24,18 @@ from databricks_mason.errors import AgentCliError from databricks_mason.render import field from databricks_mason.store_access import _databricks, apply_postgres_resources, grant_tables -from databricks_mason.tracing import TRACES_DEST_ENV, TRACES_EXPERIMENT_ENV, default_experiment +from databricks_mason.tracing import ( + TRACES_DESTINATION_ENV, + TRACES_EXPERIMENT_ENV, + TRACES_TRACKING_URI_ENV, + TRACES_WAREHOUSE_ENV, + ensure_uc_experiment, + experiment_name, + experiment_ui_url, +) + +if TYPE_CHECKING: + from databricks_mason._api_client import _MasonApiClient # TEMPORARY: the Apps build environment currently can't reach the internal pypi proxy, so builds # time out installing dependencies. Point the build at public PyPI (sanctioned interim workaround) @@ -251,21 +262,18 @@ def store_bindings(source: pathlib.Path) -> tuple[Optional[str], Optional[str]]: return memory, session -def validate_stores_and_trace_env( +def validate_stores( client, *, - app: Optional[str], memory_store: Optional[str], session_store: Optional[str], - traces_destination: Optional[str], - traces_experiment: Optional[str], -) -> dict[str, str]: - """Validate the agent's bound stores exist and build the MLFLOW_* trace env to wire in. - - Shared by `mason deploy` and `mason dev`. Stores are created by `mason memory/sessions bind` and - read from agent.toml at runtime, so this neither creates them nor writes them to app.yaml — it - only checks a bound store still exists (a typo or unbound clone fails here, not at runtime) and - returns the trace env. +) -> None: + """Validate the agent's bound stores exist. Shared by `mason deploy` and `mason dev`. + + Stores are created by `mason memory/sessions bind` and read from agent.toml at runtime, so this + neither creates them nor writes them to app.yaml — it only checks a bound store still exists (a + typo or unbound clone fails here, not at runtime). Tracing is handled separately (see + `provision_trace_experiment`). """ if memory_store and _resolve_memory_store(client, memory_store) is None: # Resolve by display name: get_memory_store looks up by resource id, not the bound name. @@ -282,18 +290,53 @@ def validate_stores_and_trace_env( hint=f"Run `mason sessions bind {session_store}` to create and bind it.", error_code=exc.error_code, ) from exc - env: dict[str, str] = {} - if traces_destination: - env[TRACES_DEST_ENV] = traces_destination - # The agent enables tracing only when BOTH a destination and an experiment are set, so - # default the experiment to this agent's per-app path (matching `mason tracing setup --app`), - # otherwise --with-traces alone would ship a half-config that silently disables tracing. - env[TRACES_EXPERIMENT_ENV] = traces_experiment or default_experiment( - client.current_user, app - ) - elif traces_experiment: - env[TRACES_EXPERIMENT_ENV] = traces_experiment - return env + + +def provision_trace_experiment( + source: pathlib.Path, app: str, client: _MasonApiClient, profile +) -> Optional[tuple[str, str]]: + """Wire Unity Catalog tracing into app.yaml if it's configured, else no-op. Shared by dev+deploy. + + Tracing is UC-only and opt-in: it's on iff the project ran ``mason tracing setup`` (a + ``catalog.schema`` bound in agent.toml). When configured, this creates+links the per-agent + UC experiment and writes ``MLFLOW_TRACKING_URI`` + ``MLFLOW_EXPERIMENT_NAME`` (the agent + runtime's enable-gate keys on the experiment) + ``MLFLOW_TRACING_DESTINATION`` (the UC schema, + which routes export to governed storage and blocks any ambient ``OTEL_EXPORTER_OTLP_*`` hijack). + Returns ``(experiment_id, catalog_schema)``, or ``None`` when tracing isn't configured - the + caller decides how to nudge the user (dev/deploy both just print a hint, never block). The + experiment id lets the caller build the MLflow experiment UI link. + """ + from databricks_mason.agent_project import AgentProject # noqa: PLC0415 - avoid import cycle + + try: + project = AgentProject.load(source) + schema, warehouse = project.trace_location, project.trace_warehouse + except AgentCliError: + schema, warehouse = None, None + if not schema: + return None + + experiment = experiment_name(client.current_user, app) + # `mlflow.create_experiment` does not create intermediate workspace folders, so create the parent + # (`/Users//mason-traces`) first - otherwise first-time provisioning on a workspace that + # doesn't have it yet fails. + client.ensure_workspace_dir(experiment.rsplit("/", 1)[0]) + experiment_id = ensure_uc_experiment(profile, experiment, schema, warehouse) + env = { + TRACES_TRACKING_URI_ENV: "databricks", + TRACES_EXPERIMENT_ENV: experiment, + TRACES_DESTINATION_ENV: schema, + } + if warehouse: + env[TRACES_WAREHOUSE_ENV] = warehouse + _upsert_manifest_env(source, env) + return experiment_id, schema + + +# One-line nudge shown by dev/deploy when tracing isn't configured (never blocks). +_TRACING_OFF_HINT = ( + "Tracing is off. Run `mason tracing setup --trace-location ` to enable it." +) def _grant_store_access( @@ -340,18 +383,6 @@ def _grant_store_access( help="Local source directory for the deployment (containing app.yaml). Defaults to the " "current directory.", ) -@click.option( - "--with-traces", - "traces_destination", - default=None, - help="UC trace destination 'catalog.schema' to wire in via MLFLOW_TRACING_DESTINATION " - "(link it first with `mason tracing setup`).", -) -@click.option( - "--traces-experiment", - default=None, - help="MLflow experiment path to wire in via MLFLOW_EXPERIMENT_NAME.", -) @click.option( "--pip-index-url", default=_DEFAULT_PIP_INDEX_URL, @@ -374,8 +405,6 @@ def deploy( obj, name, source, - traces_destination, - traces_experiment, pip_index_url, workspace_path, instances, @@ -398,26 +427,35 @@ def deploy( source_dir = pathlib.Path(source) client = obj.client() - # 1. Validate the agent's bound stores (`mason memory/sessions bind` creates them) and build any - # trace env. Stores are read from agent.toml at runtime, not wired into app.yaml; the bindings - # also drive the store access grant (step 4). + # 1. Wire UC tracing first (if configured), so its experiment is created before the app. Tracing + # is opt-in and never blocks a deploy: an unconfigured project just deploys without it. The + # experiment is keyed on the source dir name, matching `mason dev` for the same project. + traced = provision_trace_experiment(source_dir, source_dir.resolve().name, client, obj.profile) + trace_experiment_id: Optional[str] = None + trace_schema: Optional[str] = None + trace_url: Optional[str] = None + if traced: + trace_experiment_id, trace_schema = traced + trace_url = experiment_ui_url(client.host, trace_experiment_id) + + # 2. Validate the agent's bound stores (`mason memory/sessions bind` creates them). Stores are + # read from agent.toml at runtime, not wired into app.yaml; the bindings also drive the store + # access grant (step 4). memory_store, session_store = store_bindings(source_dir) with render.status("Checking stores…"): - env_updates = validate_stores_and_trace_env( - client, - app=name, - memory_store=memory_store, - session_store=session_store, - traces_destination=traces_destination, - traces_experiment=traces_experiment, - ) + validate_stores(client, memory_store=memory_store, session_store=session_store) + + env_updates: dict[str, str] = {} provisioned: dict[str, Any] = {} + if traced: + # Show the experiment id and a direct link to its MLflow traces page (not the raw name). + provisioned["Experiment"] = f"{trace_experiment_id} ({trace_schema})" + if trace_url: + provisioned["Traces"] = trace_url if memory_store: provisioned["Memory store"] = memory_store if session_store: provisioned["Session store"] = session_store - if traces_destination: - provisioned["Traces"] = traces_destination if pip_index_url: for env in _PIP_INDEX_ENVS: env_updates[env] = pip_index_url @@ -502,6 +540,9 @@ def deploy( "url": app_url, "workspace_path": ws_path, "env": env_updates, + "trace_experiment_id": trace_experiment_id, + "trace_location": trace_schema, + "trace_url": trace_url, "store_grant": "skipped" if not grants_stores else ("granted" if grant_error is None else "failed"), @@ -514,6 +555,16 @@ def deploy( (f"mason deployments get {name}", "Check its status and URL"), (f"mason deployments logs {name}", "Tail its logs"), ] + if traced: + # The UC experiment is created, but the app's SP needs write access to the schema's tables + # to actually log traces; granting it requires schema ownership, so surface it as a step. + steps.append( + f"Grant the app's service principal USE CATALOG + USE SCHEMA + MODIFY/SELECT on " + f"{trace_schema} so it can write traces." + ) + else: + # Tracing is opt-in and wasn't configured - deploy still succeeded; nudge, don't block. + steps.append(_TRACING_OFF_HINT) if app_url: steps.insert(0, (f"open {app_url}", "Open the deployed app")) if scaffolded: diff --git a/integrations/mason/src/databricks_mason/dev.py b/integrations/mason/src/databricks_mason/dev.py index f077da0d..0d5d7b27 100644 --- a/integrations/mason/src/databricks_mason/dev.py +++ b/integrations/mason/src/databricks_mason/dev.py @@ -17,12 +17,14 @@ from databricks_mason import render from databricks_mason.deploy import ( - _upsert_manifest_env, + _TRACING_OFF_HINT, + provision_trace_experiment, store_bindings, - validate_stores_and_trace_env, + validate_stores, ) from databricks_mason.errors import AgentCliError from databricks_mason.store_access import _databricks +from databricks_mason.tracing import experiment_ui_url, project_trace_location # Default local port; `databricks apps run-local` listens here unless --app-port overrides it. _DEFAULT_APP_PORT = 8000 @@ -48,25 +50,12 @@ "exists yet, and reuse it otherwise. Requires uv.", ) @click.option("--app-port", type=int, default=None, help="Port to run the app on (default 8000).") -@click.option( - "--with-traces", - "traces_destination", - default=None, - help="UC trace destination 'catalog.schema' to wire in via MLFLOW_TRACING_DESTINATION.", -) -@click.option( - "--traces-experiment", - default=None, - help="MLflow experiment path to wire in via MLFLOW_EXPERIMENT_NAME.", -) @click.pass_obj def dev( obj, source: str, prepare_environment: Optional[bool], app_port: Optional[int], - traces_destination: Optional[str], - traces_experiment: Optional[str], ) -> None: """Run a scaffolded agent locally from its app.yaml (wraps `databricks apps run-local`). @@ -75,10 +64,11 @@ def dev( ``mason deploy``. The environment is built on first run and reused after; pass ``--prepare-environment`` to force a rebuild (e.g. after changing dependencies). - Stores bound with ``mason memory/sessions bind`` are validated here and read from agent.toml at - runtime (not written to app.yaml); the ``--with-traces`` flag wires tracing env into app.yaml, - exactly as ``mason deploy`` does. Locally the store owner (you) already has access, so no - service-principal grant is needed here; that grant happens at ``mason deploy`` time. + Stores bound with ``mason memory/sessions bind`` and tracing configured with ``mason tracing + setup`` are both recorded in ``agent.toml`` and picked up here (dev validates the bound stores + exist and wires the trace experiment, exactly as ``mason deploy`` does). Locally the store owner + (you) already has access, so no service-principal grant is needed; that grant happens at ``mason + deploy`` time. """ source_dir = pathlib.Path(source) app_yaml = source_dir / "app.yaml" @@ -88,24 +78,28 @@ def dev( hint="Run from a scaffolded project, or pass --source (see `mason init`).", ) - # Validate the agent.toml store bindings exist and wire any traces into app.yaml. The stores are - # read from agent.toml at runtime, so they aren't written to app.yaml — set them (and create them) - # with `mason memory/sessions bind`. + # Stores are bound via `mason memory/sessions bind` (recorded in agent.toml) and read at runtime, + # so dev only validates that the bound stores still exist - a client/auth call made only when some + # are bound. Tracing is wired separately below (also read from agent.toml). memory_store, session_store = store_bindings(source_dir) - if memory_store or session_store or traces_destination or traces_experiment: - # The agent name defaults to the project dir name, so a per-app trace experiment here matches - # what `mason deploy ` derives. + if memory_store or session_store: with render.status("Checking stores…"): - env_updates = validate_stores_and_trace_env( - obj.client(), - app=source_dir.resolve().name, - memory_store=memory_store, - session_store=session_store, - traces_destination=traces_destination, - traces_experiment=traces_experiment, - ) - if env_updates: - _upsert_manifest_env(source_dir, env_updates) + validate_stores(obj.client(), memory_store=memory_store, session_store=session_store) + # Tracing is UC-only and opt-in: wire it into app.yaml iff `mason tracing setup` configured a + # catalog.schema for this project (the exact same path `mason deploy` takes). Check the config + # first via a cheap agent.toml read, so a plain unconfigured `mason dev` never constructs the + # workspace client or makes an auth/`me()` call - local iteration stays offline-friendly. When + # unconfigured, dev doesn't block; the startup panel nudges instead. + traced = None + trace_url = None + trace_schema, _ = project_trace_location(source) + if trace_schema: + client = obj.client() + traced = provision_trace_experiment( + source_dir, source_dir.resolve().name, client, obj.profile + ) + if traced: + trace_url = experiment_ui_url(client.host, traced[0]) # Default: prepare only when there's no venv yet, so repeat runs don't rebuild. Explicit # --prepare-environment / --no-prepare-environment overrides the auto-detect. @@ -127,7 +121,7 @@ def dev( # `run-local` prints a generic "go to http://localhost:" line that points at the chat UI — # misleading for an API-only project, which serves no page there (404). Print an accurate line up # front, keyed on whether this project actually carries the chat-app overlay. - _announce_local_url(source_dir, app_port or _DEFAULT_APP_PORT) + _announce_local_url(source_dir, app_port or _DEFAULT_APP_PORT, traced, trace_url) # Run in the project dir so run-local finds the app; stream output (no capture). _databricks( @@ -138,19 +132,38 @@ def dev( ) -def _announce_local_url(source_dir: pathlib.Path, port: int) -> None: - """Print how to reach the running app: the chat UI if present, else a sample invoke request.""" +def _announce_local_url( + source_dir: pathlib.Path, + port: int, + traced: Optional[tuple[str, str]], + trace_url: Optional[str], +) -> None: + """Print how to reach the running app: the chat UI if present, else a sample invoke request. + + Also states where traces go: the experiment id + a link to its MLflow traces page when tracing + is configured, or a one-line hint that tracing is off. + """ base = f"http://localhost:{port}" deploy_name = source_dir.resolve().name + # traced is (experiment_id, catalog_schema) when `mason tracing setup` ran, else None. + trace_field: dict[str, str] = {} + trace_step: list[str | tuple[str, str]] = [_TRACING_OFF_HINT] + if traced: + experiment_id, schema = traced + trace_field["Experiment"] = f"{experiment_id} ({schema})" + if trace_url: + trace_field["Traces"] = trace_url + trace_step = [] if (source_dir / "runtime" / "ui.py").is_file(): render.success( "Starting agent", - fields={"Chat UI": base}, + fields={"Chat UI": base, **trace_field}, next_steps=[ f"Open {base} to chat with your agent", ("mason tools add mcp ", "Give the agent a tool"), ("mason memory bind ", "Attach a memory / session store"), (f"mason deploy {deploy_name}", "Deploy it to Databricks"), + *trace_step, ], ) else: @@ -161,11 +174,12 @@ def _announce_local_url(source_dir: pathlib.Path, port: int) -> None: ) render.success( "Starting API-only agent (no chat UI — see `mason init --help`)", - fields={"Invoke": f"POST {base}/invocations"}, + fields={"Invoke": f"POST {base}/invocations", **trace_field}, next_steps=[ (sample, "Send a test request"), ("mason tools add mcp ", "Give the agent a tool"), (f"mason deploy {deploy_name}", "Deploy it to Databricks"), + *trace_step, ], ) diff --git a/integrations/mason/src/databricks_mason/help.py b/integrations/mason/src/databricks_mason/help.py index 894c82c7..96f1ba13 100644 --- a/integrations/mason/src/databricks_mason/help.py +++ b/integrations/mason/src/databricks_mason/help.py @@ -171,20 +171,23 @@ ), ), ("tracing",): ( - ("mason tracing setup --catalog main --schema agent_traces", "link a UC trace destination"), + ( + "mason tracing setup --trace-location main.agent_traces", + "configure where the agent's traces go", + ), ), ("tracing", "setup"): ( - ("mason tracing setup --catalog main --schema agent_traces", "link a UC trace destination"), + ( + "mason tracing setup --trace-location main.agent_traces --warehouse-id ", + "configure MLflow tracing to a UC schema", + ), ), ("tracing", "list"): ( - ("mason tracing list --experiment /Users/me/mason-traces/my-agent", "list recent traces"), + ("mason tracing list --trace-location main.agent_traces", "list recent traces"), ), ("tracing", "get"): (("mason tracing get ", "show one trace"),), - ("tracing", "instrument"): ( - ("mason tracing instrument --destination main.agent_traces", "print instrumentation code"), - ), ("deploy",): ( - ("mason deploy my-agent", "deploy the agent"), + ("mason deploy my-agent", "deploy the agent to Databricks Apps"), ("mason deploy my-agent --instances 2", "deploy with two instances"), ), ("deployments",): (("mason deployments list", "list agent deployments"),), diff --git a/integrations/mason/src/databricks_mason/render.py b/integrations/mason/src/databricks_mason/render.py index 764cf9f1..80239ad0 100644 --- a/integrations/mason/src/databricks_mason/render.py +++ b/integrations/mason/src/databricks_mason/render.py @@ -138,7 +138,9 @@ def _cell(value: Any) -> Any: return value if value is None: return Text("—", style=MUTED) - return str(value) + # Wrap in Text so values are shown literally: Rich treats "[...]" in a bare str as console + # markup, which would silently eat things like a "[dev]" experiment name. + return Text(str(value)) # --- detail view (aig-endpoint.png) ------------------------------------------ diff --git a/integrations/mason/src/databricks_mason/tracing.py b/integrations/mason/src/databricks_mason/tracing.py index 018dc668..5eef4c7b 100644 --- a/integrations/mason/src/databricks_mason/tracing.py +++ b/integrations/mason/src/databricks_mason/tracing.py @@ -1,19 +1,21 @@ -"""`mason tracing` — route an agent's traces to MLflow / Unity Catalog and inspect them. +"""`mason tracing` - configure where an agent's MLflow traces go, and inspect them. -Parallel to `mason memory` and `mason sessions`: `setup` provisions the trace destination -(links a UC schema to an MLflow experiment, the analog of creating a store), `list`/`get` -read traces back, and `instrument` prints the wiring snippet (the "Starter code" analog). -`mason deploy --with-traces` injects the destination into a deployment's app.yaml, exactly as -`--memory` / `--session` inject their stores. +Tracing is Unity Catalog-only and opt-in. ``mason tracing setup --trace-location `` +records a UC schema in ``agent.toml``; from then on both ``mason dev`` and ``mason deploy`` send the +agent's traces there (creating a per-app UC-linked experiment that surfaces them in the MLflow UI). +Until setup is run, tracing is simply off - neither ``dev`` nor ``deploy`` blocks on it, so a +developer without a catalog/schema is never stuck. -MLflow is an optional dependency: `setup`/`list`/`get` need `mlflow[databricks]>=3.10.1` -installed and lazily import it; `instrument` (and the deploy wiring) are pure and need nothing. +``list`` / ``get`` read traces back. MLflow is an optional dependency: ``list`` / ``get`` and the +UC experiment provisioning need ``mlflow[databricks]>=3.10.1`` and import it lazily; ``setup`` is +pure and needs nothing. """ from __future__ import annotations import os import pathlib +import re from typing import Any, Optional import click @@ -22,28 +24,47 @@ from databricks_mason.errors import AgentCliError _BREADCRUMB = "Agent Tracing" -# Fallback experiment when there's no agent context (e.g. `tracing setup` with no --app). Prefer a -# per-agent experiment via `default_experiment` so each agent's traces are isolated, not commingled. -_DEFAULT_EXPERIMENT = "/Shared/mason-agent-traces" + +# MLflow's own env vars the agent reads at runtime (wired into app.yaml by dev/deploy). +TRACES_TRACKING_URI_ENV = "MLFLOW_TRACKING_URI" +TRACES_EXPERIMENT_ENV = "MLFLOW_EXPERIMENT_NAME" +TRACES_WAREHOUSE_ENV = "MLFLOW_TRACING_SQL_WAREHOUSE_ID" +TRACES_DESTINATION_ENV = "MLFLOW_TRACING_DESTINATION" + +# Per-user workspace folder that holds each app's UC-linked tracing experiment. +_TRACES_DIR = "mason-traces" + +# A UC trace location is "catalog.schema" (no dots-in-names, no path separators, exactly one dot). +_UC_SCHEMA = re.compile(r"^[^.\s/]+\.[^.\s/]+$") + +_MIGRATE_DOCS = "https://docs.databricks.com/aws/en/mlflow3/genai/tracing/migrate-traces-to-uc" + + +def experiment_name(user: str, app: str) -> str: + """The per-app UC-linked experiment that surfaces its traces (shared by dev and deploy).""" + return f"/Users/{user}/{_TRACES_DIR}/{app}" -def default_experiment(user: str, app: Optional[str]) -> str: - """The per-agent experiment path for `app`, under the user's workspace home. +def experiment_ui_url(host: Optional[str], experiment_id: str) -> Optional[str]: + """The workspace MLflow experiment traces page for ``experiment_id`` (from the profile's host). - Keeps each agent's traces in their own experiment (and out of other users' view), instead of the - single shared bucket. `mason tracing setup`, `dev`, and `deploy` all derive the same path for a - given agent, so the experiment `setup` links is the one the agent logs to. Falls back to the - shared default only when no app name is available. + ``host`` is the workspace URL (``MasonClient.host`` / the profile's config host); returns None if + it's unavailable so callers can just omit the link. """ - if not app: - return _DEFAULT_EXPERIMENT - return f"/Users/{user}/mason-traces/{app}" + if not host or host == "unknown": + return None + return f"{host.rstrip('/')}/ml/experiments/{experiment_id}/traces" -# Env vars the deployed agent reads (see deploy.py). MLFLOW_TRACING_DESTINATION is MLflow's own -# "catalog.schema" convention; MLFLOW_EXPERIMENT_NAME is the standard MLflow experiment selector. -TRACES_DEST_ENV = "MLFLOW_TRACING_DESTINATION" -TRACES_EXPERIMENT_ENV = "MLFLOW_EXPERIMENT_NAME" +def validate_uc_schema(location: str) -> str: + """Accept a Unity Catalog 'catalog.schema'; reject anything else.""" + loc = location.strip() + if _UC_SCHEMA.match(loc): + return loc + raise AgentCliError( + f"Invalid trace location {location!r}.", + hint="Use a Unity Catalog schema in 'catalog.schema' form.", + ) # Installing the `tracing` extra (rather than a bare mlflow) is what actually resolves both the @@ -66,20 +87,12 @@ def _mlflow(): def _uc_trace_symbols(): - """Import the version-specific UC-tracing symbols, surfacing the same clean error as `_mlflow`. - - `import mlflow` succeeding doesn't guarantee these exist — they were added in the tracing API - this feature needs. Guard them so an older installed MLflow yields Mason's install hint rather - than a raw ImportError traceback. - """ + """Import the version-specific UC trace-location symbols, with a clean install hint.""" try: - from mlflow.entities import UCSchemaLocation # noqa: PLC0415 - lazy, version-specific - from mlflow.tracing import ( # noqa: PLC0415 - set_experiment_trace_location, - unset_experiment_trace_location, - ) + from mlflow.entities import UCSchemaLocation # noqa: PLC0415 - version-specific + from mlflow.tracing import set_experiment_trace_location # noqa: PLC0415 - return UCSchemaLocation, set_experiment_trace_location, unset_experiment_trace_location + return UCSchemaLocation, set_experiment_trace_location except ImportError as exc: raise AgentCliError( "This MLflow version is too old for `mason tracing setup` (UC trace destinations).", @@ -87,43 +100,87 @@ def _uc_trace_symbols(): ) from exc -def _link_trace_location(set_location, unset_location, location, exp_id: str, relink: bool) -> None: - """Link the experiment to the UC schema, handling the already-linked case. +def _configure(mlflow, profile: Optional[str], warehouse_id: Optional[str]) -> None: + """Point MLflow at the workspace (honoring mason's --profile), with a warehouse for UC ops.""" + mlflow.set_tracking_uri(f"databricks://{profile}" if profile else "databricks") + if warehouse_id: + os.environ[TRACES_WAREHOUSE_ENV] = warehouse_id + + +def _read_warehouse(explicit: Optional[str], configured: Optional[str]) -> str: + """Resolve the SQL warehouse for a UC trace *read* (`get`/`list`), or raise - reads require one. - MLflow raises if the experiment is already bound to a storage location. With `relink`, unset the - existing binding first and re-link; otherwise surface a clean error pointing at `--relink`. + Unlike dev/deploy provisioning (whose create-location call can fall back to a workspace default), + MLflow's trace-read path has no default fallback, so a warehouse must come from ``--warehouse-id``, + the project's ``agent.toml`` (`mason tracing setup`), or MLFLOW_TRACING_SQL_WAREHOUSE_ID. """ + warehouse = explicit or configured or os.getenv(TRACES_WAREHOUSE_ENV) + if not warehouse: + raise AgentCliError( + "Reading UC traces needs a SQL warehouse.", + hint="Configure one with `mason tracing setup --warehouse-id `, pass --warehouse-id, " + "or set MLFLOW_TRACING_SQL_WAREHOUSE_ID.", + ) + return warehouse + + +def ensure_uc_experiment( + profile: Optional[str], experiment_name: str, catalog_schema: str, warehouse_id: Optional[str] +) -> str: + """Create ``experiment_name`` if missing and link it to the UC ``catalog.schema`` (idempotent). + + A UC destination can only be linked to an experiment with no existing traces, so this links a + freshly created experiment; a re-deploy (experiment already linked) is a no-op, and an experiment + that already holds non-UC traces raises a clear error pointing at the migration docs. Returns the + experiment **id** (used to build the MLflow experiment UI link shown by dev/deploy). + + Creating the UC trace tables needs a SQL warehouse: ``warehouse_id`` when configured, else MLflow + falls back to the workspace default - which may not exist, so `mason tracing setup` keeps + `--warehouse-id` available for workspaces without one. + """ + mlflow = _mlflow() + _configure(mlflow, profile, warehouse_id) + catalog, _, schema = catalog_schema.partition(".") + uc_schema_location, set_location = _uc_trace_symbols() + + experiment = mlflow.get_experiment_by_name(experiment_name) + experiment_id = ( + experiment.experiment_id if experiment else mlflow.create_experiment(experiment_name) + ) try: - set_location(location=location, experiment_id=exp_id) - except Exception as exc: # noqa: BLE001 - mlflow raises a generic error when already linked - if "already" not in str(exc).lower(): - raise - if not relink: + set_location( + location=uc_schema_location(catalog_name=catalog, schema_name=schema), + experiment_id=experiment_id, + ) + except Exception as exc: # noqa: BLE001 - mlflow raises a generic error; classify by message + text = str(exc).lower() + if "contains traces" in text: raise AgentCliError( - "This experiment is already linked to a trace storage location.", - hint="Re-run with --relink to replace the existing link.", + f"Experiment {experiment_name!r} already has non-UC traces, so it can't be linked " + "to Unity Catalog.", + hint=f"Migrate the existing traces to UC: {_MIGRATE_DOCS}", ) from exc - unset_location(location=location, experiment_id=exp_id) - set_location(location=location, experiment_id=exp_id) + if "already" in text: + return experiment_id # already linked (re-deploy) - idempotent + raise AgentCliError( + f"Could not link {experiment_name!r} to {catalog_schema}: {exc}" + ) from exc + return experiment_id -def _configure(mlflow, profile: Optional[str], warehouse_id: Optional[str]) -> None: - """Point MLflow at the workspace (honoring mason's --profile) for UC-backed tracing.""" - mlflow.set_tracking_uri(f"databricks://{profile}" if profile else "databricks") - if warehouse_id: - os.environ["MLFLOW_TRACING_SQL_WAREHOUSE_ID"] = warehouse_id +def project_trace_location(source: str) -> tuple[Optional[str], Optional[str]]: + """The (UC schema, warehouse id) bound in the project's agent.toml, or (None, None). + Cheap (a small TOML read, no workspace call), so callers can decide whether tracing is even + configured before touching the client - `mason dev` uses it to stay auth-free when it isn't. + """ + from databricks_mason.agent_project import AgentProject # noqa: PLC0415 - avoid import cycle -def _ensure_experiment(mlflow, client, name: str) -> str: - experiment = mlflow.get_experiment_by_name(name) - if experiment: - return experiment.experiment_id - # create_experiment won't make the intermediate workspace folder for a nested path (e.g. - # /Users//mason-traces/), so create the parent dir first. - parent = name.rsplit("/", 1)[0] - if parent: - client.ensure_workspace_dir(parent) - return mlflow.create_experiment(name) + try: + project = AgentProject.load(source) + except AgentCliError: + return None, None + return project.trace_location, project.trace_warehouse def _attr(obj: Any, *paths: str, default: Any = None) -> Any: @@ -145,16 +202,13 @@ def _status_str(status: Any) -> Optional[str]: return getattr(status, "name", None) or str(status) -def _split_destination(destination: str) -> tuple[str, str]: - catalog, _, schema = destination.partition(".") - if not catalog or not schema: - raise AgentCliError("--destination must be 'catalog.schema'.") - return catalog, schema - - -def _experiment_url(host: str, experiment_id: str) -> str: - """Workspace URL for an experiment's Traces tab.""" - return f"{host.rstrip('/')}/ml/experiments/{experiment_id}?compareRunsMode=TRACES" +def _trace_json(trace: Any) -> dict: + return { + "trace_id": _attr(trace, "info.trace_id", "info.request_id"), + "status": _status_str(_attr(trace, "info.status", "info.state")), + "execution_time_ms": _attr(trace, "info.execution_time_ms", "info.execution_duration_ms"), + "timestamp_ms": _attr(trace, "info.timestamp_ms", "info.request_time"), + } # --- group ------------------------------------------------------------------ @@ -162,88 +216,61 @@ def _experiment_url(host: str, experiment_id: str) -> str: @click.group() def tracing() -> None: - """Set up and inspect MLflow traces (in Unity Catalog) for your agents.""" + """Configure where your agent's MLflow traces go, and inspect them.""" -# --- setup: provision the UC trace destination ------------------------------ +# --- setup: record the UC trace destination --------------------------------- @tracing.command("setup") -@click.option("--catalog", required=True, help="Unity Catalog catalog to store traces in.") -@click.option("--schema", required=True, help="Unity Catalog schema to store traces in.") @click.option( - "--app", - default=None, - help="Agent name — gives this agent its own experiment (/Users//mason-traces/) so its " - "traces aren't commingled with other agents'. Defaults to the current directory name (matching " - "`mason dev`/`deploy`). Overridden by --experiment.", -) -@click.option( - "--experiment", - default=None, - help="MLflow experiment path (overrides the per-app default derived from --app).", + "--trace-location", + "trace_location", + required=True, + help="Unity Catalog schema 'catalog.schema' where agent traces are stored.", ) @click.option( "--warehouse-id", default=None, - help="SQL warehouse id for trace queries (MLFLOW_TRACING_SQL_WAREHOUSE_ID).", + help="SQL warehouse id for provisioning the MLflow tracing experiment. Optional if " + "MLFLOW_TRACING_SQL_WAREHOUSE_ID env variable is set or if the workspace has a default " + "warehouse.", ) @click.option( - "--relink", - is_flag=True, - help="Replace an existing trace-location link on the experiment (unset, then re-link).", + "--source", + default=".", + type=click.Path(exists=True, file_okay=False), + help="Project directory containing agent.toml. Defaults to the current directory.", ) @click.pass_obj -def tracing_setup(obj, catalog, schema, app, experiment, warehouse_id, relink) -> None: - """Link a UC schema to an MLflow experiment so agent traces land in Unity Catalog.""" - mlflow = _mlflow() - _configure(mlflow, obj.profile, warehouse_id) - # Default the agent name to the current directory, matching `mason dev`/`deploy`, so running - # setup from a project dir needs no --app and still lands in that agent's own experiment. - app = app or pathlib.Path.cwd().name - client = obj.client() - exp_name = experiment or default_experiment(client.current_user, app) - exp_id = _ensure_experiment(mlflow, client, exp_name) - - UCSchemaLocation, set_location, unset_location = _uc_trace_symbols() - _link_trace_location( - set_location, - unset_location, - UCSchemaLocation(catalog_name=catalog, schema_name=schema), - exp_id, - relink, - ) - destination = f"{catalog}.{schema}" - url = _experiment_url(client.host, exp_id) +def tracing_setup(obj, trace_location, warehouse_id, source) -> None: + """Configure tracing via MLflow. + + Records the ``catalog.schema`` (and optional warehouse). From then on both ``mason dev`` and + ``mason deploy`` send the agent's traces to that schema, creating a per-app UC-linked experiment + there. Until this is run, tracing is off (neither command blocks on it). + """ + from databricks_mason.agent_project import AgentProject # noqa: PLC0415 + + location = validate_uc_schema(trace_location) + project = AgentProject.load(pathlib.Path(source)) + project.bind_trace_location(location, warehouse_id) + project.write() if obj.output == "json": - render.emit_json( - { - "experiment": exp_name, - "experiment_id": exp_id, - "destination": destination, - "url": url, - } - ) + render.emit_json({"trace_location": location, "warehouse_id": warehouse_id}) return - # dev/deploy derive this same experiment from the agent name (dev: project dir, deploy: ), - # so only spell out --traces-experiment when the experiment was set explicitly (not per-app). - exp_flag = "" if experiment is None else f" --traces-experiment {exp_name}" - deploy_name = app + fields = {"Trace location": location} + if warehouse_id: + fields["SQL warehouse"] = warehouse_id + # `list` needs a warehouse; suggest passing one only when none was just configured. + list_cmd = "mason tracing list" if warehouse_id else "mason tracing list --warehouse-id " render.success( - f"Linked traces for '{exp_name}' to {destination}", - fields={"Experiment": exp_name, "Destination": destination, "View traces": url}, + f"Configured tracing for '{location}'", + fields=fields, next_steps=[ - (f"mason dev --with-traces {destination}{exp_flag}", "Run locally with tracing on"), - ( - f"mason deploy {deploy_name} --with-traces {destination}{exp_flag}", - "Deploy with tracing on", - ), - ( - f"mason tracing instrument --destination {destination}", - "Print the code snippet to trace your own agent", - ), - (f"mason tracing list --experiment {exp_name}", "List traces once you have some"), + ("mason deploy ", "Create the UC-linked experiment and deploy"), + (list_cmd, "Read traces at this location"), ], ) @@ -253,24 +280,48 @@ def tracing_setup(obj, catalog, schema, app, experiment, warehouse_id, relink) - @tracing.command("list") @click.option( - "--experiment", default=None, help=f"MLflow experiment path (default: {_DEFAULT_EXPERIMENT})." + "--trace-location", + "trace_location", + default=None, + help="Trace location to read: a UC 'catalog.schema', or an experiment id/path. Defaults to the " + "project's configured UC schema (from `mason tracing setup`).", +) +@click.option( + "--warehouse-id", + default=None, + help="SQL warehouse id for reading UC-backed traces. Resolved from this flag, the project's " + "configured warehouse (`mason tracing setup`), or MLFLOW_TRACING_SQL_WAREHOUSE_ID; required - " + "MLflow has no workspace-default fallback for reads.", ) @click.option("--limit", type=int, default=20) +@click.option( + "--source", + default=".", + type=click.Path(file_okay=False), + help="Project directory to resolve the default trace location from (default: current dir).", +) @click.pass_obj -def tracing_list(obj, experiment, limit) -> None: - """List recent agent traces in an experiment.""" - mlflow = _mlflow() - _configure(mlflow, obj.profile, None) - exp_name = experiment or _DEFAULT_EXPERIMENT - # search_traces selects experiments by id, not name, so resolve first. A missing experiment means - # no traces have been recorded there yet — show an empty list rather than erroring. - exp = mlflow.get_experiment_by_name(exp_name) - traces = ( - mlflow.search_traces( - experiment_ids=[exp.experiment_id], max_results=limit, return_type="list" +def tracing_list(obj, trace_location, warehouse_id, limit, source) -> None: + """List recent agent traces at a trace location. + + Resolution order: ``--trace-location`` (works standalone), else the project's configured UC + schema (``mason tracing setup``). Tracing is UC-only, so a UC schema is queried through a SQL + warehouse. + """ + configured_location, configured_warehouse = project_trace_location(source) + location = trace_location or configured_location + if not location: + raise AgentCliError( + "No trace location configured for this project.", + hint="Run `mason tracing setup --trace-location ` first, or pass " + "--trace-location.", ) - if exp - else [] + + warehouse = _read_warehouse(warehouse_id, configured_warehouse) + mlflow = _mlflow() + _configure(mlflow, obj.profile, warehouse) + traces = mlflow.search_traces( + locations=[_resolve_location(mlflow, location)], max_results=limit, return_type="list" ) if obj.output == "json": @@ -286,20 +337,56 @@ def tracing_list(obj, experiment, limit) -> None: for t in traces ] render.resource_table( - f"Agent Traces · {exp_name}", + f"Agent Traces · {location}", [("Trace ID", "left"), ("Status", "left"), ("Latency (ms)", "left"), ("Created", "left")], rows, ) +def _resolve_location(mlflow, location: str) -> str: + """Turn a location spec into what search_traces wants: a UC schema or an experiment id. + + ``catalog.schema`` and bare numeric ids pass through; an experiment path (``/Users/...``) is + resolved to its id. + """ + if _UC_SCHEMA.match(location) or location.isdigit(): + return location + experiment = mlflow.get_experiment_by_name(location) + if experiment is None: + raise AgentCliError(f"No experiment found at {location!r}.") + return experiment.experiment_id + + @tracing.command("get") @click.argument("trace_id") +@click.option( + "--warehouse-id", + default=None, + help="SQL warehouse id for reading UC-backed traces. Resolved from this flag, the project's " + "configured warehouse (`mason tracing setup`), or MLFLOW_TRACING_SQL_WAREHOUSE_ID; required - " + "MLflow has no workspace-default fallback for reads.", +) +@click.option( + "--source", + default=".", + type=click.Path(file_okay=False), + help="Project directory to read the configured warehouse from (default: current dir).", +) @click.pass_obj -def tracing_get(obj, trace_id) -> None: - """Get a single trace by id (status, latency, span count, previews).""" +def tracing_get(obj, trace_id, warehouse_id, source) -> None: + """Get a single trace by id (status, latency, span count, previews). + + Needs only the id: a v4 trace id (``trace://``) is self-locating. Tracing is + UC-only, so the trace is read through a SQL warehouse - the project's configured one + (``mason tracing setup``), ``--warehouse-id``, or MLFLOW_TRACING_SQL_WAREHOUSE_ID. + """ + _, configured_warehouse = project_trace_location(source) + warehouse = _read_warehouse(warehouse_id, configured_warehouse) mlflow = _mlflow() - _configure(mlflow, obj.profile, None) + _configure(mlflow, obj.profile, warehouse) trace = mlflow.get_trace(trace_id) + if trace is None: + raise AgentCliError(f"No trace found with id {trace_id!r}.") if obj.output == "json": render.emit_json(_trace_json(trace)) return @@ -317,51 +404,3 @@ def tracing_get(obj, trace_id) -> None: }, status=_status_str(_attr(trace, "info.status", "info.state")), ) - - -def _trace_json(trace: Any) -> dict: - return { - "trace_id": _attr(trace, "info.trace_id", "info.request_id"), - "status": _status_str(_attr(trace, "info.status", "info.state")), - "execution_time_ms": _attr(trace, "info.execution_time_ms", "info.execution_duration_ms"), - "timestamp_ms": _attr(trace, "info.timestamp_ms", "info.request_time"), - } - - -# --- instrument: print the agent wiring snippet (no MLflow needed) ---------- - - -@tracing.command("instrument") -@click.option( - "--destination", - default=None, - help="UC trace destination 'catalog.schema' (from `mason tracing setup`).", -) -@click.option( - "--experiment", default=None, help=f"MLflow experiment path (default: {_DEFAULT_EXPERIMENT})." -) -@click.pass_obj -def tracing_instrument(obj, destination, experiment) -> None: - """Print the snippet that routes an OpenAI Agents SDK agent's traces to UC.""" - catalog, schema = _split_destination(destination) if destination else ("", "") - exp_name = experiment or _DEFAULT_EXPERIMENT - dest = destination or f"{catalog}.{schema}" - code = ( - "import mlflow\n" - "from mlflow.entities import UCSchemaLocation\n\n" - 'mlflow.set_tracking_uri("databricks")\n' - f'mlflow.set_experiment("{exp_name}")\n' - f'mlflow.tracing.set_destination(UCSchemaLocation(catalog_name="{catalog}", schema_name="{schema}"))\n' - "mlflow.openai.autolog() # OpenAI Agents SDK spans -> Unity Catalog traces\n" - "# NOTE: do NOT call agents.set_tracing_disabled(True) — that turns tracing off." - ) - if obj.output == "json": - render.emit_json({"destination": dest, "experiment": exp_name, "snippet": code}) - return - render.detail( - _BREADCRUMB, - dest, - {"Experiment": exp_name, "Destination": dest, "Requires": "mlflow[databricks]>=3.10.1"}, - status="ACTIVE", - snippets=[("python", "python", code)], - ) diff --git a/integrations/mason/tests/unit_tests/deploy_bugfix_test.py b/integrations/mason/tests/unit_tests/deploy_bugfix_test.py index 8ce02839..321560e9 100644 --- a/integrations/mason/tests/unit_tests/deploy_bugfix_test.py +++ b/integrations/mason/tests/unit_tests/deploy_bugfix_test.py @@ -110,13 +110,10 @@ def test_validate_stores_raises_when_session_store_missing(): "session store not found", error_code="NOT_FOUND" ) with pytest.raises(AgentCliError) as exc: - deploy_mod.validate_stores_and_trace_env( + deploy_mod.validate_stores( client, - app="a", memory_store=None, session_store="ghost", - traces_destination=None, - traces_experiment=None, ) assert "does not exist" in str(exc.value) client.get_session_store.assert_called_once_with("ghost") diff --git a/integrations/mason/tests/unit_tests/deploy_test.py b/integrations/mason/tests/unit_tests/deploy_test.py index 2b0667c4..b8720d1c 100644 --- a/integrations/mason/tests/unit_tests/deploy_test.py +++ b/integrations/mason/tests/unit_tests/deploy_test.py @@ -7,12 +7,33 @@ import types from unittest import mock +import pytest import yaml from click.testing import CliRunner from databricks_mason import deploy as deploy_mod from databricks_mason.errors import AgentCliError +# Deploy wires UC tracing when configured, creating a UC experiment (a live mlflow/workspace op). +# Tests that aren't about tracing bypass that provisioning via this autouse fixture; the +# tracing-specific tests below are listed here so the real `provision_trace_experiment` runs. +_TRACING_TESTS = { + "test_deploy_provisions_uc_experiment_when_configured", + "test_provision_trace_experiment_returns_none_without_uc", + "test_provision_trace_experiment_creates_and_wires", +} + + +@pytest.fixture(autouse=True) +def _bypass_tracing_provision(request, monkeypatch): + if request.function.__name__ in _TRACING_TESTS: + return + monkeypatch.setattr( + deploy_mod, + "provision_trace_experiment", + lambda source, app, user, profile: ("/Users/me@example.com/mason-traces/app", "cat.schema"), + ) + def test_upsert_manifest_env_scaffolds_when_missing(tmp_path: pathlib.Path): scaffolded = deploy_mod._upsert_manifest_env( @@ -69,6 +90,9 @@ class _FakeClient: host = "https://ws" current_user = "me@example.com" + def ensure_workspace_dir(self, path): + pass + def get_memory_store(self, name): return {"name": f"memory-stores/{name}"} @@ -401,10 +425,34 @@ def test_deploy_does_not_write_store_env_to_app_yaml(tmp_path: pathlib.Path, mon assert "AGENT_SESSION_ACTOR_ID" not in env -def test_deploy_with_traces_injects_tracing_env(tmp_path: pathlib.Path, monkeypatch): +_AGENT_TOML_WITH_TRACES = 'schema_version = 1\n\n[agent]\nframework = "langgraph"\n\n[trace_location]\nname = "cat.schema"\n' +_AGENT_TOML_NO_TRACES = 'schema_version = 1\n\n[agent]\nframework = "langgraph"\n' + + +def test_deploy_proceeds_without_tracing_when_unconfigured(tmp_path: pathlib.Path, monkeypatch): + # Tracing is opt-in and UC-only: an unconfigured project deploys fine (no gate) - it just gets a + # one-line hint instead of being blocked. src = tmp_path / "app" src.mkdir() (src / "app.yaml").write_text(yaml.safe_dump({"command": ["x"]})) + (src / "agent.toml").write_text(_AGENT_TOML_NO_TRACES) + # Unconfigured tracing -> provision returns None (override the autouse tuple bypass). + monkeypatch.setattr(deploy_mod, "provision_trace_experiment", lambda *a: None) + + result = _run_deploy(src, monkeypatch, []) + + assert result.exit_code == 0, result.output + assert "mason tracing setup" in result.output # soft hint, not a block + env = {e["name"]: e["value"] for e in yaml.safe_load((src / "app.yaml").read_text())["env"]} + assert not any(k.startswith("MLFLOW") for k in env) # nothing tracing-related wired + + +def test_deploy_provisions_uc_experiment_when_configured(tmp_path: pathlib.Path, monkeypatch): + # With UC configured, deploy creates+links the UC experiment and wires it into app.yaml. + src = tmp_path / "app" + src.mkdir() + (src / "app.yaml").write_text(yaml.safe_dump({"command": ["x"]})) + (src / "agent.toml").write_text(_AGENT_TOML_WITH_TRACES) monkeypatch.setattr(deploy_mod, "_deployment_exists", lambda a, p: True) monkeypatch.setattr( @@ -412,26 +460,26 @@ def test_deploy_with_traces_injects_tracing_env(tmp_path: pathlib.Path, monkeypa "_databricks", lambda args, profile, **kw: types.SimpleNamespace(returncode=0, stdout="", stderr=""), ) - - result = CliRunner().invoke( - deploy_mod.deploy, - [ - "myapp", - "--source", - str(src), - "--with-traces", - "cat.schema", - "--traces-experiment", - "/Shared/x", - ], - obj=_FakeCtx(), + linked: dict = {} + monkeypatch.setattr( + deploy_mod, + "ensure_uc_experiment", + lambda profile, experiment, schema, warehouse: linked.update( + experiment=experiment, schema=schema + ) + or experiment, ) + result = CliRunner().invoke(deploy_mod.deploy, ["myapp", "--source", str(src)], obj=_FakeCtx()) + assert result.exit_code == 0, result.output - doc = yaml.safe_load((src / "app.yaml").read_text()) - env = {e["name"]: e["value"] for e in doc["env"]} + assert linked["schema"] == "cat.schema" + env = {e["name"]: e["value"] for e in yaml.safe_load((src / "app.yaml").read_text())["env"]} + assert env["MLFLOW_TRACKING_URI"] == "databricks" + # Both wired: the experiment name (the agent runtime's tracing enable-gate keys on it) and the + # destination pinned to the UC schema (routes to governed storage + blocks ambient OTLP hijack). + assert env["MLFLOW_EXPERIMENT_NAME"] == linked["experiment"] assert env["MLFLOW_TRACING_DESTINATION"] == "cat.schema" - assert env["MLFLOW_EXPERIMENT_NAME"] == "/Shared/x" def test_resolve_memory_store_pages_at_100_and_matches_display_name(): @@ -530,32 +578,47 @@ def client(self): assert "mason memory bind ghost" in result.output -def test_with_traces_defaults_the_experiment_per_app(): - # --with-traces alone must still set the experiment, or the agent ships tracing half-configured - # (destination set, experiment missing) and silently disables it. The default is per-app, so - # each agent's traces are isolated instead of piling into one shared experiment. - env = deploy_mod.validate_stores_and_trace_env( - _FakeClient(), - app="my-agent", - memory_store=None, - session_store=None, - traces_destination="cat.schema", - traces_experiment=None, - ) - assert env["MLFLOW_TRACING_DESTINATION"] == "cat.schema" - assert env["MLFLOW_EXPERIMENT_NAME"] == "/Users/me@example.com/mason-traces/my-agent" +def test_provision_trace_experiment_returns_none_without_uc(tmp_path: pathlib.Path): + # No UC schema configured -> no-op (returns None), so the caller can proceed without tracing. + (tmp_path / "app.yaml").write_text(yaml.safe_dump({"command": ["x"]})) + (tmp_path / "agent.toml").write_text(_AGENT_TOML_NO_TRACES) + client = mock.Mock(current_user="me@example.com") + assert deploy_mod.provision_trace_experiment(tmp_path, "app", client, None) is None + client.ensure_workspace_dir.assert_not_called() # unconfigured -> no workspace side effects -def test_with_traces_explicit_experiment_wins_over_per_app(): - env = deploy_mod.validate_stores_and_trace_env( - _FakeClient(), - app="my-agent", - memory_store=None, - session_store=None, - traces_destination="cat.schema", - traces_experiment="/Shared/custom", + +def test_provision_trace_experiment_creates_and_wires(tmp_path: pathlib.Path, monkeypatch): + # With UC configured: create+link the per-agent experiment and wire it into app.yaml. + (tmp_path / "app.yaml").write_text(yaml.safe_dump({"command": ["x"]})) + (tmp_path / "agent.toml").write_text(_AGENT_TOML_WITH_TRACES) + calls: dict = {} + monkeypatch.setattr( + deploy_mod, + "ensure_uc_experiment", + lambda profile, experiment, schema, warehouse: calls.update( + experiment=experiment, schema=schema + ) + or experiment, ) - assert env["MLFLOW_EXPERIMENT_NAME"] == "/Shared/custom" + + client = mock.Mock(current_user="me@example.com") + result = deploy_mod.provision_trace_experiment(tmp_path, "myagent", client, None) + assert result is not None # configured -> wires and returns (experiment, schema) + experiment, schema = result + + assert schema == "cat.schema" + assert experiment == "/Users/me@example.com/mason-traces/myagent" + assert calls["schema"] == "cat.schema" + # the parent workspace folder is created before the (nested) experiment, else it would fail + client.ensure_workspace_dir.assert_called_once_with("/Users/me@example.com/mason-traces") + env = { + e["name"]: e["value"] for e in yaml.safe_load((tmp_path / "app.yaml").read_text())["env"] + } + # Both the experiment name (runtime enable-gate) and the UC-schema destination pin are wired. + assert env["MLFLOW_EXPERIMENT_NAME"] == experiment + assert env["MLFLOW_TRACING_DESTINATION"] == "cat.schema" + assert env["MLFLOW_TRACKING_URI"] == "databricks" def _run_deploy(src, monkeypatch, extra_args): diff --git a/integrations/mason/tests/unit_tests/dev_test.py b/integrations/mason/tests/unit_tests/dev_test.py index 81927216..98080d87 100644 --- a/integrations/mason/tests/unit_tests/dev_test.py +++ b/integrations/mason/tests/unit_tests/dev_test.py @@ -14,9 +14,13 @@ class _Ctx: def __init__(self, output: str = "text", profile=None): self.output = output self.profile = profile + self.client_calls = 0 # so tests can assert dev stays auth-free when nothing needs it - def client(self): # only used when --with-* flags are passed - return mock.Mock() + def client(self): + # Constructing the client / reading current_user is a workspace (auth) call; dev must only + # do it when tracing or a store is actually configured. + self.client_calls += 1 + return mock.Mock(current_user="me@example.com") def test_dev_prepares_when_no_venv(tmp_path: pathlib.Path): @@ -90,7 +94,9 @@ def test_dev_filters_build_index_env_via_entry_point(tmp_path: pathlib.Path): dev_yaml = tmp_path / ".mason-dev.app.yaml" assert str(dev_yaml) in cmd names = {e["name"] for e in yaml.safe_load(dev_yaml.read_text())["env"]} - assert names == {"AGENT_SESSION_STORE"} # index vars stripped, app env kept + assert "PIP_INDEX_URL" not in names and "UV_INDEX_URL" not in names # index vars stripped + assert "AGENT_SESSION_STORE" in names # app env kept + assert not any(n.startswith("MLFLOW") for n in names) # tracing not configured -> not wired def test_dev_no_entry_point_when_no_index_override(tmp_path: pathlib.Path): @@ -118,7 +124,7 @@ def test_dev_validates_bound_stores_without_writing_store_env(tmp_path: pathlib. (tmp_path / ".venv").mkdir() with ( mock.patch.object(dev_mod, "_databricks") as db, - mock.patch.object(dev_mod, "validate_stores_and_trace_env", return_value={}) as validate, + mock.patch.object(dev_mod, "validate_stores") as validate, ): result = CliRunner().invoke(dev_mod.dev, ["--source", str(tmp_path)], obj=_Ctx()) assert result.exit_code == 0, result.output @@ -129,17 +135,51 @@ def test_dev_validates_bound_stores_without_writing_store_env(tmp_path: pathlib. assert db.call_args.args[0][:2] == ["apps", "run-local"] -def test_dev_without_bindings_does_not_validate(tmp_path: pathlib.Path): - # No agent.toml store bindings (and no --with-* traces) -> nothing to validate. +def test_dev_without_bindings_stays_offline(tmp_path: pathlib.Path): + # No store bindings and no `mason tracing setup`: dev validates nothing, wires no MLFLOW/AGENT + # env, and never touches the workspace client (no auth/`me()` call), so offline local dev works. (tmp_path / "app.yaml").write_text("command: []\n") (tmp_path / ".venv").mkdir() + ctx = _Ctx() with ( mock.patch.object(dev_mod, "_databricks"), - mock.patch.object(dev_mod, "validate_stores_and_trace_env") as validate, + mock.patch.object(dev_mod, "validate_stores") as validate, + ): + result = CliRunner().invoke(dev_mod.dev, ["--source", str(tmp_path)], obj=ctx) + assert result.exit_code == 0, result.output + validate.assert_not_called() # no bindings -> nothing to validate + assert ctx.client_calls == 0 # unconfigured dev makes no workspace/auth call + import yaml + + doc = yaml.safe_load((tmp_path / "app.yaml").read_text()) or {} + env = {e["name"]: e["value"] for e in (doc.get("env") or [])} + assert not any(k.startswith("MLFLOW") for k in env) + assert not any(k.startswith("AGENT_") for k in env) + + +def test_dev_wires_tracing_through_shared_provision(tmp_path: pathlib.Path): + # When tracing IS configured (agent.toml trace_location), dev wires it via the exact same + # provision_trace_experiment path as deploy (the actual UC work is covered by deploy_test). + (tmp_path / "app.yaml").write_text("command: []\n") + (tmp_path / ".venv").mkdir() + (tmp_path / "agent.toml").write_text( + 'schema_version = 1\n\n[agent]\nframework = "openai"\n\n[trace_location]\nname = "cat.schema"\n' + ) + with ( + mock.patch.object(dev_mod, "_databricks"), + mock.patch.object( + dev_mod, + "provision_trace_experiment", + return_value=("/Users/me@example.com/mason-traces/x", "cat.schema"), + ) as prov, ): result = CliRunner().invoke(dev_mod.dev, ["--source", str(tmp_path)], obj=_Ctx()) assert result.exit_code == 0, result.output - validate.assert_not_called() + prov.assert_called_once() # configured -> tracing wired through the shared path + assert prov.call_args.args[1] == tmp_path.resolve().name # app name + assert ( + prov.call_args.args[2].current_user == "me@example.com" + ) # the client (read once configured) def test_dev_requires_app_yaml(tmp_path: pathlib.Path): diff --git a/integrations/mason/tests/unit_tests/tracing_test.py b/integrations/mason/tests/unit_tests/tracing_test.py index b69b9cfd..2133e1f8 100644 --- a/integrations/mason/tests/unit_tests/tracing_test.py +++ b/integrations/mason/tests/unit_tests/tracing_test.py @@ -1,208 +1,302 @@ -"""Unit tests for `mason tracing`: instrument snippet, destination parsing, mlflow guard. +"""Unit tests for `mason tracing`: UC-schema validation, setup persistence, list resolution, and +the deploy-time UC-experiment provisioning. -The pure-Python surface (instrument, destination parsing) is covered directly. The -mlflow-backed commands are exercised only for their "mlflow not installed" guard, since the -hermetic test env does not carry mlflow on the deptree. +MLflow-backed paths (`list`, `ensure_uc_experiment`) are exercised with a mocked `_mlflow`/UC symbols +(the hermetic test env carries no mlflow); the pure surface is tested directly. """ from __future__ import annotations import json +import pathlib +import types +from unittest import mock +import pytest from click.testing import CliRunner from databricks_mason import tracing as tracing_mod +from databricks_mason.agent_project import AgentProject from databricks_mason.errors import AgentCliError +_AGENT_TOML = 'schema_version = 1\n\n[agent]\nframework = "langgraph"\n' + class _Ctx: - """Stand-in for CliContext: tracing commands read only .profile and .output.""" + """Stand-in for CliContext: tracing reads .profile / .output, and .client() for the dev default.""" - def __init__(self, output: str = "text", profile=None): + def __init__(self, output: str = "text", profile=None, user="me@example.com"): self.output = output self.profile = profile + self._user = user + def client(self): + return mock.Mock(current_user=self._user) -def test_instrument_text_runs(): - result = CliRunner().invoke( - tracing_mod.tracing_instrument, - ["--destination", "cat.schema", "--experiment", "/Shared/x"], - obj=_Ctx(), - ) - assert result.exit_code == 0, result.output - assert "Agent Tracing" in result.output +def _project(tmp_path: pathlib.Path, *, schema: str | None = None, warehouse: str | None = None): + binding = "" + if schema: + binding = f'\n[trace_location]\nname = "{schema}"\n' + if warehouse: + binding += f'warehouse_id = "{warehouse}"\n' + (tmp_path / "agent.toml").write_text(_AGENT_TOML + binding) + return tmp_path -def test_instrument_json_snippet_contents(): - result = CliRunner().invoke( - tracing_mod.tracing_instrument, - ["--destination", "cat.schema", "--experiment", "/Shared/x"], - obj=_Ctx(output="json"), - ) - assert result.exit_code == 0, result.output - payload = json.loads(result.output) - assert payload["destination"] == "cat.schema" - assert payload["experiment"] == "/Shared/x" - snippet = payload["snippet"] - assert 'catalog_name="cat"' in snippet - assert 'schema_name="schema"' in snippet - assert "mlflow.openai.autolog" in snippet - assert 'mlflow.set_experiment("/Shared/x")' in snippet - - -def test_instrument_defaults_to_placeholders_and_default_experiment(): - result = CliRunner().invoke(tracing_mod.tracing_instrument, [], obj=_Ctx(output="json")) - assert result.exit_code == 0, result.output - payload = json.loads(result.output) - assert payload["experiment"] == tracing_mod._DEFAULT_EXPERIMENT - assert "" in payload["snippet"] and "" in payload["snippet"] +# --- validation -------------------------------------------------------------- -def test_split_destination_valid(): - assert tracing_mod._split_destination("my_cat.my_schema") == ("my_cat", "my_schema") +def test_validate_uc_schema_accepts_catalog_dot_schema(): + assert tracing_mod.validate_uc_schema("catalog.schema") == "catalog.schema" + assert tracing_mod.validate_uc_schema(" cat.sch ") == "cat.sch" -def test_split_destination_invalid_raises(): - for bad in ("nodot", ".schema", "catalog."): - try: - tracing_mod._split_destination(bad) - raise AssertionError(f"expected AgentCliError for {bad!r}") - except AgentCliError: - pass +def test_validate_uc_schema_rejects_bad_values(): + for bad in ("nodot", "cat.schema.table", "12345", "", " ", "cat/schema"): + with pytest.raises(AgentCliError): + tracing_mod.validate_uc_schema(bad) -def test_ensure_experiment_creates_parent_dir_for_nested_path(): - from unittest import mock - mlflow = mock.Mock() - mlflow.get_experiment_by_name.return_value = None # doesn't exist yet - mlflow.create_experiment.return_value = "eid-1" - client = mock.Mock() +# --- experiment naming ------------------------------------------------------- - eid = tracing_mod._ensure_experiment(mlflow, client, "/Users/me@x.com/mason-traces/demo") - assert eid == "eid-1" - # the intermediate workspace folder is created before the experiment (mlflow won't make it) - client.ensure_workspace_dir.assert_called_once_with("/Users/me@x.com/mason-traces") +def test_experiment_name(): + assert tracing_mod.experiment_name("me@x.com", "app") == "/Users/me@x.com/mason-traces/app" -def test_ensure_experiment_reuses_existing_without_mkdir(): - from unittest import mock +def test_experiment_ui_url(): + assert ( + tracing_mod.experiment_ui_url("https://x.databricks.com", "123") + == "https://x.databricks.com/ml/experiments/123/traces" + ) + # trailing slash on the host is normalized; a missing/unknown host yields no link + assert ( + tracing_mod.experiment_ui_url("https://x.databricks.com/", "123") + == "https://x.databricks.com/ml/experiments/123/traces" + ) + assert tracing_mod.experiment_ui_url(None, "123") is None + assert tracing_mod.experiment_ui_url("unknown", "123") is None - mlflow = mock.Mock() - mlflow.get_experiment_by_name.return_value = mock.Mock(experiment_id="eid-2") - client = mock.Mock() - assert tracing_mod._ensure_experiment(mlflow, client, "/Shared/x") == "eid-2" - client.ensure_workspace_dir.assert_not_called() # existing experiment -> no dir work - mlflow.create_experiment.assert_not_called() +# --- setup ------------------------------------------------------------------- -def test_default_experiment_is_per_app_under_user_home(): - assert ( - tracing_mod.default_experiment("me@x.com", "my-agent") - == "/Users/me@x.com/mason-traces/my-agent" +def test_setup_persists_schema_and_warehouse_in_agent_toml(tmp_path: pathlib.Path): + _project(tmp_path) + result = CliRunner().invoke( + tracing_mod.tracing_setup, + ["--trace-location", "cat.schema", "--warehouse-id", "wh1", "--source", str(tmp_path)], + obj=_Ctx(output="json"), ) + assert result.exit_code == 0, result.output + assert json.loads(result.output) == {"trace_location": "cat.schema", "warehouse_id": "wh1"} + project = AgentProject.load(tmp_path) + assert project.trace_location == "cat.schema" + assert project.trace_warehouse == "wh1" -def test_default_experiment_falls_back_to_shared_without_app(): - assert tracing_mod.default_experiment("me@x.com", None) == tracing_mod._DEFAULT_EXPERIMENT +def test_setup_rejects_invalid_schema(tmp_path: pathlib.Path): + _project(tmp_path) + result = CliRunner().invoke( + tracing_mod.tracing_setup, + ["--trace-location", "cat.schema.table", "--source", str(tmp_path)], + obj=_Ctx(), + ) + assert result.exit_code != 0 + assert AgentProject.load(tmp_path).trace_location is None # nothing persisted -def test_experiment_url_builds_traces_tab_link(): - url = tracing_mod._experiment_url("https://ws.databricks.com/", "123") - assert url == "https://ws.databricks.com/ml/experiments/123?compareRunsMode=TRACES" +# --- ensure_uc_experiment ---------------------------------------------------- -def test_link_trace_location_reports_already_linked_without_relink(): - def set_location(location, experiment_id): - raise RuntimeError("experiment is already linked to a storage location") +def _fake_uc(mlflow, existing=None, create_id="e1"): + mlflow.get_experiment_by_name.return_value = ( + types.SimpleNamespace(experiment_id=existing) if existing else None + ) + mlflow.create_experiment.return_value = create_id - try: - tracing_mod._link_trace_location( - set_location, lambda **k: None, object(), "e1", relink=False - ) - raise AssertionError("expected AgentCliError") - except AgentCliError as exc: - assert "--relink" in (exc.hint or "") +def test_ensure_uc_experiment_creates_and_links(): + mlflow = mock.Mock() + _fake_uc(mlflow) + set_location = mock.Mock() + with ( + mock.patch.object(tracing_mod, "_mlflow", return_value=mlflow), + mock.patch.object(tracing_mod, "_configure"), + mock.patch.object( + tracing_mod, "_uc_trace_symbols", return_value=(mock.Mock(), set_location) + ), + ): + experiment_id = tracing_mod.ensure_uc_experiment(None, "/Users/me/x", "cat.schema", "wh1") + assert experiment_id == "e1" # returns the id (for the experiment UI link), not the name + mlflow.create_experiment.assert_called_once_with("/Users/me/x") + assert set_location.call_args.kwargs["experiment_id"] == "e1" -def test_link_trace_location_relinks_when_requested(): - calls = [] - def set_location(location, experiment_id): - calls.append("set") - if calls.count("set") == 1: # first attempt fails as already-linked - raise RuntimeError("already linked") +def test_ensure_uc_experiment_idempotent_when_already_linked(): + mlflow = mock.Mock() + _fake_uc(mlflow, existing="e9") + set_location = mock.Mock(side_effect=RuntimeError("experiment is already linked to a location")) + with ( + mock.patch.object(tracing_mod, "_mlflow", return_value=mlflow), + mock.patch.object(tracing_mod, "_configure"), + mock.patch.object( + tracing_mod, "_uc_trace_symbols", return_value=(mock.Mock(), set_location) + ), + ): + # a re-deploy of an already-linked experiment is a no-op, not an error (returns the id) + assert tracing_mod.ensure_uc_experiment(None, "/Users/me/x", "cat.schema", None) == "e9" - def unset_location(location, experiment_id): - calls.append("unset") - tracing_mod._link_trace_location(set_location, unset_location, object(), "e1", relink=True) - assert calls == ["set", "unset", "set"] # try, unset existing, re-link +def test_ensure_uc_experiment_errors_on_existing_traces(): + mlflow = mock.Mock() + _fake_uc(mlflow, existing="e9") + set_location = mock.Mock(side_effect=RuntimeError("Experiment already contains traces.")) + with ( + mock.patch.object(tracing_mod, "_mlflow", return_value=mlflow), + mock.patch.object(tracing_mod, "_configure"), + mock.patch.object( + tracing_mod, "_uc_trace_symbols", return_value=(mock.Mock(), set_location) + ), + ): + with pytest.raises(AgentCliError) as exc: + tracing_mod.ensure_uc_experiment(None, "/Users/me/x", "cat.schema", None) + assert "migrate" in (exc.value.hint or "").lower() -def test_link_trace_location_propagates_unrelated_errors(): - def set_location(location, experiment_id): - raise RuntimeError("permission denied on catalog") +# --- list resolution --------------------------------------------------------- - try: - tracing_mod._link_trace_location( - set_location, lambda **k: None, object(), "e1", relink=True - ) - raise AssertionError("expected the original error") - except RuntimeError as exc: - assert "permission denied" in str(exc) # not swallowed as an already-linked case +def _fake_mlflow(traces): + fake = mock.Mock() + fake.search_traces.return_value = traces + return fake -def test_setup_requires_mlflow_when_absent(): - # mlflow is not on the hermetic test deptree, so setup should surface a clean CLI error - # (non-zero exit) rather than a traceback. - result = CliRunner().invoke( - tracing_mod.tracing_setup, ["--catalog", "c", "--schema", "s"], obj=_Ctx() - ) - assert result.exit_code != 0 +def _trace(trace_id): + return types.SimpleNamespace( + info=types.SimpleNamespace( + trace_id=trace_id, status="OK", execution_time_ms=5, timestamp_ms=1 + ) + ) -def test_list_resolves_experiment_name_to_id_for_search(): - # search_traces selects by experiment_ids, not names, so list must resolve the name first. - from unittest import mock - mlflow = mock.Mock() - mlflow.get_experiment_by_name.return_value = mock.Mock(experiment_id="eid-9") - mlflow.search_traces.return_value = [] +def test_list_uses_configured_uc_schema(tmp_path: pathlib.Path): + _project(tmp_path, schema="proj.schema", warehouse="wh-1") + fake = _fake_mlflow([_trace("tr-1")]) with ( - mock.patch.object(tracing_mod, "_mlflow", return_value=mlflow), + mock.patch.object(tracing_mod, "_mlflow", return_value=fake), mock.patch.object(tracing_mod, "_configure"), ): result = CliRunner().invoke( - tracing_mod.tracing_list, ["--experiment", "/Shared/x", "--limit", "7"], obj=_Ctx() + tracing_mod.tracing_list, ["--source", str(tmp_path)], obj=_Ctx(output="json") ) - assert result.exit_code == 0, result.output - mlflow.get_experiment_by_name.assert_called_once_with("/Shared/x") - _, kwargs = mlflow.search_traces.call_args - assert kwargs["experiment_ids"] == ["eid-9"] - assert kwargs["max_results"] == 7 - + assert fake.search_traces.call_args.kwargs["locations"] == ["proj.schema"] + assert json.loads(result.output)[0]["trace_id"] == "tr-1" -def test_list_returns_empty_when_experiment_absent(): - # A not-yet-created experiment has no traces; list should show none, not error. - from unittest import mock - mlflow = mock.Mock() - mlflow.get_experiment_by_name.return_value = None +def test_list_flag_overrides_project(tmp_path: pathlib.Path): + _project(tmp_path, schema="proj.schema", warehouse="wh-1") + fake = _fake_mlflow([]) with ( - mock.patch.object(tracing_mod, "_mlflow", return_value=mlflow), + mock.patch.object(tracing_mod, "_mlflow", return_value=fake), mock.patch.object(tracing_mod, "_configure"), ): result = CliRunner().invoke( - tracing_mod.tracing_list, ["--experiment", "/Shared/missing"], obj=_Ctx(output="json") + tracing_mod.tracing_list, + ["--trace-location", "flag.schema", "--source", str(tmp_path)], + obj=_Ctx(output="json"), ) - assert result.exit_code == 0, result.output - assert json.loads(result.output) == [] - mlflow.search_traces.assert_not_called() + assert fake.search_traces.call_args.kwargs["locations"] == ["flag.schema"] + + +def test_list_errors_without_configured_location(tmp_path: pathlib.Path): + # Tracing is UC-only + opt-in: with nothing configured and no flag, list refuses (no silent + # fallback), pointing the user at `mason tracing setup`. + _project(tmp_path) # no UC schema configured + result = CliRunner().invoke( + tracing_mod.tracing_list, ["--source", str(tmp_path)], obj=_Ctx(output="json") + ) + assert result.exit_code != 0 + assert "mason tracing setup" in result.output + + +def test_list_applies_configured_warehouse(tmp_path: pathlib.Path): + # A UC schema is queried through the project's configured SQL warehouse. + _project(tmp_path, schema="cat.schema", warehouse="wh-1") + captured: dict = {} + fake = _fake_mlflow([]) + with ( + mock.patch.object(tracing_mod, "_mlflow", return_value=fake), + mock.patch.object( + tracing_mod, "_configure", side_effect=lambda m, p, w: captured.update(warehouse=w) + ), + ): + res = CliRunner().invoke( + tracing_mod.tracing_list, ["--source", str(tmp_path)], obj=_Ctx(output="json") + ) + assert res.exit_code == 0, res.output + assert captured["warehouse"] == "wh-1" + assert fake.search_traces.call_args.kwargs["locations"] == ["cat.schema"] + + +def test_get_applies_configured_warehouse(tmp_path: pathlib.Path): + # `get` reads a UC-backed trace through the project's configured SQL warehouse, same as `list` + # (a UC trace read needs a warehouse; without one MLflow raises "SQL warehouse ID is required"). + _project(tmp_path, schema="cat.schema", warehouse="wh-1") + captured: dict = {} + fake = mock.Mock() + fake.get_trace.return_value = _trace("tr-1") + with ( + mock.patch.object(tracing_mod, "_mlflow", return_value=fake), + mock.patch.object( + tracing_mod, "_configure", side_effect=lambda m, p, w: captured.update(warehouse=w) + ), + ): + res = CliRunner().invoke( + tracing_mod.tracing_get, ["tr-1", "--source", str(tmp_path)], obj=_Ctx(output="json") + ) + assert res.exit_code == 0, res.output + assert captured["warehouse"] == "wh-1" + assert fake.get_trace.call_args.args[0] == "tr-1" + + +def test_list_errors_without_warehouse(tmp_path: pathlib.Path, monkeypatch): + # A UC read needs a SQL warehouse; with none configured/passed and no env var, `list` fails fast + # with a mason hint instead of MLflow's raw "SQL warehouse ID is required" error. + monkeypatch.delenv("MLFLOW_TRACING_SQL_WAREHOUSE_ID", raising=False) + _project(tmp_path, schema="cat.schema") # location configured, but no warehouse anywhere + result = CliRunner().invoke( + tracing_mod.tracing_list, ["--source", str(tmp_path)], obj=_Ctx(output="json") + ) + assert result.exit_code != 0 + assert "SQL warehouse" in result.output + + +def test_get_errors_without_warehouse(tmp_path: pathlib.Path, monkeypatch): + # Same requirement for `get`: no warehouse anywhere -> fail fast with the hint. + monkeypatch.delenv("MLFLOW_TRACING_SQL_WAREHOUSE_ID", raising=False) + _project(tmp_path, schema="cat.schema") + result = CliRunner().invoke( + tracing_mod.tracing_get, ["tr-1", "--source", str(tmp_path)], obj=_Ctx(output="json") + ) + assert result.exit_code != 0 + assert "SQL warehouse" in result.output + + +# --- helpers ----------------------------------------------------------------- + + +def test_project_trace_location_reads_schema_and_warehouse(tmp_path: pathlib.Path): + _project(tmp_path, schema="cat.schema", warehouse="wh1") + assert tracing_mod.project_trace_location(str(tmp_path)) == ("cat.schema", "wh1") + + +def test_project_trace_location_none_without_project(tmp_path: pathlib.Path): + assert tracing_mod.project_trace_location(str(tmp_path)) == (None, None) def test_status_str_handles_enum_like_and_none():