Partager via

MLflow Tracing pour les agents

  • Article

Important

Cette fonctionnalité est disponible en préversion publique.

Cet article décrit le suivi MLflow sur Databricks et explique comment l’utiliser pour ajouter l’observabilité à vos applications IA génératives.

Qu’est-ce que le MLflow Tracing ?

MLflow Tracing capture des informations détaillées sur l’exécution d’applications IA générative. Traçage des journaux d'entrées, des sorties et des métadonnées associées à chaque étape intermédiaire d’une requête afin de pouvoir identifier la source des bogues ou un comportement inattendu. Par exemple, si votre modèle hallucine, vous pouvez rapidement inspecter chaque étape qui a conduit à l’hallucination.

Le suivi MLflow est intégré aux outils et à l’infrastructure Databricks, ce qui vous permet de stocker et d’afficher des traces dans les notebooks Databricks ou l’interface utilisateur de l’expérience MLflow.

Le traçage en ligne capture des informations détaillées pour chaque étape d’une application IA générative

Pourquoi utiliser MLflow Tracing ?

Le suivi MLflow offre plusieurs avantages :

  • Passez en revue une visualisation de trace interactive et utilisez l’outil d’investigation pour diagnostiquer les problèmes.
  • Vérifiez que les modèles d’invite et les garde-fous produisent des résultats raisonnables.
  • Analysez la latence des différents frameworks, modèles et tailles de blocs.
  • Estimer les coûts de l’application en mesurant l’utilisation des jetons sur différents modèles.
  • Établissez des jeux de données « golden » de référence pour évaluer les performances des différentes versions.
  • Stockez les traces des points de terminaison de modèle de production pour déboguer les problèmes et effectuer une révision et une évaluation hors connexion.

Ajouter des traces à votre agent

Le suivi MLflow prend en charge trois méthodes pour ajouter des traces à vos applications IA génératives. Pour des exemples de référence API, consultez la documentation sur MLflow.

API Cas d’usage recommandé Description
Journalisation automatique MLflow Développement avec des bibliothèques GenAI intégrées La journalisation automatique enregistre automatiquement les traces pour les frameworks open source pris en charge comme LangChain, LlamaIndex et OpenAI.
API fluides Agent personnalisé avec Pyfunc API low code pour ajouter des traces sans se préoccuper de la gestion de l’arborescence de la trace. MLflow détermine automatiquement les relations d’étendue parent-enfant appropriées à l’aide de la pile Python.
API client MLflow Cas d’usage avancés tels que le multithreading MLflowClient fournit des API granulaires et thread-safe pour les cas d’usage avancés. Vous devez gérer manuellement la relation de parenté des portées. Cela vous permet de mieux contrôler le cycle de vie des traces, en particulier pour les scénarios multi-thread.

Installer MLflow Tracing

Le suivi MLflow est disponible dans les versions 2.13.0 et ultérieures de MLflow, qui est préinstallé dans <DBR< 15.4 LTS ML et versions ultérieures. Si nécessaire, installez MLflow avec le code suivant :

%pip install mlflow>=2.13.0 -qqqU
%restart_python

Vous pouvez également installer la dernière version de databricks-agents, qui inclut une version MLflow compatible :

%pip install databricks-agents

Utiliser l'autologging pour ajouter des traces à vos agents

Si votre bibliothèque GenAI prend en charge le suivi, comme LangChain ou OpenAI, activez l'autologging en ajoutant mlflow.<library>.autolog() à votre code. Par exemple:

mlflow.langchain.autolog()

Remarque

À partir de Databricks Runtime 15.4 LTS ML, le suivi MLflow est activé par défaut dans les notebooks. Pour désactiver le suivi, par exemple, avec LangChain, vous pouvez exécuter mlflow.langchain.autolog(log_traces=False) dans votre bloc-notes.

MLflow prend en charge des bibliothèques supplémentaires pour la journalisation automatique de trace. Pour obtenir la liste complète des bibliothèques intégrées, consultez la documentation de suivi MLflow.

Utiliser des API Fluent pour ajouter manuellement des traces à votre agent

Les API Fluent dans MLflow créent automatiquement des hiérarchies de trace en fonction du flux d’exécution de votre code.

Décorer votre fonction

Utilisez le décorateur @mlflow.trace pour créer une étendue pour la portée de la fonction décorée.

L’objet MLflow Span organise les étapes de suivi. Spans capture des informations sur des opérations ou des étapes individuelles, telles que les appels d’API ou les requêtes de magasin de vecteurs, au sein d’un workflow.

La portée démarre lorsque la fonction est appelée et se termine lorsqu’elle est retournée. MLflow enregistre l’entrée et la sortie de la fonction et toutes les exceptions levées à partir de la fonction.

Par exemple, le code suivant crée une étendue nommée my_function qui capture les arguments d’entrée x et y et la sortie.

