【SageMaker】TorchServeのネイティブサポートを利用したPyTorchモデルのデプロイ

以前の記事で、SageMaker Training JobsによるTabBERTモデルのFine-Tuningを行いました。
Fine-Tuning済モデルをS3にアップロードすることができたので、今回はSageMakerでモデルのデプロイをしてみました。

https://qiita.com/suzuki0430/items/738344621d79a09c8234

# TorchServe
作成した機械学習モデルを推論システムとして使うためには以下の要素が必要となります。
- 学習済みの機械学習モデル
- 学習済みモデルに入力する特徴量作成の処理(特徴量生成、トークン化)
- 外部システムから学習済みモデルを利用するためのインタフェース(I/F)

これらの機能をまとめて提供するのがモデルサービングライブラリで、TensorFlowモデルであればTensorFlow Serving、PyTorchモデルであればTorchServeといったものがあります。

TorchServeを使うと、機械学習モデルのデプロイで課題となるような部分をいろいろ解決することができます。
- 予測性能とレイテンシのバランス(パフォーマンス) ⇒ 低レイテンシになるような最適化
- 前処理や後処理用のハンドラ作成(デプロイ容易性) ⇒ 画像分類やテキスト分類などの基本的なものであればデフォルトハンドラをもつ
- 本番運用におけるモデルの管理(セキュリティ、スケーラビリティ、信頼性) ⇒ 複数モデルのサービング、モデルバージョニング(A/Bテスト)、ログ/メトリクスの監視が可能

構成としては以下のようになり、外部システム向けの推論API(Inference API)、モデル管理API(Management API)、メトリクス収集API(Metrics API)をもちます。

>![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/569054/b21da487-f8ca-cb2d-d0df-6549d4f8b31d.png)
TorchServe を使用した大規模な推論のための PyTorch モデルをデプロイする

また、TorchServeによるデプロイ方法にはいろいろあり、EC2、EKS、SageMakerを使用することができます。
現在、TorchServeはSageMakerにおけるデフォルトの推論サーバとしてネイティブにサポートされているため、以前のような、SageMakerノートブックインスタンス上でのtorch-model-archiverのインストール、モデルのアーカイブ操作からのS3アップロード、DockerイメージのビルドとECRプッシュといった手間がなくなり、デプロイまでの操作がとても楽になっています。

※以下の記事に、TorchServeのネイティブサポート前の方法によるモデルデプロイの手順が書いてあるのですが、アーカイブ作成やDockerコンテナ作成の部分が手間に感じます。

https://aws.amazon.com/jp/blogs/news/deploying-pytorch-models-for-inference-at-scale-using-torchserve/

SageMakerでデプロイした場合の構成は以下のようになります。

>![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/569054/4a25ad9a-dad1-33d1-8f4d-4fbb83d1a834.png)
Serving PyTorch models in production with the Amazon SageMaker native TorchServe integration

# 実装
以下のドキュメントやコードを元に、推論コードの実装からデプロイまで行います。

https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#create-a-pytorchmodel-object

https://aws.amazon.com/jp/blogs/machine-learning/serving-pytorch-models-in-production-with-the-amazon-sagemaker-native-torchserve-integration/

https://github.com/aws/amazon-sagemaker-examples/blob/main/frameworks/pytorch/get_started_mnist_deploy.ipynb

https://github.com/data-science-on-aws/data-science-on-aws/tree/main/09_deploy/torchserve/pytorch

## 必要ファイルとディレクトリ構造
ノートブックインスタンスのルート直下に、ファイル操作やデプロイを行うためのノートブック`deploy_tabformer.ipynb`と`code`ディレクトリを配置します。
`code`ディレクトリには、推論用のコード`inference.py`を配置します。
その他に、`inference.py`でインポートする必要があるファイルやサードパーティライブラリのバージョン指定用のファイル`requirements.txt`も配置します。
肝心のFine-Tuning済モデル`fine_tuning_model.pt`については、ノートブックからではなくS3から`model.tar.gz`として読み込みます。

