mlflow_trace_decorator_quickstart_with_automatic_function_instrumentation.py

1import mlflow
2
3# Set the experiment to organize our traces
4mlflow.set_experiment("MLflow Tracing Quickstart")
5
6# Use the @mlflow.trace decorator to automatically instrument the function
7@mlflow.trace
8def predict(input_text: str):
9    # This represents a call to an LLM or any processing logic
10    return f"Response to: {input_text}"
11
12# When the function is called, MLflow automatically starts a trace,
13# records the input and output, and logs it to the MLflow server.
14response = predict("What is MLflow Tracing?")
15
16print(f"Result: {response}")