Amazon SageMaker Studio Labで作成したノートブックをそのままバッチジョブ化する

データアナリティクス事業本部 インテグレーション部 機械学習チームの貞松です。

本記事は「クラスメソッド 機械学習チーム アドベントカレンダー 2022」の7日目です。
昨日(6日目)は じょんすみす による Vertex AI WorkbenchからBigQueryアクセスする でした。

下記の記事でも紹介されましたが、Amazon SageMaker Studio(以降SageMaker Studio)上で作成したノートブックをそのままバッチジョブとして実行できる機能が追加されました。

実は、このノートブックのバッチジョブ化機能は、Amazon SageMaker Studio Lab(以降SageMaker Studio Lab)にも追加されています。

実際にSageMaker Studio Lab上でノートブックをバッチジョブ化してみましょう。

※2022.12.08追記: 以下補足です。
SageMaker Studio Lab自体はAWSアカウントのサインアップ不要で利用可能な無料の機械学習開発環境ですが、本記事で解説するNotebook Job(ノートブックのバッチジョブ化機能)については、AWSリソース(AWSサービスとしてのSageMakerとEventBridge)を利用する為にAWSアカウントとIAMの設定が必要となります。
ご利用の際には、上記ご了承の上で、予めAWSアカウントをご準備ください。

SageMaker Studio Labとは

Amazon SageMaker Studio Labは、無料の機械学習開発環境です。
AWSアカウントにサインアップする必要はなく、Eメールアドレスさえあれば利用登録が可能です。複雑な設定やユーザー管理等も不要です。
利用用途やイメージとしてはGoogle Colabに近いサービスですね。

SageMaker Studio Lab

登録方法や使い始めの手順については、下記の記事をご参照ください。

SageMaker Studio Labで使用するIAMロールとIAMユーザーを作成する

SageMaker Studio Lab上でノートブックをバッチジョブとして実行するために、必要なIAMロールとIAMユーザーを作成します。
IAMロールは作成したバッチジョブの実行ロールとして設定するために必要であり、SageMakerサービスに対する権限が必要です。
IAMユーザーはSageMaker Studio Lab上でバッチジョブを作成するために必要であり、SageMakerとEventBridgeに対する権限が必要です。
※具体的なバッチジョブの作成方法については後述しますが、どうやら作成したバッチジョブの実行制御(スケジュール実行など)はEventBridgeが使用されるようですね。

IAMロールの作成

Trusted entity typeはAWS serviceを選択し、Use caseはSageMaker - Executionを選択します。

Permissions policyとして、AmazonSageMakerFullAccessを選択します。

任意のロール名を入力してIAMロールの作成を完了します。

作成したIAMロールのARNがバッチジョブ作成時に必要になるため控えておきます。

IAMユーザーの作成

任意のユーザー名を入力し、Access Key - Programmatic Accessがチェックされている状態にして次に進みます。

Set Permissionは、Attach existing policies directlyを選択し、AmazonSageMakerFullAccessとAmazonEventBridgeFullAccess２つのポリシーを選択します。

最終的な設定内容を確認して、IAMユーザーの作成を完了します。

SageMaker Studio Lab上でAWS CLIのConfig設定をする際に必要になるため、IAMユーザー作成完了画面で表示されるAWS Access Key IDとAWS Secret Access Keyを何かしらの形で保持しておきます(普通にCSVダウンロードで良いです)

SageMaker Studio Lab上でAWS CLIのConfigを設定する

SageMaker Studio Labから、ノートブックをバッチジョブ化したものを実行するために前述で作成したIAMユーザーの情報でAWS CLIのConfigを設定します。　　

まずはLauncherの画面(SageMaker Studio Labのトップ画面)からTerminalを開きます。

次に aws configure コマンドを実行して、必要な情報を入力していきます。
AWS Access Key IDとAWS Secret Access Keyは、作成したIAMユーザーの情報を入力します。
Default region nameとDefault output formatは必要に応じて設定します。ここではap-northeast-1とjsonで設定します。

(studiolab) studio-lab-user@default:~$ aws configure 
AWS Access Key ID []: <作成したIAMユーザーのAWS Access Key ID>
AWS Secret Access Key []: <作成したIAMユーザーのAWS Secret Access Key>
Default region name []: <任意のリージョン>
Default output format [None]: <任意のファイルフォーマット(空でも可)>

これで、AWSアカウントにアクセスできるようにSageMaker Studio Lab環境が設定されました。
実際にAWSアカウントに接続できることを確認する為に、以下のコマンドを実行します。

(studiolab) studio-lab-user@default:~$ aws sts get-caller-identity

実際に実行した様子は以下のようになります。

SageMaker Studio Lab上でノートブックを作成する

SageMaker Studio Lab上で新規にノートブックを作成する場合は、Launcherの画面からNotebookのセクションにあるdefaultPythonを選択します。

また、SageMaker Studioと同様に既存のGitリポジトリをクローンすることも出来るので、今回は amazon-sagemaker-example リポジトリをクローンして、その中にあるノートブックをバッチジョブ化します。

サイドバーのGitタブからClone a Repositoryを選択します。

Git Repository URLを入力してクローンします。

無事にクローンされました。

SageMaker Studio Lab上のノートブックをバッチジョブ化する

ノートブックをジョブ化するためには、サイドバーのファイルタブを開いて、該当のノートブックファイルを右クリックします。
右クリックメニューにあるCreate Notebook Jobを選択します。

Notebook Job作成画面が開きます。このあたりはSageMaker Studioと同様ですね。
ここで、バッチジョブを実行する際のインスタンス設定などを実施します。

Additional Optionsの中にRole ARNの設定項目があるので、ここでようやく前述で作成したIAMロールのARNの出番がやってきます。

実行パターンは即時実行とスケジュール実行があります。
このあたりも含めて、全体的にSageMaker Studioと同様です。

作成されたジョブの一覧画面です。正常にジョブが作成できていることを確認できます。

ジョブの一覧からジョブの詳細画面を開くことができます。
ジョブの設定確認だけでなく、実行中のジョブの停止やジョブ自体の削除などを実行できます。

まとめ

SageMaker Studioと同様に、SageMaker Studio Labでもノートブックをバッチジョブ化することができました。 サクッと実験できるLab環境から、実用的なジョブが生成できるのはシンプルに便利ですね。