MLflow Models

 

MLflow 모델은 다양한 다운스트림 도구에서 사용할 수 있는 머신 러닝 모델을 패키징하기 위한 표준 포맷입니다. 이 포맷은 다양한 다운스트림 도구에서 이해할 수 있는 다양한 "flavors"로 모델을 저장할 수 있는 규칙을 정의합니다.

 

Storage Format

각 MLflow 모델은 임의의 파일을 포함하는 디렉토리이며, 디렉토리 루트에는 모델을 볼 수 있는 여러 유형을 정의할 수 있는 MLmodel 파일이 있습니다. MLflow 모델의 측면에서모델이란 직렬화된 객체(pickled scikit-learn 모델)이거나 mlflow.models.set_model() API로 정의된 모델 인스턴스를 포함하는 파이썬 스크립트(또는 노트북) 일 수 있습니다. 

Flavor는 MLflow 모델의 주요 개념인데요, 이는 배포시에 라이브러리와 통합하지 않고도 모든 ML 라이브러리의 모델과 작동할 수 있도록 정해진 규칙입니다. MLflow는 모든 기본 제공 배포 도구가 지원하는 여러가지 표준 flavor를 정의합니다. 여기에는 모델을 파이썬 함수로 실행하는 python 함수 flavor도 포함됩니다.

 

MLmodel file

특정 모델이 지원하는 모든 flavor는 YAML 형식의 MLmodel 파일에 정의되어 있습니다. 예를 들어, MLflow 저장소에서 python examples/sklearn_logistic_regression/train.py를 실행하면 모델 디렉토리 아래에 다음 파일이 생성됩니다.

# Directory written by mlflow.sklearn.save_model(model, "model", input_example=...)
model/
├── MLmodel
├── model.pkl
├── conda.yaml
├── python_env.yaml
├── requirements.txt
├── input_example.json (optional, only logged when input example is provided and valid during model logging)
├── serving_input_example.json (optional, only logged when input example is provided and valid during model logging)
└── environment_variables.txt (optional, only logged when environment variables are used during model inference)

 

그리고 MLmodel 파일은 두 가지의 flavors를 설명합니다

time_created: 2018-05-25T17:28:53.35

flavors:
  sklearn:
    sklearn_version: 0.19.1
    pickled_model: model.pkl
  python_function:
    loader_module: mlflow.sklearn

 

모델 플레이버를 나열하는 flavors 필드 외에도 MLmodel YAML 형식에는 다음 필드가 포함될 수 있습니다.

• time_created: Date and time when the model was created, in UTC ISO 8601 format.
• run_id: ID of the run that created the model, if the model was saved using MLflow Tracking.
• signature: model signature in JSON format.
• input_example: reference to an artifact with input example.
• databricks_runtime: Databricks runtime version and type, if the model was trained in a Databricks notebook or job.
• mlflow_version: The version of MLflow that was used to log the model.

 

Additional Logged Files

환경 재현을 위해 모델이 로깅될 때마다 conda.yaml, python_env.yaml 및 requirements.txt 파일이 자동으로 로깅됩니다. 그런 다음 이러한 파일을 사용하여 conda 또는 virtualenv와 pip를 사용하여 종속성을 다시 설치할 수 있습니다.

모델을 로깅할 때 모델 입력 예제가 제공된 경우 두 개의 추가 파일 input_example.json 및 serving_input_example.json이 로깅됩니다.

모델을 로깅할 때 모델 메타데이터 파일(MLmodel, conda.yaml, python_env.yaml, requirements.txt)이 metadata라는 하위 디렉터리에 복사됩니다. 휠이 달린 모델의 경우 original_requirements.txt 파일도 복사됩니다.

python_env.yaml

python: 3.9.8
build_dependencies:
- pip==23.3.2
- setuptools==69.0.3
- wheel==0.42.0
dependencies:
- -r requirements.txt

 

requirements.txt

mlflow==2.9.2
scikit-learn==1.3.2
cloudpickle==3.0.0

 

conda.yaml

name: mlflow-env
channels:
  - conda-forge
dependencies:
- python=3.9.8
- pip
- pip:
  - mlflow==2.9.2
  - scikit-learn==1.3.2
  - cloudpickle==3.0.0

 