@mlflow.trace(name="agent", span_type="TYPE", attributes={"key": "value"})
def my_function(x, y):
    return x + y

Utilisez le gestionnaire de contexte de traçage

Si vous souhaitez créer une étendue pour un bloc de code arbitraire, pas seulement une fonction, vous pouvez utiliser mlflow.start_span() en tant que gestionnaire de contexte qui encapsule le bloc de code. La portée commence lorsque le contexte est entré et se termine lorsque le contexte est arrêté. Les entrées et sorties de la plage doivent être fournies manuellement à l’aide de méthodes setter de l’objet de plage généré par le gestionnaire de contexte.

with mlflow.start_span("my_span") as span:
    span.set_inputs({"x": x, "y": y})
    result = x + y
    span.set_outputs(result)
    span.set_attribute("key", "value")

Envelopper une fonction externe

Pour tracer les fonctions de bibliothèque externe, encapsulez la fonction avec mlflow.trace.

from sklearn.metrics import accuracy_score

y_pred = [0, 2, 1, 3]
y_true = [0, 1, 2, 3]

traced_accuracy_score = mlflow.trace(accuracy_score)
traced_accuracy_score(y_true, y_pred)

Exemple d’API Fluent

L’exemple suivant montre comment utiliser les API Fluent mlflow.trace et mlflow.start_span pour suivre l'quickstart-agent:

import mlflow
from mlflow.deployments import get_deploy_client

class QAChain(mlflow.pyfunc.PythonModel):
    def __init__(self):
        self.client = get_deploy_client("databricks")

    @mlflow.trace(name="quickstart-agent")
    def predict(self, model_input, system_prompt, params):
        messages = [
                {
                    "role": "system",
                    "content": system_prompt,
                },
                {
                    "role": "user",
                    "content":  model_input[0]["query"]
                }
          ]

        traced_predict = mlflow.trace(self.client.predict)
        output = traced_predict(
            endpoint=params["model_name"],
            inputs={
                "temperature": params["temperature"],
                "max_tokens": params["max_tokens"],
                "messages": messages,
            },
        )

        with mlflow.start_span(name="_final_answer") as span:
          # Initiate another span generation
            span.set_inputs({"query": model_input[0]["query"]})

            answer = output["choices"][0]["message"]["content"]

            span.set_outputs({"generated_text": answer})
            # Attributes computed at runtime can be set using the set_attributes() method.
            span.set_attributes({
              "model_name": params["model_name"],
                        "prompt_tokens": output["usage"]["prompt_tokens"],
                        "completion_tokens": output["usage"]["completion_tokens"],
                        "total_tokens": output["usage"]["total_tokens"]
                    })
              return answer

Après avoir ajouté la trace, exécutez la fonction. L’exemple suivant poursuit l’exemple avec la fonction predict() dans la section précédente. Les traces s’affichent automatiquement lorsque vous exécutez la méthode d’appel, predict().

SYSTEM_PROMPT = """
You are an assistant for Databricks users. You answer Python, coding, SQL, data engineering, spark, data science, DW and platform, API, or infrastructure administration questions related to Databricks. If the question is unrelated to one of these topics, kindly decline to answer. If you don't know the answer, say that you don't know; don't try to make up an answer. Keep the answer as concise as possible. Use the following pieces of context to answer the question at the end:
"""

model = QAChain()

prediction = model.predict(
  [
      {"query": "What is in MLflow 5.0"},
  ],
  SYSTEM_PROMPT,
  {
    # Using Databricks Foundation Model for easier testing, feel free to replace it.
    "model_name": "databricks-dbrx-instruct",
    "temperature": 0.1,
    "max_tokens": 1000,
  }
)

API client MLflow

MlflowClient expose des API granulaires et thread safe pour démarrer et terminer les traces, gérer les portées et définir des champs de portée. Il fournit un contrôle complet du cycle de vie et de la structure de trace. Ces API sont utiles lorsque les API Fluent ne suffisent pas à vos besoins, par exemple pour des applications multithreading et des rappels.

