Aller au contenu principal
Dernière mise à jour le

Traçage des modèles de fondation Databricks

MLflow Tracing offre une capacité de traçage automatique pour les modèles de base Databricks. Étant donné que les modèles de base Databricks utilisent une API compatible OpenAI, vous pouvez activer le traçage automatique en appelant la fonction mlflow.openai.autolog, et MLflow capturera les traces pour l'invocation de LLM et les Log dans l'Experimentation MLflow active.

Python
import mlflow

mlflow.openai.autolog()

MLflow Trace capture automatiquement les informations suivantes concernant les appels de modèles de fondation Databricks :

  • Prompts et réponses de complétion
  • Latences
  • Nom du modèle et Endpoint
  • Métadonnées supplémentaires telles que temperature, max_tokens, si spécifiées
  • Appel de fonction si renvoyé dans la réponse
  • Toute exception si levée
remarque

Sur les clusters de compute serverless, l'autologging n'est pas activé automatiquement. Vous devez appeler explicitement mlflow.openai.autolog() pour activer le traçage automatique de cette intégration.

Prérequis​

Pour utiliser MLflow Tracing avec les modèles de fondation Databricks, vous devez installer MLflow et le SDK OpenAI (car les modèles de fondation Databricks utilisent une API compatible avec OpenAI).

Pour les environnements de développement, installez le package MLflow complet avec les modules complémentaires Databricks et le SDK OpenAI :

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

Le package mlflow[databricks] complet inclut toutes les fonctionnalités pour le développement local et l’expérimentation sur Databricks.

remarque

MLflow 3 est fortement recommandé pour la meilleure expérience de traçage avec les modèles de base Databricks.

Avant d'exécuter les exemples, vous devrez configurer votre environnement :

Pour les utilisateurs en dehors des notebooks Databricks : Définissez vos variables d'environnement Databricks :

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

Pour les utilisateurs dans les Notebooks Databricks : ces identifiants sont définis automatiquement pour vous.

APIs prises en charge​

MLflow prend en charge le traçage automatique pour les APIs Databricks Foundation Model suivantes :

Achèvement du chat

Appel de fonction

Streaming

Asynchrone

✅

✅

✅

✅

Achèvement du chat

Appel de fonction

Streaming

Asynchrone

✅

✅

✅

✅

Pour demander la prise en charge d'APIs supplémentaires, veuillez ouvrir une demande de fonctionnalité sur GitHub.

Exemple de base​

Python
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,
)

Trace des modèles de fondation Databricks

streaming​

MLflow Tracing prend en charge l'API de streaming des modèles de fondation Databricks. Avec la même configuration de traçage automatique, MLflow trace automatiquement la réponse en streaming et affiche la sortie concaténée dans l'interface utilisateur d'étendue.

Python
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="")

Trace de streaming des modèles de fondation Databricks

Appel de fonction​

MLflow Tracing capture automatiquement la réponse d'appel de fonction des Modèles Foundation Databricks. L’instruction de fonction dans la réponse sera mise en surbrillance dans l’interface utilisateur de la trace. De plus, vous pouvez annoter la fonction d’outil avec le décorateur @mlflow.trace pour créer un intervalle pour l’exécution de l’outil.

L’exemple suivant implémente un agent d’appel de fonction simple utilisant les modèles de base de Databricks et MLflow Tracing.

Python
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)

Trace d’appel de fonction des modèles de fondation Databricks.

Modèles disponibles​

Databricks Foundation Models donne accès à une variété de modèles de pointe, notamment Llama, Anthropic et d'autres modèles de fondation de premier plan.

Pour la liste complète et la plus à jour des modèles disponibles et de leurs ID de modèle, veuillez consulter la documentation sur les modèles de fondation Databricks.

Désactiver le suivi automatique​

Le traçage automatique des modèles de fondation Databricks peut être désactivé globalement en appelant mlflow.openai.autolog(disable=True) ou mlflow.autolog(disable=True).

Sur cette page