mlflow_trace_decorator_quickstart_with_nested_functions.py

1import mlflow
2
3# Set the experiment to organize our traces
4mlflow.set_experiment("Tracing Quickstart")
5
6@mlflow.trace
7def identity(x):
8    return x
9
10@mlflow.trace
11def add(a, b):
12    return a + b
13
14@mlflow.trace
15def workflow(a, b):
16    # Traces will be nested automatically when a 
17    # traced function calls another traced function
18    return identity(add(a, b))
19
20# Execute the traced function
21result = workflow(1, 2)
22print(f"Result: {result}")