MLflow Tracing SDK

Common errors

• ModuleNotFoundError: No module named 'mlflow'

  cause The `mlflow` package, which contains the full MLflow SDK, is not installed or accessible in the current Python environment. The `mlflow-tracing` package is a lightweight subset and does not include the full `mlflow` package, leading to this error if you try to import `mlflow` directly.
  fix If full MLflow functionality is needed, install the complete package: `pip install mlflow`. If only tracing is intended, ensure `mlflow-tracing` is installed (`pip install mlflow-tracing`) and use imports specifically from `mlflow.tracing` where applicable.

• UnicodeEncodeError: 'charmap' codec can't encode character '\u03b1' in position X: character maps to <undefined>

  cause This error occurs when MLflow Tracing attempts to log data containing special characters (e.g., non-ASCII) using a system default encoding (like 'charmap' on Windows) that cannot represent those characters, typically during the process of exporting trace data to the backend.
  fix Ensure all strings passed to MLflow tracing APIs are UTF-8 encoded. A common workaround is to set the `PYTHONIOENCODING` environment variable to `utf-8` before running your application (e.g., `export PYTHONIOENCODING=utf-8` on Linux/macOS or `set PYTHONIOENCODING=utf-8` on Windows).

• WARNING mlflow.tracing.fluent: Failed to start span ... 'NoneType' object has no attribute 'set_span_type'.

  cause This warning/error indicates a failure in MLflow's autologging or manual span creation, often due to conflicts with OpenTelemetry auto-instrumentation or unexpected `None` values being passed where a valid object or span context is expected, preventing a new span from being properly initialized or linked.
  fix Verify that MLflow's tracing provider is correctly initialized and not being inadvertently reset or overridden by other OpenTelemetry configurations. Ensure that data being traced (especially in autologging scenarios) is valid and not `None` for required attributes. Upgrading `mlflow-tracing` and related AI/LLM library integrations may also resolve known compatibility issues.

• name 'MlflowClient' is not defined

  cause The `mlflow-tracing` package, being a minimalist SDK, intentionally does not include the `MlflowClient` class or other components of the full MLflow tracking client. This error occurs when code attempts to instantiate or use `mlflow.tracking.MlflowClient` while only `mlflow-tracing` is installed.
  fix If programmatic interaction with the MLflow Tracking Server (e.g., `search_runs`, `get_run`, model registry operations) is required, install the full `mlflow` package (`pip install mlflow`). If only tracing functionality is needed, refactor your code to use the `mlflow.tracing` fluent APIs and context managers without relying on `MlflowClient`.

Warnings

• breaking Do NOT co-install `mlflow-tracing` with the full `mlflow` package.
  Fix: If you need the full MLflow features, install `mlflow`. If you only need tracing in production, use `mlflow-tracing`. Uninstall one before installing the other (`pip uninstall mlflow && pip install mlflow-tracing` or vice-versa).
• gotcha Calling `mlflow.set_tracking_uri()` after OpenTelemetry auto-instrumentation can reset the global `TracerProvider`.
  Fix: To avoid silent data loss or breaking other telemetry, ensure `mlflow.set_tracking_uri()` is called early in your application's lifecycle, preferably before other OpenTelemetry instrumentation is initialized. This issue mainly applies when combining MLflow tracing with external OTel auto-instrumentation.
• gotcha Using a file-based backend store for the MLflow server can lead to poor UI/SDK performance.
  Fix: For better performance in production or with significant trace volume, configure your MLflow server to use a database-based backend store (e.g., PostgreSQL, MySQL, SQLite with `--backend-store-uri sqlite:///mlflow.db`).
• gotcha Traces might get stuck in 'in progress' and not be viewable if a process hangs or runs too long.
  Fix: Set the `MLFLOW_TRACE_TIMEOUT_SECONDS` environment variable to automatically halt and export traces that exceed a specified duration. This allows analysis even for stuck traces.
• gotcha MLflow client and server version compatibility is important for new features.
  Fix: While `mlflow-tracing` SDK and the MLflow server within the same major version are generally compatible, it's recommended to keep both client and server up-to-date to ensure new tracing features (introduced in MLflow 2.14.0 and enhanced in 3.x) are available and function correctly.
• breaking The test script failed to run due to a missing Python module (`openai`). This prevents any evaluation of the library's features.
  Fix: Install the missing 'openai' package in the test environment (`pip install openai`) or ensure the test environment is correctly configured with all necessary dependencies for the test script.
• breaking A `ModuleNotFoundError` for 'openai' indicates a missing dependency for functionalities that might integrate MLflow with OpenAI, such as tracing OpenAI API calls, or for running specific test scenarios.
  Fix: Install the `openai` package using `pip install openai`.

Install

• pip install mlflow-tracing Install MLflow Tracing SDK

Imports

• mlflow
  wrong:

  from mlflow import tracing

  correct:

  import mlflow

  While `mlflow.tracing` contains utility functions, the main entry points for autologging, manual tracing (e.g., `@mlflow.trace`), and configuration are directly under the `mlflow` namespace.
• mlflow.tracing.configure

  from mlflow import tracing
  tracing.configure(...)

  Used for advanced configuration like span processors.
• mlflow.tracing.disable

  from mlflow import tracing
  tracing.disable()

  Used to temporarily disable tracing.
• mlflow.tracing.enable

  from mlflow import tracing
  tracing.enable()

  Used to re-enable tracing if previously disabled.
• @mlflow.trace

  import mlflow
  
  @mlflow.trace
  def my_traced_function():
      pass

  Decorator for manual function instrumentation.
• mlflow.start_span

  import mlflow
  
  with mlflow.start_span('my_span_name'):
      # ... code to trace ...

  Context manager for manual code block instrumentation.

Quickstart

This quickstart demonstrates how to set up MLflow Tracing for an OpenAI call. It configures the MLflow tracking URI, sets an experiment, enables autologging for OpenAI, and then performs a simple API call. The trace for this call will be automatically logged and viewable in the MLflow UI. Make sure an MLflow server is running and `OPENAI_API_KEY` is set in your environment.

import os
import mlflow
from openai import OpenAI

# Set your MLflow Tracking URI (replace with your server, e.g., 'http://localhost:5000')
# For Databricks, use 'databricks' and ensure DATABRICKS_HOST/TOKEN are set.
mlflow.set_tracking_uri(os.environ.get('MLFLOW_TRACKING_URI', 'http://127.0.0.1:5000'))

# Set a new MLflow experiment to log traces to
mlflow.set_experiment("my_genai_app_traces")

# Ensure OpenAI API key is set for the example
if not os.environ.get("OPENAI_API_KEY"): 
    # In a real app, use a secure way to load keys (e.g., environment variable, secret manager)
    # For quick testing, you can set it directly here, but it's not recommended for production
    print("WARNING: OPENAI_API_KEY environment variable not set. Skipping OpenAI example.")
    openai_client = None
else:
    # Enable auto-tracing for OpenAI calls
    mlflow.openai.autolog()
    
    # Initialize OpenAI client
    openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

    # Make an OpenAI call - this will be automatically traced
    print("Invoking OpenAI completion...")
    response = openai_client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "You are a helpful AI assistant."},
            {"role": "user", "content": "Tell me a fun fact about Python programming."}
        ],
        max_tokens=50
    )
    print("OpenAI Response:", response.choices[0].message.content)
    print("Trace should now be visible in MLflow UI under 'my_genai_app_traces' experiment.")