Environment variables file

MLflow는 모델을 로깅할 때 모델 추론 중에 사용되는 환경 변수를 environment_variables.txt 파일에 기록합니다.

environment_variables.txt 파일에는 모델 추론 중에 사용되는 환경 변수의 이름만 포함되며 값은 저장되지 않습니다.

RECORD_ENV_VAR_ALLOWLIST = {
    # api key related
    "API_KEY",  # e.g. OPENAI_API_KEY
    "API_TOKEN",
    # databricks auth related
    "DATABRICKS_HOST",
    "DATABRICKS_USERNAME",
    "DATABRICKS_PASSWORD",
    "DATABRICKS_TOKEN",
    "DATABRICKS_INSECURE",
    "DATABRICKS_CLIENT_ID",
    "DATABRICKS_CLIENT_SECRET",
    "_DATABRICKS_WORKSPACE_HOST",
    "_DATABRICKS_WORKSPACE_ID",
}

import mlflow
import os

os.environ["TEST_API_KEY"] = "test_api_key"


class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input, params=None):
        if os.environ.get("TEST_API_KEY"):
            return model_input
        raise Exception("API key not found")


with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        "model", python_model=MyModel(), input_example="data"
    )

 

 

Model Signatures And Input Examples

MLflow에서 모델 시그니처와 입력 예제의 복잡성을 이해하는 것은 효과적인 모델 관리 및 배포에 매우 중요합니다.

• Model Signature: 모델 입력, 출력 및 추가 추론 매개변수에 대한 스키마를 정의하여 모델 상호작용을 위한 표준화된 인터페이스를 촉진합니다.
• Model Input Example: 유효한 모델 입력의 구체적인 인스턴스를 제공하여 모델 요구 사항을 이해하고 테스트하는 데 도움이 됩니다. 또한, 모델을 로깅할 때 입력 예제가 명시적으로 제공되지 않으면 모델 서명이 자동으로 추론되어 저장됩니다.
• Model Serving Payload Example: 배포된 모델 엔드포인트를 쿼리하기 위한 JSON 페이로드 예시를 제공합니다. 모델을 로깅할 때 입력 예제가 제공되면 해당 입력 예제에서 제공 페이로드 예제가 자동으로 생성되어 serving_input_example.json으로 저장됩니다.

 

 

Model API

여러 가지 방법으로 MLflow 모델을 저장하고 로드할 수 있습니다. 첫째, MLflow에는 여러 공통 라이브러리와의 통합이 포함되어 있습니다. 예를 들어, mlflow.sklearn에는 scikit-learn 모델에 대한 save_model, log_model, load_model 함수가 포함되어 있습니다. 둘째, mlflow.models.Model 클래스를 사용하여 모델을 만들고 작성할 수 있습니다. 이 클래스에는 네 가지 주요 기능이 있습니다.

• add_flavor 모델에 flavor를 추가합니다. 각 flavor에는 string name과 key-value 사전이 있으며, 값은 YAML로 직렬화할 수 있는 모든 객체가 될 수 있습니다
• save 모델을 로컬 디렉토리에 저장합니다.
• log MLflow Tracking을 사용하여 현재 실행에서 모델을 아티팩트로 기록합니다.
• load 로컬 디렉토리나 이전 실행의 아티팩트에서 모델을 로드합니다.

 

 

Models From Code

"Code에서 모델 생성" 기능은 독립 실행형 Python 스크립트에서 모델을 직접 정의하고 로깅할 수 있게 해줍니다. 이 기능은 특히 코드로 표현이 가능한 모델(훈련을 통해 최적화된 가중치가 필요하지 않은 모델)이나 외부 서비스를 활용하는 애플리케이션(예: LangChain 체인)을 로깅할 때 유용합니다. 또 다른 장점으로는 이 접근 방식을 통해 Python의 pickle 또는 cloudpickle 모듈 사용을 완전히 피할 수 있다는 점인데, 이는 신뢰할 수 없는 모델을 로드할 때 발생할 수 있는 보안 위험을 줄이는 데 기여합니다.

이 기능은 LangChain, LlamaIndex 및 PythonModel 모델에서만 지원됩니다.

