Skip to main content
Last updated on

Get started: MLflow Tracing for GenAI in a local IDE

This quickstart helps you integrate your GenAI app with MLflow Tracing if you use a local development environment such as an IDE (VS Code, PyCharm, Cursor, or others) or a locally hosted notebook environment, such as Jupyter. If you use a Databricks notebook, see the Databricks notebook quickstart instead.

This tutorial walks you through:

  • Creating an MLflow experiment in your Databricks workspace
  • Connecting your local development environment to the experiment
  • Running a simple GenAI application locally on your machine
  • Viewing the trace data uploaded to your Databricks workspace

trace

Prerequisites​

  • Access to a Databricks workspace.

Step 1: Install MLflow​

Install MLflow with Databricks connectivity:

Bash
pip install --upgrade "mlflow[databricks]>=3.1" openai

Step 2: Create a new MLflow experiment​

An MLflow experiment is the container for your GenAI application. For more information, see Experiments.

  1. Open your Databricks workspace.
  2. In the left sidebar, under AI/ML, click Experiments.
  3. At the top of the Experiments page, click GenAI apps & agents.
  4. To get the experiment ID and path, click the information icon Info icon. in the upper-left. You use these values in later steps.

create experiment

Step 3: Connect your environment to MLflow​

The following code snippets show how to set up authentication using a Databricks personal access token (PAT). MLflow also works with the other Databricks-supported authentication methods.

  1. In your MLflow Experiment, click the kebab menu icon Kebab menu icon. > Log traces locally > Generate API Key.

  2. Copy and run the generated code in your terminal.

    Bash
    export DATABRICKS_TOKEN=<databricks-personal-access-token>
    export DATABRICKS_HOST=https://<workspace-name>.cloud.databricks.com
    export MLFLOW_TRACKING_URI=databricks
    export MLFLOW_REGISTRY_URI=databricks-uc
    export MLFLOW_EXPERIMENT_ID=<experiment-id>

The steps above use the managed experiment backend, which is the fastest path to your first trace but caps traces at 100,000 per experiment. For scalable, governed storage with no trace limit, Databricks recommends binding your experiment to a Unity Catalog trace location.

Add the following to your terminal environment (alongside the variables from Step 3):

Bash
export MLFLOW_TRACING_SQL_WAREHOUSE_ID=<warehouse-id>

When you create app.py in Step 4, add the following at the top before any tracing calls:

Python
import mlflow
from mlflow.entities.trace_location import UnityCatalog

mlflow.set_experiment(
    experiment_id="<experiment-id>",  # the ID you copied in Step 2
    trace_location=UnityCatalog(
        catalog_name="<UC_CATALOG_NAME>",
        schema_name="<UC_SCHEMA_NAME>",
        table_prefix="<UC_TABLE_PREFIX>",
    ),
)

For full setup details, see Store OpenTelemetry traces in Unity Catalog.

Step 4: Create and instrument your application​

Create your GenAI app with tracing enabled. The code runs entirely on your local machine and uploads traces to your MLflow experiment in the Databricks workspace.

  1. Create a Python file named app.py in your project directory.

  2. Initialize an OpenAI client to connect to either Databricks-hosted LLMs or LLMs hosted by OpenAI.

    Use databricks-openai to get an OpenAI client that connects to Databricks-hosted LLMs. Select a model from the available foundation models.

    Python
    import mlflow
    from databricks_openai import DatabricksOpenAI

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

    # Traces are written to the experiment set by MLFLOW_EXPERIMENT_ID and the
    # tracking server set by MLFLOW_TRACKING_URI, both configured in the previous step

    # Create an OpenAI client that is connected to Databricks-hosted LLMs
    client = DatabricksOpenAI()

    # Select an LLM
    model_name = "databricks-claude-sonnet-5"

  3. Define and run your application:

    Use the @mlflow.trace decorator, which makes it easy to trace any Python function, combined with OpenAI automatic instrumentation to capture the details of the call to the OpenAI SDK.

    Python
    # 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(
            # Uses a Databricks-hosted LLM by default. To use an AI Gateway, Model Serving endpoint, or your own OpenAI credentials, replace `model_name` with a valid model such as `gpt-5`.
            model=model_name,
            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)

  4. Run the application:

    Bash
    python app.py

    This code runs locally: the app.py file lives on your machine, and the environment variables from Step 3 tell the MLflow client where to upload trace data. To change the prompt or agent logic, edit the local file and re-run it.

For details on adding tracing to apps, see Add traces to applications: automatic and manual tracing and MLflow Tracing Integrations (more than 20 library integrations).

Step 5: View the trace in MLflow​

When a trace appears in the Databricks UI, your code ran on your laptop and the trace was shipped over HTTPS.

  1. Return to the MLflow experiment you selected in Step 2.
  2. The generated trace appears in the Traces tab.
  3. Click the trace to view its details.

Trace Details

Understand the trace​

The new trace 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

Even this minimal trace surfaces useful information about your application's behavior, including:

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

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.

Guides and references​

For details on concepts and features in this guide, see:

On this page