```
.
├─ code
│   ├─ inference.py (推論)
│   ├─ common.py (モデル)
│   ├─ preprocessing.py (特徴量生成)
│   ├─ config.json (設定ファイル)
│   ├─ pytorch_model.bin (事前学習済モデル)
│   ├─ summary.encoder_fit.pkl (エンコーダー)
│   ├─ vocab_token2id.bin (辞書ファイル)
│   └─ requirements.txt (ライブラリバージョン指定)
└─ deploy_tabformer.ipynb (デプロイ用コード)
```

今回、`code`内のファイル群については、`deploy_tabformer.ipynb`上でファイル操作をして配置します(`common.py` `inference.py` `preprocessing.py` `requirements.txt`は手動で配置)。

`pre_trained_model.tar.gz`にすべてのファイルが含まれるのですが、不要なファイルもあるため、一旦`model`ディレクトリに展開して、必要なファイルのみ`code`にコピーします。

ノートブックのストレージ確保のため、不要となった`pre_trained_model.tar.gz`や`model`は削除します。

```py:deploy_tabformer.ipynb
import boto3
import tarfile
import os
import shutil

# 必要ファイルのダウンロード
s3 = boto3.resource('s3')
bucket = s3.Bucket('tabformer-opt')
bucket.download_file('<path>/model.tar.gz', 'pre_trained_model.tar.gz')

# ファイル展開
with tarfile.open(name='pre_trained_model.tar.gz', mode="r:gz") as mytar:
    mytar.extractall('model')

# tar.gzの削除
os.remove('./pre_trained_model.tar.gz')

# 必要ファイルのコピー
shutil.copy('./model/vocab_token2id.bin', './code')
shutil.copy('./model/summary.3.2022-10-01_2022-11-30.encoder_fit.pkl', './code')
shutil.copy('./model/checkpoint-500/config.json', './code')
shutil.copy('./model/checkpoint-500/pytorch_model.bin', './code')

# modelディレクトリの削除
shutil.rmtree('./model')
```

## inference.pyの作成
`inference.py`では、以下の関数を記述します。

- `model_fn`：モデルの読み込み
- `predict_fn`：推論処理
- `input_fn`：リクエストデータのデコード(`predict_fn`へのデータ受け渡し)
- `output_fn`：レスポンスデータの生成(`predict_fn`からデータ受け取り)

```py:inference.py
import json
import sys
import logging
import pickle
import torch
from torch import nn
from os import path
from preprocessing import ActionHistoryPreprocessing
from common import CommonModel

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
logger.addHandler(logging.StreamHandler(sys.stdout))

MODEL_NAME = 'fine_tuning_model.pt'
JSON_CONTENT_TYPE = 'application/json'

def model_fn(model_dir):
    model_path = '{}/{}'.format(model_dir, MODEL_NAME)
    
    try:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')    

        pretrained_model = path.join("/opt/ml/model/code/", f"pytorch_model.bin")
        pretrained_config = path.join("/opt/ml/model/code/", f"config.json")
        model = CommonModel(pretrained_config, pretrained_model)
        model.to(device)
        model.load_state_dict(torch.load(model_path))
        model.eval()

        return model
    except Exception as e:
        logger.exception(f"Exception in model fn {e}")
        return None

    
def predict_fn(input_data, model):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    input_records = input_data['records']

    token2id_file = path.join("/opt/ml/model/code/", f"vocab_token2id.bin")
    encoder_fname = path.join("/opt/ml/model/code/", f"summary.encoder_fit.pkl")
    encoder_fit = pickle.load(open(encoder_fname, "rb"))

    preprocessed_data = ActionHistoryPreprocessing(
        input_data=input_records,
        token2id_file=token2id_file,
        encoder_fit=encoder_fit)

    model_input = torch.tensor([preprocessed_data.getitem()], dtype=torch.long)
    model_output = model(model_input.to(device))
    
    pred = torch.argmax(model_output, 1)
    pred_list = pred.tolist()

    result = {
        "reaction": pred_list[0]
    }

    return result


def input_fn(serialized_input_data, content_type=JSON_CONTENT_TYPE):  
    if content_type == JSON_CONTENT_TYPE:
        data = json.loads(serialized_input_data)
        return data
    else:
        pass


def output_fn(prediction_output, accept=JSON_CONTENT_TYPE):
    if accept == JSON_CONTENT_TYPE:
        return json.dumps(prediction_output), accept
    raise Exception('Requested unsupported ContentType in Accept: ' + accept)
```

