Ray Serve 모델 레지스트리 통합

Ray Serve는 Python 네이티브라서 MLflow 같은 모델 레지스트리와 특별한 설정 없이 자연스럽게 통합돼요. 복잡한 어댑터나 glue code 없이 모델 레지스트리에 등록된 모델을 직접 로드해서 프로덕션 워크로드로 서빙할 수 있어요. 이 가이드에선 Ray Serve와 모델 레지스트리를 연동해 end-to-end ML 서빙 파이프라인을 만드는 방법을 다뤄요.

출처: Model Registry Integration — Ray Docs

왜 Python 네이티브 통합이 중요한가

프레임워크 특화 서빙 솔루션은 커스텀 어댑터나 복잡한 설정이 필요한 반면, Ray Serve는 임의의 Python 코드를 실행해요. 따라서:

표준 Python 클라이언트로 어떤 모델 레지스트리에서든 모델을 직접 로드할 수 있고
모델 로딩과 추론을 하나의 deployment에서 결합할 수 있으며
YAML 설정이나 커스텀 직렬화 형식과 씨름하지 않고 빠르게 반복할 수 있어요.

MLflow와 통합

MLflow는 ML 수명주기를 관리하는 인기 있는 오픈소스 플랫폼이에요. Ray Serve는 MLflow Model Registry에서 모델을 로드해 프로덕션에서 서빙하기 쉽게 만들어줘요.

MLflow 모델 서빙 모범 사례

모델 시그니처와 입력 스키마 검증 사용: mlflow.models.infer_signature로 모델 시그니처를 항상 로그해서 MLflow가 입력을 검증하게 하세요. 업스트림 코드가 바뀌어도 조용한 실패(silent failure)를 막고, 서빙 중 자동 스키마 강제를 가능하게 해요.
의존성을 명시적으로 패키징: 모델을 로그할 때 pip_requirements를 쓰고 핵심 라이브러리 버전을 고정(pin)하세요. 그러면 학습·평가·서빙 환경에서 모델이 동일하게 동작해요.
전처리 파이프라인 영속화: scikit-learn을 쓴다면 전처리 단계를 포함한 완전한 Pipeline 객체를 로그하세요. 그러면 학습과 서빙 변환이 일관되게 유지돼요.
LLM·디퓨전 모델은 Hugging Face Hub나 Weights & Biases 사용: MLflow 내장 REST 서버는 고동시성(high-concurrency) GPU 워크로드에 최적화돼 있지 않아요. 대규모 언어 모델, 디퓨전 모델, 기타 무거운 트랜스포머 기반 아키텍처는 모델 레지스트리로 Hugging Face Hub나 Weights & Biases를 사용하세요. 이 플랫폼들은 대형 모델 아티팩트에 더 나은 도구를 제공하고, Ray Serve가 GPU 배칭·자동 스케일링·스케줄링을 효율적으로 처리해요.

모델 학습 및 등록

다음 예시는 scikit-learn 모델을 모범 사례에 맞춰 학습하고 MLflow에 등록하는 방법이에요:

from sklearn.datasets import make_regression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
import mlflow
import mlflow.sklearn
import mlflow.pyfunc
from mlflow.entities import LoggedModelStatus
from mlflow.models import infer_signature
import numpy as np


