Traçage de Haystack
Haystack est un framework d'orchestration d'IA open source pour la création d'applications LLM prêtes pour la production, de systèmes de recherche sémantique et de systèmes de questions-réponses.
MLflow Tracing fournit une capacité de traçage automatique pour Haystack. Vous pouvez activer le traçage pour Haystack en appelant la fonction mlflow.haystack.autolog, et les traces sont automatiquement consignées dans l'Experiment MLflow actif lors de l'invocation de pipelines et de composants.
import mlflow
mlflow.haystack.autolog()
La trace MLflow capture automatiquement les informations suivantes concernant les exécutions de pipeline Haystack :
- Pipelines et composants
- Latences
- Métadonnées des composants.
- Utilisation et coût des jetons
- Informations sur les accès au cache
- Toute exception si levée
Sur les clusters de compute serverless, l'autologging n'est pas activé automatiquement. Vous devez appeler explicitement mlflow.haystack.autolog() pour activer le traçage automatique de cette intégration.
Prérequis
Pour utiliser MLflow Tracing avec Haystack, vous devez installer MLflow et les packages Haystack.
- Development
- Production
Pour les environnements de développement, installez le package MLflow complet avec les modules complémentaires Databricks et les packages Haystack :
pip install --upgrade "mlflow[databricks]>=3.1" haystack-ai
Le package mlflow[databricks] complet inclut toutes les fonctionnalités pour le développement local et l’expérimentation sur Databricks.
Pour les déploiements de production, installez les packages mlflow-tracing et Haystack :
pip install --upgrade mlflow-tracing haystack-ai
Le mlflow-tracing package est optimisé pour une utilisation en production.
MLflow 3 est fortement recommandé pour une meilleure expérience de traçage avec Haystack.
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 :
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.
Clés API : Assurez-vous que les clés API de votre fournisseur LLM sont configurées. Pour les environnements de production, utilisez AI Gateway ou les secrets Databricks au lieu de valeurs codées en dur pour une gestion sécurisée des clés API.
export OPENAI_API_KEY="your-openai-api-key"
# Add other provider keys as needed
Exemple d'utilisation
L’exemple suivant montre comment utiliser Haystack avec le suivi MLflow. Cet exemple crée un pipeline RAG (Génération augmentée de récupération) simple avec un récupérateur, un générateur de prompts et un LLM.
import mlflow
import os
from haystack import Pipeline
from haystack.components.builders import PromptBuilder
from haystack.components.generators import OpenAIGenerator
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack import Document
# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "your-openai-api-key" # Uncomment and set if not globally configured
# Enable auto tracing for Haystack
mlflow.haystack.autolog()
# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/haystack-tracing-demo")
# Create a simple document store with sample documents
document_store = InMemoryDocumentStore()
document_store.write_documents([
Document(content="Paris is the capital of France."),
Document(content="Berlin is the capital of Germany."),
Document(content="Rome is the capital of Italy."),
])
# Build a simple RAG pipeline
template = """
Given the following documents, answer the question.
Documents:
{% for doc in documents %}
{{ doc.content }}
{% endfor %}
Question: {{ question }}
Answer:
"""
pipe = Pipeline()
pipe.add_component("retriever", InMemoryBM25Retriever(document_store=document_store))
pipe.add_component("prompt_builder", PromptBuilder(template=template))
pipe.add_component("llm", OpenAIGenerator(model="gpt-4o-mini"))
pipe.connect("retriever", "prompt_builder.documents")
pipe.connect("prompt_builder", "llm")
# Run the pipeline - trace will be automatically logged
result = pipe.run({
"retriever": {"query": "What is the capital of France?"},
"prompt_builder": {"question": "What is the capital of France?"}
})
print(result["llm"]["replies"][0])
Pour les environnements de production, utilisez AI Gateway ou les secrets Databricks au lieu de valeurs codées en dur pour la gestion sécurisée des clés API.
Suivi de l'utilisation des jetons
MLflow suit automatiquement l'utilisation des jetons pour les pipelines Haystack lors de l'utilisation de MLflow version 3.4.0. ou ultérieure. Les informations d'utilisation des jetons incluent les jetons d'entrée, les jetons de sortie et les jetons totaux consommés pendant l'exécution du pipeline.
import mlflow
mlflow.haystack.autolog()
from haystack import Pipeline
from haystack.components.generators import OpenAIGenerator
# Create and run a pipeline
pipe = Pipeline()
pipe.add_component("llm", OpenAIGenerator(model="gpt-4o-mini"))
# Run the pipeline and retrieve trace information
with mlflow.start_span(name="haystack_pipeline_run") as span:
result = pipe.run({"llm": {"prompt": "What is the capital of France?"}})
print(result["llm"]["replies"][0])
# Token usage is automatically logged and visible in the MLflow UI
trace_info = mlflow.get_last_active_trace()
print(f"Trace ID: {trace_info.request_id}")
Les détails d'utilisation des jetons sont affichés dans l'interface utilisateur de suivi MLflow, vous permettant de surveiller et d'optimiser les performances et les coûts de votre pipeline.
Désactiver le suivi automatique
Le traçage automatique pour Haystack peut être désactivé globalement en appelant mlflow.haystack.autolog(disable=True) ou mlflow.autolog(disable=True).