Implementation:Mlflow Mlflow Trace Decorator

Knowledge Sources	MLflow MLflow Tracing API
Domains	ML_Ops, LLM_Observability
Last Updated	2026-02-13 20:00 GMT

Overview

Concrete tool for manually instrumenting Python functions and code blocks to create MLflow trace spans provided by the MLflow library.

Description

MLflow provides two complementary APIs for manual code instrumentation: the @mlflow.trace decorator and the mlflow.start_span context manager.

The @mlflow.trace decorator wraps a function so that each invocation automatically creates a span. The span captures the function's input arguments and return value, records timing, and sets the span status to ERROR if an exception is raised (including the exception message and stack trace in the span attributes). The decorator supports synchronous functions, async functions, generators, async generators, class methods, and static methods. It can also be used as a direct wrapper around external library functions via mlflow.trace(math.factorial)(5).

The mlflow.start_span context manager provides imperative, block-level control. Within the with block, the yielded LiveSpan object exposes methods to set inputs, outputs, and custom attributes. Child spans can be created by nesting additional start_span calls or @mlflow.trace decorators within the block. When the context manager is used at the top level (not inside another span), it creates a root span and the complete trace is logged when the root span ends.

Both APIs support span typing via SpanType (LLM, RETRIEVER, CHAIN, TOOL, AGENT, etc.), custom attributes for metadata, and optional trace destination routing. The @mlflow.trace decorator also supports an output_reducer parameter for generator functions to collapse streamed outputs into a single span output value, and a sampling_ratio_override to control per-function sampling rates.

Usage

Use @mlflow.trace for straightforward function-level tracing where automatic input/output capture is sufficient. Use mlflow.start_span when you need to set inputs and outputs conditionally, attach dynamic attributes, or instrument a code block that is not a standalone function. Combine both in the same application to build rich, nested trace trees.

Code Reference

Source Location

• Repository: mlflow
• File: mlflow/tracing/fluent.py
• Lines (@mlflow.trace): L112-269
• Lines (mlflow.start_span): L490-587

Signature

# Decorator
@mlflow.trace(
    func=None,
    name=None,
    span_type=SpanType.UNKNOWN,
    attributes=None,
    output_reducer=None,
    trace_destination=None,
    sampling_ratio_override=None,
) -> Callable

# Context manager
mlflow.start_span(
    name="span",
    span_type=SpanType.UNKNOWN,
    attributes=None,
    trace_destination=None,
) -> Generator[LiveSpan, None, None]

Import

import mlflow
# or for explicit imports:
from mlflow import trace, start_span

I/O Contract

Inputs

Name	Type	Required	Description
func	Callable	No	(Decorator only) The function to decorate. Must not be provided when using as a decorator with parentheses.
name	str	No	Display name for the span. Defaults to the function name (decorator) or "span" (context manager).
span_type	str or SpanType	No	The type of span: LLM, RETRIEVER, CHAIN, TOOL, AGENT, EMBEDDING, UNKNOWN, etc. Default UNKNOWN.
attributes	dict[str, Any]	No	Custom metadata dictionary to attach to the span.
output_reducer	Callable[[list[Any]], Any]	No	(Decorator only) Reduces generator outputs into a single span output value.
trace_destination	TraceLocationBase	No	Override destination for the trace. Only effective on root spans.
sampling_ratio_override	float	No	(Decorator only) Override global sampling ratio (0.0 to 1.0). Only applies to root spans.

Outputs

Name	Type	Description
(decorator)	Callable	The decorated function that auto-creates spans on each invocation.
(context manager)	Generator[LiveSpan, None, None]	Yields a LiveSpan object for setting inputs, outputs, and attributes within the block.

Usage Examples

Basic Usage

import mlflow


@mlflow.trace
def my_function(x, y):
    return x + y


# Calling the function creates a span named "my_function"
result = my_function(3, 4)

Decorator with Custom Span Type

import mlflow
from mlflow.entities import SpanType


@mlflow.trace(name="retrieve_docs", span_type=SpanType.RETRIEVER)
def retrieve(query: str) -> list[str]:
    # retrieval logic
    return ["doc1", "doc2"]

Context Manager Usage

import mlflow

with mlflow.start_span("my_span") as span:
    x = 1
    y = 2
    span.set_inputs({"x": x, "y": y})

    z = x + y

    span.set_outputs(z)
    span.set_attribute("key", "value")

Nested Spans

import mlflow


@mlflow.trace(name="pipeline", span_type="CHAIN")
def pipeline(query):
    docs = retrieve(query)
    return generate(query, docs)


@mlflow.trace(name="retrieve", span_type="RETRIEVER")
def retrieve(query):
    return ["doc1", "doc2"]


@mlflow.trace(name="generate", span_type="LLM")
def generate(query, docs):
    return f"Answer based on {len(docs)} docs"

Related Pages

Implements Principle

• Principle:Mlflow_Mlflow_Code_Instrumentation