def train_and_register_model():
    # Initialize model in PENDING state
    logged_model = mlflow.initialize_logged_model(
        name="my-random-forest-reg-model",
        model_type="sklearn",
        tags={"model_type": "random_forest"},
    )

    try:
        with mlflow.start_run() as run:
            X, y = make_regression(n_features=4, n_informative=2, random_state=0, shuffle=False)
            X_train, X_test, y_train, y_test = train_test_split(
                X, y, test_size=0.2, random_state=42
            )

            params = {"max_depth": 2, "random_state": 42}

            # Best Practice: Use sklearn Pipeline to persist preprocessing
            # This ensures training and serving transformations stay aligned
            pipeline = Pipeline([
                ("scaler", StandardScaler()),
                ("regressor", RandomForestRegressor(**params))
            ])
            pipeline.fit(X_train, y_train)

            # Log parameters and metrics
            mlflow.log_params(params)

            y_pred = pipeline.predict(X_test)
            mlflow.log_metrics({"mse": mean_squared_error(y_test, y_pred)})

            # Best Practice: Infer model signature for input validation
            # Prevents silent failures from mismatched feature order or missing columns
            signature = infer_signature(X_train, y_pred)

            # Best Practice: Pin dependency versions explicitly
            # Ensures identical behavior across training, evaluation, and serving
            pip_requirements = [
                f"scikit-learn=={__import__('sklearn').__version__}",
                f"numpy=={np.__version__}",
            ]

            # Log the sklearn pipeline with signature and dependencies
            mlflow.sklearn.log_model(
                sk_model=pipeline,
                name="sklearn-model",
                input_example=X_train[:1],
                signature=signature,
                pip_requirements=pip_requirements,
                registered_model_name="my-random-forest-reg-model",
                model_id=logged_model.model_id,
            )

        # Finalize model as READY
        mlflow.finalize_logged_model(logged_model.model_id, LoggedModelStatus.READY)

        mlflow.set_logged_model_tags(
            logged_model.model_id,
            tags={"production": "true"},
        )

    except Exception as e:
        # Mark model as FAILED if issues occur
        mlflow.finalize_logged_model(logged_model.model_id, LoggedModelStatus.FAILED)
        raise

    # Retrieve and work with the logged model
    final_model = mlflow.get_logged_model(logged_model.model_id)
    print(f"Model {final_model.name} is {final_model.status}")

이 함수는 전처리를 담은 Pipeline으로 감싼 RandomForestRegressor를 학습하고, 시그니처와 고정된 의존성과 함께 모델을 로그해 my-random-forest-reg-model이라는 이름으로 MLflow Model Registry에 등록해요.

모델 로드 및 서빙

MLflow에 모델을 등록했다면 Ray Serve로 로드해서 서빙할 수 있어요. 다음 예시는 warm-start 초기화로 MLflow Model Registry에서 모델을 로드하는 deployment를 만드는 방법이에요:

from ray import serve
import mlflow.pyfunc
import numpy as np


@serve.deployment
class MLflowModelDeployment:
    def __init__(self):
        # Search for models with production tag
        models = mlflow.search_logged_models(
            filter_string="tags.production='true' AND name='my-random-forest-reg-model'",
            order_by=[{"field_name": "creation_time", "ascending": False}],
        )
        if models.empty:
            raise ValueError("No model with production tag found")

        # Get the most recent production model
        model_row = models.iloc[0]
        artifact_location = model_row["artifact_location"]

        # Best Practice: Load model once during initialization (warm-start)
        # This eliminates first-request latency spikes
        self.model = mlflow.pyfunc.load_model(artifact_location)

        # Pre-warm the model with a dummy prediction
        dummy_input = np.zeros((1, 4))
        _ = self.model.predict(dummy_input)

    async def __call__(self, request):
        data = await request.json()
        features = np.array(data["features"])

        # MLflow validates input against the logged signature automatically
        prediction = self.model.predict(features)
        return {"prediction": prediction.tolist()}


app = MLflowModelDeployment.bind()

핵심 포인트는 이렇습니다. 모델을 __init__에서 한 번만 로드(warm-start)해 첫 요청의 지연 스파이크를 없애고, MLflow가 로그된 시그니처로 입력을 자동 검증한다는 점이에요. 이 접근은 어떤 모델 레지스트리에서든 동작하며, 학습·서빙 환경 간 변환이 일관되게 유지되도록 전처리 파이프라인과 의존성 버전을 함께 영속화하는 게 관건이에요.

Ray Serve 모델 레지스트리 통합