[MLFlow] MLFlow의 LLMOps-2 (LLMOps 기능 알아보기)

 

들어가며

이전 포스팅에 이어, 실제 MLflow 환경을 Docker로 구성해서 모델 실행 결과를 추적하고, 모델 레지스트리와 프롬프트 레지스트리를 관리하고, 이를 평가하는 기능을 구현해보도록 하겠습니다. 

1. 로컬 MLFlow 서버 올리기 (docker)

MLflow를 Docker를 사용해서 로컬 서버(내 컴퓨터) 환경에 올려보도록 하겠습니다. (Docker 환경이 구성이 되어있다고 가정하고 진행하도록 하곘습니다. Docker를 잘 모르시는 분들은 다음 포스팅을 참고해주세요)

•  다음 docker pull 명령어를 통해 이미지를 받습니다 (3.x.x 버전 이상을 받아주세요)

docker pull ghcr.io/mlflow/mlflow:v3.3.0

• 컨테이너 서버를 올립니다
  
  • Window 유저가 아니시라면 "MSYS_NO_PATHCONV=1'는 제거해주세요

mkdir -p ./mlflow
MSYS_NO_PATHCONV=1 docker run -it\
    --name mlflow-server\
    -p 5000:5000 -v "$(pwd)/mlflow:/mlruns"\
    ghcr.io/mlflow/mlflow:v3.3.0\
    mlflow server --host 0.0.0.0\
    --backend-store-uri /mlruns\
    --default-artifact-root /mlruns

잘 동작한다면, 다음과 같이 로그가 뜨게 됩니다.

• 다음은, 컨테이너 서버를 내리는 방법입니다 (실습 이후, 컨테이너 종료를 권장드립니다)

docker stop mlflow-server

 

컨테이너가 잘 올라갔다면 http://localhost:5000에 접속하셔서, 다음과 같은 MLFlow 대시보드를 확인하실 수 있습니다.

MLflow 대시보드

2. MLFlow 핵심 기능 설명

MLFLow 환경이 구성되었다면, 이제 MLFlow를 구성하는 Experiments, Models, Prompts에 대해 간단히 설명드리겠습니다. 그리고 이를 활용해서 MLFlow로 LLMOps 환경을 구축하는 코드를 만들어 보도록 하겠습니다.

2-1. Experiments (실험 관리)

• 모델을 개발하며 수행하는 모든 실험 과정을 기록하고 추적하는 기능입니다. 같은 모델이라도 파라미터나 프롬프트를 바꾸거나 다른 데이터를 사용하며 달라지는 모든 시도를 체계적으로 관리할 수 있게 해줍니다.

2-2. Models (모델 관리 및 배포)

• 서비스에 사용할 모델을 표준화된 형식으로 패키징하고, 버전 관리하며, 쉽게 배포 및 롤백할 수 있도록 돕는 기능입니다. 다양한 환경에서 모델을 일관된 방식으로 사용할 수 있게 해주는 것이 핵심입니다.

2-3. Prompts (프롬프트 관리)

• LLM(대규모 언어 모델) 애플리케이션 개발에서 핵심적인 프롬프트를 체계적으로 관리하기 위한 최신 기능입니다. LLM의 성능이 프롬프트에 따라 크게 달라지기 때문에, 프롬프트를 코드처럼 버전 관리하고 재사용하는 것이 중요해졌기에 생긴 기능입니다.

 

3. 간단한 LLMOps 환경 구축하기

1.  Experiments 만들기

Experiments 프로젝트를 만들어서, 모델 관리를 위한 환경을 구축합니다. 

1. Experiments - "Create"를 눌러 Experiments를 만듭니다. (Expeiment 이름은 임의로 설정합니다)
2. Experiments 이름 옆에 정보(ℹ️) 탭을 클릭해서 Experiment ID를 확인합니다

Experiments를 만들었으면, 이제 해당 Experiment ID를 사용해서 모델이나 평가 결과들을 기록할 수 있습니다.

