將追蹤新增至您的代理程式

發行項
11/08/2024

重要

這項功能處於公開預覽狀態。

本文章說明如何使用 MLflow 追蹤提供的 Fluent 和 MLflowClient API，將追蹤新增至您的代理程式。

注意

如需 MLflow 追蹤的詳細 API 參考和程式碼範例，請參閱 MLflow 文件。

需求

MLflow 2.13.1

使用自動記錄將追蹤新增至您的代理程式

如果您使用支援追蹤的 GenAI 程式庫 (例如 LangChain、LlamaIndex 或 OpenAI)，您可以針對程式庫整合啟用 MLflow 自動記錄以啟用追蹤。例如，使用 mlflow.langchain.autolog() 自動將追蹤新增至以 LangChain 為基礎的代理程式。

注意

自 Databricks Runtime 15.4 LTS ML 起，預設將會在筆記本中啟用 MLflow。若要停用追蹤，例如使用 LangChain，您可以在筆記本中執行 mlflow.langchain.autolog(log_traces=False)。

mlflow.langchain.autolog()

MLflow 支援其他程式庫進行追蹤自動記錄。如需整合程式庫的完整清單，請參閱 MLflow 追蹤文件。

使用 Fluent API 手動將追蹤新增至您的代理程式

以下是使用 Fluent API：mlflow.trace 和 mlflow.start_span 將追蹤新增至 quickstart-agent 的簡易範例。建議將此用於 PyFunc 模型。


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

執行推斷

檢測程式碼之後，您可以像平常一樣執行函數。以下範例延續在前一章節中使用 predict() 函數的範例。您執行引動方法 predict() 時，會自動顯示追蹤。


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 API

MLflow 中的 Fluent API 會根據執行程式碼的位置和時間，自動建構追蹤階層。下列章節說明使用 MLflow 追蹤 Fluent API 支援的工作。

裝飾函數

您可以使用 @mlflow.trace 裝飾項目來裝飾函數，為裝飾之函數的範圍建立跨度。跨度會在叫用函數時開始，並在函數傳回時結束。 MLflow 會自動記錄函數的輸入和輸出，以及從函數引發的任何例外狀況。例如，執行下列程式碼會使用名稱「my_function」來建立跨度，從而擷取輸入引述 x 和 y，以及函數的輸出。

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

使用追蹤內容管理員

如果您想要為任意程式碼區塊建立跨度，而不只是函數，您可以使用 mlflow.start_span() 作為包裝程式碼區塊的內容管理員。跨度會在輸入內容時開始，並在內容結束時結束。跨度輸入和輸出應透過從內容管理員產生之跨度物件的 setter 方法來手動提供。

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

包裝外部函數

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)

MLflow 用戶端 API

MlflowClient 會公開細微、安全執行緒 API 來開始和結束追蹤、管理跨度，以及設定跨度欄位。它提供追蹤生命週期和結構的完整控制。 Fluent API 不足以滿足您的需求時，這些 API 很實用，例如多執行緒應用程式和回呼。

以下是使用 MLflow 用戶端建立完整追蹤的步驟。

透過 client = MlflowClient() 建立 MLflowClient 的執行個體。
使用 client.start_trace() 方法啟動追蹤。這樣會開始追蹤內容，並啟動絕對根跨度，且傳回根跨度物件。這個方法必須在 start_span() API 之前執行。
1. 在 client.start_trace() 設定追蹤的屬性、輸入和輸出。
注意

Fluent API 中的 start_trace() 方法沒有對等項目。這是因為 Fluent API 會自動初始化追蹤內容，並判斷它是否為以受控狀態為基礎的根跨度。
start_trace() API 會傳回跨度。取得要求 ID、追蹤的唯一識別碼 (也稱為 trace_id)，以及使用 span.request_id 和 span.span_id 傳回之跨度的 ID。
使用 client.start_span(request_id, parent_id=span_id) 啟動子系跨度，以設定跨度的屬性、輸入和輸出。
1. 此方法需要 request_id 和 parent_id，才能將跨度與追蹤階層中的正確位置建立關聯。它會傳回另一個跨度物件。
透過呼叫 client.end_span(request_id, span_id) 結束子系跨度。
針對您想要建立的任何子系跨度重複 3 - 5。
所有子系跨度都結束之後，呼叫 client.end_trace(request_id) 即可關閉整個追蹤並加以記錄。

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

共用方式為