Hinzufügen von Ablaufverfolgungen zu Ihren Agents
Wichtig
Dieses Feature befindet sich in der Public Preview.
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.
Hinweis
Die detaillierte API-Referenz sowie Codebeispiele zur MLflow-Ablaufverfolgung finden Sie in der MLflow-Dokumentation.
- MLflow 2.13.1
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.
Hinweis
Ab Databricks Runtime 15.4 LTS ML ist die MLflow-Ablaufverfolgung standardmäßig in Notebooks aktiviert. Um die Ablaufverfolgung zu deaktivieren, z. B. mit LangChain, können Sie in Ihrem Notebook mlflow.langchain.autolog(log_traces=False)
ausführen.
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.
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
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,
}
)
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.
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
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")
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)
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.
Erstellen Sie eine Instanz von MLflowClient mit
client = MlflowClient()
.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 derstart_span()
API ausgeführt werden.- 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.- Legen Sie Ihre Attribute, Eingaben und Ausgaben für die Ablaufverfolgung in
Die start_trace() API gibt einen Span zurück. Rufen Sie mithilfe von
span.request_id
undspan.span_id
die Anforderungs-ID, einen eindeutigen Bezeichner der Ablaufverfolgung, der auch alstrace_id
bezeichnet wird, und die ID des zurückgegebenen Span ab.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.- Diese Methode erfordert
request_id
undparent_id
, um den Span der richtigen Position in der Ablaufverfolgungshierarchie zuzuordnen. Sie gibt ein weiteres Span-Objekt zurück.
- Diese Methode erfordert
Beenden Sie den untergeordneten Span durch Aufrufen von
client.end_span(request_id, span_id)
.Wiederholen Sie die Schritte 3 bis 5 für alle untergeordneten Spans, die Sie erstellen möchten.
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})