Instrumenting With Mlflow Tracing

MLflow Tracing Instrumentation Guide

Language-Specific Guides

Based on the user's project, load the appropriate guide:

• Python projects: Read references/python.md
• TypeScript/JavaScript projects: Read references/typescript.md

If unclear, check for package.json (TypeScript) or requirements.txt/pyproject.toml (Python) in the project.

---

What to Trace

Trace these operations (high debugging/observability value):

Skip tracing these (too granular, adds noise):

• Simple data transformations (dict/list manipulation)
• String formatting, parsing, validation
• Configuration loading, environment setup
• Logging or metric emission
• Pure utility functions (math, sorting, filtering)

Rule of thumb: Trace operations that are important for debugging and identifying issues in your application.

---

Verification

After instrumenting the code, always verify that tracing is working.

Planning to evaluate your agent? Tracing must be working before you run agent-evaluation. Complete verification below first.

1. Run the instrumented code — execute the application or agent so that at least one traced operation fires 2. Confirm traces are logged — use mlflow.search_traces() or MlflowClient().search_traces() to check that traces appear in the experiment:

import mlflow

traces = mlflow.search_traces(experiment_ids=["<experiment_id>"])
print(f"Found {len(traces)} trace(s)")
assert len(traces) > 0, "No traces were logged — check tracking URI and experiment settings"

3. Verify spans were captured — confirm the trace contains the expected spans, not just an empty shell:

trace = traces.iloc[0]
spans = mlflow.get_trace(trace.trace_id).data.spans
print(f"Trace has {len(spans)} span(s)")
for span in spans:
    print(f"  - {span.name} ({span.span_type})")

4. Report the result — tell the user how many traces and spans were found and confirm tracing is working

If no traces appear

Check these in order:

• Tracking URI not set — is mlflow.set_tracking_uri(...) called before the agent run? Without this, traces go to a local ./mlruns directory instead of the configured server.
• Autolog warnings — did mlflow.autolog() or framework-specific mlflow.<framework>.autolog() raise any warnings during setup? Check stderr for patching failures.
• Wrong experiment ID — verify the experiment ID passed to search_traces() matches the experiment active when the code ran (mlflow.get_experiment_by_name(...) to confirm).
• Network/auth issues — can the process reach the tracking server? Check for connection errors or 401/403 responses in logs.

For automated validation, use agent-evaluation/scripts/validate_tracing_runtime.py.

---

Feedback Collection

Log user feedback on traces for evaluation, debugging, and fine-tuning. Essential for identifying quality issues in production.

See references/feedback-collection.md for:

• Recording user ratings and comments with mlflow.log_feedback()
• Capturing trace IDs to return to clients
• LLM-as-judge automated evaluation

---

Reference Documentation

Production Deployment

See references/production.md for:

• Environment variable configuration
• Async logging for low-latency applications
• Sampling configuration (MLFLOW_TRACE_SAMPLING_RATIO)
• Lightweight SDK (mlflow-tracing)
• Docker/Kubernetes deployment

Advanced Patterns

See references/advanced-patterns.md for:

• Async function tracing
• Multi-threading with context propagation
• PII redaction with span processors

Distributed Tracing

See references/distributed-tracing.md for:

• Propagating trace context across services
• Client/server header APIs

Advanced Tracing Patterns

Contents

• Async Function Support
• Multi-Threading with Context Propagation
• PII Redaction with Span Processors
• Custom Span Attributes
• Error Handling and Status
• Trace Linking

For feedback collection, see feedback-collection.md.

---

Async Function Support

Trace async functions using the same decorator:

import mlflow
import asyncio
from mlflow.entities import SpanType

@mlflow.trace(span_type=SpanType.LLM)
async def async_llm_call(prompt: str) -> str:
    response = await llm_client.chat(prompt)
    return response.content

@mlflow.trace(name="parallel_retrieval", span_type=SpanType.RETRIEVER)
async def retrieve_parallel(queries: list[str]) -> list[list[str]]:
    tasks = [retrieve_documents(q) for q in queries]
    return await asyncio.gather(*tasks)

---

Multi-Threading with Context Propagation

