wehub-resource-sync
59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
"""
|
|
This example demonstrates how to create a trace with multiple spans using the high-level MLflow fluent APIs.
|
|
"""
|
|
|
|
import mlflow
|
|
|
|
mlflow.set_experiment("mlflow-tracing-example")
|
|
|
|
|
|
# Decorating the function with `@mlflow.trace` decorator is the easiest way to trace your function.
|
|
# MLflow will create a trace for function calls and automatically
|
|
# captures function name, inputs, output, and more.
|
|
@mlflow.trace
|
|
def f1(x: int) -> int:
|
|
return x + 1
|
|
|
|
|
|
# You can also specify additional metadata for the trace
|
|
@mlflow.trace(
|
|
span_type="math",
|
|
attributes={"operation": "addition"},
|
|
)
|
|
def f2(x: int) -> int:
|
|
# MLflow keeps track of the call hierarchy. Calling `f1` inside
|
|
# `f2` will create a child span `f1` under the `f2` span.
|
|
x = f1(x) + 2
|
|
|
|
# You can also create a span for an arbitrary block of code using `with mlflow.start_span` context manager.
|
|
with mlflow.start_span(name="leaf", attributes={"operation": "exponentiation"}) as span:
|
|
# Inputs and outputs need to be set explicitly for manually created spans.
|
|
span.set_inputs({"x": x})
|
|
x = x**2
|
|
span.set_outputs({"x": x})
|
|
|
|
return x
|
|
|
|
|
|
assert f2(1) == 16
|
|
|
|
# You can access the last trace via get_last_active_trace_id API.
|
|
trace_id = mlflow.get_last_active_trace_id()
|
|
trace = mlflow.get_trace(trace_id)
|
|
|
|
# Alternatively, you can use `search_traces` API to retrieve
|
|
# traces that meet certain criteria.
|
|
traces = mlflow.search_traces(
|
|
filter_string="timestamp > 0",
|
|
max_results=1,
|
|
)
|
|
|
|
# Print the trace in JSON format
|
|
print(trace.to_json(pretty=True))
|
|
|
|
print(
|
|
"\033[92m"
|
|
+ "🤖Now run `mlflow server` and open MLflow UI to see the trace visualization!"
|
|
+ "\033[0m"
|
|
)
|