MLflow-Ablaufverfolgung für Agents

Dieser Artikel beschreibt die MLflow-Ablaufverfolgung auf Databricks und wie Sie damit die Observierbarkeit Ihrer generativen KI-Anwendungen erhöhen können.

Was ist die MLflow-Ablaufverfolgung?

Die MLflow-Ablaufverfolgung erfasst detaillierte Informationen zur Ausführung von generativen KI-Anwendungen. Die Ablaufverfolgung protokolliert Eingaben, Ausgaben und Metadaten, die den einzelnen Zwischenschritten einer Anforderung zugeordnet sind, damit Sie die Quelle von Fehlern und unerwartetem Verhalten anheften können. Wenn Ihr Modell zum Beispiel halluziniert, können Sie jeden Schritt, der zu der Halluzination geführt hat, schnell überprüfen.

Die MLflow-Ablaufverfolgung ist in Databricks-Tools und -Infrastruktur integriert, sodass Sie Ablaufverfolgungen in Databricks-Notizbüchern oder der MLflow Experiment UI speichern und anzeigen können.

Gründe für die Verwendung der MLflow-Ablaufverfolgung

Die MLflow-Ablaufverfolgung bietet mehrere Vorteile:

• Überprüfen Sie eine interaktive Ablaufverfolgungsvisualisierung, und verwenden Sie das Untersuchungstool, um Probleme zu diagnostizieren.
• Überprüfen Sie, dass Promptvorlagen und Schutzmaßnahmen angemessene Ergebnisse erzielen.
• Analysieren Sie die Latenz verschiedener Frameworks, Modelle und Blockgrößen.
• Schätzen Sie die Anwendungskosten, indem Sie die Tokennutzung über verschiedene Modelle hinweg messen.
• Richten Sie Benchmark-"goldene" Datasets ein, um die Leistung verschiedener Versionen zu bewerten.
• Speichern Sie Ablaufverfolgungen von Produktionsmodellendpunkten, um Probleme zu debuggen und Offlineüberprüfungen und -auswertungen durchzuführen.

Hinzufügen von Ablaufverfolgungen zu Ihrem Agent

MLflow Tracing unterstützt drei Methoden zur Einbindung von Traces in Ihre generierenden KI-Anwendungen. Details zur API-Referenz finden Sie in der MLflow-Dokumentation.

Installieren der MLflow-Ablaufverfolgung

Die MLflow-Tracing-Funktion ist in MLflow-Versionen 2.13.0 und höher verfügbar, die in <DBR< 15.4 LTS ML und höher vorinstalliert sind. Installieren Sie bei Bedarf MLflow mit dem folgenden Code:

%pip install mlflow>=2.13.0 -qqqU
%restart_python

Alternativ können Sie die neueste Version von databricks-agentsinstallieren, die eine kompatible MLflow-Version enthält:

%pip install databricks-agents

Verwenden der automatischen Protokollierung zum Hinzufügen von Ablaufverfolgungen zu Ihren Agents

Wenn Ihre GenAI-Bibliothek die Ablaufverfolgung unterstützt, z. B. LangChain oder OpenAI, aktivieren Sie die automatische Protokollierung, indem Sie Ihrem Code mlflow.<library>.autolog() hinzufügen. Zum Beispiel:

mlflow.langchain.autolog()

MLflow unterstützt zusätzliche Bibliotheken für die automatische Ablaufverfolgungsprotokollierung. Eine vollständige Liste der integrierten Bibliotheken finden Sie in der Dokumentation zur MLflow-Ablaufverfolgung.

Verwenden von Fluent-APIs zum manuellen Hinzufügen von Ablaufverfolgungen zu Ihrem Agent

Fluent-APIs in MLflow erstellen automatisch Ablaufverfolgungshierarchien basierend auf dem Ausführungsfluss ihres Codes.

Dekorieren Sie Ihre Funktion

Verwenden Sie das @mlflow.trace-Decorator-Element, um einen Span für den Umfang der ergänzten Funktion zu erstellen.

Das MLflow Span-Objekt organisiert die Schritte der Ablaufverfolgung. Spans erfasst Informationen zu einzelnen Vorgängen oder Schritten, z. B. API-Aufrufe oder Vektorspeicherabfragen, innerhalb eines Workflows.

Die Spanne beginnt, wenn die Funktion aufgerufen wird und endet, wenn sie zurückgegeben wird. MLflow zeichnet die Eingabe und Ausgabe der Funktion und alle Ausnahmen auf, die von der Funktion ausgelöst werden.

Der folgende Code erstellt beispielsweise einen Bereich namens my_function, der Eingabeargumente x und y und die Ausgabe erfasst.

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

Verwenden des Ablaufverfolgungskontext-Managers

Wenn Sie einen Bereich für einen beliebigen Codeblock erstellen möchten, nicht nur eine Funktion, können Sie mlflow.start_span() als Kontext-Manager verwenden, der den Codeblock umschließt. Die Spanne beginnt, wenn der Kontext eingegeben wird und endet, wenn der Kontext beendet wird. Die Ein- und Ausgabe des Span sollten manuell über Settermethoden des Span-Objekts bereitgestellt werden, das vom Kontext-Manager zurückgegeben wird.

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

Umschließen einer externen Funktion

Um externe Bibliotheksfunktionen zu verfolgen, schließen Sie die Funktion mit mlflow.traceum.

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)

Fluent-API-Beispiel

Das folgende Beispiel zeigt, wie Sie die Fluent-APIs mlflow.trace und mlflow.start_span verwenden, um die quickstart-agentnachzuverfolgen:

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