MLflow uses contextvars to track the active trace. When spawning threads, explicitly copy the context:

import mlflow
from mlflow.entities import SpanType
from concurrent.futures import ThreadPoolExecutor
import contextvars

@mlflow.trace(name="parallel_processing", span_type=SpanType.CHAIN)
def process_items(items: list) -> list:
    results = []

    def process_one(item):
        with mlflow.start_span(name=f"process_{item}") as span:
            span.set_inputs({"item": item})
            result = heavy_computation(item)
            span.set_outputs({"result": result})
            return result

    with ThreadPoolExecutor(max_workers=4) as executor:
        # Copy context to each thread
        ctx = contextvars.copy_context()
        futures = [executor.submit(ctx.run, process_one, item) for item in items]
        results = [f.result() for f in futures]

    return results

Using `run_in_executor` with asyncio:

import asyncio
import contextvars

async def async_with_thread_work():
    loop = asyncio.get_event_loop()
    ctx = contextvars.copy_context()

    with mlflow.start_span(name="parent") as span:
        # Run blocking code in thread while preserving trace context
        result = await loop.run_in_executor(
            None,
            ctx.run,
            blocking_function,
            arg1, arg2
        )
        return result

---

PII Redaction with Span Processors

Redact sensitive data before traces are logged:

import mlflow
from mlflow.tracing.processor import SpanProcessor
import re

class PIIRedactionProcessor(SpanProcessor):
    """Redact PII from span inputs and outputs."""

    PATTERNS = {
        "email": r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}",
        "phone": r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b",
        "ssn": r"\b\d{3}-\d{2}-\d{4}\b",
        "credit_card": r"\b\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}\b",
    }

    def on_start(self, span, parent_context):
        pass  # Nothing to do on start

    def on_end(self, span):
        # Redact inputs
        if span.inputs:
            span._inputs = self._redact_dict(span.inputs)

        # Redact outputs
        if span.outputs:
            span._outputs = self._redact_dict(span.outputs)

    def _redact_dict(self, data: dict) -> dict:
        result = {}
        for key, value in data.items():
            if isinstance(value, str):
                result[key] = self._redact_string(value)
            elif isinstance(value, dict):
                result[key] = self._redact_dict(value)
            elif isinstance(value, list):
                result[key] = [
                    self._redact_string(v) if isinstance(v, str) else v
                    for v in value
                ]
            else:
                result[key] = value
        return result

    def _redact_string(self, text: str) -> str:
        for pii_type, pattern in self.PATTERNS.items():
            text = re.sub(pattern, f"[REDACTED_{pii_type.upper()}]", text)
        return text


# Register the processor
mlflow.tracing.add_span_processor(PIIRedactionProcessor())

Selective field redaction:

class SelectiveRedactionProcessor(SpanProcessor):
    """Redact specific fields by name."""

    SENSITIVE_FIELDS = {"password", "api_key", "token", "secret", "credentials"}

    def on_end(self, span):
        if span.inputs:
            span._inputs = self._redact_fields(span.inputs)
        if span.outputs:
            span._outputs = self._redact_fields(span.outputs)

    def _redact_fields(self, data: dict) -> dict:
        result = {}
        for key, value in data.items():
            if key.lower() in self.SENSITIVE_FIELDS:
                result[key] = "[REDACTED]"
            elif isinstance(value, dict):
                result[key] = self._redact_fields(value)
            else:
                result[key] = value
        return result

---

Custom Span Attributes

Add custom metadata to spans for filtering and analysis:

from mlflow.entities import SpanType

@mlflow.trace(span_type=SpanType.CHAIN)
def process_request(request_type: str, priority: int) -> str:
    span = mlflow.get_current_active_span()

    # Add custom attributes
    span.set_attributes({
        "request_type": request_type,
        "priority": priority,
        "environment": "production",
        "version": "1.2.3",
    })

    result = handle_request(request_type)
    return result

Query by custom attributes:

client = mlflow.MlflowClient()
traces = client.search_traces(
    experiment_ids=["1"],
    filter_string="attribute.request_type = 'urgent' AND attribute.priority > 5"
)

---

Error Handling and Status

