다음을 통해 공유

에이전트에 추적 추가

  • 아티클

Important

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

이 문서에서는 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 추적은 기본적으로 Notebooks 내에서 사용하도록 설정됩니다. 예를 들어, LangChain에서 추적을 사용하지 않도록 설정하기 위해 Notebook에서 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,
  }
)

흐름 API

MLflow의 Fluent API는 코드가 실행되는 위치와 시기에 따라 추적 계층을 자동으로 생성합니다. 다음 섹션에서는 MLflow 추적 흐름 API를 사용하여 지원되는 작업에 대해 설명합니다.

함수 데코레이트

@mlflow.trace decorator를 통해 함수를 데코레이팅하여 데코레이팅된 함수의 범위에 대한 범위를 만들 수 있습니다. 범위는 함수가 호출될 때 시작되고 반환될 때 끝납니다. 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는 다중 스레드 애플리케이션 및 콜백과 같은 요구 사항에 흐름 API가 충분하지 않은 경우에 유용합니다.

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

  1. client = MlflowClient()에 의해 MLflowClient의 인스턴스를 만듭니다.

  2. client.start_trace() 메서드를 사용하여 추적을 시작합니다. 그러면 추적 컨텍스트가 시작되고 절대 루트 범위가 시작되며 루트 범위 개체가 반환됩니다. 이 메서드는 start_span() API 전에 실행해야 합니다.

    1. client.start_trace()에서 추적에 대한 특성, 입력 및 출력을 설정합니다.

    참고 항목

    흐름 API의 start_trace() 메서드와 동일한 메서드는 없습니다. 흐름 API는 추적 컨텍스트를 자동으로 시작하고 관리되는 상태에 따라 루트 범위인지 여부를 결정하기 때문입니다.

  3. start_trace() API는 범위를 반환합니다. 요청 ID, 추적의 고유 식별자(trace_id라고도 함), span.request_id 및 span.span_id를 사용하여 반환된 범위의 ID를 가져옵니다.

  4. 범위에 대한 특성, 입력 및 출력을 설정하는 데 client.start_span(request_id, parent_id=span_id)을 사용하여 자식 범위를 시작합니다.

    1. 이 메서드는 추적 계층 구조의 올바른 위치와 범위를 연결하기 위해 request_id 및 parent_id가 필요합니다. 다른 범위 개체는 반환합니다.

  5. client.end_span(request_id, span_id)을 호출하여 자식 범위를 종료합니다.

  6. 만들려는 모든 자식 범위에 대해 3~5를 반복합니다.

  7. 모든 자식 범위가 종료되면 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})