MLflow Tracing-Integrationen

MLflow Tracing ist in eine breite Palette beliebter generativer KI-Bibliotheken und Frameworks integriert und bietet eine einzeilige automatische Ablaufverfolgung für alle. Dies ermöglicht es Ihnen, sofortige Beobachtbarkeit in Ihren GenAI-Anwendungen mit minimalem Setup zu erlangen.

Diese umfassende Unterstützung bedeutet, dass Sie die Observability ohne erhebliche Codeänderungen erhalten können, indem Sie die Tools einsetzen, die Sie bereits im Einsatz haben. Für benutzerdefinierte Komponenten oder nicht unterstützte Bibliotheken bietet MLflow auch leistungsstarke manuelle Ablaufverfolgungs-APIs.

Die automatische Ablaufverfolgung erfasst die Logik und Zwischenschritte Ihrer Anwendung, wie LLM-Aufrufe, Toolnutzung und Agenteninteraktionen, je nachdem, wie Sie die spezifische Bibliothek oder das SDK implementiert haben.

Top-Integrationen auf einen Blick

Hier sind Schnellstartbeispiele für einige der am häufigsten verwendeten Integrationen. Klicken Sie auf eine Registerkarte, um ein einfaches Verwendungsbeispiel anzuzeigen. Detaillierte Voraussetzungen und erweiterte Szenarien finden Sie auf ihren dedizierten Integrationsseiten (verlinkt über die Registerkarten oder die nachstehende Liste).

OpenAI

import mlflow
import openai

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# import os
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

# Enable auto-tracing for OpenAI
mlflow.openai.autolog()

# Set up MLflow tracking
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/openai-tracing-demo")

openai_client = openai.OpenAI()

messages = [
    {
        "role": "user",
        "content": "What is the capital of France?",
    }
]

response = openai_client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    temperature=0.1,
    max_tokens=100,
)
# View trace in MLflow UI

Vollständiger OpenAI-Integrationsleitfaden

LangChain

import mlflow
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# import os
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.langchain.autolog()

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/langchain-tracing-demo")

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7, max_tokens=1000)
prompt = PromptTemplate.from_template("Tell me a joke about {topic}.")
chain = prompt | llm | StrOutputParser()

chain.invoke({"topic": "artificial intelligence"})
# View trace in MLflow UI

Vollständiges LangChain-Integrationshandbuch

LangGraph

import mlflow
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# import os
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.langchain.autolog() # LangGraph uses LangChain's autolog

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/langgraph-tracing-demo")

@tool
def get_weather(city: str):
    """Use this to get weather information."""
    return f"It might be cloudy in {city}"

llm = ChatOpenAI(model="gpt-4o-mini")
graph = create_react_agent(llm, [get_weather])
result = graph.invoke({"messages": [("user", "what is the weather in sf?")]})
# View trace in MLflow UI

Vollständiges LangGraph-Integrationshandbuch

Anthropic

import mlflow
import anthropic
import os

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.anthropic.autolog()

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/anthropic-tracing-demo")

client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
message = client.messages.create(
    model="claude-3-5-sonnet-20241022",
    max_tokens=1024,
    messages=[{"role": "user", "content": "Hello, Claude"}],
)
# View trace in MLflow UI

Vollständiges anthropisches Integrationshandbuch

DSPy

import mlflow
import dspy

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# import os
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.dspy.autolog()

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/dspy-tracing-demo")

lm = dspy.LM("openai/gpt-4o-mini") # Assumes OPENAI_API_KEY is set
dspy.configure(lm=lm)

class SimpleSignature(dspy.Signature):
    input_text: str = dspy.InputField()
    output_text: str = dspy.OutputField()

program = dspy.Predict(SimpleSignature)
result = program(input_text="Summarize MLflow Tracing.")
# View trace in MLflow UI

Vollständiger DSPy-Integrationsleitfaden

Databricks, eine Plattform für Datenanalyse

import mlflow
import os
from openai import OpenAI # Databricks FMAPI uses OpenAI client

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.openai.autolog() # Traces Databricks FMAPI using OpenAI client

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/databricks-fmapi-tracing")

client = OpenAI(
    api_key=os.environ.get("DATABRICKS_TOKEN"),
    base_url=f"{os.environ.get('DATABRICKS_HOST')}/serving-endpoints"
)
response = client.chat.completions.create(
    model="databricks-llama-4-maverick",
    messages=[{"role": "user", "content": "Key features of MLflow?"}],
)
# View trace in MLflow UI

Vollständiges Databricks-Integrationshandbuch

Bedrock

import mlflow
import boto3

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# import os
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.bedrock.autolog()

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/bedrock-tracing-demo")

bedrock = boto3.client(
    service_name="bedrock-runtime",
    region_name="us-east-1" # Replace with your region
)
response = bedrock.converse(
    modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
    messages=[{"role": "user", "content": "Hello World in one line."}]
)
# View trace in MLflow UI

Vollständiger Bedrock-Integrationsleitfaden

AutoGen

import mlflow
from autogen import ConversableAgent
import os

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.autogen.autolog()

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/autogen-tracing-demo")

config_list = [{"model": "gpt-4o-mini", "api_key": os.environ.get("OPENAI_API_KEY")}]
assistant = ConversableAgent("assistant", llm_config={"config_list": config_list})
user_proxy = ConversableAgent("user_proxy", human_input_mode="NEVER", code_execution_config=False)

user_proxy.initiate_chat(assistant, message="What is 2+2?")
# View trace in MLflow UI