코드에서 모델을 로깅하려면 mlflow.models.set_model() API를 활용할 수 있습니다. 이 API를 사용하면 모델 클래스의 인스턴스를 직접 정의하여 모델이 정의된 파일 내에서 모델을 지정할 수 있습니다. 이러한 모델을 로깅할 때에는 객체 대신 파일 경로를 지정하며, 이 경로는 모델 클래스 정의와 set_model API가 적용된 사용자 정의 모델 인스턴스를 포함하는 Python 파일을 가리킵니다.

아래 그림은 표준 모델 로깅 프로세스와 "Code에서 모델 생성" 기능을 활용해 저장할 수 있는 모델의 로깅 프로세스를 비교하여 보여줍니다. 이는 "Code에서 모델 생성" 기능을 사용할 수 있는 모델에 해당합니다.

 

 

 

예를 들어, my_model.py라는 별도의 파일에 모델을 정의합니다.

import mlflow
from mlflow.models import set_model


class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        return model_input


# Define the custom PythonModel instance that will be used for inference
set_model(MyModel())

 

그런 다음 다른 Python 스크립트에서 파일 경로에서 모델을 로깅합니다.

import mlflow

model_path = "my_model.py"

with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        python_model=model_path,  # Define the model as the path to the Python file
        artifact_path="my_model",
    )

# Loading the model behaves exactly as if an instance of MyModel had been logged
my_model = mlflow.pyfunc.load_model(model_info.model_uri)

 

"Code에서 모델 생성" 기능은 외부 파일 참조에서 가져온 import 문을 캡처하는 것을 지원하지 않습니다. 만약 pip install로 해결되지 않는 의존성이 있다면, 해당 의존성은 code_paths 기능을 사용하여 절대 경로를 통한 import로 포함하고 해결해야 합니다. 하지만 코드 경로 의존성 처리의 제한 사항을 고려하면, 코드에서 정의된 모델에 필요한 모든 로컬 의존성을 동일한 Python 스크립트 파일 내에 캡슐화하는 것을 권장합니다. 이를 통해 구현을 단순화하고 의존성 문제를 최소화할 수 있습니다.

코드에서 모델을 정의하고 mlflow.models.set_model() API를 사용할 때, 로깅되는 스크립트 내에 정의된 코드는 내부적으로 실행되어 코드의 유효성을 확인합니다. 만약 스크립트에 외부 서비스와의 연결이 포함되어 있다면(예: LangChain에서 생성형 AI 서비스에 연결), 모델 로깅 시 해당 서비스로 연결 요청이 발생할 수 있다는 점을 유의해야 합니다. 이러한 동작을 고려하여 외부 서비스 연결이 필요한 경우, 로깅 프로세스 중 발생할 수 있는 요청 비용이나 시간 지연을 사전에 계획하는 것이 중요합니다.

 

 

Built-In Model Flavors

MLflow는 애플리케이션에서 유용할 수 있는 여러 가지 표준 flavors를 제공합니다. 특히, 많은 배포 도구가 이러한 flavors를 지원하므로 이러한 flavors 중 하나로 자체 모델을 내보내 모든 도구의 이점을 얻을 수 있습니다. 아래는 인기있는 flavors들이며, 더 다양한 flavors는 여기서 확인하시기 바랍니다.

• Python Function (python_function)
• PyTorch (pytorch)
• Scikit-learn (sklearn)
• TensorFlow (tensorflow)
• ONNX (onnx)
• OpenAI (openai) (Experimental)
• LangChain (langchain) (Experimental)
• Transformers (transformers) (Experimental)
• ...

 

Python Function

python_function 모델 플레이버는 MLflow Python 모델의 기본 인터페이스로 작동합니다. 모든 MLflow Python 모델은 python_function 모델로 로드 가능해야 하며, 이를 통해 모델이 어떤 영속성 모듈이나 프레임워크를 사용해 생성되었는지와 상관없이 다른 MLflow 도구와 함께 동작할 수 있습니다. 이러한 상호운용성은 매우 강력하며, 어떤 Python 모델이든 다양한 환경에서 프로덕션화할 수 있도록 합니다.

