MLflow 추적을 사용하여 에이전트 관찰 가능

이 문서에서는 Databricks에서 MLflow 추적을 사용하여 생성 AI 애플리케이션에 가시성을 추가하는 방법을 설명합니다.

MLflow 추적이란 무엇입니까?

MLflow 추적 개발에서 배포에 이르는 생성 AI 애플리케이션에 대한 엔드투엔드 관찰 가능성을 제공합니다. 추적은 Databricks의 Gen AI 도구 집합과 완전히 통합되어 전체 개발 및 프로덕션 수명 주기에 걸쳐 자세한 인사이트를 캡처합니다.

다음은 Gen AI 애플리케이션에서 추적을 위한 주요 사용 사례입니다.

• 간소화된 디버깅: 추적은 Gen AI 애플리케이션의 각 단계에 대한 가시성을 제공하여 문제를 보다 쉽게 진단하고 해결할 수 있도록 합니다.

• 오프라인 평가: 추적은 에이전트 평가에 중요한 데이터를 생성하므로 시간이 지남에 따라 에이전트의 품질을 측정하고 개선할 수 있습니다.

• 프로덕션 모니터링: 추적은 에이전트 동작 및 자세한 실행 단계에 대한 가시성을 제공하므로 프로덕션 환경에서 에이전트 성능을 모니터링하고 최적화할 수 있습니다.

• 감사 로그: MLflow 추적은 에이전트 작업 및 결정에 대한 포괄적인 감사 로그를 생성합니다. 이는 예기치 않은 문제가 발생할 때 규정 준수를 보장하고 디버깅을 지원하는 데 중요합니다.

요구 사항

MLflow 추적은 MLflow 버전 2.13.0 이상에서 사용할 수 있습니다. Databricks는 최신 버전의 MLflow를 설치하여 최신 기능 및 향상된 기능에 액세스하는 것이 좋습니다.

%pip install mlflow>=2.13.0 -qqqU
%restart_python

자동 추적

MLflow 자동 로깅 사용하면 코드 mlflow.<library>.autolog()한 줄을 추가하여 에이전트를 신속하게 계측할 수 있습니다.

MLflow는 가장 인기 있는 에이전트 작성 라이브러리에 대한 자동 로깅을 지원합니다. 각 작성 라이브러리에 대한 자세한 내용은 MLflow 자동 로깅 설명서참조하세요.

자동 로깅 사용 안 함

자동 로깅 추적은 다음 라이브러리에 대해 Databricks Runtime 15.4 ML 이상에서 기본적으로 사용하도록 설정됩니다.

• LangChain
• 랭그래프(Langgraph)
• OpenAI
• 라마인덱스

이러한 라이브러리에 대한 자동 로깅 추적을 사용하지 않도록 설정하려면 Notebook에서 다음 명령을 실행합니다.

`mlflow.<library>.autolog(log_traces=False)`

수동으로 추적 추가

자동 로깅은 에이전트를 계측하는 편리한 방법을 제공하지만 에이전트를 보다 세밀하게 계측하거나 자동 로깅에서 캡처하지 않는 추가 추적을 추가할 수 있습니다. 이러한 경우 MLflow 추적 API를 사용하여 추적을 수동으로 추가합니다.

MLflow 추적 API 추적의 트리 구조를 관리할 걱정 없이 추적을 추가하기 위한 하위 코드 API입니다. MLflow는 Python 스택을 사용하여 적절한 부모-자식 범위 관계를 자동으로 결정합니다.

자동 로깅 및 수동 추적 결합

수동 추적 API는 자동 로깅과 함께 사용할 수 있습니다. MLflow는 자동 로깅 및 수동 추적을 통해 만든 범위를 결합하여 에이전트 실행의 전체 추적을 만듭니다. 자동 로깅 및 수동 추적을 결합하는 예제는 MLflow 추적사용하여 도구 호출 에이전트 계측을 참조하세요.

@mlflow.trace 데코레이터를 사용하여 함수 추적

코드를 수동으로 계측하는 가장 간단한 방법은 @mlflow.trace 데코레이터를 사용하여 함수를 데코레이트하는 것입니다. MLflow 추적 데코레이터 데코레이팅된 함수의 범위가 있는 "범위"를 만듭니다. 이 범위는 추적의 실행 단위를 나타내며 추적 시각화에 단일 행으로 표시됩니다. 범위는 함수, 대기 시간 및 함수에서 발생한 예외의 입력 및 출력을 캡처합니다.

