Share via

Tracing a GenAI App (Notebook)

This quickstart helps you integrate your GenAI app with MLflow Tracing if you use a Databricks Notebook as your development environment. If you use a local IDE, please use the IDE quickstart instead.

What you'll achieve

By the end of this tutorial, you will have:

  • A Databricks Notebook with a linked MLflow Experiment for your GenAI app
  • A simple GenAI application instrumented with MLflow Tracing
  • A trace from that app in your MLflow Experiment

trace

Prerequisites

  • Databricks Workspace: Access to a Databricks workspace.

Step 1: Create a Databricks Notebook

Note

Creating a Databricks Notebook will create an MLflow Experiment that is the container for your GenAI application. Learn more about the Experiment and what it contains in the data model section.

  1. Open your Databricks workspace
  2. Go to New at the top of the left sidebar
  3. Click Notebook

Notebook Creation

Databricks runtimes include MLflow. However, for the best experience with GenAI capabilities, including the most comprehensive tracing features and robust support, it is highly recommended to use the latest version of MLflow.

You can update MLflow in your notebook by running:

%pip install --upgrade "mlflow[databricks]>=3.1" openai
dbutils.library.restartPython()
  • mlflow[databricks]>=3.1: This command ensures you have MLflow 3.1 or a more recent version, along with the databricks extra for seamless connectivity and functionality within Databricks.
  • dbutils.library.restartPython(): This is crucial to ensure the newly installed version is used by the Python kernel.

Warning

While tracing features 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 and robust support.

Step 3: Instrument your application

Tip

Databricks provided out of the box access to popular frontier and open source foundational LLMs. To run this quickstart, you can:

  1. Use the Databricks hosted LLMs
  2. Directly use your own API key from an LLM provider
  3. Create an external model to enable governed access to your LLM provider's API keys

If you prefer to use one of the other 20+ LLM SDKs (Anthropic, Bedrock, etc) or GenAI authoring frameworks (LangGraph, etc) that MLflow supports, follow the instructions that appear when you create a new MLflow Experiment.

Select the appropriate integration for your application:

  1. Create a cell in your notebook with the below code

    Here, we use the @mlflow.trace decorator that makes it easy to trace any Python application combined with the OpenAI automatic instrumentation to captures the details of the call to the OpenAI SDK.

    The code snippet below uses Anthropic's Claude Sonnet LLM. You can select another LLM from the list of supported foundational models.

    import mlflow
from databricks.sdk import WorkspaceClient

# Enable MLflow's autologging to instrument your application with Tracing
mlflow.openai.autolog()

# Create an OpenAI client that is connected to Databricks hosted LLMs
mlflow_creds = mlflow.utils.databricks_utils.get_databricks_host_creds()
client = OpenAI(
    api_key=mlflow_creds.token,
    base_url=f"{mlflow_creds.host}/serving-endpoints"
)

# Use the trace decorator to capture the application's entry point
@mlflow.trace
def my_app(input: str):
  # This call is automatically instrumented by `mlflow.openai.autolog()`
  response = client.chat.completions.create(
    model="databricks-claude-sonnet-4",  # This example uses a Databricks hosted LLM - you can replace this with any AI Gateway or Model Serving endpoint. If you have an external model endpoint configured, you can use that model name here.
    messages=[
      {
        "role": "system",
        "content": "You are a helpful assistant.",
      },
      {
        "role": "user",
        "content": input,
      },
    ],
  )
  return response.choices[0].message.content

result = my_app(input="What is MLflow?")
print(result)

  2. Run the Notebook cell to execute the code

Step 4: View the Trace in MLflow

  1. The Trace will appear below the Notebook cell
  2. Optionally, you can go to the MLflow Experiment UI to see the Trace
    1. Click on the Experiment icon on the right side of your screen
    2. Click the open icon next to Experiment Runs
    3. You will now see the generated trace in the Traces tab
    4. Click on the trace to view its details

Understanding the Trace

The trace you've just created shows:

  • Root Span: Represents the inputs to the my_app(...) function
    • Child Span: Represents the OpenAI completion request
  • Attributes: Contains metadata like model name, token counts, and timing information
  • Inputs: The messages sent to the model
  • Outputs: The response received from the model

This simple trace already provides valuable insights into your application's behavior, such as:

  • What was asked
  • What response was generated
  • How long the request took
  • How many tokens were used (affecting cost)

Tip

For more complex applications like RAG systems or multi-step agents, MLflow Tracing provides even more value by revealing the inner workings of each component and step.

Next steps

Continue your journey with these recommended actions and tutorials.

Reference guides

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