02-experiment-tracing-and-uc-storage

name	02-experiment-tracing-and-uc-storage
description	Use when setting up MLflow experiments, tracing, or UC OTEL trace storage for a GenAI agent. Covers structured experiment paths, tracing decorators, manual spans, tags, connection pooling, and Unity Catalog OTEL storage for SQL-queryable trace retention. Foundation Step 2. Consumes MLflow environment from Step 1.
license	Apache-2.0
clients	["ide_cli","genie_code"]
bundle_resource	none
deploy_verb	none
deploy_note	Experiment + tracing + UC OTEL trace storage configured via the MLflow SDK; OTEL trace tables land in the per-user prefixed schema. No bundle resource. Identical on both clients; on Genie Code use its serverless runtime + runDatabricksCli for any CLI step. See `skills/genie-code-environment`.
coverage	full
metadata	{"last_verified":"2026-08-30","volatility":"high","upstream_sources":[],"author":"prashanth-subrahmanyam","version":"3.6.1","domain":"genai-agents","pipeline_position":"F2","consumes":"mlflow_environment","produces":"experiment_paths, tracing_config, connection_pool, f2_grants_complete, otel_table_prefix, mlflow_tracing_sql_warehouse_id, app_service_principal_grants","grounded_in":"docs.databricks.com/aws/en/mlflow3/genai/tracing/trace-unity-catalog, docs.databricks.com/aws/en/mlflow3/genai, docs.databricks.com/aws/en/mlflow3/genai/tracing/app-instrumentation, docs.databricks.com/aws/en/mlflow3/genai/tracing/app-instrumentation/automatic, docs.databricks.com/aws/en/mlflow3/genai/tracing/prod-tracing, docs.databricks.com/aws/en/mlflow3/genai/tracing/add-context-to-traces"}

Experiment tracing setup

When to Use

Use this skill when you need to:

• Organize MLflow experiments so runs are discoverable by space, domain, and lifecycle stage
• Add tracing to GenAI agents (decorators, nested spans, inputs/outputs)
• Configure MLflow for multi-stage pipelines (development, evaluation, deployment) with consistent paths and UC prompt registry visibility
• Tune HTTP client behavior before high-throughput tracing or evaluation workloads

Prerequisite: complete Foundation Step 1 (MLflow foundation) so tracking URI and authentication are already correct. See MLflow GenAI Foundation (Foundation Step 1).

TypeScript / Node agents: this skill is the Python instrumentation reference. For the official mlflow-tracing + mlflow-openai npm path (Node-native mlflow.init, tracedOpenAI, mlflow.trace, withSpan, session grouping), see the sibling skill 02b-typescript-tracing. Use OTLP (via custom OpenTelemetry instrumentation when the TypeScript SDK does not fit) only as a fallback when you need vendor-neutral spans or already run an OpenTelemetry collector.

Production deployment: the env-var matrix for deployed agents (ENABLE_MLFLOW_TRACING, MLFLOW_EXPERIMENT_ID, SP CAN_EDIT on the experiment, the Git-folder caveat, Production Monitoring → Delta) lives in references/prod-tracing-deployment.md. Track A and Track C deployment skills link there.

User / session / environment context: the canonical reference for attributing traces to a user (mlflow.trace.user), grouping multi-turn conversations (mlflow.trace.session), and overriding mlflow.source.type from APP_ENVIRONMENT lives in 02c-trace-context-and-environments. The "Trace tags and metadata" section below shows the call-site shape; F2c is the long form (tags vs metadata, auto-populated fields, search examples, deployment overrides).

Which approach: automatic vs manual vs combined

Before writing tracing code, pick the right approach. Source: Add traces to applications (overview).

Scenario	Recommended approach
You use one GenAI library (LangChain, LlamaIndex, DSPy, …)	Automatic tracing only — mlflow.<library>.autolog().
You call an LLM SDK directly (OpenAI, Anthropic, Mistral, …)	Automatic for the SDK + a thin @mlflow.trace wrapper around your run() / orchestration function so all calls roll up into one trace.
You use multiple frameworks / SDKs in one workflow	Enable autolog() for each framework + use @mlflow.trace to combine them into a single root trace.
All other scenarios (custom logic, tool routing, complex retry/fallback, framework-less)	Manual with @mlflow.trace decorators first; drop down to mlflow.start_span only when you need finer-grained control.

Start with automatic. It's the fastest way to get traces working. Add manual tracing later if you need more control. Both approaches feed the same trace tree — @mlflow.trace parent spans naturally nest auto-traced child spans.

