[アップデート] Amazon SageMaker Studio の新機能として、ノートブックコードを本番環境に対応したジョブに自動変換する機能を発表！ #reinvent

はじめに

こんちには。

データアナリティクス事業本部 機械学習チームの中村です。

本記事では、re:Invent 2022中に発表された機能で、ノートブックコードを本番環境に対応したジョブに自動変換する機能について紹介します。

公式アナウンスはこちらです。

セッションとしては、アナウンス後に開催された以下の機械学習のリーダーシップセッションで言及がありました。

• Innovate with AI/ML to transform your business (AIM217-L)

機能の概要

これまで、データサイエンティストや開発者がノートブックを実運用に移す場合、以下のような作業が必要でした。

• ノートブックのセルからコードを手作業でコピー
• スクリプトとすべての依存関係のあるパッケージをコンテナ環境に構築
• コンテナをジョブとして実行するようにスケジュール

また、これらのCI/CDを構築するためには、これを実行するパイプラインの構築・管理が必要でした。 この作業は、機械学習のコア部分の開発ではない時間として、時間を消費しています。

今回の機能では、開発者がノートブックを選択し、Studio ビジュアルインタフェースから数回クリックするだけで本番環境でジョブとして実行するよう自動化することができるようになりました。

ジョブの実行がスケジュールされると、SageMaker Studio は自動的に以下を行います。

• ノートブック全体のスナップショットを取得
• 依存関係をコンテナにパッケージ化し、インフラストラクチャを構築
• ノートブックを自動ジョブとして実行し、ジョブ完了時にインフラストラクチャのプロビジョニングを解除

詳細は以下のAWS公式ブログにも説明があります。

使ってみた

本機能では、SageMaker Studio Notebookで使用できるので、そちらをまず立ち上げます。

次に、本機能を使用するにはSageMakerの実行ロールに以下の権限が必要のため、追加します。

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Service": "sagemaker.amazonaws.com"
            },
            "Action": "sts:AssumeRole"
        },
        {
            "Effect": "Allow",
            "Principal": {
                "Service": "events.amazonaws.com"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}

加えて、以下のポリシーを実験のために追加しておきます。

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "VisualEditor0",
            "Effect": "Allow",
            "Action": [
                "events:TagResource",
                "events:DeleteRule",
                "events:PutTargets",
                "events:DescribeRule",
                "events:PutRule",
                "events:RemoveTargets",
                "events:DisableRule",
                "events:EnableRule"
            ],
            "Resource": "*",
            "Condition": {
              "StringEquals": {
                "aws:ResourceTag/sagemaker:is-scheduling-notebook-job": "true"
              }
            }
        },
        {
            "Sid": "VisualEditor1",
            "Effect": "Allow",
            "Action": "iam:PassRole",
            "Resource": "arn:aws:iam::*:role/*",
            "Condition": {
                "StringLike": {
                    "iam:PassedToService": "events.amazonaws.com"
                }
            }
        },
        {
            "Sid": "VisualEditor2",
            "Effect": "Allow",
            "Action": "sagemaker:ListTags",
            "Resource": "arn:aws:sagemaker:*:*:user-profile/*/*"
        }
    ]
}

実行ロールの状態は、Studioユーザの実行ロールの内容を確認してみてください。

ジョブ用のノートブックを準備

SageMaker StudioはUIが一新されてますのでご注意ください。初期状態ではHomeタブが開いていますので、「Open Launcher」をクリックし、

Create NotebookでまずはNotebookを起動します。

以下のチュートリアルに従って、通信会社の解約予測のデータセットを使った一連の処理を記述します。

以下のようにコードを記述します。

最初にはパラメータ化するものを書いておくと、分かりやすいと思われます。

このパラメータは、ジョブ実行時にオプションとして与えることで変更することが可能です。

# These parameters can be modified when executing this notebook as a job
number_rf_estimators = 100

次に前処理です。

# You can download additional assets for your use in your notebook jobs
!aws s3 cp s3://sagemaker-sample-files/datasets/scripts/sagemaker-studio-scheduling/synthetic_data.py ./