2.  모델 실행 결과 Trace 해보기 

간단한 Langgrpah AI 모델을 만들어서 해당 모델의 실행 결과를 experiments에 등록해보겠습니다. LangChain이나 LangGraph를 사용해보신 분들이라면, LangSmith의 기능과 같이 모델 실행 결과를 로그로 남기는 기능입니다.

1. 간단한 Gen AI 서비스를 만듭니다. (간단한 Gen AI 서비스 샘플 코드는 Github을 참고해주세요)
2. `set_mlflow_experiments`함수를 실행해 mlflow Experiments 환경을 설정해줍니다.
3. mlflow Experiemnts - Traces 탭을 확인합니다.

from langgraph_sample import get_llm, build_graph, stream_answers
import mlflow

# mlflow experiments setting
def set_mlflow_experiments():
    mlflow.set_tracking_uri("http://localhost:5000")
    mlflow.set_experiment(experiment_id="417156008591960350")

    # langchain log를 할려면 autolog만 하면 됩니다.
    mlflow.langchain.autolog()

if __name__ =="__main__":
    set_mlflow_experiments()
    question = "서울 가락시장역 근처 맛집 3개를 알려줘"
    
    # 랭그래프 모델은 github 참고
    llm = get_llm()
    graph = build_graph(llm)
    for text in stream_answers(graph, question):
        print(text)

잘 동작했다면, 모델 실행 결과가 로그로 남게 됩니다.

Trace를 살펴보시면 다음과 같이 실험이 기록된 것을 확인하실 수 있습니다.

3.  모델 등록해서 사용하기

3-1. 모델 run 실행

이제 모델을 Experiments에 등록해서 Trace를 진행해보도록 하겠습니다. 모델을 Logging한다고 표현하며, 크게 어렵지 않습니다.

1. experiments를 설정합니다.
2. logging할 모델을 가져와서 mlflow.start_run으로 모델을 로깅합니다 (run_name이 None이면, 임의의 run_name 생성)
3. Experiments - Runs에서 실행된 run을 확인합니다.
4. Experiments - Models에서 Run에 의해 실행된 모델을 확인합니다. (Model은 Run 없이도 실행이 가능합니다)

모델을 logging할 때는 log하는 모델의 객체를 정의해서 "mlflow.models.set_model"을 설정해둬야 합니다.

#simple_langgraph_agent.py

# agent 소스코드 하단에 다음과 같이 그래프를 정의해야 합니다.
llm = get_llm()
graph_model = build_graph(llm)
mlflow.models.set_model(graph_model)

def set_mlflow_experiments(experiments_id):
    mlflow.set_tracking_uri("http://localhost:5000")
    mlflow.set_experiment(experiment_id=experiments_id)

    # langchain log를 할려면 autolog만 하면 된다.
    mlflow.langchain.autolog()
# default experiments log


def log_model(run_name=None):
    with mlflow.start_run(run_name=run_name):
        logged_agent_info = mlflow.langchain.log_model(
            lc_model="simple_langgraph_agent.py",
            name="simple_model"  # graph workflow
        )
    return logged_agent_info

    
if __name__ == '__main__':
    set_mlflow_experiments()
    logged_agent_info = log_model(graph)
    model_uri = logged_agent_info.model_uri

다음과 같이 모델이 Experiments에 기록되는 것을 확인할 수 있습니다.

run_name을 설정하면, 다음과 같이 설정을 할 수 있습니다.

3-2. 모델 Registry 등록

1. "Models" 탭에서 "Create Model"로 모델 레지스트리를 생성합니다.
2. Experiments - Models에서 Run에 의해 실행된 모델을 확인합니다. (이때, Model name은 모델 log시 설정한 name입니다)
3. Moddel Registry로 등록할 모델을 선택합니다.
4. 상단에 "Register Model" 중 앞서 생성한 모델 레지스트리를 선택합니다..
5. "Details - Model versions"에 등록된 모델 버전을 확인합니다.