`model_fn`では、モデルや`config.json`などのファイルをコンテナのディレクトリから読み込みます。
例えば、モデル`fine_tuning_model.pt`については、S3から`model_dir`内にコピーされ、ノートブックの`code`ディレクトリに配置したコードは、デプロイ時にコンテナの`/opt/ml/model/code/`にコピーされます。
なお、特徴量を生成する`ActionHistoryPreprocessing`や`CommonModel`などのモジュールについては、以下のようにインポートします。
```py
from preprocessing import ActionHistoryPreprocessing
from common import CommonModel
```

## モデルデプロイ
`PyTorchModel`オブジェクトを作成して`model.deploy()`を行うと、SageMakerのTorchServe上でコンテナが起動します。

まず、SageMakerのセッションやロールを準備し、`PyTorchModel`オブジェクトを作成します。
- `model_data`：`model.tar.gz`のS3パス
- `name`：SageMakerに登録するモデル名
- `role`：SageMakerを利用するためのロール
- `source_dir`：コンテナの`/opt/ml/model/code/`にコピーするノートブック内のディレクトリ名
- `entry_point`：`source_dir`内でエントリーポイントとなるファイル
- `framework_version`：コンテナのPyTorchバージョン
- `py_version`：コンテナのPythonバージョン

`framework_version`と`py_version`が指定されている場合、デプロイ時に適当なDockerイメージをSageMaker側で選択して実行してくれます。
例えば、framework_version=1.9.0, py_version=38を指定すると、`763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-inference:1.9.0-gpu-py38`が選択されます。
カスタムしたDockerイメージを使いたい場合は、`image_uri`でDockerイメージのURLを指定します。
また、`framework_version`については、学習時の環境と揃えないと、モデル読み込み時にエラーが発生する場合があります。

```py:deploy_tabformer.ipynb
import boto3
import sagemaker
import pandas as pd

from sagemaker import get_execution_role
from sagemaker.utils import name_from_base
from sagemaker.pytorch.model import PyTorchModel

from sagemaker.predictor import Predictor
from sagemaker.serializers import JSONSerializer
from sagemaker.deserializers import JSONDeserializer

sess = sagemaker.Session()
bucket = sess.default_bucket()
role = sagemaker.get_execution_role()
region = boto3.Session().region_name

sm = boto3.Session().client(service_name='sagemaker', region_name=region)

class ReactionPrediction(Predictor):
    def __init__(self, endpoint_name, sagemaker_session):
        super().__init__(endpoint_name, 
                         sagemaker_session=sagemaker_session, 
                         serializer=JSONSerializer(), 
                         deserializer=JSONDeserializer()
        )

model = PyTorchModel(model_data='s3://tabformer-opt/<path>/model.tar.gz',
                     name='tabformer-opt',
                     role=role,
                     entry_point='inference.py',
                     source_dir='code',
                     framework_version='1.9.0',
                     py_version='py38',
                     predictor_cls=ReactionPrediction)
```
`model.deploy()`でデプロイを行います。
`instance_type`でインスタンスタイプを指定したり、`endpoint_name`でエンドポイント名を指定することができます。

