palantir_models における Spark ML モデルのサポート

Apache Spark™ は、Foundry の計算処理を支える主要なエンジンの1つであり、幅広い機械学習機能 ↗を提供しています。Foundry は Spark MLlib をサポートしていますが、Spark はもともと分散型の機械学習フレームワークであるため、その特性に起因する注意点があります。

たとえば、Foundry の一部の機能では、モデルをコンテナイメージにカプセル化し、Rubix ↗ などの基盤インフラストラクチャを利用して、モデルを実行する複数のコンテナをプロビジョニングおよび管理します。Foundry は、Spark モデル用のこのようなコンテナイメージのビルドをまだサポートしていません。

そのため、scikit-learn、xgboost、keras など、単一ノードでのトレーニングと推論をネイティブにサポートするフレームワークを優先することを推奨します。特に、モデルを REST API として利用する場合は、Spark ML の使用を推奨しません。

Spark モデルのトレーニングとバッチ推論

Foundry では、Code Workspaces は単一ノードで実行され、Code Workbooks はレガシーとみなされているため、コードリポジトリで Spark モデルをトレーニングすることを推奨します。ただし、Foundry での Spark ML モデルの開発は、ほかのフレームワークでの開発と基本的に変わりません。

• Model API は、Spark DataFrame を入力としてサポートしています。
• MLlib パイプラインモデル ↗ は、palantir_models のバージョン 0.1599.0 で導入された SparkMLAutoSerializer クラスを使用して自動シリアライズできます。

以下は、Spark ML を使用してオープンソースの Iris データセット ↗で多クラス分類モデルをトレーニングするコードの例です。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
# 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()
    # 列は sepallength, sepalwidth, petallength, petalwidth, variety であると仮定します
    feature_cols = iris_data.columns[:-1]
    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)

    # 学習済みモデルを ModelAdapter でラップします
    foundry_model = SparkModelAdapter(model)

    predictions = foundry_model.transform(test_data).df_out
    # 学習済みモデルを実験とともに Foundry に公開して書き込みます
    # モデルと実験が公開されると、モデルページにすぐに表示されます
    inference_data_out.write_dataframe(
        predictions
    )
    model_output.publish(
        model_adapter=foundry_model,
    )

対応するアダプターコードは以下のとおりです。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
# src/main/model_adapters/adapter.py
import palantir_models as pm
from palantir_models.serializers import SparkMLAutoSerializer
from pyspark.ml.functions import vector_to_array
from pyspark.sql import functions as F


class SparkModelAdapter(pm.ModelAdapter):
    NUM_CLASSES = 3

    @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),
            # 入力の学習用データセットにもこの列がありますが、
            # 実際の推論時には存在しないため、アダプターに追加する必要はありません。
            # ("variety", str),
        ]
        output_cols = input_cols + [
            ("label", int)
        ]
        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)
        # 簡略化のため、確率ベクトルとラベルベクトルは返さず、
        # API で指定されているとおりにラベルのみを返します。
        return predictions.drop("features", "rawPrediction", "probability", "prob_array")

ライブ推論

Modeling Objectives では試験的機能として技術的にサポートされていますが、ライブ推論に Spark モデルを使用することは、一般に推奨しません。ライブ推論用のモデルを設計する際は、ほかの単一ノード用ライブラリで構築したモデルを優先することを推奨します。