【MLflow】LLM 트레이스 관리 화면을 자연어로 커스터마이징하기 (Custom Trace View)

이 글에 대하여

MLflow 3.16.0에 추가된 Custom Trace Views 기능은 트레이스 관리 UI를 자연어로 자유롭게 커스터마이징할 수 있는 기능입니다.

MLflow AI Assistant를 사용함으로써 UI가 구성됩니다.

Custom Trace Views란 무엇인가요?

트레이드 상세 화면은 일반적으로 왼쪽의 스팬(span) 트리와 오른쪽의 선택된 스팬의 입력 및 출력이라는 정해진 형태를 가지고 있습니다.

RAG나 에이전트의 트레이스를 검토할 때, 실제로 보고 싶은 것은 "사용자의 질문과 최종 답변", "각 도구 호출 결과", "평가 입력란" 정도인 경우가 많습니다. 그런데도 스팬 트리를 하나하나 열어 돌아다녀야 합니다.

Custom Trace Views는 이 화면을 대체하는 자유로운 레이아웃을 직접 준비할 수 있는 기능입니다. 물론 사용자가 UI를 코드로 작성하는 것이 아니라, "질문과 답변을 위에 나란히 배치하고, 아래에 도구 출력을 카드로"와 같이 문장으로 요청하면 MLflow AI Assistant가 보기 좋게 구성해 줍니다.

사전에 다음 두 가지 사항에 유의하는 것이 좋습니다.

1. 뷰는 특정 트레이스의 복사가 아닌 레이아웃입니다.

'루트 스팬의 입력', '첫 번째 도구 호출의 출력'과 같이 역할을 지정하여 데이터를 가리키기 때문에, 같은 Experiment의 다른 트레이스를 열어도 그 트레이스 값으로 동일한 레이아웃이 그려집니다.

※ 뷰는 Experiment 단위로 저장되며, 같은 Experiment를 여는 모든 사람이 사용할 수 있습니다. 하나의 Experiment에 저장할 수 있는 뷰는 최대 50개입니다.

2. 구성 요소(Component)는 고정되어 있습니다.

MLflow AI Assistant가 임의의 HTML이나 React를 생성하는 것이 아니라, 정해진 14가지 종류의 구성 요소를 조합한 JSON을 반환합니다. 구현할 수 있는 범위는 이 구성 요소로 결정됩니다.

준비 사항

필요한 것은 MLflow 3.16 이상의 트래킹 서버, MLflow AI Assistant 사용이 가능한 상태, 그리고 함수 호출(function calling)이 가능한 모델 세 가지입니다.

공식 MLflow AI Assistant 페이지에서 설정을 진행합니다.

AI Assistant의 프로바이더로 Ollama 또는 MLflow AI Gateway를 사용하는 경우, 도구를 호출할 수 있는 모델을 이용해야 합니다.

사용해 보기

트레이스 열기

Experiment의 Traces 목록에서 트레이스를 하나 엽니다.

헤더 오른쪽의 'Default view'가 뷰 전환 버튼입니다. 이것을 누르면 'Create custom view'를 선택할 수 있습니다.

아래는 데모 데이터로 만든 Custom Trace View의 예시입니다. 예를 들어 이런 식으로 원하는 레이아웃의 뷰를 만들 수 있습니다.

뷰 만들기

'Create custom view'를 선택하면 설명문을 입력하는 화면이 나타납니다. 다음과 같이 작성해 보겠습니다.

사용자의 질문과 최종 답변을 위에 나란히 배치하고, 그 아래에 각 도구 호출 출력을 카드로 표시해 주세요. 마지막으로, 답변의 정확도를 1~5로 고르는 라디오 버튼과 자유 기술 코멘트란, 전송 버튼을 추가해 주세요.

'Build with Assistant'를 누르면 오른쪽 영역에 Assistant 패널이 열리고, 트레이스 화면은 'Building this view...'가 됩니다.

1분 남짓 만에 초안이 완성되었습니다. 질문 카드, 도구 출력 카드, 평가란까지 요청한 대로 구성되어 있습니다.

눈에 띄는 점은 '최종 답변' 카드가 비어 있다는 것과, 도구 호출이 3개인데 카드가 2장밖에 없다는 것입니다.

