Aller au contenu principal
Dernière mise à jour le

Traçage automatique

Ajoutez une ligne de code, mlflow.<library>.autolog(), pour tracer automatiquement votre application d'IA générative. Le traçage automatique fonctionne avec plus de 20 bibliothèques et frameworks pris en charge prêts à l'emploi.

remarque

Sur les clusters de compute Serverless, la journalisation automatique pour les cadres de traçage genAI n'est pas activée automatiquement. Vous devez activer explicitement l'autologging en appelant la fonction mlflow.<library>.autolog() appropriée pour les intégrations spécifiques que vous souhaitez suivre.

Prérequis​

Databricks recommande MLflow 3 pour les dernières fonctionnalités de traçage GenAI.

Exécutez ce qui suit dans un Notebook Databricks pour installer le package mlflow et le package d'intégration que vous souhaitez utiliser. Cet exemple utilise OpenAI :

  • mlflow[databricks] >=3.1 : Fonctionnalités MLflow de base avec fonctionnalités GenAI et connectivité Databricks.
  • openai>=1.0.0 : Nécessaire uniquement pour exécuter l'exemple de traçage automatique de base sur cette page (si vous utilisez d'autres fournisseurs de LLM, installez plutôt leurs SDK respectifs).
  • Bibliothèques supplémentaires : installez les bibliothèques spécifiques pour les intégrations que vous souhaitez utiliser.

Installer les exigences de base :

Python
%pip install --upgrade "mlflow[databricks]>=3.1" "openai>=1.0.0"
# Also install libraries you want to trace (langchain, anthropic, etc.)

dbutils.library.restartPython()

Configurer les identifiants​

Dans un Notebook Databricks, définissez toutes les clés d'API LLM nécessaires :

Python
import os
os.environ["OPENAI_API_KEY"] = "your-api-key"

# Add other provider keys as needed
# os.environ["ANTHROPIC_API_KEY"] = "your-api-key"
# os.environ["MISTRAL_API_KEY"] = "your-api-key"

Exemple de traçage automatique​

MLflow trace automatiquement plus de 20 frameworks pris en charge. Voici comment activer le traçage automatique pour les agents OpenAI se connectant aux APIs de modèle de fondation Databricks.

Python
import mlflow
import os
from openai import OpenAI

# Databricks Foundation Model APIs use Databricks authentication.

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/databricks-sdk-autolog-example")

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

# 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"
)

# Query Llama 4 Maverick using OpenAI client
response = client.chat.completions.create(
    model="databricks-llama-4-maverick",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What are the key features of MLflow Tracing?"}
    ],
    max_tokens=150,
    temperature=0.7
)

print(response.choices[0].message.content)
# Your calls to Databricks Foundation Model APIs are automatically traced!

Traçage automatique de plusieurs frameworks​

Vous pouvez utiliser le traçage automatique pour plusieurs frameworks dans le même agent. Découvrez tous les frameworks pris en charge.

Le code suivant combine les appels d'API directs OpenAI, les chaînes LangChain et la logique personnalisée dans une seule trace pour faciliter le debugging et le monitoring.

Python
%pip install --upgrade "langchain==0.3.7" "langchain-openai==0.2.8"
Python
import mlflow
import openai
from mlflow.entities import SpanType
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

# Enable auto-tracing for both OpenAI and LangChain
mlflow.openai.autolog()
mlflow.langchain.autolog()

# Create OpenAI client
client = openai.OpenAI()

@mlflow.trace(span_type=SpanType.CHAIN)
def multi_provider_workflow(query: str):
    # First, use OpenAI directly for initial processing
    analysis = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "Analyze the query and extract key topics."},
            {"role": "user", "content": query}
        ]
    )
    topics = analysis.choices[0].message.content

    # Then use LangChain for structured processing
    llm = ChatOpenAI(model="gpt-4o-mini")
    prompt = ChatPromptTemplate.from_template(
        "Based on these topics: {topics}\nGenerate a detailed response to: {query}"
    )
    chain = prompt | llm
    response = chain.invoke({"topics": topics, "query": query})

    return response

# Run the function
result = multi_provider_workflow("Explain quantum computing")

Combiner le traçage manuel et automatique​

Utilisez @mlflow.trace avec le traçage automatique pour créer des traces unifiées pour les scénarios suivants :

  • Plusieurs appels LLM dans un seul workflow
  • Systèmes multi-agents avec différents fournisseurs
  • Logique personnalisée entre les appels LLM
Python
import mlflow
import openai
from mlflow.entities import SpanType

mlflow.openai.autolog()

# Create OpenAI client
client = openai.OpenAI()

@mlflow.trace(span_type=SpanType.CHAIN)
def run(question):
    messages = build_messages(question)
    # MLflow automatically generates a span for OpenAI invocation
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        max_tokens=100,
        messages=messages,
    )
    return parse_response(response)


@mlflow.trace
def build_messages(question):
    return [
        {"role": "system", "content": "You are a helpful chatbot."},
        {"role": "user", "content": question},
    ]


@mlflow.trace
def parse_response(response):
    return response.choices[0].message.content


run("What is MLflow?")

L'exécution de ce code génère une seule trace qui combine les spans manuels avec le traçage automatique OpenAI :

Mélange de traçage automatique et manuel

Exemple avancé : appels LLM multiples​

L'exemple suivant montre comment le traçage automatique capture plusieurs appels LLM dans un seul workflow.

Ce modèle est utile lorsque votre application doit prendre des décisions séquentielles, telles que l'analyse d'une query en premier lieu, puis la génération d'une réponse basée sur cette analyse.

Python
import mlflow
import openai
from mlflow.entities import SpanType

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

# Create OpenAI client
client = openai.OpenAI()

@mlflow.trace(span_type=SpanType.CHAIN)
def process_user_query(query: str):
    # First LLM call: Analyze the query
    analysis = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "Analyze the user's query and determine if it requires factual information or creative writing."},
            {"role": "user", "content": query}
        ]
    )
    analysis_result = analysis.choices[0].message.content

    # Second LLM call: Generate response based on analysis
    if "factual" in analysis_result.lower():
        # Use a different model for factual queries
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "Provide a factual, well-researched response."},
                {"role": "user", "content": query}
            ]
        )
    else:
        # Use a different model for creative queries
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "Provide a creative, engaging response."},
                {"role": "user", "content": query}
            ]
        )

    return response.choices[0].message.content

# Run the function
result = process_user_query("Tell me about the history of artificial intelligence")

Cela crée une trace avec :

  • Portée parente pour process_user_query
  • Deux spans enfants pour les appels OpenAI

Ressources supplémentaires​

Sur cette page