예를 들어 다음 코드는 입력 인수 my_function 및 x 출력을 캡처하는 y이라는 범위를 만듭니다.

import mlflow

@mlflow.trace
def add(x: int, y: int) -> int:
  return x + y

범위 이름, 범위 유형을 사용자 지정하고 범위에 사용자 지정 특성을 추가할 수도 있습니다.

from mlflow.entities import SpanType

@mlflow.trace(
  # By default, the function name is used as the span name. You can override it with the `name` parameter.
  name="my_add_function",
  # Specify the span type using the `span_type` parameter.
  span_type=SpanType.TOOL,
  # Add custom attributes to the span using the `attributes` parameter. By default, MLflow only captures input and output.
  attributes={"key": "value"}
)
def add(x: int, y: int) -> int:
  return x + y

컨텍스트 관리자를 사용하여 임의의 코드 블록 추적

함수뿐만 아니라 임의의 코드 블록에 대한 범위를 만들려면 mlflow.start_span() 코드 블록을 래핑하는 컨텍스트 관리자로 사용합니다. 범위는 컨텍스트가 입력될 때 시작되고 컨텍스트가 종료될 때 종료됩니다. 컨텍스트 관리자가 생성한 범위 개체의 setter 메서드를 사용하여 범위 입력 및 출력을 수동으로 제공해야 합니다. 자세한 내용은 MLflow 설명서 - 컨텍스트 처리기참조하세요.

with mlflow.start_span(name="my_span") as span:
  span.set_inputs({"x": x, "y": y})
  result = x + y
  span.set_outputs(result)
  span.set_attribute("key", "value")

저수준 추적 라이브러리

MLflow는 추적 트리 구조를 명시적으로 제어하기 위한 하위 수준 API도 제공합니다. MLflow 설명서 - 수동 계측참조하세요.

추적 예제: 자동 로깅 및 수동 추적 결합

다음 예제에서는 OpenAI 자동 로깅 및 수동 추적을 결합하여 도구 호출 에이전트를 완전히 계측합니다.

import json
from openai import OpenAI
import mlflow
from mlflow.entities import SpanType

client = OpenAI()

# Enable OpenAI autologging to capture LLM API calls
# (*Not necessary if you are using the Databricks Runtime 15.4 ML and above, where OpenAI autologging is enabled by default)
mlflow.openai.autolog()

# Define the tool function. Decorate it with `@mlflow.trace` to create a span for its execution.
@mlflow.trace(span_type=SpanType.TOOL)
def get_weather(city: str) -> str:
  if city == "Tokyo":
    return "sunny"
  elif city == "Paris":
    return "rainy"
  return "unknown"


tools = [
  {
    "type": "function",
    "function": {
      "name": "get_weather",
      "parameters": {
        "type": "object",
        "properties": {"city": {"type": "string"}},
      },
    },
  }
]

_tool_functions = {"get_weather": get_weather}

# Define a simple tool-calling agent
@mlflow.trace(span_type=SpanType.AGENT)
def run_tool_agent(question: str):
  messages = [{"role": "user", "content": question}]

  # Invoke the model with the given question and available tools
  response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    tools=tools,
  )
  ai_msg = response.choices[0].message
  messages.append(ai_msg)

  # If the model requests tool calls, invoke the function(s) with the specified arguments
  if tool_calls := ai_msg.tool_calls:
    for tool_call in tool_calls:
      function_name = tool_call.function.name
      if tool_func := _tool_functions.get(function_name):
        args = json.loads(tool_call.function.arguments)
        tool_result = tool_func(**args)
      else:
        raise RuntimeError("An invalid tool is returned from the assistant!")

      messages.append(
        {
          "role": "tool",
          "tool_call_id": tool_call.id,
          "content": tool_result,
        }
      )

    # Send the tool results to the model and get a new response
    response = client.chat.completions.create(
      model="gpt-4o-mini", messages=messages
    )

  return response.choices[0].message.content

# Run the tool calling agent
question = "What's the weather like in Paris today?"
answer = run_tool_agent(question)

태그를 사용하여 추적에 주석 달기

MLflow 추적 태그는 대화 ID, 사용자 ID, Git 커밋 해시 등과 같은 추적에 사용자 지정 메타데이터를 추가할 수 있는 키-값 쌍입니다. 추적을 필터링하고 검색하기 위해 MLflow UI에 태그가 표시됩니다.