Spans automatically capture exceptions, but you can add custom error handling:

from mlflow.entities import SpanType

@mlflow.trace(span_type=SpanType.CHAIN)
def risky_operation(data: dict) -> str:
    span = mlflow.get_current_active_span()

    try:
        result = process(data)
        span.set_status("OK")
        return result
    except ValidationError as e:
        span.set_status("ERROR")
        span.set_attributes({
            "error_type": "validation",
            "error_message": str(e),
        })
        raise
    except Exception as e:
        span.set_status("ERROR")
        span.set_attributes({
            "error_type": "unexpected",
            "error_message": str(e),
        })
        # Log but don't re-raise for graceful degradation
        return fallback_response()

---

Trace Linking

Link related traces (e.g., retry attempts, follow-up requests):

from mlflow.entities import SpanType

@mlflow.trace(span_type=SpanType.CHAIN)
def process_with_retry(data: dict, parent_trace_id: str = None) -> str:
    span = mlflow.get_current_active_span()

    if parent_trace_id:
        span.set_attributes({
            "linked_trace_id": parent_trace_id,
            "link_type": "retry",
        })

    try:
        return process(data)
    except RetryableError:
        current_trace_id = span.request_id
        return process_with_retry(data, parent_trace_id=current_trace_id)

Distributed Tracing

Connect spans across multiple services into a single trace by propagating trace context over HTTP.

Overview

MLflow provides two helper functions for distributed tracing:

• Client: mlflow.tracing.get_tracing_context_headers_for_http_request() - fetches headers to propagate
• Server: mlflow.tracing.set_tracing_context_from_http_request_headers() - extracts trace context from headers

---

Client Example

import requests
import mlflow
from mlflow.tracing import get_tracing_context_headers_for_http_request

with mlflow.start_span("client-root"):
    headers = get_tracing_context_headers_for_http_request()
    requests.post(
        "https://your.service/handle", headers=headers, json={"input": "hello"}
    )

---

Server Example (Flask)

import mlflow
from flask import Flask, request
from mlflow.tracing import set_tracing_context_from_http_request_headers

app = Flask(__name__)

@app.post("/handle")
def handle():
    headers = dict(request.headers)
    with set_tracing_context_from_http_request_headers(headers):
        with mlflow.start_span("server-handler") as span:
            # Your logic here
            span.set_attribute("status", "ok")
    return {"ok": True}

---

Server Example (FastAPI)

import mlflow
from fastapi import FastAPI, Request
from mlflow.tracing import set_tracing_context_from_http_request_headers

app = FastAPI()

@app.post("/handle")
async def handle(request: Request):
    headers = dict(request.headers)
    with set_tracing_context_from_http_request_headers(headers):
        with mlflow.start_span("server-handler") as span:
            # Your logic here
            span.set_attribute("status", "ok")
    return {"ok": True}

---

Result

Spans from client and server appear as a single connected trace in MLflow UI, showing end-to-end execution across services.

Feedback Collection

Log user feedback on traces for evaluation, debugging, and fine-tuning. Essential for identifying quality issues in production.

---

Recording User Feedback

Python:

import mlflow

def record_feedback(trace_id: str, rating: int):
    """Record user feedback for a trace."""
    mlflow.log_feedback(
        trace_id=trace_id,
        name="user_rating",
        value=rating,
        source=mlflow.entities.feedback.FeedbackSource(
            source_type="HUMAN",
            source_id="web_ui"
        )
    )

TypeScript:

import * as mlflow from "mlflow-tracing";

async function recordFeedback(traceId: string, rating: number) {
    await mlflow.logFeedback({
        traceId,
        name: "user_rating",
        value: rating,
        source: { sourceType: "HUMAN", sourceId: "web_ui" },
    });
}

---

Capturing Trace ID for Feedback

Return the trace ID to the client so they can submit feedback later.

from mlflow.entities import SpanType

@mlflow.trace(span_type=SpanType.CHAIN)
def chat(message: str) -> dict:
    response = generate_response(message)

    # Get trace ID to return to client for later feedback
    trace_id = mlflow.get_current_active_span().trace_id

    return {
        "response": response,
        "trace_id": trace_id  # Client uses this to submit feedback
    }