Vollständiges AutoGen-Integrationshandbuch

Sichere API-Schlüsselverwaltung

Für Produktionsumgebungen empfiehlt Databricks die Verwendung von AI Gateway oder Databricks Secrets zur Verwaltung von API-Schlüsseln. AI Gateway ist die bevorzugte Methode und bietet zusätzliche Governance-Features.

AI-Gateway (empfohlen)

Databricks empfiehlt Mosaik AI Gateway für die Steuerung und Überwachung des Zugriffs auf generative KI-Modelle.

Erstellen Sie einen Foundation Model-Endpunkt, der mit AI-Gateway konfiguriert ist:

1. Gehen Sie in Ihrem Databricks-Arbeitsbereich zu Serving und >.
2. Wählen Sie einen Endpunkttyp und einen Anbieter aus.
3. Konfigurieren Sie den Endpunkt mit Ihrem API-Schlüssel.
4. Aktivieren Sie während der Endpunktkonfiguration AI-Gateway , und konfigurieren Sie bei Bedarf Geschwindigkeitsbeschränkungen, Fallbacks und Schutzläufe.
5. Sie können automatisch generierten Code abrufen, um schnell mit der Abfrage des Endpunkts zu beginnen. Gehen Sie zu Serving> Ihr Endpunkt >Verwendung>Abfrage. Stellen Sie sicher, den Tracing-Code hinzuzufügen:

import mlflow
from openai import OpenAI
import os

# How to get your Databricks token: https://docs.databricks.com/en/dev-tools/auth/pat.html
# DATABRICKS_TOKEN = os.environ.get('DATABRICKS_TOKEN')
# Alternatively in a Databricks notebook you can use this:
DATABRICKS_TOKEN = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()

# Enable auto-tracing for OpenAI
mlflow.openai.autolog()

# Set up MLflow tracking (if running outside Databricks)
# If running in a Databricks notebook, these are not needed.
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/my-genai-app")

client = OpenAI(
  api_key=DATABRICKS_TOKEN,
  base_url="<YOUR_HOST_URL>/serving-endpoints"
)

chat_completion = client.chat.completions.create(
  messages=[
  {
    "role": "system",
    "content": "You are an AI assistant"
  },
  {
    "role": "user",
    "content": "What is MLflow?"
  }
  ],
  model="<YOUR_ENDPOINT_NAME>",
  max_tokens=256
)

print(chat_completion.choices[0].message.content)

Databricks Geheimnisse

Verwenden Sie geheime Databricks-Schlüssel zum Verwalten von API-Schlüsseln:

1. Erstellen Sie einen geheimen Bereich, und speichern Sie Ihren API-Schlüssel:

   from databricks.sdk import WorkspaceClient
   
   # Set your secret scope and key names
   secret_scope_name = "llm-secrets"  # Choose an appropriate scope name
   secret_key_name = "api-key"        # Choose an appropriate key name
   
   # Create the secret scope and store your API key
   w = WorkspaceClient()
   w.secrets.create_scope(scope=secret_scope_name)
   w.secrets.put_secret(
       scope=secret_scope_name,
       key=secret_key_name,
       string_value="your-api-key-here"  # Replace with your actual API key
   )
   

2. Abrufen und Verwenden des geheimen Schlüssels in Ihrem Code:

   import mlflow
   import openai
   import os
   
   # Configure your secret scope and key names
   secret_scope_name = "llm-secrets"
   secret_key_name = "api-key"
   
   # Retrieve the API key from Databricks secrets
   os.environ["OPENAI_API_KEY"] = dbutils.secrets.get(
       scope=secret_scope_name,
       key=secret_key_name
   )
   
   # Enable automatic tracing
   mlflow.openai.autolog()
   
   # Use OpenAI client with securely managed API key
   client = openai.OpenAI()
   response = client.chat.completions.create(
       model="gpt-4o-mini",
       messages=[{"role": "user", "content": "Explain MLflow Tracing"}],
       max_tokens=100
   )
   

Aktivierung mehrerer Integrationen für die automatische Rückverfolgung

Da GenAI-Anwendungen häufig mehrere Bibliotheken kombinieren, können Sie mit der MLflow-Ablaufverfolgung die automatische Ablaufverfolgung für mehrere Integrationen gleichzeitig aktivieren, was eine einheitliche Ablaufverfolgung ermöglicht.

Um zum Beispiel sowohl LangChain als auch direktes OpenAI-Tracking zu aktivieren:

import mlflow

# Enable MLflow Tracing for both LangChain and OpenAI
mlflow.langchain.autolog()
mlflow.openai.autolog()

# Your code using both LangChain and OpenAI directly...
# ... an example can be found on the Automatic Tracing page ...

MLflow generiert eine einzelne, zusammenhängende Ablaufverfolgung, die Schritte aus LangChain- und direct OpenAI LLM-Aufrufen kombiniert, sodass Sie den vollständigen Ablauf überprüfen können. Weitere Beispiele für die Kombination von Integrationen finden Sie auf der Seite "Automatische Ablaufverfolgung ".

Automatische Ablaufverfolgung deaktivieren

Die automatische Ablaufverfolgung für eine bestimmte Bibliothek kann durch das Aufrufen von mlflow.<library>.autolog(disable=True) deaktiviert werden. Um alle Integrationen für die automatische Protokollierung auf einmal zu deaktivieren, verwenden Sie mlflow.autolog(disable=True).

import mlflow

# Disable for a specific library
mlflow.openai.autolog(disable=True)

# Disable all autologging
mlflow.autolog(disable=True)