MLflow API 또는 MLflow UI를 사용하여 태그를 진행 중이거나 완료된 추적으로 설정할 수 있습니다. 다음 예제에서는 mlflow.update_current_trace() API를 사용하여 진행 중인 추적에 태그를 추가하는 방법을 보여 줍니다.

@mlflow.trace
def my_func(x):
    mlflow.update_current_trace(tags={"fruit": "apple"})
    return x + 1

추적 태그 지정 및 추적을 사용하여 추적을 필터링하고 검색하는 방법에 대한 자세한 내용은 MLflow 설명서 - 추적 태그설정을 참조하세요.

추적 데이터 검토

에이전트를 실행한 후 추적을 검토하려면 다음 옵션 중 하나를 사용합니다.

• 인라인 시각화: Databricks 노트북에서는 추적 데이터가 셀 출력에서 인라인으로 렌더링됩니다.
• MLflow 실험: Databricks에서 실험으로 이동> 실험 >추적 선택하여 실험의 모든 추적을 보고 검색합니다.
• MLflow 실행: 에이전트가 활성 MLflow 런에서 작동하면, MLflow UI의 실행 페이지에 추적이 표시됩니다.
• 에이전트 평가 UI: Mosaic AI 에이전트 평가에서 평가 결과의 "자세한 추적 보기"를 클릭하여 각 에이전트 실행에 대한 추적을 검토할 수 있습니다.
• 추적 검색 API: 프로그래밍 방식으로 추적을 검색하려면 추적 검색 API사용합니다.

추적을 사용하여 에이전트를 평가하다

추적 데이터는 에이전트를 평가하는 데 중요한 리소스 역할을 합니다. MLflow 추적은 모델 실행에 대한 자세한 정보를 캡처하여 오프라인 평가에 중요한 역할을 합니다. 추적 데이터를 사용하여 골든 데이터 세트에 대한 에이전트의 성능을 평가하고, 문제를 식별하고, 에이전트의 성능을 향상시킬 수 있습니다.

%pip install -U mlflow databricks-agents
%restart_python

import mlflow

# Get the recent 50 successful traces from the experiment
traces = mlflow.search_traces(
    max_results=50,
    filter_string="status = 'OK'",
)

traces.drop_duplicates("request", inplace=True) # Drop duplicate requests.
traces["trace"] = traces["trace"].apply(lambda x: x.to_json()) # Convert the trace to JSON format.

# Evaluate the agent with the trace data
mlflow.evaluate(data=traces, model_type="databricks-agent")

에이전트 평가에 대한 자세한 내용은 평가를 실행하고결과를 확인하세요.

유추 테이블을 사용하여 배포된 에이전트 모니터링

에이전트가 Mosaic AI Model Serving에 배포된 후 유추 테이블을 사용하여 에이전트를 모니터링할 수 있습니다. 유추 테이블에는 검토 앱의 요청, 응답, 에이전트 추적 및 에이전트 피드백에 대한 자세한 로그가 포함됩니다. 이 정보를 통해 문제를 디버그하고, 성능을 모니터링하고, 오프라인 평가를 위한 골든 데이터 세트를 만들 수 있습니다.

에이전트 배포에 유추 테이블을 사용하도록 설정하려면 AI 에이전트에 대한 유추 테이블 사용참조하세요.

온라인 흔적 조회

Notebook을 사용하여 유추 테이블을 쿼리하고 결과를 분석합니다.

추적을 시각화하려면 display(<the request logs table>) 실행하고 검사할 행을 선택합니다.

# Query the inference table
df = spark.sql("SELECT * FROM <catalog.schema.my-inference-table-name>")
display(df)

에이전트 모니터링

Gen AI 앱모니터링하는 방법을 참조하세요.

추적 부하 대기 시간

추적은 성능 영향을 최소화하기 위해 비동기적으로 작성됩니다. 그러나 추적은 특히 각 유추 요청에 대한 추적 크기가 큰 경우 엔드포인트 응답 속도에 대기 시간을 추가합니다. Databricks는 프로덕션에 배포하기 전에 추적 대기 시간 영향을 이해하기 위해 엔드포인트를 테스트하는 것이 좋습니다.

다음 표에서는 추적 크기별 대기 시간 영향에 대한 대략적인 예상을 제공합니다.

문제 해결

문제 해결 및 일반적인 질문에 대해서는 MLflow 설명서: 추적 방법 가이드 및 MLflow 설명서: FAQ를 참조하세요.