이런 경우에는 Assistant에게 수정해 달라고 합니다. 패널이 열려 있다면 그대로 채팅창에, 닫혀 있다면 'Edit with Assistant'를 눌러 연 다음, "『최종 답변』 카드가 비어 있습니다. chat_agent 스팬의 출력인 choices[0].message.content를 Markdown 형식으로 표시해 주세요. 그 외는 변경하지 말아 주세요"라고 보냅니다. 그러면 해당 부분만 수정된 뷰가 돌아왔습니다.

한 번에 완벽하게 나오는 것을 기대하기보다는, 초안을 보고 수정한다는 전제로 사용하는 것이 좋을 것 같습니다.

완료되면 'Save'를 눌러 이름을 지정합니다. 여기서는 'QA review'로 했습니다.

저장하면 헤더 표시가 "Custom view"에서 이 이름으로 바뀌고, 뷰 전환 메뉴에도 나열됩니다.

다른 트레이스에서 열기

저장된 뷰를 도구 호출이 하나밖에 없는 다른 트레이스(날씨를 물어본 것)로 열면 다음과 같습니다.

도구 카드는 get_weather 한 개만 나타나고, 두 번째는 대응하는 스팬이 없기 때문에 그려지지 않습니다. 이는 레이아웃이 "몇 번의 도구 호출"과 같이 데이터를 지칭하기 때문입니다.

반면에 "최종 답변"은 다시 비어 있습니다. 수정할 때 "chat_agent 스팬의 출력"으로 지정했기 때문에, 뷰는 그 이름의 스팬을 찾지만, 이 트레이스의 루트 스팬은 agent라는 다른 이름이었습니다. 만약 뷰를 다른 트레이스에서도 사용할 계획이라면, 특정 스팬 이름보다는 "루트 스팬의 출력"처럼 역할로 요청하는 것이 재사용성이 높습니다. 이것도 Edit with Assistant에서 수정할 수 있습니다.

뷰에서 피드백 보내기

이번 뷰에는 평가란을 추가했습니다. 4를 선택하고 코멘트를 작성한 후 전송합니다.

전송하면 버튼이 약 2초 동안 "Feedback submitted"로 바뀌었다가 칸이 비어질 뿐, 화면의 다른 곳은 변하지 않습니다.

저장되었는지 여부는 "Default view"로 돌아가서 "Assess"를 열면 알 수 있습니다. Feedback에 Accuracy = 4가 사람(HUMAN) 평가로 붙어 있었습니다.

API에서 읽으면, 일반적인 log_feedback()을 실행한 것과 구별할 수 없습니다.

import mlflow
mlflow.set_tracking_uri("http://localhost:5050")
trace = mlflow.get_trace("tr-ad394dd4623643f93da70de8593e855a")
...

즉, 뷰의 평가란은 리뷰 화면을 원하는 형태로 만드는 수단일 뿐이며, 기록 장소는 기존 메커니즘입니다. 평가 이름(여기서는 Accuracy)은 뷰를 만들 때 Assistant가 결정하므로, 이름을 지정하고 싶다면 설명문에 작성해야 합니다.

이름 변경 및 삭제

뷰를 표시하는 동안 "Edit with Assistant" 옆의 세 점 메뉴를 열면 "Rename view"와 "Delete view"가 있습니다. 삭제하면 같은 Experiment를 보는 모든 사람에게서 사라지므로 주의해야 합니다.

구성 요소 및 조합

Assistant가 사용할 수 있는 구성 요소는 다음 14가지입니다 (3.16.1 기준).

구성 요소가 고정되어 있기 때문에, "표로 정렬하고 싶다", "그래프를 만들고 싶다"와 같은 요청은 아직 어렵습니다. 반대로 말하면, 이 표에 있는 것들의 조합이라면 설명문을 잘 작성하면 대부분 만들 수 있습니다.

