Spark ML 모델 학습과 배치 추론

Foundry는 Spark MLlib를 지원하며 모델 API에 Spark DataFrame을 선언할 수 있습니다. 다만 Spark 모델을 실행할 전용 컨테이너 이미지를 만드는 경로는 지원하지 않아 일부 배포 기능에 제약이 있습니다. 일반적인 모델, 특히 REST로 제공할 모델에는 scikit-learn, XGBoost 또는 Keras처럼 단일 노드에서 학습하고 추론할 수 있는 프레임워크를 권장합니다.

Spark 모델 학습은 Code Repositories를 권장합니다. Code Workspaces는 단일 노드에서 실행하며 Code Workbook은 기존 방식입니다. Spark MLlib PipelineModel은 palantir_models 0.1599.0부터 SparkMLAutoSerializer로 자동 직렬화할 수 있습니다.

Iris 분류 예제

입력 데이터셋에는 sepallength, sepalwidth, petallength, petalwidth 수치 특성과 variety 품종 정답이 있어야 합니다. 아래 경로를 실제 입력, 모델과 추론 출력으로 바꿉니다. CBAC를 사용하면 출력 리소스를 미리 생성합니다.

src/main/model_training/model_training.py에 다음 코드를 둡니다. 학습 및 시험 데이터를 seed 42로 70%와 30%로 나누고 VectorAssembler, StringIndexer 및 RandomForestClassifier를 파이프라인으로 학습합니다. 트리 수는 10, 최대 깊이는 5입니다.

# src/main/model_training/model_training.py
from pyspark.ml import Pipeline
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.feature import StringIndexer, VectorAssembler

from transforms.api import transform, Input, Output
from palantir_models.transforms import ModelOutput

from main.model_adapters.adapter import SparkModelAdapter


@transform(
    iris_data_in=Input("<PATH_TO_FOLDER>/iris_data"),
    model_output=ModelOutput("<PATH_TO_FOLDER>/spark_model"),
    inference_data_out=Output("<PATH_TO_FOLDER>/inference_data_out")
)
def compute(iris_data_in, model_output, inference_data_out):
    iris_data = iris_data_in.dataframe()
    # assuming columns sepallength, sepalwidth, petallength, petalwidth, variety
    feature_cols = ["sepallength", "sepalwidth", "petallength", "petalwidth"]
    train_data, test_data = iris_data.randomSplit([0.7, 0.3], seed=42)

    assembler = VectorAssembler(inputCols=feature_cols, outputCol="features")
    label_indexer = StringIndexer(inputCol="variety", outputCol="label", handleInvalid="keep")

    rf = RandomForestClassifier(
        labelCol="label",
        featuresCol="features",
        numTrees=10,
        maxDepth=5,
        seed=42
    )
    pipeline = Pipeline(stages=[assembler, label_indexer, rf])
    model = pipeline.fit(train_data)

    # Wrap the trained model in a ModelAdapter
    foundry_model = SparkModelAdapter(model)

    predictions = foundry_model.transform(test_data).df_out
    inference_data_out.write_dataframe(
        predictions
    )
    model_output.publish(
        model_adapter=foundry_model,
    )

어댑터

src/main/model_adapters/adapter.py의 어댑터는 학습된 PipelineModel을 저장합니다. 실제 추론 입력에는 variety 정답이 없을 수 있으므로 API의 필수 입력에 포함하지 않습니다. label은 학습 과정에서 만든 정답 인덱스이고 prediction이 모델의 실제 예측값입니다. 클래스 인덱스와 품종의 대응은 학습한 StringIndexer에서 정합니다.

# src/main/model_adapters/adapter.py
import palantir_models as pm
from palantir_models.serializers import SparkMLAutoSerializer


class SparkModelAdapter(pm.ModelAdapter):
    @pm.auto_serialize(model=SparkMLAutoSerializer())
    def __init__(self, model):
        self.model = model

    @classmethod
    def api(cls):
        input_cols = [
            ("sepallength", float),
            ("sepalwidth", float),
            ("petallength", float),
            ("petalwidth", float),
            # The input training dataset also has this column, but it doesn't need to be
            # added to the adapter since it will not be present at actual inference time.
            # ("variety", str),
        ]
        output_cols = input_cols + [
            ("prediction", float)
        ]
        inputs = {
            "df_in": pm.Spark(columns=input_cols),
        }
        outputs = {"df_out": pm.Spark(columns=output_cols)}
        return inputs, outputs

    def predict(self, df_in):
        predictions = self.model.transform(df_in)
        # Remove intermediate vector columns from the returned DataFrame.
        return predictions.drop("features", "rawPrediction", "probability", "prob_array")

출력에서 features, rawPrediction, probability 및 prob_array 중간 벡터 열을 제거합니다. prediction은 예측한 클래스 인덱스를 담는 수치 열입니다. 시험 입력에 있던 정답 관련 열은 추가 열로 남을 수 있습니다.

온라인 추론의 제한

Spark 모델의 온라인 추론은 Modeling Objectives의 실험적 지원 경로입니다. 해당 배포에는 JDK와 Spark 배포본이 있지만 로컬 Spark만 실행하므로 모든 처리가 단일 JVM에서 이뤄집니다. 지원 Spark 버전이 변경될 수 있으며 이전 버전과의 호환을 보장하지 않으므로 현재 버전과 맞지 않으면 모델을 다시 빌드해야 합니다.

요청은 JSON으로 들어오므로 Spark DataFrame 입력의 변환 문제가 생길 수 있습니다. 전처리 또는 어댑터 변환에서 필요한 형식으로 변환합니다. 예를 들어 Pandas 표는 SparkSession.builder.getOrCreate().createDataFrame(df)로 Spark 표로 만든 뒤 모델에 전달할 수 있습니다. 온라인 모델을 새로 설계한다면 단일 노드 라이브러리를 우선 검토합니다.

Spark 환경에서 별도 라이브러리의 분산 학습을 연결하려면 PySpark XGBoost 학습을 사용할 수 있습니다.