Hinzufügen von Ablaufverfolgungen zu Ihren Agents

In diesem Artikel wird gezeigt, wie Sie Ihren Agents mithilfe der Fluent- und MLflowClient-APIs die Ablaufverfolgungen hinzufügen, die mit der MLflow-Ablaufverfolgung zur Verfügung gestellt wurden.

Anforderungen

• MLflow 2.13.1

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

Wenn Sie eine GenAI-Bibliothek verwenden, die Unterstützung für die Ablaufverfolgung hat (z. B. LangChain, LlamaIndex oder OpenAI), können Sie die automatische MLflow-Protokollierung für die Bibliotheksintegration aktivieren, um die Ablaufverfolgung zu aktivieren. Verwenden Sie z. B mlflow.langchain.autolog() zum automatischen Hinzufügen von Ablaufverfolgungen zu Ihrem LangChain-basierten Agent.

mlflow.langchain.autolog()

MLflow unterstützt zusätzliche Bibliotheken für die automatische Ablaufprotokollierung. 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

Im folgenden kurzen Beispiel werden die Fluent-APIs: mlflow.trace und mlflow.start_span verwendet, um Ablaufverfolgungen zu quickstart-agent hinzuzufügen. Dies wird für PyFunc-Modelle empfohlen.

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

Durchführen von Rückschlüssen

Nachdem Sie Ihren Code instrumentiert haben, können Sie Ihre Funktion wie gewohnt ausführen. Im Folgenden wird das Beispiel mit der Funktion predict() im vorherigen Abschnitt fortgesetzt. Die Ablaufverfolgungen werden automatisch angezeigt, wenn Sie die Aufrufmethode predict() ausführen.

SYSTEM_PROMPT = """
You are an assistant for Databricks users. You are answering python, coding, SQL, data engineering, spark, data science, DW and platform, API or infrastructure administration question related to Databricks. If the question is not related to one of these topics, kindly decline to answer. If you don't know the answer, just 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,
  }
)

Fluent APIs

Die Fluent APIs in MLflow erstellen automatisch die Ablaufverfolgungshierarchie basierend darauf wo und wann der Code ausgeführt wird. In den folgenden Abschnitten werden die unterstützten Aufgaben mithilfe der Fluent APIs für die MLflow-Ablaufverfolgung beschrieben.

Dekorieren Sie Ihre Funktion

Sie können Ihre Funktion mit dem @mlflow.trace-Dekorierer dekorieren, um ein Span für den Geltungsbereich der dekorierten Funktion zu schaffen. Der Span beginnt, wenn die Funktion aufgerufen wird und endet, wenn sie zurückgegeben wird. MLflow zeichnet automatisch die Eingabe und Ausgabe der Funktion sowie alle Ausnahmen auf, die von der Funktion ausgelöst werden. Wenn Sie z. B. den folgenden Code ausführen, wird ein Span mit dem Namen „my_function“ erstellt, wobei die Eingabeargumente x und y sowie die Ausgabe der Funktion erfasst werden.

@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. Der Span 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

Die Funktion mlflow.trace kann als Wrapper verwendet werden, um eine Funktion Ihrer Wahl nachzuverfolgen. Dies ist nützlich, wenn Sie Funktionen nachverfolgen möchten, die aus externen Bibliotheken importiert wurden. Sie generiert den gleichen Span, den Sie erreichen würden, indem Sie diese Funktion dekorieren.

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)

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 nicht für Ihre Anforderungen ausreichen, z. B. für Multithreadanwendungen und Rückrufe.

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

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

2. Starten Sie eine Ablaufverfolgung mithilfe der Methode client.start_trace(). Dadurch wird der Ablaufverfolgungskontext initiiert und ein absoluter Stammspan gestartet und ein Stammspan-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

   In den Fluent APIs gibt es keine Entsprechung für die Methode start_trace(). Dies liegt daran, dass die Fluent APIs den Ablaufverfolgungskontext automatisch initialisieren und bestimmen basierend auf dem verwalteten Zustand, ob es sich um den Stammspan handelt.

3. Die start_trace() API gibt einen Span 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 Span 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. Nachdem alle untergeordneten Spans beendet wurden, rufen Sie client.end_trace(request_id) auf, um die gesamte 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"]

# Span has to be ended explicitly
mlflow_client.end_span(request_id, span_id=span_ss.span_id, outputs=retrieved)

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