※ 이미지를 표시하는 TraceImage가 최근 PR (#26297)에서 추가되었기 때문에, 다음 릴리스부터 사용할 수 있을 것입니다.

역할로 나누면 3개 그룹으로 나뉩니다.

뼈대는 Row

、Column

、Card

로, 레이아웃은 이 요소들의 중첩으로 결정됩니다.

내용물은 Text
、Markdown
、KeyValueViewer
、StatCard
、Icon
으로 구성되어 있으며, 트레이스 데이터(스팬의 입력・출력・속성, 트레이스의 상태・지연 시간・토큰 수)를 표시하는 부품입니다.

평가 관련 부분은 AssessmentCard
과 AssessmentBoard
이 이미 존재하는 평가를 표시하고, RadioGroup
과 Feedback*
이 입력을 담당합니다. 입력 중 FeedbackThumbsUpDownButtons
는 누르면 즉시 기록되며, 그 외의 것은 FeedbackSubmit
을 눌렀을 때 기록됩니다.

이번에 만든 뷰를 부품으로 분해해 보겠습니다. 뷰의 정의는 Experiment의 태그에 JSON 형태로 들어가 있으며, 다음과 같은 코드로 추출할 수 있습니다.

import json
import mlflow
mlflow.set_tracking_uri("http://localhost:5050")
...

추출한 JSON을 읽어 트리 구조로 만들면 다음과 같습니다(id는 생략).

Column
├── Row
│ ├── Card > Column
...

부품의 내용은 값 그 자체가 아니라 '어떤 스팬의 어떤 필드인지'로 작성되어 있습니다. 예를 들어 질문 칸은 다음과 같은 형태입니다.

{"$source": "spanField", "spanRef": "root", "field": "inputs"}

어려움에 부딪히는 지점 (ハマりどころ)

모델이 도구를 호출하지 않으면, 화면이 멈춘 상태로 있게 됩니다.

직접 가지고 있는 Ollama의 llama3.1 8B 모델로 테스트했을 때, Assistant는 뷰의 JSON을 반환하는 대신 무관한 Python 코드를 문장으로 반환했습니다. 트레이스 화면은 'Building this view...' 상태에서 10분 이상 멈췄고 오류도 발생하지 않았습니다. Ollama와 MLflow AI Gateway의 프로바이더에서는 뷰가 도구 호출로 생성되므로, 도구 호출에 대응하는 모델을 사용해야 합니다. 멈춘 경우에는 페이지를 새로 고침하면 돌아옵니다.

또한, Docker 등으로 서버가 다른 호스트에 있는 경우 Assistant가 거부됩니다. 초기 설정에서는 서버 프로세스와 같은 머신의 브라우저에서의 조작만 허용하기 때문입니다. 서버에 MLFLOW_ENABLE_REMOTE_ASSISTANT=true
을 설정하고 프로바이더를 MLflow AI Gateway로 하면 사용할 수 있습니다. 이 설정을 하지 않은 상태에서 Assistant API is only accessible from the same host
라는 메시지가 나오면, 서버가 다른 호스트에 있다는 것이 원인입니다.

원하는 레이아웃이 한 번에 만들어지지 않는 경우도 많습니다. Edit with Assistant로 지적하면 수정되므로, 초안을 보고 고치고 저장하는 흐름으로 사용하면 성공할 것 같습니다.

사용 팁 (使い方のコツ)

현재로서는 다음 두 가지를 의식하면 원하는 뷰에 가까워지기 쉽습니다.

한 번에 만들려고 하지 말고, Edit을 반복하기

처음 설명문에서 모든 것을 말하려 하기보다는, **먼저 골격만 만들게 하고, 초안을 보면서

※ 구성 요소 목록과 역할은 '구성 요소 조합' 표와 같습니다. 어떤 스팬을 가리키는지는 이름이 아니라 '루트 스팬', '몇 번째 TOOL 스팬'과 같은 역할로 작성해야, 다른 트레이스에서 열었을 때도 동일하게 표시됩니다.

맺음말

Custom Trace Views는 트레이스를 보는 방식을 문장으로 요청하여 만드는 기능입니다.

14가지 종류의 구성 요소 조합이 가능하며, 뷰는 Experiment별로 저장됩니다. 평가 필드의 전송지는 평소의 assessment와 동일하며, 뷰는 화면을 만들기 위한 수단일 뿐 저장 장소가 아닙니다. 도구 호출이 가능한 모델이 필요하며, 작은 로컬 모델의 경우 오류를 내지 않고 멈춥니다.

현재로서는 막 출시된 기능이라 아직 할 수 있는 것이 많지는 않지만, LLMOps의 범위를 크게 넓혀줄 잠재력 있는 기능이라고 생각합니다!

덧붙여 말씀드리자면, 공식 문서에 Custom Trace View 기능 페이지가 존재하지 않아 제가 문서를 작성했습니다. 무사히 병합되었으므로, 다음 릴리스(3.16.1 이후)에서 공식 문서에서도 상세 사양을 확인할 수 있을 것 같습니다.

토론

'뷰의 평가 필드가 log_feedback과 같은 기록지가 된다'는 내용이 있어 API로 보아도 구별할 수 없다는 것을 알고 참고가 되었습니다.