MLflow Custom Flavor 만들기 — 새로운 프레임워크를 MLflow에 통합하기

MLflow는 이미 scikit-learn, TensorFlow, PyTorch 등 다양한 머신러닝 프레임워크를 기본적으로 지원하지만,
직접 만든 모델이나 새로운 프레임워크는 기본적으로 통합되어 있지 않습니다.
이런 경우, Custom Flavor(사용자 정의 플래버) 기능을 사용하면 MLflow의 생태계 안에 손쉽게 통합할 수 있습니다.

 

 

1. Custom Flavor란?

 

MLflow의 Custom Flavor는

“MLflow가 기본적으로 지원하지 않는 모델을 표준 인터페이스로 통합하기 위한 구조”

를 의미합니다.

 

즉, 새로운 머신러닝 프레임워크나 사내 모델을
MLflow의 저장(save), 로드(load), 서빙(serve) 흐름 속에 쉽게 포함시킬 수 있도록 해줍니다.

 

2. Custom Flavor를 만드는 이유

MLflow는 sklearn, xgboost, pytorch 같은 내장 flavor를 제공합니다.
하지만 새로운 프레임워크나 맞춤형 모델을 관리하려면 직접 flavor를 정의해야 합니다.

예를 들어:

• 사내 전용 모델 형식을 MLflow에 연결하고 싶을 때
• MLflow가 공식적으로 지원하지 않는 프레임워크(JAX, CatBoost 등)를 추가하고 싶을 때
• 특정 전처리·후처리 로직을 포함한 모델을 MLflow로 관리하고 싶을 때

이때 “Custom Flavor”가 필요합니다.

 

 

3. Custom Flavor의 구성 요소

하나의 Custom Flavor는 다음 세 부분으로 이루어집니다:

1. Python 클래스
   • save_model(), load_model(), predict() 메서드를 구현해야 합니다.
   • 이 클래스가 Flavor의 핵심 동작을 정의합니다.
2. 모델 디렉토리 구조
   • MLmodel: 메타데이터 파일
   • model.pkl: 실제 모델 파일
   • Flavor 로직이 정의된 Python 파일
3. MLflow Flavor 메타데이터
   • 모델의 flavor 정보(flavors),
   • Python 버전,
   • 로더 모듈 정보(loader_module) 등이 포함됩니다.

 

4. Custom Flavor 디렉토리 구조

custom_flavor_model/
├── MLmodel
├── model.pkl
└── custom_flavor.py

 

• MLmodel: 모델을 로드할 때 참고하는 메타데이터
• model.pkl: 학습된 모델이 저장된 파일
• custom_flavor.py: Flavor의 저장·로딩 로직이 정의된 스크립트

 

5. 예제 코드 — Custom Flavor 구현하기

def save_model(path, model):
    os.makedirs(path, exist_ok=True)

    # 1️⃣ 모델 저장 (Pickle 형식)
    model_path = os.path.join(path, "model.pkl")
    with open(model_path, "wb") as f:
        pickle.dump(model, f)

    # 2️⃣ MLmodel 메타데이터 정의
    mlmodel_data = {
        "flavors": {
            "python_function": {
                "loader_module": __name__,  # 현재 파일을 로더로 지정
                "python_version": "3.10"
            }
        }
    }

    # 3️⃣ MLmodel 파일 저장
    with open(os.path.join(path, "MLmodel"), "w") as f:
        json.dump(mlmodel_data, f, indent=4)

# Custom Flavor로 저장
save_model("custom_flavor_model", model)

 

이 구조는 MLflow가 해당 모델을 로드할 때 MLmodel 파일을 참고하여
어떤 Python 모듈(loader_module)을 통해 모델을 불러올지 결정하게 해줍니다.

 

 

정리하면 

 

MLflow의 Custom Flavor 기능은
단순히 모델을 저장하는 수준을 넘어,
새로운 머신러닝 프레임워크나 사내 모델을 완전히 MLflow 생태계에 통합할 수 있는 강력한 확장 포인트입니다.

즉, “MLflow가 아직 모르는 모델도 MLflow답게 관리할 수 있게 만드는 기능”입니다.