Databricksにおけるカスタムモデルのトレーニング、SageMakerエンドポイントへのデプロイメント

Databricks MLflowのモデルサービングは、モデルのテスト目的でクイックにモデルの動作確認するという目的には適していますが、大量の同時アクセスが予期される、モデルで使用するデータをS3から取得するなど複雑な前処理を組み込みたいといった要件があるプロダクション目的でモデルをサービングする際にはSageMakerのエンドポイントなどを使用することをお勧めしています。

特に、機械学習モデルの推論処理の前処理として「`モデルを呼び出す際、IDを指定し推論に用いるデータをS3から取得する`」という要件がある場合、`PyFunc`としてモデルを実装・トレーニングし、モデルをSageMakerのエンドポイントにデプロイするアプローチを取る必要があります。

> **注意**

> - 要件が`モデルを呼び出す際に前処理を追加したい`のみであり、S3へのアクセスが不要であれば、`PyFunc`による実装のみで十分です。トレーニングしたモデルをMLflowのモデルサービングで運用することが可能です。
- 上記要件に`S3へのアクセス`が追加される場合、SageMakerのエンドポイントにモデルのデプロイを行う必要があります。S3へのアクセスを必要とするモデルをMLflowのモデルサービングでデプロイしても、モデルを呼び出す際にエラーとなります。
![Screen Shot 2022-02-04 at 11.46.15.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/1024f2f6-aa75-d88a-dd0e-ed1063e5230a.png)


本ノートブックでは、ワインの品質を予測するモデルを構築しますが、前処理を追加したカスタムモデルをSageMakerのエンドポイントにデプロイし、推論に用いるデータのIDを指定して推論を実行します。

