機械学習プロジェクトにおけるSageMaker Processingの使い所

みなさんこんにちは。機械学習チームのたかぱい（@takapy0210）です。

2021年もあと1ヶ月となりましたね。皆様いかがお過ごしでしょうか。

...さて12月といえば、毎年恒例のアドベントカレンダーの季節ですね！
というわけで、2021年もコネヒト Advent Calendarが始まります！🎉

初日となる本エントリでは、機械学習チームで使用しているSageMaker*1の機能である、Processing*2について、活用事例とともにご紹介しようと思います。

SageMaker Processingとは?
- SKLearnProcessor / PySparkProcessor
- Processor / ScriptProcessor
これまでの課題感
SageMaker Processorの活用方法
- SageMakerからECRにあるコンテナを指定してProcessor Jobを起動する
SageMaker Processorを使うメリット
Processorのちょっとイマイチな点
最後に

SageMaker Processingとは?

機械学習の前処理、後処理、モデル評価といったワークフローをSageMaker上で簡単に行うためのPython SDKです。
独自のコンテナ上で任意のpythonスクリプトを実行でき、処理が終了するとインスタンスが自動的に停止するというシンプルな機能から、機械学習関連のデータ処理だけでなく、学習や推論にも使うことができるサービスとなっています。

このPython SDKには主に4つのクラスが用意されていて、それぞれで使用方法が異なります。

sagemaker.sklearn.processing.SKLearnProcessor
sagemaker.processing.PySparkProcessor
sagemaker.processing.Processor
sagemaker.processing.ScriptProcessor

それぞれについて簡単にご紹介します。
ちなみに、コネヒトでよく用いているのは sagemaker.processing.Processorです。
（以下で紹介するサンプルコードは、notebook上から実行していると思ってください）

SKLearnProcessor / PySparkProcessor

あらかじめ用意されたscikit-learnやPySparkの実行環境（コンテナ）を使用するクラスです。
frameworkのバージョンやインスタンスタイプ、インスタンスの台数を指定し、実行したいpythonスクリプトをrunメソッドの引数に与えて実行します。

from sagemaker.sklearn.processing import SKLearnProcessor
from sagemaker.processing import ProcessingInput, ProcessingOutput

# SKLearnProcessorクラスの実行例
sklearn_processor = SKLearnProcessor(framework_version='0.20.0',
                                     role=role,
                                     instance_count=1,
                                     instance_type='ml.m5.xlarge')

sklearn_processor.run(
    code='processing.py',  # ここで実行したいスクリプトを指定（このスクリプトはSageMaker上に存在する必要がある）
    inputs=[ProcessingInput(
        source='dataset.csv',
        destination='/opt/ml/processing/input')],
    outputs=[ProcessingOutput(source='/opt/ml/processing/output/train'),
        ProcessingOutput(source='/opt/ml/processing/output/validation'),
        ProcessingOutput(source='/opt/ml/processing/output/test')]
)

aws.amazon.com

Processor / ScriptProcessor

独自の実行環境（コンテナ）で処理を行いたい時に使用するクラスです。
Processorクラスはrunメソッドにcode引数（pythonスクリプト）を必要としないクラスで、実行したい処理をDocker image内に内包させる必要があります。
ScriptProcessorクラスはコンテナの作成と処理コードを分けて実行するクラスで、runメソッドにpythonスクリプトを受け渡し、処理を実行することができます。

from sagemaker.processing import Processor

# Processorクラスの実行例
processor = Processor(
    image_uri=image,  # コンテナイメージを指定
    entrypoint=["python3", "/opt/program/processing.py"],  # 実行したいスクリプトを指定（コンテナに内包しておく必要がある）
    env={"PYTHON_ENV": "hoge"},
    role=role,
    instance_count=1,
    instance_type="ml.t3.2xlarge"
)

processor.run()

docs.aws.amazon.com

これまでの課題感

インタラクティブな分析やモデリングなどは以前のブログでご紹介したように、分析用コンテナイメージを用意し、それをローカルマシン上で動かす or SageMaker上（クラウド上）で動かすことで、各人が業務を進めていました。