Supported Value Types

# Numeric rating
mlflow.log_feedback(trace_id, name="rating", value=5, source=source)

# Text comment
mlflow.log_feedback(trace_id, name="comment", value="Very helpful!", source=source)

# Boolean thumbs up/down
mlflow.log_feedback(trace_id, name="thumbs_up", value=True, source=source)

---

Feedback Source Types

Production Configuration

This guide covers configuration and optimization for logging traces in production environments (environment variables, async logging, sampling).

For how to construct traces (instrumentation methods, what to trace, decorators, manual spans), refer to:

• python.md - Python instrumentation guide
• typescript.md - TypeScript instrumentation guide

---

Contents

• Environment Variables
• Async Logging
• Sampling

---

Environment Variables

Configure MLflow Tracing via environment variables for production deployments:

# Required: Tracking server URI
export MLFLOW_TRACKING_URI="http://mlflow-server:5000"

# Optional: Default experiment
export MLFLOW_EXPERIMENT_NAME="production-agent"

# Optional: Authentication
export MLFLOW_TRACKING_USERNAME="user"
export MLFLOW_TRACKING_PASSWORD="password"
# Or use token-based auth
export MLFLOW_TRACKING_TOKEN="your-token"

---

Async Logging

For latency-sensitive applications, enable async logging to avoid blocking on trace uploads:

import mlflow

mlflow.config.enable_async_logging(True)

Or via environment variable:

export MLFLOW_ENABLE_ASYNC_LOGGING=true

Behavior:

• Traces are queued and uploaded in a background thread
• Function returns immediately after local span creation
• Failed uploads are retried automatically

Flush on shutdown:

import atexit
import mlflow

atexit.register(mlflow.flush_trace_async_logging)

---

Sampling

Reduce tracing overhead by sampling a fraction of requests:

# Trace 10% of requests
export MLFLOW_TRACE_SAMPLING_RATIO=0.1

Note: Sampling is random per-trace. All spans within a sampled trace are captured.

MLflow Tracing - Python Guide

Contents

• Quick Start
• Instrumentation Methods (AutoLogging, Decorator, Manual Spans)
• User/Session Tracking
• Combining AutoLogging with Custom Tracing
• Common Issues

---

Quick Start

Install and Configure

pip install mlflow>=3.8.0

Check if `MLFLOW_TRACKING_URI` and `MLFLOW_EXPERIMENT_ID` are already set in the environment. If both are set, skip the configuration below — MLflow will use them automatically.

Only call these if the environment is NOT pre-configured:

import mlflow

mlflow.set_tracking_uri("http://localhost:5000")  # skip if MLFLOW_TRACKING_URI is set
mlflow.set_experiment("my-agent")                 # skip if MLFLOW_EXPERIMENT_ID is set

Enable Tracing

For supported frameworks (LangChain, LangGraph, OpenAI, etc.):

mlflow.langchain.autolog()  # or openai, anthropic, litellm, etc.

For custom code:

from mlflow.entities import SpanType

@mlflow.trace(span_type=SpanType.CHAIN)
def my_function(query: str) -> str:
    # Your code here
    return result

---

Instrumentation Methods

Method 1: AutoLogging (Recommended for Frameworks)

Zero-code instrumentation for supported libraries. See the Integrations page for the complete list.

import mlflow

# Enable before importing/using the library
mlflow.langchain.autolog()    # LangChain, LangGraph
mlflow.openai.autolog()       # OpenAI SDK
mlflow.anthropic.autolog()    # Anthropic SDK
mlflow.gemini.autolog()       # Google Gemini (google-genai SDK)
mlflow.litellm.autolog()      # LiteLLM
mlflow.dspy.autolog()         # DSPy
mlflow.autogen.autolog()      # AutoGen
mlflow.crewai.autolog()       # CrewAI

Method 2: Decorator (Recommended for Custom Code)

Prefer decorator over manual spans - it auto-captures function name, inputs, and outputs.

from mlflow.entities import SpanType