https://aws.amazon.com/jp/sagemaker/pricing/

```py:deploy_tabformer.ipynb
predictor = model.deploy(initial_instance_count=1, 
                         instance_type='ml.g4dn.2xlarge',
                         endpoint_name='tabformer-opt', 
                         wait=False)
```

デプロイに成功すると、SageMakerの`推論 > モデル`でモデル設定やコンテナの詳細、`推論 > エンドポイント`でエンドポイント設定などを確認することができます(ステータスも`InService`となります)。
![スクリーンショット 2023-01-15 18.17.13.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/569054/b9e28d6e-0e07-df1b-5b9c-60f038e6f7bd.png)

![スクリーンショット 2023-01-15 18.16.39.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/569054/5a86e65f-fe3d-046f-7dee-aa50abf47056.png)

デプロイモデルが正常にレスポンスを返すかどうかを確認するために、リクエストデータ(2連続の時系列データ)を入力してみます。

```py
test_data = {"records": [['2022-10-01 04:06:01', 1844060, 1643750, None, 'transition', None, None, None, 76, 'HP', 3, 'HP', 'smartphone', None, None, None, None, None, None, 'https://aaa.com/', 14.0], ['2022-10-01 08:05:31', 1845599, 1645078, None, 'transition', None, None, None, 76, 'HP', 3, 'HP', 'smartphone', None, None, None, None, None, None, 'https://bbb.com/', 19.0]]}
prediction = predictor.predict(test_data)
```
正常にレスポンスが返りました。

```
{'reaction': 0}
```

## 遭遇したエラーとその解決方法
エラーにハマりすぎて年末年始が消えたので、メモとして残します。

### モデルの保存形式による読み込みエラー

```
ModelError: An error occurred (ModelError) when calling the InvokeEndpoint operation: Received server error (500) from primary with message "'NoneType' object is not callable
```

```
RuntimeError: version_ <= kMaxSupportedFileFormatVersionINTERNAL ASSERT FAILED at "../caffe2/serialize/inline_container.cc":139, please report a bug to PyTorch. Attempted to read a PyTorch file with version 10, but the maximum supported version for reading is 6. Your PyTorch installation may be too old.
```

学習時と推論時の`torch`バージョンの差分によってモデルをうまくロードできないのかと思い、`torch`と`torchvision`のバージョンを両環境で揃えたのですが、解決できませんでした。
そこで、Fine-Tuning後の保存方法を`torch.jit.save()`(TorchScript)から`torch.save()`(モデル全体)に変えたところ別のエラーが発生しました。

```
/opt/conda/lib/python3.8/site-packages/sklearn/base.py:310: UserWarning: Trying to unpickle estimator LabelEncoder from version 1.1.3 when using version 0.24.2. This might lead to breaking code or invalid results. Use at your own risk.
```
`torch.save(model.state_dict(), ~)`(モデルパラメータ)に変えたところ、解決しました。

https://stackoverflow.com/questions/62592915/unable-to-load-model-trained-on-pytotch-version-1-6-0-in-pytorch-version-1-5-1

https://github.com/pytorch/pytorch/issues/18325

https://take-tech-engineer.com/pytorch-model-save-load/

### サードパーティライブラリの未インポートによるエラー
```
ModuleNotFoundError: No module named 'transformers'
```
`common.py`でインポートしている`transformers`がコンテナにインストールされていないことでエラーが発生しました。
`code`に`requirements.txt`を作成し、ライブラリ名とバージョンを指定することで、デプロイ時にインストールされるようにしました。

```txt:requirements.txt
transformers==3.2.0
```

# 参考資料

https://takaherox.hatenablog.com/entry/2021/01/02/144216

https://note.panicyusuke.com/TorchServe

https://aws.amazon.com/jp/blogs/news/sagemaker-custom-containers-pattern-inference/

https://qiita.com/mshr-ito/items/f456f60fc5c557a19cf5