For the full 20+ supported autolog integrations (LLM SDKs, orchestrators, agent frameworks, embedding libraries) plus the multi-framework combine pattern and the serverless-compute caveat, see references/autolog-integrations.md.

Experiment path organization

CRITICAL: consume the experiment path from state — do not invent one

The workshop pins MLflow experiment paths to the same user-and-use-case identity that backs APP_NAME (e.g. jane-d-stayfinder) so concurrent attendees on a shared workspace cannot collide on a single experiment, and so the leaf in the MLflow UI is never a generic word like Tracing, traces, Default, or my-agent.

The canonical derivation lives in vibecoding-state migrate_canonical and is captured in state at the prompt that first resolves $APP_NAME / $AGENT_NAME:

State field	Derivation	Example
mlflow_experiment_path	/Users/<user_email>/mlflow/<APP_NAME or AGENT_NAME>-agent	/Users/[e-mail]/mlflow/jane-d-stayfinder-agent
mlflow_feedback_experiment_path	/Users/<user_email>/mlflow/<APP_NAME>-feedback	/Users/[e-mail]/mlflow/jane-d-stayfinder-feedback

This skill consumes those values from state://Resources.mlflow_experiment_path rather than constructing its own. If state shows <pending> for the path, halt and route back to vibecoding-state migrate_canonical — do not paper over it with a hand-rolled /Shared/... default.

Path template (for projects that do not run on top of vibecoding-state)

If your project does not use the vibecoding-state skill, define a template that still pins identity onto the leaf:

EXPERIMENT_PATH_TEMPLATE = "/Users/{{ user_email }}/mlflow/{{ app_name }}-{{ stage }}"

Where app_name is the user-prefixed, use-case-suffixed identity (e.g. jane-d-stayfinder) and stage ∈ {agent, eval, feedback, deploy}.

Three-experiment lifecycle pattern

For multi-stage pipelines, use separate experiments (one leaf per stage under the same app_name):

Stage	Leaf	Purpose
agent / dev	<app_name>-agent	Interactive debugging, short runs, permissive logging — the default tracing destination
eval	<app_name>-eval	Benchmarks, mlflow.genai.evaluate, regression gates
feedback	<app_name>-feedback	End-user thumbs / human assessments persisted from the AppKit feedback skill
deploy	<app_name>-deploy	Production or promotion runs, stricter tags and retention

The leaf must always carry <app_name> so that browsing MLflow experiments lists jane-d-stayfinder-agent, jane-d-stayfinder-eval, etc. — never a bare agent / eval / Tracing.

Setting the experiment

When running inside the workshop, read the path from state:

import mlflow

# state://Resources.mlflow_experiment_path is already pinned to
# /Users/<user_email>/mlflow/<APP_NAME>-agent by vibecoding-state.migrate_canonical.
experiment_path = state["Resources"]["mlflow_experiment_path"]
mlflow.set_experiment(experiment_path)

Stand-alone projects build the path from the same identity inputs:

import mlflow

user_email = "[e-mail]"
app_name = "jane-d-stayfinder"  # ${FIRSTNAME}-${LASTINITIAL}-${use_case_slug}
experiment_path = f"/Users/{user_email}/mlflow/{app_name}-agent"
mlflow.set_experiment(experiment_path)

Set the experiment early in your entrypoint — before enabling autolog and making any LLM calls. Never use a literal leaf like traces, Tracing, or my-agent; the leaf is the only thing surfacing in the MLflow UI search column and a generic value defeats per-attendee isolation.

For complete experiment organization patterns including ExperimentManager, search, cleanup, and decision tables, see: references/experiment-organization.md.

CRITICAL: Prompt registry linkage

Prompts registered in Unity Catalog must be linked to the experiment or they will not surface correctly in the Experiment UI for prompt-aware workflows.

After set_experiment, set the experiment tag:

mlflow.set_experiment_tags({
    "mlflow.promptRegistryLocation": f"{catalog}.{schema}",
})

Use your UC catalog and schema where prompts are registered. Without mlflow.promptRegistryLocation, UC-registered prompts may not appear as expected in the UI.

Tracing with decorators

Use @mlflow.trace for automatic span creation around functions. Pick a name and span_type that match how you want traces grouped in the UI.

import mlflow


@mlflow.trace(name="classify_intent", span_type="AGENT")
def classify_intent(query: str) -> dict:
    ...


@mlflow.trace(name="call_llm", span_type="LLM")
def call_llm(prompt: str) -> str:
    ...


@mlflow.trace(name="evaluate_response", span_type="JUDGE")
def evaluate_response(response: str) -> float:
    ...