Les étapes suivantes permettent de créer une trace complète à l’aide du client MLflow.

  1. Créez une instance de MLflowClient par client = MlflowClient().

  2. Démarrez une trace à l’aide de la méthode client.start_trace(). Cette opération lance le contexte de trace et démarre une étendue racine absolue et retourne un objet d’étendue racine. Cette méthode doit être exécutée avant l’API start_span().

    1. Définissez vos attributs, entrées et sorties pour la trace dans client.start_trace().

    Remarque

    Il n’existe pas d’équivalent à la méthode start_trace() dans les API Fluent. Cela est dû au fait que les API Fluent initialisent automatiquement le contexte de trace et déterminent s’il s’agit de l’étendue racine en fonction de l’état managé.

  3. L’API start_trace() retourne une portée. Obtenez l’ID de requête, un identificateur unique de la trace également appelé trace_idet l’ID de l’étendue retournée à l’aide de span.request_id et de span.span_id.

  4. Démarrez une portée enfant en utilisant client.start_span(request_id, parent_id=span_id) pour définir vos attributs, entrées et sorties pour la portée.

    1. Cette méthode nécessite request_id et parent_id d’associer l’étendue à la position correcte dans la hiérarchie de trace. Elle retourne un autre objet de portée.

  5. Terminez la portée enfant en appelant client.end_span(request_id, span_id).

  6. Répétez les étapes 3 à 5 pour créer toutes les portées enfant que vous souhaitez.

  7. Une fois toutes les sous-étendues terminées, appelez client.end_trace(request_id) pour fermer la trace et l’enregistrer.

from mlflow.client import MlflowClient

mlflow_client = MlflowClient()

root_span = mlflow_client.start_trace(
  name="simple-rag-agent",
  inputs={
          "query": "Demo",
          "model_name": "DBRX",
          "temperature": 0,
          "max_tokens": 200
         }
  )

request_id = root_span.request_id

# Retrieve documents that are similar to the query
similarity_search_input = dict(query_text="demo", num_results=3)

span_ss = mlflow_client.start_span(
      "search",
      # Specify request_id and parent_id to create the span at the right position in the trace
        request_id=request_id,
        parent_id=root_span.span_id,
        inputs=similarity_search_input
  )
retrieved = ["Test Result"]

# You must explicitly end the span
mlflow_client.end_span(request_id, span_id=span_ss.span_id, outputs=retrieved)

root_span.end_trace(request_id, outputs={"output": retrieved})

Examen des traces

Pour passer en revue les traces après l’exécution de l’agent, utilisez l’une des options suivantes :

  • La visualisation des traces est affichée en ligne dans la sortie de la cellule.
  • Les traces sont consignées dans votre expérience MLflow. Vous pouvez consulter et rechercher la liste complète des traces historiques dans l’onglet Traces de la page Expérience. Lorsque l’agent s’exécute dans un Run MLflow actif, les traces s’affichent sur la page Run.
  • Récupérez les traces par programmation à l’aide de l’API search_traces().

Utiliser MLflow Tracing en production

MLflow Tracing est également intégré à Mosaic AI Model Serving, ce qui vous permet de déboguer efficacement les problèmes, de surveiller les performances et de créer un jeu de données golden pour une évaluation hors ligne. Lorsque le suivi MLflow est activé pour votre point de terminaison de service, les traces sont enregistrées dans unetable d’inférence sous la colonne response.

Vous pouvez visualiser les traces journalisées dans les tables d’inférence en interrogeant la table et en affichant les résultats dans un bloc-notes. Utilisez display(<the request logs table>) dans votre bloc-notes et sélectionnez les lignes individuelles des traces que vous souhaitez visualiser.

Pour activer le suivi MLflow pour votre point de terminaison de service, vous devez définir la variable d’environnement ENABLE_MLFLOW_TRACING sur True dans la configuration du point de terminaison. Pour savoir comment déployer un point de terminaison avec des variables d’environnement personnalisées, consultez Ajouter des variables d’environnement de texte brut. Si vous avez déployé votre agent à l’aide de l’API deploy(), les traces sont automatiquement enregistrées dans une table d’inférence. Consultez Déployer un agent pour une application d’IA générative.

Remarque

L’écriture de traces dans une table d’inférence s’effectue de façon asynchrone. Elle n’ajoute donc pas la même surcharge que dans l’environnement de notebook pendant le développement. Cependant, cela peut encore entraîner une certaine surcharge au niveau de la vitesse de réponse du point de terminaison, en particulier lorsque la taille de la trace pour chaque demande d’inférence est importante. Databricks ne garantit aucun contrat de niveau de service (SLA) pour l’impact réel de la latence sur votre point de terminaison de modèle, car il dépend fortement de l’environnement et de l’implémentation du modèle. Databricks vous recommande de tester les performances de votre point de terminaison et d’obtenir des insights sur la surcharge de suivi avant le déploiement sur une application de production.

Le tableau suivant fournit une indication approximative de l’impact sur la latence d’inférence pour différentes tailles de trace.

Taille de trace par requête Impact sur la latence (ms)
~10 Ko ~ 1 ms
~ 1 Mo 50 ~ 100 ms
10 Mo 150 ms ~

Limites

  • Le suivi MLflow est disponible dans les notebooks Databricks et les tâches de notebooks et la mise en service de modèle.

L’enregistrement automatique LangChain peut ne pas prendre en charge toutes les API de prédiction LangChain. Pour obtenir la liste complète des API prises en charge, consultez documentation MLflow.