@mlflow.trace(span_type=SpanType.RETRIEVER)
def retrieve_documents(query: str) -> list[str]:
    return documents

@mlflow.trace(span_type=SpanType.TOOL)
def search_database(sql: str) -> dict:
    return results

Span types: LLM, CHAIN, TOOL, AGENT, RETRIEVER, EMBEDDING, RERANKER, PARSER, UNKNOWN

Method 3: Manual Spans (When Decorator Not Possible)

Use only when you can't use a decorator:

• Tracing code not wrapped in a function (e.g., script-level code, loop bodies)
• Dynamic span names computed at runtime

with mlflow.start_span(name=f"process_{item_id}") as span:
    span.set_inputs({"query": query})  # Must set manually
    result = process(query)
    span.set_outputs({"result": result})  # Must set manually

---

User/Session Tracking

For multi-turn applications, use standard metadata fields mlflow.trace.user and mlflow.trace.session.

from fastapi import Request
from mlflow.entities import SpanType

@app.post("/chat")
def handle_chat(request: Request, body: ChatRequest):
    user_id = request.headers.get("X-User-ID", "anonymous")
    session_id = request.headers.get("X-Session-ID", "default")
    return chat(body.message, user_id, session_id)

@mlflow.trace(span_type=SpanType.CHAIN)
def chat(message: str, user_id: str, session_id: str) -> str:
    mlflow.update_current_trace(
        metadata={
            "mlflow.trace.user": user_id,
            "mlflow.trace.session": session_id,
        }
    )
    return response

Query traces by user:

traces = mlflow.search_traces(
    filter_string="metadata.`mlflow.trace.user` = 'user123'"
)

---

Combining AutoLogging with Custom Tracing

import mlflow
from mlflow.entities import SpanType
from langchain_openai import ChatOpenAI

mlflow.langchain.autolog()

@mlflow.trace(name="rag_pipeline", span_type=SpanType.CHAIN)
def rag_query(question: str) -> str:
    docs = retrieve_documents(question)  # Custom function

    llm = ChatOpenAI()  # Auto-traced by autolog
    response = llm.invoke(format_prompt(docs, question))

    return response.content

---

Common Issues

Traces not appearing? 1. Verify tracking URI is correct (MLFLOW_TRACKING_URI env var or mlflow.set_tracking_uri()) 2. Ensure autolog is called before framework imports 3. Check experiment is configured (MLFLOW_EXPERIMENT_ID env var or mlflow.set_experiment())

Nested spans not connected?

• Use @mlflow.trace or context managers consistently
• For threading, see advanced-patterns.md

MLflow Tracing - TypeScript Guide

Contents

• Quick Start
• API Reference (Core APIs, Span Manipulation, Auto-Tracing Wrappers, Span Types)
• Instrumentation Methods (Function Wrapper, Manual Spans)
• User/Session Tracking
• Combining Auto-Tracing with Custom Tracing
• Common Issues

---

Quick Start

Install and Configure

npm install mlflow-tracing

import * as mlflow from "mlflow-tracing";

mlflow.init({
    trackingUri: "http://localhost:5000",
    experimentId: "my-agent",
});

Enable Tracing

const myFunction = mlflow.trace(
    (query: string) => {
        // Your code here
        return result;
    },
    { name: "my_function", spanType: mlflow.SpanType.CHAIN }
);

---

API Reference

Core Tracing APIs

import * as mlflow from "mlflow-tracing";

// Initialize (required before tracing)
mlflow.init({
    trackingUri: "http://localhost:5000",
    experimentId: "my-experiment",
});

// Function wrapper - creates traced version of a function
const tracedFn = mlflow.trace(
    (arg: string) => { return result; },
    { name: "my_function", spanType: mlflow.SpanType.CHAIN }
);

// Manual span - for dynamic names or non-function code
const result = await mlflow.withSpan(
    { name: "dynamic_span", spanType: mlflow.SpanType.TOOL },
    async (span) => {
        span.setInputs({ key: "value" });
        const result = await doWork();
        span.setOutputs({ result });
        return result;
    }
);

// Flush traces before process exit
await mlflow.flushTraces();

Span Manipulation