Common span_type values: AGENT, TOOL, LLM, RETRIEVER, JUDGE, EMBEDDING. Align names with your team's conventions so traces stay searchable across services.

For complete decorator and async tracing examples, see: references/tracing-patterns.md.

For the 20+ mlflow.<library>.autolog() integrations (OpenAI, Anthropic, Mistral, LangChain, LangGraph, LlamaIndex, DSPy, LiteLLM, etc.), the multi-framework combine snippet, and the serverless-compute caveat (autolog is not auto-enabled), see references/autolog-integrations.md.

Manual span creation

For fine-grained control (nested work units, partial inputs/outputs, retries), use mlflow.start_span. This pattern matches how the optimizer wraps LLM calls.

For complex tracing, open a span with span_type=SpanType.CHAIN, set inputs before the call, record token usage, and set outputs on success or failure — including retry events via SpanEvent.

Illustrative nested pattern (same structural idea: parent span, child LLM span, explicit inputs/outputs):

import mlflow


def run_optimization_step(query, context):
    with mlflow.start_span(name="optimization_step") as span:
        span.set_inputs({"query": query})

        with mlflow.start_span(name="strategist_call", span_type="LLM") as llm_span:
            llm_span.set_inputs({"prompt": formatted_prompt})
            result = call_llm(formatted_prompt)
            llm_span.set_outputs({"response": result})

        span.set_outputs({"result": result})
        return result

In production code you may prefer from mlflow.entities import SpanType and types such as SpanType.CHAIN for LLM orchestration spans, consistent with _traced_llm_call.

For the full _traced_llm_call implementation, error handling, token logging, and a multi-step agent example with nested AGENT/LLM/TOOL/JUDGE spans, see: references/tracing-patterns.md.

Trace tags and metadata

Enrich the current trace with session, user, and deployment context so runs are filterable and attributable. Reserved identity fields belong under metadata= (immutable, MLflow-recognized for UI filter / group); mutable routing dimensions belong under tags=.

import os

mlflow.update_current_trace(
    metadata={
        "mlflow.trace.user":    user_id,
        "mlflow.trace.session": session_id,
        "mlflow.source.type":   os.getenv("APP_ENVIRONMENT", "development"),
        "agent_version":        "1.2.0",
        "space_id":             space_id,
    },
    tags={
        "domain":   domain,
        "sla_tier": "gold",
    },
)

Call this from code that runs inside an active trace (for example after mlflow.start_run / autolog / @mlflow.trace has established trace context). Setting mlflow.trace.user / mlflow.trace.session under tags= still works for read-back but loses the immutability guarantee and the UI's first-class user / session facets — prefer metadata.

For the full tag taxonomy, metadata patterns, trace search queries, and monitoring dashboard integration, see: references/trace-context-patterns.md. For the canonical reference on user / session / environment context (auto-populated metadata, APP_ENVIRONMENT override, search by metadata), see 02c-trace-context-and-environments.

Connection pool configuration

Reduce flaky failures under load by setting MLflow HTTP client defaults before heavy tracing or evaluation traffic:

import os

os.environ.setdefault("MLFLOW_HTTP_REQUEST_MAX_RETRIES", "5")
os.environ.setdefault("MLFLOW_HTTP_REQUEST_TIMEOUT", "120")

Set these as early as possible in the job or app entrypoint (alongside other MLflow env vars from Foundation Step 1). Adjust retries and timeout for your workspace network and batch sizes.

For connection pool tuning in high-throughput serving scenarios and async tracing performance tips, see: references/tracing-patterns.md § 8.

DO / DON'T examples

Experiment organization

DO — Pin the experiment leaf to the user-and-use-case identity, and prefer reading from vibecoding-state:

# In a workshop-managed project, read the pre-derived path from state.
experiment_path = state["Resources"]["mlflow_experiment_path"]
# e.g. "/Users/[e-mail]/mlflow/jane-d-stayfinder-agent"
mlflow.set_experiment(experiment_path)

# Stand-alone project — build the path from the same identity inputs.
user_email = "[e-mail]"
app_name = "jane-d-stayfinder"  # ${FIRSTNAME}-${LASTINITIAL}-${use_case_slug}
experiment_path = f"/Users/{user_email}/mlflow/{app_name}-agent"
mlflow.set_experiment(experiment_path)

DON'T — Use a generic leaf, a hand-rolled /Shared/... default, or a hard-coded workspace path. The leaf is what shows up in the MLflow UI experiment list, and traces / Tracing / my-agent give every attendee on a shared workspace the same name:

# WRONG: generic leaf — collides across attendees, useless in the UI