Führen Sie nach dem Hinzufügen der Ablaufverfolgung die Funktion aus. Das folgende setzt das Beispiel mit der Funktion predict() im vorherigen Abschnitt fort. Die Ablaufverfolgungen werden automatisch angezeigt, wenn Sie die Aufrufmethode predict() ausführen.

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,
  }
)

MLflow-Client-APIs

MlflowClient macht präzise, threadsichere APIs verfügbar, um Ablaufverfolgungen zu starten und zu beenden, Spans zu verwalten und Span-Felder festzulegen. Sie bietet vollständige Kontrolle über den Ablaufverfolgungslebenszyklus und die Struktur. Diese APIs sind nützlich, wenn die Fluent-APIs für Ihre Anforderungen nicht ausreichen, z. B. bei Multithread-Anwendungen und Callback-Funktionen.

Im Folgenden werden Schritte zum Erstellen einer vollständigen Ablaufverfolgung mithilfe des MLflow-Clients beschrieben.

1. Erstellen Sie eine Instanz von MLflowClient durch client = MlflowClient().

2. Starten Sie eine Nachverfolgung mithilfe der client.start_trace()-Methode. Dadurch wird der Ablaufverfolgungskontext initiiert, ein absoluter Stammspan gestartet und ein Stamm-Span-Objekt zurückgegeben. Diese Methode muss vor der start_span()-API ausgeführt werden.

   1. Legen Sie Ihre Attribute, Eingaben und Ausgaben für die Ablaufverfolgung in client.start_trace() fest.

   Hinweis

   Es gibt keine Entsprechung mit der start_trace()-Methode in den Fluent-APIs. Dies liegt daran, dass die Fluent-APIs den Tracing-Kontext automatisch initialisieren und bestimmen, ob es sich basierend auf dem verwalteten Zustand um die Wurzelspanne handelt.

3. Die start_trace()-API gibt eine Spanne zurück. Rufen Sie mithilfe von span.request_id und span.span_id die Anforderungs-ID, einen eindeutigen Bezeichner der Ablaufverfolgung, der auch als trace_id bezeichnet wird, und die ID des zurückgegebenen Spans ab.

4. Beginnen Sie einen untergeordneten Span indem Sie mithilfe von client.start_span(request_id, parent_id=span_id) Ihre Attribute, Eingaben und Ausgaben für den Span festlegen.

   1. Diese Methode erfordert request_id und parent_id, um den Span der richtigen Position in der Ablaufverfolgungshierarchie zuzuordnen. Sie gibt ein weiteres Span-Objekt zurück.

5. Beenden Sie den untergeordneten Span durch Aufrufen von client.end_span(request_id, span_id).

6. Wiederholen Sie die Schritte 3 bis 5 für alle untergeordneten Spans, die Sie erstellen möchten.

7. Rufen Sie client.end_trace(request_id) auf, nachdem alle untergeordneten Spans beendet sind, um die Ablaufverfolgung zu schließen und aufzuzeichnen.

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})

Prüfen von Ablaufverfolgungen

Verwenden Sie eine der folgenden Optionen, um Protokolle nach dem Ausführen des Agents zu überprüfen:

• Die Visualisierung von Ablaufverfolgungen wird inline in der Zellenausgabe gerendert.
• Die Ablaufverfolgungen werden in Ihrem MLflow-Experiment protokolliert. Sie können die vollständige Liste des Ablaufverfolgungsverlaufs auf der Registerkarte Ablaufverfolgungen auf der Seite Experiment ansehen und durchsuchen. Wenn der Agent bei einer aktiven MLflow-Ausführung ausgeführt wird, erscheinen Ablaufverfolgungen auf der Seite Ausführung.
• Rufen Sie Ablaufverfolgungen programmgesteuert über die search_traces()-API ab.

Verwenden der MLflow-Ablaufverfolgung in der Produktion

Die MLflow-Ablaufverfolgung ist auch in Mosaik AI Model Serving integriert, sodass Sie Probleme effizient debuggen, die Leistung überwachen und ein goldenes Dataset für die Offlineauswertung erstellen können. Wenn die MLflow-Ablaufverfolgung für Ihren Bereitstellungsendpunkt aktiviert ist, werden Ablaufverfolgungen in einer Rückschlusstabelle unter der Spalte response aufgezeichnet.

Sie können die Ablaufverfolgungen visualisieren, die in Ableitungstabellen protokolliert werden, indem Sie die Tabelle abfragen und die Ergebnisse in einem Notizbuch anzeigen. Verwenden Sie display(<the request logs table>) in Ihrem Notizbuch und wählen Sie die einzelnen Zeilen der Spuren aus, die Sie visualisieren möchten.

Zum Aktivieren der MLflow-Ablaufverfolgung für Ihren Bereitstellungsendpunkt müssen Sie die Umgebungsvariable ENABLE_MLFLOW_TRACING in der Endpunktkonfiguration auf True festlegen. Informationen zum Bereitstellen eines Endpunkts mit benutzerdefinierten Umgebungsvariablen finden Sie unter Hinzufügen von Nur-Text-Umgebungsvariablen. Wenn Sie Ihren Agent mithilfe der deploy() API bereitgestellt haben, werden Ablaufverfolgungen automatisch in einer Rückschlusstabelle protokolliert. Siehe Bereitstellen eines Agents für generative KI-Anwendung.

Begrenzungen

• Die MLflow-Ablaufverfolgung ist in Databricks-Notebooks, Notebookaufträgen und Modellbereitstellungen verfügbar.

LangChain-Autoprotokollierung unterstützt unter Umständen nicht alle LangChain-Vorhersage-APIs. Die vollständige Liste der unterstützten APIs finden Sie in MLflow-Dokumentation.