// Get current active span (inside a traced function)
const span = mlflow.getCurrentActiveSpan();
if (span) {
    span.setInputs({ query });
    span.setOutputs({ result });
    span.setAttribute("model", "gpt-4o-mini");
}

// Update current trace metadata/tags
mlflow.updateCurrentTrace({
    tags: { environment: "production" },
    metadata: {
        "mlflow.trace.user": userId,
        "mlflow.trace.session": sessionId,
    },
    requestPreview: "Custom request summary...",
    responsePreview: "Custom response summary...",
});

Auto-Tracing Wrappers

import { tracedOpenAI } from "mlflow-openai";
import { OpenAI } from "openai";

// Wrap OpenAI client - all calls auto-traced
const openai = tracedOpenAI(new OpenAI());

Span Types

mlflow.SpanType.LLM        // LLM inference calls
mlflow.SpanType.CHAIN      // Multi-step pipelines
mlflow.SpanType.TOOL       // Tool/function calls
mlflow.SpanType.AGENT      // Agent orchestration
mlflow.SpanType.RETRIEVER  // Document retrieval
mlflow.SpanType.EMBEDDING  // Embedding generation
mlflow.SpanType.RERANKER   // Result reranking
mlflow.SpanType.PARSER     // Output parsing
mlflow.SpanType.UNKNOWN    // Default/other

---

Instrumentation Methods

Method 1: Function Wrapper (Recommended)

Prefer wrapper over manual spans - it auto-captures function name, inputs, and outputs.

const retrieveDocuments = mlflow.trace(
    (query: string): string[] => {
        return documents;
    },
    { name: "retrieve_documents", spanType: mlflow.SpanType.RETRIEVER }
);

const searchDatabase = mlflow.trace(
    (sql: string): Record<string, unknown> => {
        return results;
    },
    { name: "search_database", spanType: mlflow.SpanType.TOOL }
);

Method 2: Manual Spans with withSpan

Use when you can't use a function wrapper (dynamic span names, non-function code):

const result = await mlflow.withSpan(
    { name: `process_${itemId}` },
    async (span) => {
        span.setInputs({ query });
        const result = await process(query);
        span.setOutputs({ result });
        return result;
    }
);

---

User/Session Tracking

For multi-turn applications, use standard metadata fields mlflow.trace.user and mlflow.trace.session.

app.post('/chat', async (req, res) => {
    const userId = req.header('X-User-ID') || 'anonymous';
    const sessionId = req.header('X-Session-ID') || 'default';
    const response = await chat(req.body.message, userId, sessionId);
    res.json({ response });
});

const chat = mlflow.trace(
    async (message: string, userId: string, sessionId: string) => {
        await mlflow.updateCurrentTrace({
            metadata: {
                "mlflow.trace.user": userId,
                "mlflow.trace.session": sessionId,
            },
        });
        return response;
    },
    { name: "chat", spanType: mlflow.SpanType.CHAIN }
);

---

Combining Auto-Tracing with Custom Tracing

import * as mlflow from "mlflow-tracing";
import { tracedOpenAI } from "mlflow-openai";
import { OpenAI } from "openai";

// Wrap OpenAI client for auto-tracing
const openai = tracedOpenAI(new OpenAI());

const retrieveDocuments = mlflow.trace(
    async (query: string): Promise<string[]> => {
        // Your retrieval logic
        return docs;
    },
    { name: "retrieve_documents", spanType: mlflow.SpanType.RETRIEVER }
);

const ragQuery = mlflow.trace(
    async (question: string): Promise<string> => {
        const docs = await retrieveDocuments(question);  // Custom traced function

        const response = await openai.chat.completions.create({  // Auto-traced
            model: "gpt-4o-mini",
            messages: [{ role: "user", content: formatPrompt(docs, question) }],
        });

        return response.choices[0].message.content || "";
    },
    { name: "rag_pipeline", spanType: mlflow.SpanType.CHAIN }
);

---

Common Issues

Traces not appearing? 1. Verify mlflow.init() is called with correct tracking URI 2. Check experiment ID is set

Traces not sent before exit?

• Call await mlflow.flushTraces() before process exit to ensure all spans are sent