추가로, python_function 모델 플레이버는 Python 모델을 위한 범용 파일 시스템 모델 형식을 정의하고, 이를 저장하거나 로드하는 유틸리티를 제공합니다. 이 형식은 모델을 로드하고 사용하는 데 필요한 모든 정보를 포함하므로 자체 완결성을 갖추고 있습니다. 의존성은 모델과 함께 직접 저장되거나 Conda 환경을 통해 참조됩니다. 이러한 모델 형식은 다른 도구가 MLflow와 통합하여 자신들의 모델을 활용할 수 있도록 지원합니다.

 

Function-based Model

추가적인 종속 메서드 없이 간단한 Python 함수를 직렬화하려는 경우 키워드 인수 python_model을 통해 predict 메서드를 간단히 로깅하면 됩니다.

import mlflow
import pandas as pd


# Define a simple function to log
def predict(model_input):
    return model_input.apply(lambda x: x * 2)


# Save the function as a model
with mlflow.start_run():
    mlflow.pyfunc.log_model("model", python_model=predict, pip_requirements=["pandas"])
    run_id = mlflow.active_run().info.run_id

# Load the model from the tracking server and perform inference
model = mlflow.pyfunc.load_model(f"runs:/{run_id}/model")
x_new = pd.Series([1, 2, 3])

prediction = model.predict(x_new)
print(prediction)

 

Class-based Model

예를 들어 전처리, 복잡한 예측 논리 또는 사용자 정의 직렬화를 처리하는 클래스와 같이 보다 복잡한 객체를 직렬화하려는 경우 PythonModel 클래스를 서브클래싱해야 합니다. PythonModel 구현은 지나치게 복잡하므로 간단한 사례에는 함수 기반 모델을 대신 사용하는 것이 좋습니다.

import mlflow
import pandas as pd


class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input, params=None):
        return [x * 2 for x in model_input]


# Save the function as a model
with mlflow.start_run():
    mlflow.pyfunc.log_model(
        "model", python_model=MyModel(), pip_requirements=["pandas"]
    )
    run_id = mlflow.active_run().info.run_id

# Load the model from the tracking server and perform inference
model = mlflow.pyfunc.load_model(f"runs:/{run_id}/model")
x_new = pd.Series([1, 2, 3])

