커스텀 모델 기초: 시그니처와 PyFunc

import mlflow
from mlflow.models import ModelSignature, infer_signature
from mlflow.types.schema import Schema, ColSpec, ParamSchema, ParamSpec

# 방법 1: 학습 데이터에서 자동 추론 (권장)
import pandas as pd
import numpy as np

X_sample = pd.DataFrame({
    "amount": [50000.0, 25.0],
    "merchant_category": ["online_retail", "grocery"],
    "hour": [3, 12],
    "is_international": [1, 0]
})
y_sample = pd.DataFrame({
    "fraud_probability": [0.85, 0.02],
    "is_fraud": [1, 0]
})

# 입력/출력 데이터로 시그니처 자동 추론
signature = infer_signature(X_sample, y_sample)
print(signature)
# inputs: [amount: double, merchant_category: string, hour: long, is_international: long]
# outputs: [fraud_probability: double, is_fraud: long]

# 방법 2: 수동으로 시그니처 정의 (정밀한 제어가 필요할 때)
signature = ModelSignature(
    inputs=Schema([
        ColSpec("double", "amount"),
        ColSpec("string", "merchant_category"),
        ColSpec("long", "hour"),
        ColSpec("integer", "is_international")
    ]),
    outputs=Schema([
        ColSpec("double", "fraud_probability"),
        ColSpec("integer", "is_fraud"),
        ColSpec("string", "risk_level")
    ]),
    # params: 추론 시 동적으로 변경할 수 있는 파라미터 (선택)
    params=ParamSchema([
        ParamSpec("threshold", "double", 0.7)
    ])
)

# 방법 3: params를 활용한 동적 추론 파라미터
class FraudModelWithParams(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        import joblib
        self.model = joblib.load(context.artifacts["model_path"])
        self.scaler = joblib.load(context.artifacts["scaler_path"])

    def predict(self, context, model_input, params=None):
        # params에서 임계값을 런타임에 받음 (기본값 0.7)
        threshold = (params or {}).get("threshold", 0.7)

        scaled_input = self.scaler.transform(model_input)
        probabilities = self.model.predict_proba(scaled_input)[:, 1]

        return pd.DataFrame({
            "fraud_probability": probabilities,
            "is_fraud": (probabilities >= threshold).astype(int)
        })

# params가 포함된 시그니처 생성
X_sample = X_test[:5]
predictions = FraudModelWithParams().predict(None, X_sample)
signature = infer_signature(
    X_sample,
    predictions,
    params={"threshold": 0.7}   # 기본 파라미터 값
)

import mlflow.pyfunc
import pandas as pd
import numpy as np

class FraudDetectionModel(mlflow.pyfunc.PythonModel):
    """전처리 + 모델 추론 + 후처리를 하나로 묶는 커스텀 모델"""

    def load_context(self, context):
        """모델 로드 시 한 번 실행 (무거운 초기화 작업)"""
        import joblib

        # artifacts에서 모델과 전처리기 로드
        self.model = joblib.load(context.artifacts["model_path"])
        self.scaler = joblib.load(context.artifacts["scaler_path"])
        self.threshold = 0.7  # 커스텀 임계값

    def predict(self, context, model_input, params=None):
        """추론 요청마다 실행되는 메서드"""
        # 1. 전처리
        scaled_input = self.scaler.transform(model_input)

        # 2. 예측 확률 계산
        probabilities = self.model.predict_proba(scaled_input)[:, 1]

        # 3. 후처리 (커스텀 임계값 + 리스크 등급)
        results = pd.DataFrame({
            "fraud_probability": probabilities,
            "is_fraud": (probabilities >= self.threshold).astype(int),
            "risk_level": pd.cut(
                probabilities,
                bins=[0, 0.3, 0.7, 1.0],
                labels=["LOW", "MEDIUM", "HIGH"]
            )
        })

        return results

import joblib
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.preprocessing import StandardScaler

# 모델과 스케일러 학습
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
model = GradientBoostingClassifier(n_estimators=200)
model.fit(X_train_scaled, y_train)

# 아티팩트 저장
joblib.dump(model, "/tmp/model.joblib")
joblib.dump(scaler, "/tmp/scaler.joblib")

# 시그니처 자동 추론
sample_input = pd.DataFrame(X_test[:3], columns=feature_cols)
sample_output = FraudDetectionModel().predict(None, sample_input)
signature = infer_signature(sample_input, sample_output)

# 커스텀 모델을 MLflow에 로깅
with mlflow.start_run(run_name="fraud-custom-pyfunc"):
    model_info = mlflow.pyfunc.log_model(
        artifact_path="fraud_model",
        python_model=FraudDetectionModel(),
        artifacts={
            "model_path": "/tmp/model.joblib",
            "scaler_path": "/tmp/scaler.joblib"
        },
        conda_env={
            "dependencies": [
                "python=3.10",
                {"pip": ["scikit-learn==1.4.0", "pandas>=2.0", "numpy>=1.24"]}
            ]
        },
        signature=signature,
        input_example=sample_input,
        registered_model_name="catalog.schema.fraud_detection_custom"
    )