에이전트에 추적 추가

아티클
06/25/2024

Important

이 기능은 공개 미리 보기 상태입니다.

이 문서에서는 MLflow 추적에서 사용할 수 있는 Fluent 및 MLflowClient API를 사용하여 에이전트에 추적을 추가하는 방법을 보여 줍니다.

참고 항목

MLflow 추적에 대한 자세한 API 참조 및 코드 예제는 MLflow 설명서를 참조 하세요.

요구 사항

MLflow 2.13.1

Fluent API를 사용하여 에이전트에 추적 추가

다음은 Fluent API mlflow.start_span mlflow.trace 를 사용하고 추적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

유추 수행

코드를 계측한 후에는 정상적으로 함수를 실행할 수 있습니다. 다음은 이전 섹션의 함수를 사용하여 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를 노출합니다. 추적 수명 주기 및 구조를 완전히 제어할 수 있습니다. 이러한 API는 다중 스레드 애플리케이션 및 콜백과 같은 요구 사항에 Fluent API가 충분하지 않은 경우에 유용합니다.

다음은 MLflow 클라이언트를 사용하여 전체 추적을 만드는 단계입니다.

에 의해 client = MlflowClient()MLflowClient의 인스턴스를 만듭니다.
메서드를 사용하여 추적을 시작합니다 client.start_trace() . 그러면 추적 컨텍스트가 시작되고 절대 루트 범위가 시작되고 루트 범위 개체가 반환됩니다. 이 메서드는 API 전에 start_span() 실행해야 합니다.
1. 에서 추적 client.start_trace()에 대한 특성, 입력 및 출력을 설정합니다.
참고 항목

Fluent API의 start_trace() 메서드와 동일한 값은 없습니다. Fluent API는 추적 컨텍스트를 자동으로 초기화하고 관리되는 상태에 따라 루트 범위인지 여부를 결정하기 때문입니다.
start_trace() API는 범위를 반환합니다. 요청 ID, 추적의 고유 식별자, trace_id반환된 범위의 ID(사용 span.request_id 및 span.span_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})

다음을 통해 공유