推論に用いるデータは[databricks-datasets](https://qiita.com/taka_yayoi/items/dcf77d0b007fae774ce5)のワインデータに`lot_id`という連番の列を追加したものとします。この`lot_id`を指定して推論を行います。

![Screen Shot 2022-02-02 at 18.15.28.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/c9bb8140-2b1b-0b66-4f00-97a51780adbd.png)


# 要件
- 本書の処理を行うクラスターには、SageMakerおよびS3にアクセスできるインスタンスプロファイルをアタッチします。
- クラスターライブラリとしてPyPIから以下をインストールします。
  - s3fs
  - sagemaker
  - botocore

# AWSの設定

以下の手順に従ってSageMakerにアクセスするためのロールを作成します。

[機械学習モデルをSageMakerにデプロイするのためのAWS認証設定のセットアップ \- Qiita](https://qiita.com/taka_yayoi/items/e26e4546cdaf13069aad)

モデルがアクセスするデータが格納されているS3にアクセスするためのインラインポリシーを追加します。

```json:JSON
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "s3:ListBucket"
            ],
            "Resource": [
                "arn:aws:s3:::<S3バケット名>"
            ]
        },
        {
            "Effect": "Allow",
            "Action": [
                "s3:PutObject",
                "s3:GetObject",
                "s3:DeleteObject",
                "s3:PutObjectAcl"
            ],
            "Resource": [
                "arn:aws:s3:::<S3バケット名>/*"
            ]
        }
    ]
}
```

以下の手順に従って、上記SageMaker用ロールのインスタンスプロファイルARNを用いてDatabricksでインスタンスプロファイルを作成します。

[Databricksにおけるインスタンスプロファイルを用いたS3バケットへのセキュアなアクセス \- Qiita](https://qiita.com/taka_yayoi/items/446c7971be354f88c679#%E3%82%B9%E3%83%86%E3%83%83%E3%83%975-databricks%E3%81%AB%E3%82%A4%E3%83%B3%E3%82%B9%E3%82%BF%E3%83%B3%E3%82%B9%E3%83%97%E3%83%AD%E3%83%95%E3%82%A1%E3%82%A4%E3%83%AB%E3%82%92%E8%BF%BD%E5%8A%A0%E3%81%99%E3%82%8B)

Databricksのデプロイに使用したロールに以下のポリシーを追加します。

```json:JSON
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": [
                "iam:PassRole"
            ],
            "Resource": [
                "arn:aws:iam::<AWSアカウントID>:role/<上で作成したSageMaker用ロール>"
            ],
            "Effect": "Allow"
        }
    ]
}
```

# モデルのトレーニング & ロギング

## トレーニング、テストデータセットの準備

`databrikcs-datasets`のワインデータを用いてトレーニングデータを準備します。

```py:Python
import pandas as pd

white_wine = pd.read_csv("/dbfs/databricks-datasets/wine-quality/winequality-white.csv", sep=";")
red_wine = pd.read_csv("/dbfs/databricks-datasets/wine-quality/winequality-red.csv", sep=";")

red_wine['is_red'] = 1
white_wine['is_red'] = 0

data = pd.concat([red_wine, white_wine], axis=0)

# カラム名から空白を削除
data.rename(columns=lambda x: x.replace(' ', '_'), inplace=True)

# qualityを2値に変換します
high_quality = (data.quality >= 7).astype(int)
data.quality = high_quality
display(data)
```

```py:Python
# トレーニング、テストデータセットに分割します
from sklearn.model_selection import train_test_split

train, test = train_test_split(data, random_state=123)
X_train = train.drop(["quality"], axis=1)
X_test = test.drop(["quality"], axis=1)
y_train = train.quality
y_test = test.quality
```

## PyFuncカスタムモデルの定義

```py:Python
import mlflow
import mlflow.pyfunc
import mlflow.sklearn
import numpy as np
import sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score
from mlflow.models.signature import infer_signature
from mlflow.utils.environment import _mlflow_conda_env
import cloudpickle
import time

# カスタムモデルに必要なライブラリ
import boto3
import fsspec
import s3fs

# PyFuncのカスタムモデルとして定義
class SklearnModelWrapper(mlflow.pyfunc.PythonModel):
  """
  sklearnのモデルをトレーニング、使用するクラス
  """
  
  def __init__(self, model):
    self.model = model
    
  def predict(self, context, model_input):
    """これは抽象化関数です。sklearnモデルを取り出すためにメソッドをカスタマイズします。
        Args:
            context ([type]): モデルのアーティファクトが格納されるMLflowコンテキスト
            model_input ([type]): データを検索するキー
        Returns:
            [type]: ロードされたモデルアーティファクト
    """
    
    import csv
    import pandas as pd
    import boto3
    import s3fs
    
    # S3バケット名
    bucket_name = '<前処理で読み込むデータを格納しているS3バケット>'
    # ファイルパス
    savepath = 'wine_quality_w_id.csv'
    
    # S3からのデータ取得
    # 本処理を実行するクラスターにS3にアクセス可能なインスタンスプロファイルをアタッチすること    
    s3_client = boto3.client('s3')
    response = s3_client.get_object(Bucket=bucket_name, Key=savepath)
    file = response["Body"]
    df = pd.read_csv(file, header=0, delimiter=",", low_memory=False)
    #print(df)
    
    # SageMakerエンドポイント経由の場合、引数はDataFrameとして渡されるのでDataFrameから値を取り出す     
    if isinstance(model_input, pd.DataFrame):
      row_id = int(model_input['data'].iloc[0])
    else:
      row_id = int(model_input)
    
    print("model_input:", model_input)
    
    target_data = df.drop("quality", axis=1)    
    target_data = target_data[target_data['lot_id'] == row_id].drop("lot_id", axis=1)
    print("target_data:", target_data)
    
    return self.model.predict_proba(target_data)[:,1]    
```

## モデルのトレーニング

シグネチャの実装はTODOです。

```py:Python
# mlflow.start_runは、このモデルのパフォーマンスを追跡するための新規MLflowランを生成します。
# コンテキスト内で、使用されたパラメーターを追跡するためにmlflow.log_param、精度のようなメトリクスを追跡するために
# mlflow.log_metricを呼び出します。
with mlflow.start_run(run_name='untuned_random_forest') as run:
  n_estimators = 10
  model = RandomForestClassifier(n_estimators=n_estimators, random_state=np.random.RandomState(123))
  model.fit(X_train, y_train)

  # predict_probaは[prob_negative, prob_positive]を返却するので、出力を[:, 1]でスライスします。
  predictions_test = model.predict_proba(X_test)[:,1]
  auc_score = roc_auc_score(y_test, predictions_test)
  mlflow.log_param('n_estimators', n_estimators)
  # メトリックとしてROC曲線のAUCを使用します。
  mlflow.log_metric('auc', auc_score)
  wrappedModel = SklearnModelWrapper(model)

  # モデルの入出力スキーマを定義するシグネチャをモデルとともに記録します。
  # モデルがデプロイされた際に、入力を検証するためにシグネチャが用いられます。
  # TODO:シグネチャの定義
  #signature = infer_signature(np.array([0]), wrappedModel.predict(None, np.array([0])))
  #signature = infer_signature(0, wrappedModel.predict(None, 0))
  #print(signature)
  
  # MLflowにはモデルをサービングする際に用いられるconda環境を作成するユーティリティが含まれています。
  # 必要な依存関係がconda.yamlに保存され、モデルとともに記録されます。
  conda_env =  _mlflow_conda_env(
        additional_conda_deps=None,
        additional_pip_deps=["cloudpickle=={}".format(cloudpickle.__version__), 
                             "scikit-learn=={}".format(sklearn.__version__),
                             "boto3==1.20.1",
                             "fsspec=={}".format(fsspec.__version__),
                             "s3fs=={}".format(s3fs.__version__),
                             "botocore==1.23.24"
                            ],
        additional_conda_channels=None,
    )
  #mlflow.pyfunc.log_model("random_forest_model", python_model=wrappedModel, conda_env=conda_env, signature=signature)
  mlflow.pyfunc.log_model("random_forest_model", python_model=wrappedModel, conda_env=conda_env)
```

## モデルの動作確認

```py:Python
# ラン(トレーニング)のIDを取得
previous_run_id = run.info.run_id

# トラッキングされたモデルのロード
loaded_model = mlflow.pyfunc.load_model(f"runs:/{previous_run_id}/random_forest_model")
```

```py:Python
# IDを指定して予測を実行
out_predict = loaded_model.predict("8")
print(out_predict)
```

![Screen Shot 2022-02-02 at 18.04.29.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/9f770533-ce0d-7a9c-8537-f2f49d5a7904.png)

# カスタムモデルのサービング

MLflowで記録されたモデルをSageMakerのエンドポイントにデプロイします。

## モデルのデプロイ

SageMakerエンドポイントにモデルをデプロイする際には、AmazonのElastic Container Registry (ECR)のDockerイメージを指定します。SageMakerはモデルをサービングする際にこのイメージを使用します。

ローカルマシンでMLflow CLI:`mlflow sagemaker build-and-push-container`を実行し、`mlflow-pyfunc`イメージをビルドし、イメージをECRのリポジトリにアップロードします。ECRリポジトリに`mlflow-pyfunc`イメージが作成されます。アップロード完了後、DockerイメージのURLを取得することができます。

上記CLIコマンドを実行する前に以下の設定を行なっていることを確認してください。
- ローカルマシンでAWS CLIのインストール、設定を行います。
- ローカルマシンでPyPIを使って`mlflow`、`boto3`をインストールしておきます。

![](https://sajpstorage.blob.core.windows.net/demo20220129-pytorch-mlflow/ECR.png)

```py:Python
region = "ap-northeast-1"
run_id1 = previous_run_id
model_uri = "runs:/" + run_id1 + "/random_forest_model"
```

```py:Python
# 以下のECR Dockerイメージに対するURLを<ECR-URL>で置き換えてください。
# ECR URLは以下のフォーマットである必要があります: {account_id}.dkr.ecr.{region}.amazonaws.com/{repo_name}:{tag}
image_ecr_url = "<AWSアカウントID>.dkr.ecr.ap-northeast-1.amazonaws.com/mlflow-pyfunc:1.23.1"
```

```py:Python
import mlflow.sagemaker as mfs
app_name = "wine-quality" # SageMakerエンドポイント名
```

以下のAPIを用いてSageMakerエンドポイントにモデルをデプロイします。

```py:Python
# 4-5分かかります
mfs.deploy(app_name=app_name, model_uri=model_uri, image_url=image_ecr_url, region_name=region, mode="replace") # mode create/replace
```

![Screen Shot 2022-02-02 at 18.06.44.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/0660342d-b64f-57b7-4550-9bf311fdd01f.png)

以下のセルを実行することで新規SageMakerエンドポイントのステータスを確認することができます。

**注意**: アプリケーションのステータスは**Creating**であるべきです。ステータスが**InService**になるまで待ってください。それまではクエリーのリクエストは失敗します。

![](https://sajpstorage.blob.core.windows.net/demo20220129-pytorch-mlflow/sagemaker_endpoint_inservice.png)

```py:Python
# エンドポイントのステータスを確認するヘルパー関数
def check_status(app_name):
  sage_client = boto3.client('sagemaker', region_name=region)
  endpoint_description = sage_client.describe_endpoint(EndpointName=app_name)
  endpoint_status = endpoint_description["EndpointStatus"]
  return endpoint_status

print("Application status is: {}".format(check_status(app_name)))
```

![Screen Shot 2022-02-02 at 18.08.44.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/6f9ea313-4b1a-17e1-f867-ed7e27ade0d1.png)

## エンドポイントからモデルの呼び出し

HTTPリクエストを送信してサンプルインプットを評価します。

`boto3`で提供される`sagemaker-runtime` APIを用いてSageMakerエンドポイントREST APIにクエリーを実行します。

```py:Python
import json
import numpy as np

# パラメータ指定
data = [{"data": "8"}] # 推論に用いるデータのIDを指定します
input_json = json.dumps(data)
print(input_json)

def query_endpoint(app_name, input_json):
  client = boto3.session.Session().client("sagemaker-runtime", region)
  
  response = client.invoke_endpoint(
      EndpointName=app_name,
      Body=input_json,
      ContentType='application/json',
  )
  
  preds = response['Body'].read().decode("ascii")
  preds = json.loads(preds)
  print("Received response: {}".format(preds))
  return preds

print("Sending batch prediction request with input dataframe json: {}".format(input_json))

# デプロイされたモデルにポストすることで入力を評価します
prediction1 = query_endpoint(app_name=app_name, input_json=input_json)
```

![Screen Shot 2022-02-02 at 18.09.36.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/7b218232-9a36-68db-0a41-a9bba4a8d637.png)

# サンプルノートブック

ノートブックで参照しているCSVファイルも格納されています。

https://github.com/taka-yayoi/public_repo/tree/main/custom_model_sagemaker

# 参考資料

- [機械学習モデルをSageMakerにデプロイするのためのAWS認証設定のセットアップ \- Qiita](https://qiita.com/taka_yayoi/items/e26e4546cdaf13069aad)
- [SageMakerへのscikit\-learnモデルのデプロイメント \- Qiita](https://qiita.com/taka_yayoi/items/79f024874075ea9e1cd3)
- [Databricksにおけるインスタンスプロファイルを用いたS3バケットへのセキュアなアクセス \- Qiita](https://qiita.com/taka_yayoi/items/446c7971be354f88c679)
- [Preprocess input data before making predictions using Amazon SageMaker inference pipelines and Scikit\-learn \| AWS Machine Learning Blog](https://aws.amazon.com/blogs/machine-learning/preprocess-input-data-before-making-predictions-using-amazon-sagemaker-inference-pipelines-and-scikit-learn/)

### Databricks 無料トライアル

[Databricks 無料トライアル](https://databricks.com/jp/try-databricks)