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.
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).
- OpenAI
- LangChain
- LangGraph
- Anthropic
- DSPy
- Databricks
- Bedrock
- AutoGen
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
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
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
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
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
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
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
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
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.
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.
- AI Gateway (Recommended)
- Databricks secrets
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 :
- Dans votre workspace Databricks, accédez à Service > Créer un nouvel endpoint .
- Choisissez un type d'Endpoint et un fournisseur.
- Configurez l’Endpoint avec votre clé API.
- 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.
- 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 :
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)
Utilisez les secrets Databricks pour gérer les clés API :
-
Créez un Secret Scope et stockez votre API clé :
Pythonfrom 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
) -
Récupérez et utilisez le secret dans votre code :
Pythonimport 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
)
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 :
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).
import mlflow
# Disable for a specific library
mlflow.openai.autolog(disable=True)
# Disable all autologging
mlflow.autolog(disable=True)