SageMakerとServerlessを組み合わせて、お手軽にscikit-learnの機械学習APIを作る

SageMakerとServerlessを使ってscikit-learnの機械学習APIを作る方法を紹介します。

公式ドキュメントやその他の記事の多くはコンソール操作やnotebook上での操作が多く含んでいて、そのコードのまま本番運用に使うのは難しいと感じたので、この記事では コンソール操作やnotebook上での操作なしでスクリプトだけで完結 できるようにしています。カスタマイズすれば本番運用で使えるはずです。

また公式ドキュメントにもExampleがいくつかあるのですが、色々な処理を含んでいて、自分には理解し辛い部分がありました。今回、SageMakerを理解するためにもっとシンプルなToy Exampleを作ってみました。

作るもの

テストデータとしてよく使われる iris データをRandomForestで予測するAPIを作成します。 最終的な結果として下記のように特徴量を投げると判別結果を返すようなものになります。

curl -X POST -H "Content-Type: application/json" -d "{\"data\": [5.9, 3.0, 5.1, 1.8]}" https://xxxxxxxx.execute-api.ap-northeast-1.amazonaws.com/dev/invocations
# => {"y": 2}

環境

• Python: 3.7
• Serverless Framework: 1.41.1
• SageMaker SDK: 1.19.0
• boto3:1.9.140

アーキテクチャ構成

全体のアーキテクチャは以下のようになります。SageMakerの呼び出し元が「学習実行/予測エンドポイントデプロイ時」と「予測時」で異なるので注意してください。別々にしている理由は、「学習実行/予測エンドポイントデプロイ」の実行時間が、学習させるデータやモデルによって異なり、Lambdaの現在の実行上限である10分を超える可能性があるためです。

フォルダ構成

最終的なプロジェクト構成は下記のようになります。

 tree -L 2
# ├── layer_requirements
# │   ├── Pipfile
# │   ├── Pipfile.lock
# │   └── serverless.yml
# ├── script
# │   ├── src/iris.py
# │   ├── train.py
# └── serverless
#     ├── predict.py
#     └── serverless.yml

事前準備

それでは事前準備からしていきましょう。

Layer作成

今回はLambda上でSageMaker SDKを使うのですが、依存関係にnumpyなど容量が大きいものがあり、毎回デプロイするのは大変なので、Layer化しておきます。 Serverless上でLambda Layerを使う方法については以前まとめたので、下記の記事を参照ください。

blog.ikedaosushi.com

以下のようにyamlを書きます。

layer_requirements/serverless.yml

service: serverless-sagemaker-layer

plugins:
  - serverless-python-requirements
custom:
  region: ap-northeast-1
  pythonRequirements:
    usePipenv: true
    dockerizePip: true
    slim: true
    layer: true
    noDeploy: [pytest, jmespath, docutils, pip, python-dateutil, setuptools, s3transfer, six]

provider:
  name: aws
  runtime: python3.7
  region: ${self:custom.region}

resources:
  Outputs:
    ServerlessSageMakerLayerExport:
      Value:
        Ref: PythonRequirementsLambdaLayer

あとはデプロイするだけです。

cd layer_requirements
sls plugin install -n serverless-python-requirements
pipenv install sagemaker
sls deploy

コンソールからデプロイできていることを確認できます。

以下は補足なので気になった方だけ読んでください。

補足ポイント1: なぜわざわざSageMaker SDKをLambda上で使うか

「ここまでするのならLambda上ではboto3を使えばいいのでは」と感じた方もいるかと思います。理由があって、SageMaker上で作成したscikit-learnモデルの推論ではnumpyを使ったencodingが必要で(自分が調べた限り)、結局Lambda上でnumpyが必要になりencodingのコードも煩わしく、SageMaker SDKを使うとシンプルに書けるので、こちらを採用しました。

補足ポイント2: noDeploy 追加(=>boto3最新版をデプロイ)

(ちょっとしたハマりどころなのですが、)Lambdaでデフォルトでインストールされているboto3はSageMaker SDKの依存関係と互換性がないので最新版をデプロイする必要があります。

今回使っている serverless-python-requirements プラグインではboto3をデフォルトで除外するようになっているので、 noDeploy 設定を追加してboto3がデプロイされるようにしています。このプラグインを使っていなかったとしても、とにかくboto3の最新版を使うようにすればOKです。

Resource作成/環境変数定義

アプリケーション部分を書く前にS3 BucketやIAM RoleなどのResourceの作成や環境変数の定義などを行っておきます。

serverless/serverless.yml

service: sagemaker-serverless-example