import pandas as pd
from synthetic_data import generate_data

previous_two_weeks_data = generate_data(5000, label_known=True)
todays_data = generate_data(300, label_known=False)

previous_two_weeks_data.head()

def process_data(df, label_known):
    """
    This function represents typical data munging that needs to be performed for machine learning workflows.
    We'll cast data types, drop columns that are irrelevant as features to the model, and convert categorical
    features into indicator variables. If the data is not from the ground truth dataset, we'll drop it's churn label.
    """
    df = df.drop("Phone", axis=1)
    df["Area Code"] = df["Area Code"].astype(object)
    df = df.drop(["Day Charge", "Eve Charge", "Night Charge", "Intl Charge"], axis=1)
    model_data = pd.get_dummies(df)

    if label_known:
        model_data = pd.concat(
            [
                model_data["Churn?_True."],
                model_data.drop(["Churn?_False.", "Churn?_True."], axis=1),
            ],
            axis=1,
        )
        model_data = model_data.rename(columns={"Churn?_True.": "Churn"})

    return model_data

processed_prior_data = process_data(previous_two_weeks_data, label_known=True)
processed_todays_data = process_data(todays_data, label_known=False)

次にモデルを定義して、学習する処理を記述します。

import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import f1_score, confusion_matrix, ConfusionMatrixDisplay

y = np.ravel(processed_prior_data[["Churn"]])
x = processed_prior_data.drop(["Churn"], axis=1)

x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.25)

clf = RandomForestClassifier(n_estimators=int(number_rf_estimators), criterion="gini")
clf.fit(x_train, y_train)

ここからは推論処理を実行に、性能を確認しています。

# Run Inference on test set
from IPython.display import Markdown as md

y_pred = clf.predict(x_test)
f1 = f1_score(y_test, y_pred)
md("## Refreshed Model's F1 Score on validation data: {:.2f}".format(f1))

# Plot test set results
cm = confusion_matrix(y_test, y_pred)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=["Churn", "Didn't Churn"])
disp.plot()
plt.show()

実際に1日分のデータを推論したものを表示しています。

todays_predictions = clf.predict(processed_todays_data)
accounts_likely_to_churn = pd.concat(
    [pd.DataFrame(todays_predictions, columns=["Predicted Churn"]), processed_todays_data], axis=1
)
accounts_likely_to_churn = accounts_likely_to_churn[
    accounts_likely_to_churn["Predicted Churn"] == 1
]

accounts_likely_to_churn.head()

本番環境では、これらの予測結果を DynamoDB などのデータベースに格納したり、学習したモデルを SageMaker モデルレジストリに登録することなどが考えられます。

ジョブ作成

以下にある、「Create notebook job」をクリックします。

インスタンスなどを選択します。notebook jobに選択できるのは、Trainingで使用可能なインスタンスのみとなっているようです。

それぞれのインスタンスの価格については以下を参照ください。

次にパラメータを以下のように指定してみます。

最後に、「Run now」を選択し「Create」を押下してジョブを実行します。

実行結果は、以下に表示されます。

Output filesの、Notebookからは、各セルの出力が含まれたNotebookにアクセス可能で、Output Logからは実行時のログにアクセスできるようです。

Input fileとOutput filesの実体はS3にあり、Notebook JobsのJob Detailに記載の以下で確認できます。

なお、ジョブ実行時のインスタンスは自動的に削除されるようです。

さいごに（念のため）

さいごに手元のnotebook作成時に立ち上がっているINSTANCEを忘れずに止めるか、Studio自体のシャットダウンを忘れないように実施しましょう。

まとめ

いかがでしたでしょうか。Notebookをジョブに変換できる機能はとても便利そうで、スケジュール実行などにも対応しておりとても良いと感じました。Notebook自体をジョブ実行するPapermillなどのPythonモジュールは確かに存在していたのですが、それをStudioのGUIからポチポチすることにより、実現可能となったので、楽になったかなと思います。

今後実務で、活用していけたらと思いました。