Experiments - Models - Registerd models

상단에 "Register Model" 중 앞서 생성한 모델 레지스트리를 선택

모델 버전 생성

3-2. 모델 Registry에서 로드해서 사용 (Alias로 로드하기)

1. Models Registry에서 방금 등록한 모델을 확인합니다.
2. 특정 Version의 모델에 Alias(Champion)를 태깅합니다 (Champion or Challenger)
3. 다음 코드를 사용해 모델을 Load하고, 추론을 진행합니다. (로드한 모델은 기존 모델의 추론 방법과 동일합니다)
4. 이렇게 model registry에 등록한 모델로 추론을 진행할 경우, Trace에 모델 버전을 기록할 수 있습니다.

client = MlflowClient()

model_name= 'sample-model'
model_alias ="champion"
model_version = '1'


question = '가락시장 근처 맛집 3군데를 알려줘'
model_uri = f"models:/{model_name}@{model_alias}"
loaded_model = load_model(model_uri)
for text in stream_answers(loaded_model, question):
    print(text)

모델을 alias로 설정해서 다음과 같이 로드할 수 있습니다.

 

Registry에 등록된 모델로 추론을 진행할 시, Trace에 모델 Version을 기록할 수 있습니다.

 

4. Prompt 등록해서 사용하기

프롬프트 레지스트리를 활용해서, 프롬프트를 관리할 수 있습니다. 

1. "Prompts탭 - Create prompt"를 사용해 Prompt 프로젝트를 생성합니다.
2. Prompt "Name"과 첫 번째 버전의 "Prompt"를 작성합니다. (이때, Name은 수정이 불가합니다)
3. 생성한 프롬프트를 확인하고, alias를 태깅합니다 (version으로도 프롬프트를 로드할 수 있습니다)
4. 생성한 프롬프트를 다음 코드로 로드할 수 있습니다.
5. 이렇게 생성한 프롬프트를 llm 서비스에서 로드해 모델 추론을 진행합니다.

def load_prompts_by_version(version=1):
    prompt = mlflow.genai.load_prompt(f"prompts:/sample-test/{version}")
    prompt_str = prompt.template
    return prompt_str

def load_prompts_by_alias(alias='champion'):
    prompt = mlflow.genai.load_prompt(f"prompts:/sample-test@{alias}")
    prompt_str = prompt.template
    return prompt_str
    

# 프롬프트 사용 예시 (챗봇 정의 부분)
def prompt_chatbot(state: State, llm: ChatOpenAI):
    """LangGraph 노드: 누적 메시지와 추가 프롬프트를 기반으로 LLM을 호출한다."""
    prompt_str = load_prompts_by_alias()
    messages = [("system", prompt_str)] + state["messages"]
    print(messages)
    return  {"messages": [llm.invoke(messages)]}

프롬프트 등록 (Version 1) - Dict 형태 결과 반환 요청

 

프롬프트(Version1) 사용해서 실행한 추론 결과

 

프롬프트 등록 (Version 2) - 지역명 제외 요청

프롬프트 (Version2) 사용해서 실행한 추론 결과

정리하며

이번 포스팅까지 MLFlow를 사용한 LLMOps 기능을 간단하게 정리해봤습니다. 회사에서는 Databricks Managed MLFlow를 사용하다보니, 로컬에서 환경을 구성하고, 예시 샘플과 코드를 만드느라 시간이 꽤 소요된 것 같습니다. 해당 포스팅은 정말 기본적인 등록하고, 로드해서 사용하는 기능들만 소개시켜드렸기에 관심이 있으신 분들은 MLFlow나 Databricks 공식 Docuemnts를 참고해보셔도 좋을 것 같습니다. 언제가 될 지 모르겠지만, 추후 포스팅에서는 해당 기능을 활용한 파이프라인 구축이나 심화 기능 등 조금 더 기능적인 부분을 다뤄보도록 하겠습니다.