分析にある程度目処が立ったタイミングで実際のプロダクションコードを書いて、それをプロダクション用のコンテナ環境に乗せてAWS上で動かす基盤を整える、というフローで開発を進めることが多いのですが、このプロダクションコードをどの環境でどのように検証すれば良いのか、という部分に課題を感じていました。

例えば...

プロダクション用コンテナ環境を使ってローカルで検証したいが、計算リソースの関係でローカルマシンでは本番相当のデータ量で検証できない
クラウド上にある分析用コンテナ環境を使って検証しても良いが、実際に動く環境とは異なるので、プロダクション環境で実行したときに違う挙動をする可能性がある

などなど。

このような課題を解消すべく、現在使っているのがSageMaker Processorです。

SageMaker Processorの活用方法

冒頭で述べたように、弊社でよく利用しているのは sagemaker.processing.Processor クラスです。

下記図は、弊社の機械学習基盤の一部を抽象的に表したものですが、緑色の部分がいわゆる分析環境で、青色の部分が今回紹介するSageMaker Processorを活用したプロダクション動作検証環境になっています。
（図にある分析用コンテナイメージは全プロジェクト共通で、本番用コンテナイメージはプロジェクトごとに異なっています）

f:id:taxa_program:20211129183719p:plain — 分析用コンテナと本番用コンテナの使い分けイメージ

実際にSageMaker Processorを利用するのに必要な手順は以下の3つです。

実行したいpythonスクリプトを内包したコンテナイメージを作成する
1.のコンテナイメージをECRにPUSHする
SageMaker notebook上からECRにあるコンテナを指定してProcessor Jobを起動する

3.について、コードと共に簡単に解説していきます。

SageMakerからECRにあるコンテナを指定してProcessor Jobを起動する

notebook上からは以下のようにして利用することができます。
コードを見ていただければ分かるように、notebookを起動しているインスタンスとは別のインスタンスを指定して実行できるので、処理内容に合わせて様々なリソース上で実行することができます。

以下の例では processing.py を実行しています。

import boto3
import sagemaker
from sagemaker import get_execution_role
from sagemaker.processing import Processor

role = get_execution_role()
sess = sagemaker.Session()
region = boto3.session.Session().region_name
account = sess.boto_session.client('sts').get_caller_identity()['Account']

# 実行したいコンテナを指定（{}の中は任意の文字列に変更してください）
image = f'{account}.dkr.ecr.{region}.amazonaws.com/{image_name}:{tag_name}'

# データ整形JOB
processor = Processor(
    image_uri=image,  # 実行するコンテナイメージ
    entrypoint=["python3", "/opt/program/processing.py"],  # 実行するスクリプト 
    env={"PYTHON_ENV": "hoge"},  # 環境変数
    role=role,  # 実行ロール
    instance_count=1,  # インスタンスの数
    instance_type="ml.t3.2xlarge"  # インスタンスの種類
)

# 処理の実行
processor.run()

例えば、今回実行したprocessing.pyの処理内容が「データを加工して、train.csvというファイル名でS3に保存する」というものだったとします。

この時、生成されたデータが想定通りのデータになっているか確認したい時もあると思います。
そんな時はnotebook上から以下のような処理を実行することでインタラクティブにデータをチェックこともできます。

# 該当データをS3からSageMakerのvolume上にダウンロード
sess.download_data('./', bucket='hoge', key_prefix='dir/train.csv')

# pandasで読み込む
train = pd.read_csv('train.csv')

この辺のことがシームレスに行えるのもSageMakerの良さだと思います。

SageMaker Processorを使うメリット

今までも述べてきましたが、メリットについて改めてまとめると、

本番環境のコンテナ&本番データを使って、（ある程度）お手軽に動作検証ができる。
notebookとProcessor Jobで 異なるインスタンスを使えるので、Processer Jobだけに強めのインスタンスを割り当てて実行する、みたいなことができる。
- Processer Jobで起動したインスタンスは処理終了時に自動的にシャットダウンされるので、「起動しっぱなしでコストが...汗」ということも防ぐことができる。
S3を媒介として、データのやりとりも簡単にできる。

という点が挙げられるかな、と思います。