Ablaufverfolgung von Databricks Foundation-Modellen

MLflow Tracing bietet automatische Tracing-Fähigkeiten für Databricks Foundation Models. Da Databricks Foundation Models eine openAI-kompatible API verwenden, können Sie die automatische Ablaufverfolgung aktivieren, indem Sie die mlflow.openai.autolog Funktion aufrufen, und MLflow erfasst Ablaufverfolgungen für LLM-Aufrufe und protokolliert sie beim aktiven MLflow Experiment.

import mlflow

mlflow.openai.autolog()

Die MLflow-Ablaufverfolgung erfasst automatisch die folgenden Informationen zu Databricks Foundation Model-Aufrufen:

• Eingabeaufforderungen und Vervollständigungsantworten
• Wartezeiten
• Modellname und Endpunkt
• Zusätzliche Metadaten wie z. B. temperature, max_tokens, wenn angegeben
• Funktionsaufrufe, wenn sie in der Antwort enthalten sind
• Jede ausgelöste Ausnahme

Voraussetzungen

Um die MLflow-Ablaufverfolgung mit Databricks Foundation Models zu verwenden, müssen Sie MLflow und das OpenAI SDK installieren (da Databricks Foundation Models eine openAI-kompatible API verwenden).

Entwicklung

Installieren Sie für Entwicklungsumgebungen das vollständige MLflow-Paket mit Databricks-Extras und dem OpenAI SDK:

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

Das vollständige mlflow[databricks] Paket enthält alle Features für die lokale Entwicklung und Das Experimentieren mit Databricks.

Produktion

Installieren Sie mlflow-tracing und das OpenAI SDK für Produktionsbereitstellungen.

pip install --upgrade mlflow-tracing openai

Das Paket ist für den mlflow-tracing Produktionseinsatz optimiert.

Bevor Sie die Beispiele ausführen, müssen Sie Ihre Umgebung konfigurieren:

Für Benutzer außerhalb von Databricks-Notizbüchern: Legen Sie Ihre Databricks-Umgebungsvariablen fest:

export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="your-personal-access-token"

Für Benutzer innerhalb von Databricks-Notizbüchern: Diese Anmeldeinformationen werden automatisch für Sie festgelegt.

Unterstützte APIs

MLflow unterstützt die automatische Ablaufverfolgung für die folgenden Databricks Foundation Model-APIs:

Um Support für zusätzliche APIs anzufordern, öffnen Sie eine Featureanfrage auf GitHub.

Einfaches Beispiel

import mlflow
import os
from openai import OpenAI

# Databricks Foundation Model APIs use Databricks authentication.

# Enable auto-tracing for OpenAI (which will trace Databricks Foundation Model API calls)
mlflow.openai.autolog()

# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/databricks-foundation-models-demo")

# Create OpenAI client configured for Databricks
client = OpenAI(
    api_key=os.environ.get("DATABRICKS_TOKEN"),
    base_url=f"{os.environ.get('DATABRICKS_HOST')}/serving-endpoints"
)

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

response = client.chat.completions.create(
    model="databricks-llama-4-maverick",
    messages=messages,
    temperature=0.1,
    max_tokens=100,
)

Streamen

Die MLflow-Ablaufverfolgung unterstützt die Streaming-API von Databricks Foundation Models. Mit dem gleichen Setup der automatischen Ablaufverfolgung verfolgt MLflow die Streamingantwort automatisch und rendert die verkettete Ausgabe in der Span-UI.

import mlflow
import os
from openai import OpenAI

# Enable auto-tracing for OpenAI (which will trace Databricks Foundation Model API calls)
mlflow.openai.autolog()

# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/databricks-streaming-demo")

# Create OpenAI client configured for Databricks
client = OpenAI(
    api_key=os.environ.get("DATABRICKS_TOKEN"),
    base_url=f"{os.environ.get('DATABRICKS_HOST')}/serving-endpoints"
)

stream = client.chat.completions.create(
    model="databricks-llama-4-maverick",
    messages=[
        {"role": "user", "content": "Explain the benefits of using Databricks Foundation Models"}
    ],
    stream=True,  # Enable streaming response
)

for chunk in stream:
    print(chunk.choices[0].delta.content or "", end="")

Funktionsaufruf

MLflow Tracing erfasst automatisch die Antwort auf Funktionsaufrufe von Databricks Foundation Models. Die Funktionsanweisung in der Antwort wird in der Ablaufverfolgungsschnittstelle hervorgehoben. Darüber hinaus können Sie die Toolfunktion mit dem @mlflow.trace Dekorator annotieren, um einen Bereich für die Toolausführung zu erstellen.

Im folgenden Beispiel wird ein einfacher Funktionsaufruf-Agent mit Databricks Foundation Models und MLflow Tracing implementiert.

import json
import os
from openai import OpenAI
import mlflow
from mlflow.entities import SpanType

# Enable auto-tracing for OpenAI (which will trace Databricks Foundation Model API calls)
mlflow.openai.autolog()

# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/databricks-function-agent-demo")

# Create OpenAI client configured for Databricks
client = OpenAI(
    api_key=os.environ.get("DATABRICKS_TOKEN"),
    base_url=f"{os.environ.get('DATABRICKS_HOST')}/serving-endpoints"
)

# Define the tool function. Decorate it with `@mlflow.trace` to create a span for its execution.
@mlflow.trace(span_type=SpanType.TOOL)
def get_weather(city: str) -> str:
    if city == "Tokyo":
        return "sunny"
    elif city == "Paris":
        return "rainy"
    return "unknown"

tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "parameters": {
                "type": "object",
                "properties": {"city": {"type": "string"}},
            },
        },
    }
]

_tool_functions = {"get_weather": get_weather}

# Define a simple tool calling agent
@mlflow.trace(span_type=SpanType.AGENT)
def run_tool_agent(question: str):
    messages = [{"role": "user", "content": question}]

    # Invoke the model with the given question and available tools
    response = client.chat.completions.create(
        model="databricks-llama-4-maverick",
        messages=messages,
        tools=tools,
    )
    ai_msg = response.choices[0].message

    # If the model requests tool call(s), invoke the function with the specified arguments
    if tool_calls := ai_msg.tool_calls:
        for tool_call in tool_calls:
            function_name = tool_call.function.name
            if tool_func := _tool_functions.get(function_name):
                args = json.loads(tool_call.function.arguments)
                tool_result = tool_func(**args)
            else:
                raise RuntimeError("An invalid tool is returned from the assistant!")

            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": tool_result,
                }
            )

        # Send the tool results to the model and get a new response
        response = client.chat.completions.create(
            model="databricks-llama-4-maverick", messages=messages
        )

    return response.choices[0].message.content

# Run the tool calling agent
question = "What's the weather like in Paris today?"
answer = run_tool_agent(question)

Verfügbare Modelle

Databricks Foundation Models bietet Zugriff auf eine Vielzahl modernster Modelle, darunter Llama, Anthropic und andere führende Foundation-Modelle.

Die vollständige und up-to-aktuelle Liste der verfügbaren Modelle mit ihren Modell-IDs finden Sie in der Dokumentation zu den Databricks Foundation Models.

Automatische Ablaufverfolgung deaktivieren

Die automatische Ablaufverfolgung für Databricks Foundation Models kann global durch Aufrufen von mlflow.openai.autolog(disable=True) oder mlflow.autolog(disable=True) deaktiviert werden.