Skip to main content

Manual Tracing

While MLflow's automatic tracing provides instant observability for supported frameworks, manual tracing gives you complete control over how your GenAI applications are instrumented. This flexibility is essential for building production-ready applications that require detailed monitoring and debugging capabilities.

Prerequisites​

This guide requires the following packages:

  • mlflow[databricks]>=3.1: Core MLflow functionality with GenAI features and Databricks connectivity.
  • openai>=1.0.0: (Optional) Only if your custom code interacts with OpenAI; replace with other SDKs if needed.

Install the basic requirements:

Python
%pip install --upgrade "mlflow[databricks]>=3.1"
# %pip install --upgrade openai>=1.0.0 # Install if needed
MLflow Version Recommendation

While manual tracing capabilities are available in MLflow 2.15.0+, it is strongly recommended to install MLflow 3 (specifically 3.1 or newer if using mlflow[databricks]) for the latest GenAI capabilities, including expanded tracing features, refined span types, improved context propagation, and robust support.

tip

Running in a Databricks notebook? MLflow is pre-installed. You only need to install additional SDKs if your manually traced code uses them.

Running locally? You'll need to install mlflow[databricks] and any other SDKs your code calls.

When to Use Manual Tracing​

Manual tracing is the right choice when you need:

Custom Trace Structure

  • Define exactly which parts of your code to trace
  • Create custom hierarchies of spans
  • Control span boundaries and relationships

Example Use Case: Tracing specific business logic within a RAG pipeline where you want to measure retrieval vs. generation latency separately.

Manual Tracing Approaches​

MLflow provides three levels of abstraction for manual tracing, each suited to different use cases:

The high-level APIs provide an intuitive way to add tracing with minimal code changes. They automatically handle trace lifecycle, exception tracking, and parent-child relationships.

Best for: Function-level tracing with minimal code changes

Python
import mlflow
from mlflow.entities import SpanType

@mlflow.trace(span_type=SpanType.CHAIN)
def process_request(query: str) -> str:
    # Your code here - automatically traced!
    result = generate_response(query)
    return result

@mlflow.trace(span_type=SpanType.LLM)
def generate_response(query: str) -> str:
    # Nested function - parent-child relationship handled automatically
    return llm.invoke(query)

Key Benefits:

  • One-line instrumentation
  • Automatic exception handling
  • Works with async/generator functions
  • Compatible with auto-tracing

Learn more about decorators →

2. Low-Level Client APIs (Advanced)​

For scenarios requiring complete control over trace lifecycle, the client APIs provide direct access to MLflow's tracing backend.

Python
from mlflow import MlflowClient

client = MlflowClient()

# Start a trace
root_span = client.start_trace("complex_workflow")

# Create child spans with explicit parent relationships
child_span = client.start_span(
    name="data_retrieval",
    request_id=root_span.request_id,
    parent_id=root_span.span_id,
    inputs={"query": query}
)

# End spans explicitly
client.end_span(
    request_id=child_span.request_id,
    span_id=child_span.span_id,
    outputs={"documents": documents}
)

# End the trace
client.end_trace(request_id=root_span.request_id)

When to Use Client APIs:

  • Custom trace ID management
  • Integration with existing observability systems
  • Complex trace lifecycle requirements
  • Non-standard tracing patterns
warning

Client APIs require manual management of:

  • Parent-child relationships
  • Span lifecycle (start/end)
  • Exception handling
  • Thread safety

Learn more about client APIs →

API Comparison​

Feature

Decorator

Context Manager

Client APIs

Automatic Parent-Child

Yes

Yes

No - manual management

Exception Handling

Automatic

Automatic

Manual

Works with Auto-trace

Yes

Yes

No

Thread Safety

Automatic

Automatic

Manual

Custom Trace IDs

No

No

Yes

Best For

Function tracing

Code block tracing

Advanced control

Common Patterns​

Combining with Auto-Tracing​

Manual tracing seamlessly integrates with MLflow's auto-tracing capabilities:

Python
import mlflow
import openai

# Enable auto-tracing for OpenAI
mlflow.openai.autolog()

@mlflow.trace(span_type="CHAIN")
def rag_pipeline(query: str):
    # Manual span for retrieval
    with mlflow.start_span(name="retrieval") as span:
        docs = retrieve_documents(query)
        span.set_outputs({"doc_count": len(docs)})

    # Auto-traced OpenAI call
    response = openai.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": f"Answer based on: {docs}\n\nQuery: {query}"}]
    )

    return response.choices[0].message.content

Complex Workflow Tracing​

For complex workflows with multiple steps, use nested spans to capture detailed execution flow:

Python
@mlflow.trace(name="data_pipeline")
def process_data_pipeline(data_source: str):
    # Extract phase
    with mlflow.start_span(name="extract") as extract_span:
        raw_data = extract_from_source(data_source)
        extract_span.set_outputs({"record_count": len(raw_data)})

    # Transform phase
    with mlflow.start_span(name="transform") as transform_span:
        transformed = apply_transformations(raw_data)
        transform_span.set_outputs({"transformed_count": len(transformed)})

    # Load phase
    with mlflow.start_span(name="load") as load_span:
        result = load_to_destination(transformed)
        load_span.set_outputs({"status": "success"})

    return result

Customizing Request and Response Previews in the UI​

The MLflow UI provides Request and Response columns in the Traces tab that show a preview of the overall trace's input and output. By default, these are truncated. When using manual tracing, especially with the @mlflow.trace decorator or context managers that create the root span of a trace, you can customize these previews using mlflow.update_current_trace().

This is useful for complex data structures where the default preview might not be informative.

Python
import mlflow
import openai # Assuming openai is used, replace if not

# This example assumes you have an OpenAI client initialized and API key set up.
# client = openai.OpenAI()

@mlflow.trace
def predict(messages: list[dict]) -> str:
    # Customize the request preview to show the first and last messages
    custom_preview = f'{messages[0]["content"][:10]} ... {messages[-1]["content"][:10]}'
    mlflow.update_current_trace(request_preview=custom_preview)

    # Call the model
    # response = openai.chat.completions.create(
    #     model="gpt-4o-mini",
    #     messages=messages,
    # )
    # return response.choices[0].message.content
    return f"Response based on {len(messages)} messages."

messages = [
    {"role": "user", "content": "Hi, how are you?"},
    {"role": "assistant", "content": "I'm good, thank you!"},
    {"role": "user", "content": "What's your name?"},
    # ... (long message history)
    {"role": "assistant", "content": "Bye!"},
]
predict(messages)

This allows you to tailor the preview to be more informative for your specific data structures.

Next steps​

Continue your journey with these recommended actions and tutorials.

Reference guides​

Explore detailed documentation for concepts and features mentioned in this guide.

tip

Most users should start with the high-level APIs (decorators and context managers). They provide the best balance of ease-of-use and functionality while maintaining compatibility with MLflow's ecosystem.

Last updated on