機械学習プロジェクトのフォーマットを用意してSageMakerで実行する

こんにちは、株式会社LIFULLの二宮です。

機械学習のプロジェクトでは、分析やモデルの精度評価でうまくいった後、それをうまく既存のシステムで利用できるようにしなければいけません。その際に、自分たちのチームでは実装担当のエンジニアの役割分担に苦労していました。

「データサイエンティストがこのフォーマットで作ってくれれば簡単に組み込めるよ！」という状態を目指して、[Amazon SageMaker](https://aws.amazon.com/jp/sagemaker/)をラップして、ある程度汎用的に使える開発フォーマットとツールを用意しました。


# Amazon SageMakerとは

> Amazon SageMaker は、すべての開発者とデータサイエンティストに機械学習モデルの構築、トレーニング、デプロイ手段を提供します。Amazon SageMaker は、機械学習のワークフロー全体をカバーする完全マネージド型サービスです。データをラベル付けして準備し、アルゴリズムを選択して、モデルのトレーニングを行い、デプロイのための調整と最適化を行い、予測を行い、実行します。モデルをより少ない労力と費用で、本番稼働させることができます。

主な機能としては、特定の仕様に合わせたDockerイメージを用意すれば、以下のような機能を利用できます。

* [S3にあるファイルを読み込んで、機械学習モデルのトレーニングを行う](https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/how-it-works-training.html)
* [推論APIのエンドポイントをデプロイする](https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/how-it-works-hosting.html)
* [バッチ変換（推論のバッチ処理）を行う](https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/how-it-works-batch.html)

SageMakerで自作のDockerイメージを用意する場合の仕様については、公式ドキュメントや @taniyam さん（私と同じチームです）の記事を読んでください。

* [Amazon SageMaker で独自のアルゴリズムやモデルを使用する](https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/your-algorithms.html)
* [SageMakerの汎用的なdockerイメージ](https://qiita.com/taniyam/items/cf398f2462b4f2ebc792)

# 機械学習プロジェクトのフォーマット

まず、データサイエンティストには次のようなディレクトリ構成で用意してもらうようにしました。

```bash
.
├── README.md
├── Dockerfile
├── config.yml
├── pyproject.toml (poetryの設定ファイル)
├── script
│   └── __init__.py
└── tests
    └── __init__.py
```

`script/__init__.py` にメインの処理が書かれていて、以下のようなスクリプトになっています。 `simple_sagemaker_manager` が今回用意したライブラリです。

```python
import pandas as pd
from typing import List
from pathlib import Path
from sklearn import tree
from simple_sagemaker_manager.image_utils import AbstractModel


def train(training_path: Path) -> AbstractModel:
    """学習を行う。

    Args:
        training_path (Path): csvファイルのあるディレクトリ
    
    Returns:
        Model: AbstractModelを継承したモデルオブジェクト
        
    """
    train_data = pd.concat([pd.read_csv(fname, header=None) for fname in training_path.iterdir()])
    train_y = train_data.iloc[:, 0]
    train_X = train_data.iloc[:, 1:] 

    # Now use scikit-learn's decision tree classifier to train the model.
    clf = tree.DecisionTreeClassifier(max_leaf_nodes=None)
    clf = clf.fit(train_X, train_y)
    return Model(clf)


class Model(AbstractModel):
    """AbstractModelにシリアライズの方法などが記述されている。
    """

    def predict(self, matrix: List[List[float]]) -> List[List[str]]:
        """推論処理。

        Args:
            matrix (List[List[float]]): テーブルデータ

        Returns:
            list: 推論結果

        """
        # ここで返した結果が推論APIのレスポンスになります
        return [[x] for x in self.model.predict(pd.DataFrame(matrix))]
```

`AbstractModel` は次のような定義で、 `save` メソッドが呼び出された結果（pickleでシリアライズされた結果）が保存され、学習バッチの実行時は（SageMakerのシステムで利用して）これがモデルとしてS3に保存されます。また、 `save` と `load` をオーバーライドすることでシリアライズ方式を切り替えられるようにしています。

```python
import pickle
from abc import ABC, abstractmethod
from dataclasses import dataclass


@dataclass
class AbstractModel(ABC):
    model: object

    @classmethod
    def load(cls, model_path):
        # 学習バッチ時にモデルを保存する
        with open(model_path / 'model.pkl', 'rb') as f:
            model = pickle.load(f)
        return cls(model)

    def save(self, model_path):
        # 推論時にモデルを読み込む
        with open(model_path / 'model.pkl', 'wb') as f:
            pickle.dump(self.model, f)

    @abstractmethod
    def predict(self, json):
        pass
```

Pythonの[poetry](https://poetry.eustace.io/)などのプロジェクトを参考にして、cliで操作を行うようにしています。SageMakerのDockerイメージの開発フローは以下のようになります。

* プロジェクトの雛形を作成する ( `smcli new プロジェクト名` )
* 雛形を編集する
* イメージをビルドする ( `smcli build` )
* イメージをECRにpushする ( `smcli push` )

また、Dockerfileを編集できるようにしたのは、一部の機械学習ライブラリがAnacondaでしかインストールできないため、「Python3の公式イメージ以外にも挿し替えられるようにしてほしい」という要望を受けたためです。

# SageMakerの実行管理

`boto3` を直接実行するのもきついので、こちらもラップしたライブラリを用意しました。たくさんの操作がありますが、多くのプロジェクトで我々がやりたいことは「モデルを学習する」「推論APIを立てる OR バッチ変換ジョブを実行する」の3つなので、それがわかるようなインターフェイスにしています。

```python
from simple_sagemaker_manager.executor import SageMakerExecutor
from simple_sagemaker_manager.executor.classes import TrainInstance, TrainSpotInstance, Image


client = SageMakerExecutor()

# 通常インスタンスで学習する場合
model = client.execute_batch_training(
    instance=TrainInstance(
        instance_type='ml.m4.xlarge',
        instance_count=1,
        volume_size_in_gb=10,
        max_run=100
    ),
    image=Image(
        name="decision-trees-sample",
        uri="xxxxxxxxxx.dkr.ecr.ap-northeast-1.amazonaws.com/decision-trees-sample:latest"
    ),
    input_path="s3://xxxxxxxxxx/DEMO-scikit-byo-iris",
    output_path="s3://xxxxxxxxxx/output",
    role="arn:aws:iam::xxxxxxxxxx"
)


# スポットインスタンスで学習する場合
model = client.execute_batch_training(
    instance=TrainSpotInstance(
        instance_type='ml.m4.xlarge',
        instance_count=1,
        volume_size_in_gb=10,
        max_run=100,
        max_wait=1000
    ),
    image=Image(
        name="decision-trees-sample",
        uri="xxxxxxxxxx.dkr.ecr.ap-northeast-1.amazonaws.com/decision-trees-sample:latest"
    ),
    input_path="s3://xxxxxxxxxx/DEMO-scikit-byo-iris",
    output_path="s3://xxxxxxxxxxx/output",
    role="arn:aws:iam::xxxxxxxxxxxxx"
)
```

推論APIは以下のように作っています。工夫した点としては以下の通りです。

* 指定した名前(name)のエンドポイントが存在しない場合はエンドポイントを新規作成する
* 存在する場合はupdateを行う。Updatingの間もリクエストは受け付ける。
* modelsをリストで受け取れるようにした。複数のモデルを指定した場合、[Pipelineモデル](https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/inference-pipeline-create-console.html)を作成した上でデプロイを行う。

```python
from simple_sagemaker_manager.executor import SageMakerExecutor
from simple_sagemaker_manager.executor.classes import EndpointInstance, Model

client = SageMakerExecutor()


# 特定のモデルをデプロイする場合
# modelsに複数のモデルを指定すると、Pipelineモデルを作成して利用します
client.deploy_endpoint(
    instance=EndpointInstance(
        instance_type='ml.m4.xlarge',
        initial_count=1,
        initial_variant_wright=1
    ),
    models=[
        Model(
            name='decision-trees-sample-191028-111309-538454',
            model_arn='arn:aws:sagemaker:ap-northeast-1:xxxxxxxxxx',
            image_uri='xxxxxxxxxx.dkr.ecr.ap-northeast-1.amazonaws.com/decision-trees-sample:latest',
            model_data_url='s3://xxxxxxxxxx/model.tar.gz'
        )
    ],
    name='sample-endpoint',
    role="arn:aws:iam::xxxxxxxxxx"
)

# execute_batch_trainingの結果を渡すこともできます
model = client.execute_batch_training(
    # 引数は略
) 

client.deploy_endpoint(
    instance=EndpointInstance(
        instance_type='ml.m4.xlarge',
        initial_count=1,
        initial_variant_wright=1
    ),
    models=[model],
    name='sample-endpoint',
    role="arn:aws:iam::xxxxxxxxxx"
)
```

エンドポイント以外（学習バッチジョブなど）の名前は、自動で現在時刻の文字列を足して重複を避ける実装にしています。ただしエンドポイントだけは「同じ名前のものがあればアップデートする」という挙動にして利便性を高めています。

また、省略しますが、バッチ変換ジョブのメソッドも同様に実装しています。


# これからの課題

このように実装して、今は一部のプロジェクトの実装で実際に利用してもらっています。ただ未実装の課題もいくつかあり、チーム内の他の課題もまだまだあります。

* テーブルデータ以外(画像など)のコンテナイメージの雛形は未実装です
* 学習バッチのSpotInstanceで、処理の中断時に中間状態をS3に保存する機能
* ローカルで推論APIを立ち上げるなど、APIの実装やテストを支援する機能が無い
* 学習時のメトリクスの取得。これは別メンバーが[MLflow](https://qiita.com/fam_taro/items/155912068ff475a53e44)に注目して検証しています。
* データの分析やモデルの精度評価についてもまだまだ課題が多そう

また、チーム内で実際に利用してもらうと、使い勝手の悪い部分もそれなりに出てきてしまっているので、それらの問題を解決してより機械学習プロジェクトを効率化していこうと思います。