MLflow Custom Flavor — 코드 흐름과 동작 순서 분석

MLflow는 기본적으로 scikit-learn, pytorch, tensorflow 등의 모델은 바로 관리할 수 있지만,
사용자가 직접 만든 모델(Custom Model)은 MLflow가 어떻게 예측해야 할지 모릅니다.

이럴 때 사용하는 것이 바로 Custom Flavor입니다.
이번 글에서는 다음 예제 코드를 기준으로 MLflow의 Custom Flavor 동작 순서를 완전히 해부합니다. 👇

 

 

1. 전체 코드 요약

import mlflow.pyfunc
import json, os, pickle
import pandas as pd

# 1️⃣ 사용자 정의 모델
class CustomMLModel:
    def predict(self, input_data):
        return input_data ** 2  # 입력값 제곱

# 2️⃣ MLflow용 래퍼 (PythonModel 상속)
class CustomMLflowFlavor(mlflow.pyfunc.PythonModel):
    def __init__(self, model=None):
        self.model = model

    def load_context(self, context):
        with open(context.artifacts["model"], "rb") as f:
            self.model = pickle.load(f)

    def predict(self, model_input):
        return self.model.predict(model_input)

# 3️⃣ MLflow 로딩용 함수
def _load_pyfunc(model_path):
    with open(os.path.join(model_path, "model.pkl"), "rb") as f:
        model = pickle.load(f)
    return CustomMLflowFlavor(model)

# 4️⃣ 모델 저장 함수
def save_model(path, model):
    os.makedirs(path, exist_ok=True)

    # 모델을 Pickle로 저장
    with open(os.path.join(path, "model.pkl"), "wb") as f:
        pickle.dump(model, f)

    # MLmodel 메타데이터 작성
    mlmodel_data = {
        "flavors": {
            "python_function": {
                "loader_module": __name__,
                "python_version": "3.10"
            }
        }
    }
    with open(os.path.join(path, "MLmodel"), "w") as f:
        json.dump(mlmodel_data, f, indent=4)

# 5️⃣ 저장 및 예측
model = CustomMLModel()
save_model("custom_flavor_model", model)
loaded_model = mlflow.pyfunc.load_model("custom_flavor_model")
data = pd.DataFrame([2, 3, 4])
print(loaded_model.predict(data))  # [4, 9, 16]

 

 

2. 단계별 상세 해석

 1️⃣ CustomMLModel — 실제 예측 로직을 가진 모델

class CustomMLModel:
    def predict(self, input_data):
        return input_data ** 2

• MLflow가 모르는 모델입니다.
• 단순히 입력값을 제곱해서 반환하는 로직만 가집니다.
• scikit-learn이나 tensorflow가 아닌 완전 커스텀 로직입니다.

➡️ 이 모델은 MLflow 입장에서는 “낯선 객체”입니다.
그래서 바로는 저장할 수 없습니다.

 

 2️⃣ CustomMLflowFlavor — MLflow가 이해할 수 있게 감싸주는 Wrapper

class CustomMLflowFlavor(mlflow.pyfunc.PythonModel):
    def __init__(self, model=None):
        self.model = model

    def load_context(self, context):
        with open(context.artifacts["model"], "rb") as f:
            self.model = pickle.load(f)

    def predict(self, model_input):
        return self.model.predict(model_input)

 

• mlflow.pyfunc.PythonModel을 상속합니다.
• MLflow는 이 클래스를 보면 “아, 이건 내가 실행할 수 있는 모델이구나!” 하고 인식합니다.

메서드	설명
load_context()	MLflow Tracking Server나 로컬 저장소에서 모델을 불러올 때 실행됨
predict()	MLflow가 실제로 예측을 요청할 때 호출되는 함수

 

 

 

 3️⃣ _load_pyfunc() — MLflow의 로딩 진입점

def _load_pyfunc(model_path):
    with open(os.path.join(model_path, "model.pkl"), "rb") as f:
        model = pickle.load(f)
    return CustomMLflowFlavor(model)

