Aller au contenu principal

Intégrations MLflow Tracing

MLflow Tracing est intégré à un large éventail de bibliothèques et de frameworks d'IA générative populaires, offrant une expérience de **traçage automatique en une seule ligne** pour tous. Ceci vous permet d'obtenir une observabilité immédiate de vos applications GenAI avec une configuration minimale.

Cette prise en charge étendue signifie que vous pouvez obtenir une observabilité sans modifications significatives du code, en tirant parti des outils que vous utilisez déjà. Pour les composants personnalisés ou les bibliothèques non prises en charge, MLflow fournit également de puissantes API de traçage manuel.

Le traçage automatique capture la logique et les étapes intermédiaires de votre application, telles que les appels LLM, l'utilisation d'outils et les interactions avec les agents, en fonction de votre implémentation de la bibliothèque ou du SDK spécifique.

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.

Intégrations principales en un coup d'œil

Voici des exemples de quick-start pour certaines des intégrations les plus couramment utilisées. Cliquez sur un tab pour voir un exemple d'utilisation de base. Pour des prérequis détaillés et des scénarios plus avancés pour chacun, veuillez visiter leurs pages d'intégration dédiées (liées depuis les tab ou la liste ci-dessous).

Python
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

Guide d'intégration complet d'OpenAI

Gestion sécurisée des clés API

Pour les environnements de production, Databricks vous recommande d'utiliser AI Gateway ou les secrets Databricks pour gérer les clés API. AI Gateway est la méthode préférée et offre des fonctionnalités de gouvernance supplémentaires.

attention

Ne commit jamais les clés API directement dans votre code ou vos Notebooks. Veuillez toujours utiliser la Passerelle AI ou les secrets Databricks pour les informations d'identification sensibles.

Databricks recommande AI Gateway pour la gouvernance et le monitoring de l’accès aux modèles d’IA générative.

Créez un endpoint de modèle de fondation configuré avec AI Gateway :

  1. Dans votre workspace Databricks, accédez à Service > Créer un nouvel endpoint .
  2. Choisissez un type d'Endpoint et un fournisseur.
  3. Configurez l’Endpoint avec votre clé API.
  4. Lors de la configuration de l'endpoint, activez **AI Gateway** et configurez la limitation de débit, les fallbacks et les garde-fous selon les besoins.
  5. Vous pouvez obtenir du code auto-généré pour rapidement start l'Endpoint et exécuter des requêtes. Accédez à Serving > votre Endpoint > Use > Query . Veillez à ajouter le code de traçage :
Python
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)

Activation de multiples intégrations de traçage automatique

Étant donné que les applications GenAI combinent souvent plusieurs bibliothèques, MLflow Tracing vous permet d'activer le traçage automatique pour plusieurs intégrations simultanément, offrant une expérience de traçage unifiée.

Par exemple, pour activer le traçage LangChain et direct OpenAI :

Python
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 générera une trace unique et cohérente qui combine les étapes de LangChain et les appels directs de LLM OpenAI, vous permettant d'inspecter le flux complet. D'autres exemples de combinaison d'intégrations peuvent être trouvés sur la page Traçage automatique.

Désactivation du traçage automatique

Le traçage automatique pour toute bibliothèque spécifique peut être désactivé en appelant mlflow.<library>.autolog(disable=True). Pour désactiver toutes les intégrations de journalisation automatique en une seule fois, utilisez mlflow.autolog(disable=True).

Python
import mlflow

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

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