custom:
  default_stage: dev
  region: ap-northeast-1
  stage: ${opt:stage, self:custom.default_stage}
  s3:
    bucket: ${self:service}-${self:custom.stage}
    train_base_key: train
    train_file_name: train.csv
    model_base_key: artifacts
  sagemaker:
    resource: arn:aws:sagemaker:${self:custom.region}:*:*
    endpoint_name: ${self:service}-endpoint-${self:custom.stage}
  logs:
    resource: arn:aws:logs:${self:custom.region}:*:log-group:/aws/sagemaker/*
  layer:
    service: serverless-sagemaker-layer
    export: ServerlessSagemakerLayerExport
    layer: ${cf:${self:custom.layer.service}-${self:custom.stage}.${self:custom.layer.export}}

provider:
  name: aws
  runtime: python3.7
  stage: ${self:custom.stage}
  region: ${self:custom.region}
  environment:
    S3_BUCKET: ${self:custom.s3.bucket}
    SM_ROLE_ARN: !GetAtt SageMakerServiceRole.Arn
    SM_ENDPOINT_NAME: ${self:custom.sagemaker.endpoint_name}
  iamRoleStatements:
    - Effect: Allow
      Action: sagemaker:*
      Resource: ${self:custom.sagemaker.resource}

resources:
  Resources:
    S3Bucket:
      Type: AWS::S3::Bucket
      Properties:
        BucketName: ${self:custom.s3.bucket}
        AccessControl: Private
    SageMakerServiceRole:
      Type: AWS::IAM::Role
      Properties:
        AssumeRolePolicyDocument:
          Version: '2012-10-17'
          Statement:
          - Effect: Allow
            Principal:
              Service:
              - sagemaker.amazonaws.com
            Action:
            - sts:AssumeRole
        ManagedPolicyArns:
          - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
  # 学習/デプロイスクリプトから参照できるようにOutputsを書いておきます。
  Outputs:
    S3Bucket:
      Value: ${self:custom.s3.bucket}
    S3TrainBaseKey:
      Value: ${self:custom.s3.train_base_key}
    S3TrainFileName:
      Value: ${self:custom.s3.train_file_name}
    S3ModelBaseKey:
      Value: ${self:custom.s3.model_base_key}
    SmRoleArn:
      Value: !GetAtt SageMakerServiceRole.Arn
    SmEndpointName:
      Value: ${self:custom.sagemaker.endpoint_name}

あとはコマンドでデプロイするだけです。

cd serverless
sls deploy

コンソールからそれぞれのResourceが作成されていることが確認できます。

S3 Bucket

IAM Role

学習用のデータをS3にアップロード

s3://${S3Bucket}/${S3TrainBaseKey}/${S3TrainFileName} に学習したいデータをアップロードしてください。 ※ S3Bucket 、 S3TrainBaseKey 、 S3TrainFileName はserverles.ymlで設定した値で、そのまま実行していれば s3://sagemaker-serverless-example-dev/train/train.csv になります。

アップロード方法がよくわからない方はアップロードスクリプトを書いたのでこちらを使ってください。

=> serverless-sagemaker-example/upload_iris.py at master · ikedaosushi/serverless-sagemaker-example · GitHub

これで準備完了です。

モデル学習&デプロイ

さて、それではアプリケーション側を実装していきます。

学習スクリプトの準備

SageMakerの学習用インスタンスで実行されるスクリプトを書きます。

script/src/iris.py

# -*- coding: utf-8 -*-

import os
from pathlib import Path

import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.externals import joblib

TRAIN_FILE_NAME = "train.csv"
MODEL_FILE_NAME = "model.joblib"

def train(train_dir: Path, model_dir: Path):
    # S3からダウンロードされたファイルを読み込み
    train_file = train_dir/TRAIN_FILE_NAME # /opt/ml/input/data/train/train.csv
    df = pd.read_csv(train_file, engine='python')

    # 説明変数と目的変数に分ける
    X = df.drop('y', axis=1)
    y = df['y']

    # 学習
    clf = RandomForestClassifier()
    clf.fit(X, y)

    # 書き出し
    joblib.dump(clf, model_dir/MODEL_FILE_NAME)

if __name__ == '__main__':
    model_dir = os.environ['SM_MODEL_DIR'] # /opt/ml/model
    train_dir = os.environ['SM_CHANNEL_TRAIN'] # /opt/ml/input/data/train 
    train_dir = Path(train_dir)
    model_dir = Path(model_dir)
    train(train_dir, model_dir)

def model_fn(model_dir: str):
    """Predictで使う用の関数。学習されたモデルを返す"""
    model_path = Path(model_dir)/MODEL_FILE_NAME
    clf = joblib.load(model_path)

    return clf

ポイントは次の通りです。

データ読み込み&学習

環境変数 SM_CHANNEL_TRAIN (ディレクトリのパス)に、次の fit で指定するS3のPathからダウンロードされたファイルが配置されるので、それを読み込んで学習します。 学習したモデルは SM_MODEL_DIR (ディレクトリのパス) に任意の名前で保存します。

予測

model_fn メソッドで上で保存したモデルを読み込んで返します。これがAPIに使われます。

学習&デプロイ

それではSageMakerで学習とデプロイをスクリプトで行います。

script/train.py

import os
import json
import argparse
from pathlib import Path

import boto3
import sagemaker
from sagemaker.sklearn.estimator import SKLearn

# CloudFormationから環境変数を読み出し
## CFのStack設定
SERVICE_NAME = "sagemaker-serverless-example"
ENV = os.environ.get("ENV", "dev")
STACK_NAME = f"{SERVICE_NAME}-{ENV}"

## Outputsを{Key: Valueの形で読み出し}
stack = boto3.resource('cloudformation').Stack(STACK_NAME)
outputs = {o["OutputKey"]: o["OutputValue"] for o in stack.outputs}

S3_BUCKET = outputs["S3Bucket"]
S3_TRAIN_BASE_KEY = outputs["S3TrainBaseKey"]
S3_MODEL_BASE_KEY = outputs["S3ModelBaseKey"]

SM_ROLE_ARN = outputs["SmRoleArn"]
SM_ENDPOINT_NAME = outputs["SmEndpointName"]

INPUT_PATH = f"s3://{S3_BUCKET}/{S3_TRAIN_BASE_KEY}"
OUTPUT_PATH = f's3://{S3_BUCKET}/{S3_MODEL_BASE_KEY}'


def main(update_endpoint=False):
    script_path = str(Path(__file__).parent/"src/iris.py")
    train_instance_type = "ml.m5.large"
    initial_instance_count = 1
    hosting_instance_type = "ml.t2.medium"

    sagemaker_session = sagemaker.Session()
    # 学習
    sklearn = SKLearn(
        entry_point=script_path,
        train_instance_type=train_instance_type,
        role=SM_ROLE_ARN,
        sagemaker_session=sagemaker_session,
        output_path=OUTPUT_PATH
    )
    sklearn.fit({'train': INPUT_PATH})

    # デプロイ
    sklearn.deploy(
        initial_instance_count=initial_instance_count,
        instance_type=hosting_instance_type,
        endpoint_name=SM_ENDPOINT_NAME,
        update_endpoint=update_endpoint
    )

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--update-endpoint', action='store_true')
    args = parser.parse_args()
    update_endpoint = args.update_endpoint
    main(update_endpoint)

python script/train.py
# => 学習&デプロイされる

問題なく実行されていれば、モデルとエンドポイントが作成されます。

モデル

エンドポイント

補足

既にエンドポイントがデプロイされている場合は update_endpoint オプションを付ける必要があるのでスクリプトも引数を受け取れるようにしてあります。

python script/train.py --update-endpoint
# => エンドポイントを更新

予測API作成

それでは作成されたエンドポイントを使ってAPIを作成してみましょう。 次のようにserverless.ymlとPythonスクリプトを実装します。

serverless/serverless.yml

functions:
  predict:
    handler: predict.predict
    layers:
      - ${self:custom.layer.layer}
    events:
      - http:
          path: invocations
          method: post

serverless/predict.py

import os
import json

from sagemaker.sklearn.model import SKLearnPredictor

SM_ENDPOINT_NAME =  os.environ.get("SM_ENDPOINT_NAME")

def predict(event, contect):
    data = json.loads(event["body"])['data']
    predictor = SKLearnPredictor(endpoint_name=SM_ENDPOINT_NAME)
    y = predictor.predict([data]).tolist()[0]
    body = {
        "y": y
    }

    response = {
        "statusCode": 200,
        "body": json.dumps(body)
    }

    return response

出来たらデプロイします。

sls deploy
# ...略
# endpoints:
#   POST - https://xxxxxxxxx.execute-api.ap-northeast-1.amazonaws.com/dev/invocations
# ...略

デプロイできたらAPIをチェックしてみます。

curl -X POST -H "Content-Type: application/json" -d "{\"data\": [5.9, 3.0, 5.1, 1.8]}" https://xxxxxxxx.execute-api.ap-northeast-1.amazonaws.com/dev/invocations
# => {"y": 2}

無事APIを完成させることができました。

さいごに

SageMakerとServerlessを使ってscikit-learnの機械学習APIを使う例を紹介しました。 この例をカスタマイズすれば、PyTorchやTensorflowのモデルにも適用可能だと思います。 というわけで、次回はPyTorchを使った画像認識APIを作ってみる予定です。 何か不明点や改善点などありましたら Twitter まで教えてもらえたら幸いです。

今回使ったコードはすべて次のリポジトリに置いてあります。

github.com