print(f"Prediction:
    {model.predict(x_new)}")

 

 

Example: Creating a custom “add n” model

다음은 Pandas DataFrame 입력의 모든 열에 지정된 숫자 값 nnn을 더하는 사용자 정의 모델을 정의하는 예제입니다. 이 예제에서는 n=5 로 설정된 모델 인스턴스를 MLflow 모델 형식으로 저장하고, 이후 python_function 형식으로 모델을 로드하여 샘플 입력 데이터를 평가하는 과정을 보여줍니다.

import mlflow.pyfunc


# Define the model class
class AddN(mlflow.pyfunc.PythonModel):
    def __init__(self, n):
        self.n = n

    def predict(self, context, model_input, params=None):
        return model_input.apply(lambda column: column + self.n)


# Construct and save the model
model_path = "add_n_model"
add5_model = AddN(n=5)
mlflow.pyfunc.save_model(path=model_path, python_model=add5_model)

# Load the model in `python_function` format
loaded_model = mlflow.pyfunc.load_model(model_path)

# Evaluate the model
import pandas as pd

model_input = pd.DataFrame([range(10)])
model_output = loaded_model.predict(model_input)
assert model_output.equals(pd.DataFrame([range(5, 15)]))

 

 

Example: Saving an XGBoost model in MLflow format

아래는 XGBoost를 사용하여 Gradient Boosted Tree 모델을 학습하고 이를 MLflow와 통합하여 모델 추론을 수행하는 예제입니다. 이 과정에서는 XGBoost 모델을 MLflow의 python_function 추론 API에 맞게 래핑하고, 이를 활용해 MLflow 모델로 저장한 후, 테스트 데이터를 평가하는 데 사용합니다.

# Load training and test datasets
from sys import version_info
import xgboost as xgb
from sklearn import datasets
from sklearn.model_selection import train_test_split

PYTHON_VERSION = f"{version_info.major}.{version_info.minor}.{version_info.micro}"
iris = datasets.load_iris()
x = iris.data[:, 2:]
y = iris.target
x_train, x_test, y_train, _ = train_test_split(x, y, test_size=0.2, random_state=42)
dtrain = xgb.DMatrix(x_train, label=y_train)

# Train and save an XGBoost model
xgb_model = xgb.train(params={"max_depth": 10}, dtrain=dtrain, num_boost_round=10)
xgb_model_path = "xgb_model.pth"
xgb_model.save_model(xgb_model_path)

# Create an `artifacts` dictionary that assigns a unique name to the saved XGBoost model file.
# This dictionary will be passed to `mlflow.pyfunc.save_model`, which will copy the model file
# into the new MLflow Model's directory.
artifacts = {"xgb_model": xgb_model_path}

# Define the model class
import mlflow.pyfunc


class XGBWrapper(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        import xgboost as xgb

        self.xgb_model = xgb.Booster()
        self.xgb_model.load_model(context.artifacts["xgb_model"])

    def predict(self, context, model_input, params=None):
        input_matrix = xgb.DMatrix(model_input.values)
        return self.xgb_model.predict(input_matrix)


# Create a Conda environment for the new MLflow Model that contains all necessary dependencies.
import cloudpickle

conda_env = {
    "channels": ["defaults"],
    "dependencies": [
        f"python={PYTHON_VERSION}",
        "pip",
        {
            "pip": [
                f"mlflow=={mlflow.__version__}",
                f"xgboost=={xgb.__version__}",
                f"cloudpickle=={cloudpickle.__version__}",
            ],
        },
    ],
    "name": "xgb_env",
}

# Save the MLflow Model
mlflow_pyfunc_model_path = "xgb_mlflow_pyfunc"
mlflow.pyfunc.save_model(
    path=mlflow_pyfunc_model_path,
    python_model=XGBWrapper(),
    artifacts=artifacts,
    conda_env=conda_env,
)

# Load the model in `python_function` format
loaded_model = mlflow.pyfunc.load_model(mlflow_pyfunc_model_path)

# Evaluate the model
import pandas as pd

test_predictions = loaded_model.predict(pd.DataFrame(x_test))
print(test_predictions)

 

 

Example: Logging a transformers model with hf:/ schema to avoid copying large files

이 예제는 Hugging Face Hub에서 제공하는 Transformers 모델을 직접 로깅하기 위해 특별한 스키마 hf:/를 사용하는 방법을 보여줍니다. 이 방식은 모델이 너무 크거나, 특히 모델을 직접 서빙하려는 경우에 유용합니다. 다만, 모델을 로컬에서 다운로드하고 테스트하려는 경우에는 추가적인 저장 공간 절약이 되지는 않습니다.

import mlflow
from mlflow.models import infer_signature
import numpy as np
import transformers


# Define a custom PythonModel
class QAModel(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        """
        This method initializes the tokenizer and language model
        using the specified snapshot location from model context.
        """
        snapshot_location = context.artifacts["bert-tiny-model"]
        # Initialize tokenizer and language model
        tokenizer = transformers.AutoTokenizer.from_pretrained(snapshot_location)
        model = transformers.BertForQuestionAnswering.from_pretrained(snapshot_location)
        self.pipeline = transformers.pipeline(
            task="question-answering", model=model, tokenizer=tokenizer
        )

    def predict(self, context, model_input, params=None):
        question = model_input["question"][0]
        if isinstance(question, np.ndarray):
            question = question.item()
        ctx = model_input["context"][0]
        if isinstance(ctx, np.ndarray):
            ctx = ctx.item()
        return self.pipeline(question=question, context=ctx)


# Log the model
data = {"question": "Who's house?", "context": "The house is owned by Run."}
pyfunc_artifact_path = "question_answering_model"
with mlflow.start_run() as run:
    model_info = mlflow.pyfunc.log_model(
        artifact_path=pyfunc_artifact_path,
        python_model=QAModel(),
        artifacts={"bert-tiny-model": "hf:/prajjwal1/bert-tiny"},
        input_example=data,
        signature=infer_signature(data, ["Run"]),
        extra_pip_requirements=["torch", "accelerate", "transformers", "numpy"],
    )