• MLflow는 모델을 로드할 때 MLmodel 파일을 먼저 읽고,
  "loader_module": "__main__"을 확인합니다.
• 그 모듈에서 _load_pyfunc() 함수를 자동으로 호출합니다.

➡️ 이 함수는 저장된 Pickle 파일을 불러와 CustomMLflowFlavor 객체로 감싸서 반환합니다.

 

 4️⃣ save_model() — MLflow가 이해할 수 있는 구조로 저장

def save_model(path, model):
    # model.pkl 저장
    # MLmodel 파일 생성 (loader_module 정보 포함)

MLflow 모델 저장 구조는 아래와 같습니다 👇

custom_flavor_model/
├── model.pkl      ← 실제 모델 객체 (Pickle)
└── MLmodel        ← MLflow 메타데이터 파일

{
    "flavors": {
        "python_function": {
            "loader_module": "__main__",
            "python_version": "3.10"
        }
    }
}

• "loader_module" → MLflow가 이 모듈 안에서 _load_pyfunc()를 찾습니다.
• "python_function" → PyFunc 형태의 모델임을 의미합니다.

 

 5️⃣ load_model() — MLflow가 모델을 불러오는 과정

loaded_model = mlflow.pyfunc.load_model("custom_flavor_model")

이 한 줄이 내부적으로는 아래처럼 동작합니다:

1. MLmodel 파일 읽기
2. "loader_module" = __main__ 확인
3. __main__._load_pyfunc(model_path) 호출
4. CustomMLflowFlavor(model) 객체 생성
5. CustomMLflowFlavor.load_context() 실행 (필요시)
6. 준비 완료 → 예측 가능 상태

 6️⃣ predict() 호출 과정

predictions = loaded_model.predict(data)

예측 요청이 들어오면 흐름은 이렇게 이어집니다 👇

MLflow (pyfunc)
   ↓
CustomMLflowFlavor.predict()
   ↓
CustomMLModel.predict()
   ↓
결과 반환 ([4, 9, 16])

즉, CustomMLflowFlavor는 MLflow 입장에서
**“내가 예측을 요청할 수 있는 모델”**로 보이지만,
실제 계산은 내부의 CustomMLModel이 수행합니다.

이건 오버라이딩(override) 이 아니라
위임(delegation) 구조입니다.

 

 

3. 전체 실행 순서도

┌───────────────────────────────────────────┐
│               save_model()                │
│───────────────────────────────────────────│
│ ① model.pkl 저장                         │
│ ② MLmodel 메타데이터 생성                │
└───────────────────────────────────────────┘
                    ↓
┌───────────────────────────────────────────┐
│         mlflow.pyfunc.load_model()        │
│───────────────────────────────────────────│
│ ③ MLmodel 읽기                            │
│ ④ loader_module에서 _load_pyfunc() 호출  │
│ ⑤ CustomMLflowFlavor 객체 생성            │
│ ⑥ load_context() (필요 시 실행)           │
└───────────────────────────────────────────┘
                    ↓
┌───────────────────────────────────────────┐
│           predict(model_input)            │
│───────────────────────────────────────────│
│ ⑦ MLflow → CustomMLflowFlavor.predict()   │
│ ⑧ 내부 모델의 predict() 위임 호출         │
│ ⑨ 결과 반환 ([4, 9, 16])                 │
└───────────────────────────────────────────┘

 

 

정리하면

 

이 구조는 “MLflow가 전혀 모르는 모델”이라도
PythonModel 상속 + _load_pyfunc() 정의만으로
MLflow 생태계 안에서 저장·로드·예측이 가능한 형태로 통합하는 방식입니다.

즉,
MLflow의 Custom Flavor는 단순히 모델을 저장하는 기능이 아니라,

MLflow가 이해할 수 없는 새로운 모델을 표준 규격(PyFunc) 으로 변환하는 “언어 통역기” 역할을 합니다.