AWS Step Functionsを使用してオンプレミスの機械学習コードを Amazon SageMaker に取り込む

ローカルデスクトップなどオンプレミス環境で開発した機械学習のコードがあり、非効率であることは分かりつつも修正の時間を取ることができないケースも多くあるのではないでしょうか。その課題の解決策として、オンプレミス環境で作成された機械学習コードを、最小限のリファクタリングでAWS 環境に移行し、またAWS内の機械学習サービス Amazon SageMaker の機能を最大限活用する方法について考えたいと思います。

本記事でのオンプレミスのコードがさしているものは、ローカルデスクトップやお客様のサーバーなどで手動で実行するように開発されたコードで、AWS SDK for Python (Boto3) や Amazon SageMaker Python SDK などのクラウド対応の SDK でビルドされていないコードです。つまりこれらのオンプレミスのコードはクラウド展開に最適化されていません。移行のベストプラクティスは、Amazon SageMaker API または SageMaker Python SDK を使用してこれらのオンプレミスのコードをリファクタリングすることです。ただし、オンプレミスで運用されているモデルが多数存在する組織では、それらのモデルをすべて書き直す時間やリソースがない場合があります。

本記事では、Amazon SageMaker と AWS Step Functions を使用して、最小限のコードリファクタリングでオンプレミスの推論用 ML コードを AWS 環境に移行する、スケーラブルで実装が簡単なアプローチを紹介します。このソリューションを拡張してさらに機能を追加することも可能です。またデータサイエンティストと MLOps エンジニアという 2 人の異なるペルソナがオンプレミスのコードをどのようにリフトアンドシフトするのかをお伝えしたいと思います。

ソリューション概要

このフレームワークでは、オンプレミスのコードを SageMaker Processing Jobとしてコンテナ内で実行します。SageMaker はオンプレミスから持ってきたスクリプトを処理コンテナ内で実行します。処理コンテナイメージは、SageMaker の組み込みイメージでもカスタムイメージでもかまいません。またProcessing Jobの基盤となるインフラストラクチャはSageMaker によって完全に管理されます。オンプレミスのコードを変更する必要はなく、SageMaker Processing Jobの作成に慣れていれば問題ありません。

ここでは、データサイエンティストと MLOps エンジニアという 2 人のペルソナが作業することを想定しています。データサイエンティストは、手動または AWS CodeCommit などのコードリポジトリからコードを複製して、コードの SageMaker 移行を担当します。Amazon SageMaker Studio には、ML ライフサイクルのさまざまなステップを実装するための統合開発環境 (IDE) が用意されています。データサイエンティストは、これを使用して、デプロイに必要なコードアーティファクトを含むカスタムコンテナを手動で構築します。コンテナはデプロイのために、Amazon Elastic Container Registry ( Amazon ECR ) などのコンテナレジストリに登録されます。

MLOps エンジニアは Step Functions ワークフローを担当します。このワークフローを再利用して、データサイエンティストが適切なパラメータを使用して開発したカスタムコンテナをデプロイできます。Step Functions ワークフローは、必要なユースケースに合わせてモジュール化することができたり、もしくは1 つのステップで 1 つのプロセス実行のみを構成することもできます。コードの移行に必要な労力を最小限に抑えるため、完全に機能するデプロイプロセスを構築するための3つのモジュール要素に絞りました。

前処理
推論
後処理

次の図は、本記事のソリューションアーキテクチャとワークフローを示しています。

このソリューションには、次の手順が含まれます。

データサイエンティストは、Studioを使用してコードリポジトリから複製することでオンプレミスのコードをインポートし、そのコードを ML ライフサイクルのステップ（前処理、推論、後処理）に従って別々のコンポーネントにモジュール化します。
データサイエンティストは特に SageMaker が提供する Studio イメージビルド CLI ツールを使用して Docker イメージを構築します。この CLI ツールを使用すると、データサイエンティストは Studio 内で直接イメージを構築し、そのイメージを Amazon ECR に自動的に登録できます。
MLOps エンジニアは、登録されたコンテナイメージをStep Functions を使用して特定のユースケース用にデプロイします。Step Functions は、Amazon States Languageを使用して SageMaker API を直接制御できるサーバーレスワークフローサービスです。

SageMaker Processing Job

SageMaker Processing Jobの実行方法を理解していきましょう。次の図は、SageMaker がどのようにProcessing Jobを起動するかを示しています。

SageMaker はスクリプトを受け取り、Amazon Simple Storage Service (Amazon S3) からデータをコピーし、処理コンテナを取り出します。処理コンテナイメージは、SageMaker の組み込みイメージでも、ユーザーが提供するカスタムイメージでもかまいません。Processing Jobの基盤となるインフラストラクチャは、SageMaker によって完全に管理されます。クラスターリソースはJobの実行中にプロビジョニングされ、Jobが完了すると削除されます。Processing Jobの出力は、指定した S3 バケットに保存されます。独自のコンテナを構築する方法の詳細については、「独自の処理コンテナを構築する (高度なシナリオ)」を参照してください。

SageMaker Processing Job は、Docker コンテナエントリポイントスクリプトを使用して処理イメージを設定します。AppSpecification API の ContainerEntryPoint と ContainerArguments パラメータを使用して、独自のカスタムエントリポイントを提供することもできます。独自のカスタムエントリポイントを使用する場合は、イメージを再構築せずにスタンドアロンスクリプトとして実行できる柔軟性が高まります。

この例では、カスタムコンテナを構築し、推論用に SageMaker Processing Job を使用します。前処理用の Processing Job と後処理用の Processing Job は、ビルド済みの scikit-learn コンテナを備えたスクリプトモードを利用します。

前提条件

この記事を最後まで進めるには、以下の前提条件を満たす手順を実行してください。

Studio ドメインを作成します。手順については、「クイックセットアップを使用して Amazon SageMaker ドメインにオンボードする」を参照してください。
S3 バケットを作成します。（サンプルコードではSageMaker Studio環境作成時に自動で作成されるデフォルトバケットを使用しているため、バケット作成をしない場合でも実行可能です）
提供された GitHub リポジトリを Studio にクローンします。

GitHub リポジトリは、ML ライフサイクルの各段階に対応するフォルダーに整理されているため、ナビゲーションや管理が簡単になります。

オンプレミスのコードの移行

このステップでは、データサイエンティストの担当するオンプレミスのコードの移行について解説します。

まず、build_and_push.ipynb ノートブックを開くことから始めます。

ノートブックの最初のセルには、Studio Image Build CLI をインストールする手順が示されます。この CLI は、完結なコマンド操作により再利用可能なビルド環境を自動的に作成することができるため、セットアッププロセスの簡略化が可能です。CLI を使用すると、イメージの作成はビルドを指示するのと同じくらい簡単で、その結果が Amazon ECR 内のイメージの場所へのリンクになります。このアプローチにより、CLIによって調整された複雑な基盤となるワークフローを管理する必要がなくなり、イメージ構築プロセスが合理化されます。

build コマンドを実行する前に、GitHub内の readme または関連記事で指定されているように、コマンドを実行するロールに必要な権限があることを確認することが重要です。必要な権限を付与していない場合、ビルドプロセス中にエラーが発生する可能性があります。

次のコードを参照してください。

#Install sagemaker_studio_image_build utility 
import sys 
!{sys.executable} -m pip install sagemaker_studio_image_build

オンプレミスのコードを効率化するには、preprocessing.py、predict.py、postprocessing.py という名前の 3 つの Python スクリプトに分割します。プログラミングのベストプラクティスに順守する形で、コードをメイン関数から呼び出される関数に変更してください。、requirements.txt ファイルがすべてのカスタムライブラリを含んでおり、必要なすべてのライブラリがインポートされることを確認してください。

コードを整理したら、Docker コンテナと一緒にrequirements.txtファイルをパッケージ化します。以下のコマンドを使用して、Studio 内からコンテナを簡単にビルドできます。

sm-docker build .

デフォルトでは、イメージは「latest」というタグが付与され、Studioから呼び出された ECR リポジトリにプッシュされます。さらに、Studio アプリケーションの実行ロールがデフォルトの SageMaker Python SDK S3 バケットとともに使用されます。ただし、これらの設定は適切な CLI オプションを使用して簡単に変更できます。次のコードを参照してください。

sm-docker build . --repository mynewrepo:1.0 --role SampleDockerBuildRole --bucket sagemaker-us-east-1-0123456789999 --vpc-id vpc-0c70e76ef1c603b94 --subnet-ids subnet-0d984f080338960bb,subnet-0ac3e96808c8092f2 --security-group-ids sg-0d31b4042f2902cd0

コンテナがビルドされ ECR リポジトリに登録されたので、今度はそれを使って predict.py を実行する方法をさらに詳しく見ていきましょう。また、ビルド済みの scikit-learn コンテナを使用して preprocessing.py と postprocessing.py を実行するプロセスについても説明します。

コンテナの本番実装

このステップでは、前のステップで作成したコンテナを本番実装するMLOpsエンジニアのアクションについて説明します。

Step Functionsを使用してワークフローを調整します。Step Functions を使用すると、オンプレミスのシステムに存在する可能性のある既存の依存関係に対応しながら、さまざまなサービスをワークフローに柔軟に統合できます。このアプローチにより、必要なすべてのコンポーネントがシームレスに統合され、目的の順序で実行されるため、効率的なワークフローソリューションが実現します。

Step Functions では、特定の AWS サービスをAmazon States Languageから直接制御できます。Step Functions の操作と SageMaker との統合について詳しくは、「Step Functions で SageMaker を管理する」を参照してください。SageMaker の Step Functions 統合機能を使用して、スクリプトモードで SageMaker Processing Jobを使用して前処理スクリプトと後処理スクリプトを実行し、カスタムコンテナを使用して SageMaker Processing Jobとして推論を実行します。またそのためにAWS SDK for Python (Boto3) のCreateProcessingJob API を使用します。

前処理

SageMaker には、カスタムコードを実行するためのオプションがいくつか用意されています。カスタムの依存関係がないスクリプトしかない場合は、そのスクリプトを独自スクリプトの持ち込み (BYOS) として実行できます。そのためには、スクリプトをビルド済みの scikit-learn フレームワークコンテナに渡し、AppSpecification API の ContainerArguments パラメータと ContainerEntryPoint パラメータを使用して、スクリプトモードで SageMaker Processing Jobを実行するだけです。こちらの方が簡単なスクリプトを実行するのには適した方法です。

CreateProcessingJob API を呼び出してカスタムスクリプトを実行する方法を理解するには、サンプル Step Functions ワークフローの「前処理スクリプトモード」の状態設定を確認してください。

推論

独自の処理コンテナを構築するアプローチを使用してカスタムコンテナを実行できます。SageMaker Processing Job は /opt/ml ローカルパスで動作し、設定からProcessing Inputsをそのローカルパスに指定できます。

次に、Processing Jobは指定された入力データをローカルコンテナにコピーし、Jobを開始します。Jobが完了すると、Processing Outputs のローカルパスで指定された出力データが、指定されたストレージサービス（本記事の場合S3）にコピーされます。

CreateProcessingJob API を呼び出してカスタムコンテナを実行する方法を理解するには、Step Functions ワークフローサンプルの「推論カスタムコンテナ」の状態設定を確認してください。

後処理

Step Functionsの CreateProcessingJob ステップを使用すると、後処理スクリプトを前処理スクリプトと同じように実行できます。後処理スクリプトを実行すると、推論ジョブの完了後にカスタム処理タスクを実行できます。

Step Functionsワークフローの作成

プロトタイプをすばやく作成するには、Step Functions の Amazon States Language を使用します。Step Functions の定義は、States Language を使用して直接編集できます。Step Functions ワークフローのサンプルを参照してください。

Step Functions コンソールを開き、ステートマシンを選択。ステートマシンの作成から「ワークフローをコードで作成」を選択すると、新しい Step Functions ステートマシンを作成できます。

Step Functions では、使用しているリソースを確認してロールを作成できます。ただし、次のメッセージが表示される場合があります。

これに対処するには、Step Functions の AWS Identity and Access Management (IAM) ロールを作成する必要があります。手順については、「ステートマシン用の IAM ロールの作成」を参照してください。次に、以下の IAM ポリシーをアタッチして、ワークフローを実行するために必要な権限を指定します。

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "sagemaker:createProcessingJob",
                "sagemaker:ListTags",
                "sagemaker:AddTags"
            ],
            "Resource": "*"
        },
        {
            "Effect": "Allow",
            "Action": [
                "iam:PassRole"
            ],
            "Resource": "*",
            "Condition": {
                "StringEquals": {
                    "iam:PassedToService": "sagemaker.amazonaws.com"
                }
            }
        }
    ]
}

参考: ロールの作成手順

IAMコンソールのロールから「ロールを作成」を選択
「AWSのサービス」を選択後、「他の AWS のサービスのユースケース」からStep Functionsを選択し次へ
その他はデフォルトのまま作成
作成したロールを選択し以下の画像のように「ポリシーをアタッチ」を選択し、その他の許可ポリシーからCloudWatchEventsFullAccessを検索し追加

再度作成したロールを選択し以下の画像のように「インラインポリシーを作成」を選択し上記のようなjson形式でポリシーを作成する

最後に再度ロール詳細画面の概要にあるARNをコピー
Step Functionsの画面で上記でコピーしたロールARNを貼り付けて実行

次の図は、Step Functions ワークフローの各ステップへのデータとコンテナイメージのフローを示しています。

以下は、Step Functions で初期化するために最低限必要なパラメータのリストです。前述のノートブックで実行した時のパラメータなどを参考に、この入力パラメータ JSON のサンプルを編集頂ければと思います。

input_uri — 入力ファイルの S3 URI
output_uri — 出力ファイルの S3 URI
code_uri — スクリプトファイルの S3 URI
custom_image_uri — ビルドしたカスタムコンテナのコンテナ URI
scikit_image_uri — ビルド済みの scikit-learn フレームワークのコンテナ URI
role — Jobを実行するための実行ロール
instance_type — コンテナを実行するために使用する必要があるインスタンスタイプ
volume_size — コンテナに必要なストレージボリュームサイズ
max_runtime — コンテナの最大ランタイム。デフォルト値は 1 時間

以下の画像のようにStepFunctinosを実行するときに入力パラメータをJSONフォーマットで指定できます

ワークフローを実行する

オンプレミスのコードを管理しやすい部分、つまり前処理、推論、後処理に分割しました。また推論のニーズをサポートするために、必要なライブラリ依存関係を備えたカスタムコンテナを構築しました。またStep Functions で SageMaker API を呼び出すことでその機能をより活かすようにしました。そしてSageMaker API を使用してカスタムコードを実行する方法を 2 つ示しました。1 つは、ビルド済みのイメージを使用して実行時にカスタムスクリプトを使用する SageMaker Processing Jobで、もう 1 つは推論を実行するために必要なアーティファクトがパッケージされたカスタムコンテナを使用する SageMaker Processing Jobです。

次の図は、Step Functions ワークフローの実行を示しています。

サマリー

この記事では、ローカル開発環境からオンプレミスのな ML Python コードを移行し、標準化された MLOps 手順を実装するプロセスについて説明しました。このアプローチでは、何百ものモデルを簡単に移行し、各企業に適した移行方法に組み込むことができます。またSageMaker でカスタムコードを実行する方法を 2 種類紹介しましたので、お客様のニーズに最も合った方法を選択頂くことができるかと思います。

細かいレベルでカスタマイズ可能なソリューションが必要な場合は、カスタムコンテナアプローチを使用することをお勧めします。前述の前処理ステップで説明したように、基本的なスクリプトがあって、カスタムコンテナを作成する必要がない場合は、ビルド済みのイメージを使用してカスタムスクリプトを実行する方が適しているケースもあります。さらに必要に応じてこの記事の推論ステップをコンテナ化する方法と同様に、オンプレミスのモデルのトレーニングと評価ステップをコンテナ化することもできます。

著者について

Bhavana Chirumamilla は AWS のシニアレジデントアーキテクトで、データと機械学習の運用に強い情熱を持っています。豊富な経験と熱意をもって、企業が効果的なデータと ML 戦略を構築できるよう支援しています。休日には、家族と過ごしたり、旅行、ハイキング、ガーデニング、ドキュメンタリー鑑賞など、さまざまなアクティビティを楽しんでいます。

Shyam Namavaram は、アマゾンウェブサービス (AWS) で人工知能 (AI) と機械学習 (ML) を専門とするシニアソリューションアーキテクトです。技術ガイダンスを提供し、AWS での安全なクラウドソリューションの革新と構築を支援することで、お客様の AI と ML の運用を加速させることに熱心に取り組んでいます。専門はAI、ML、コンテナ、分析技術です。仕事以外では、スポーツをしたり、トレッキングで自然を感じるのが大好きです。

Qingwei Li は、アマゾンウェブサービスの機械学習スペシャリストです。彼はリサーチアドバイザーの助成金口座を破棄し、約束されたノーベル賞を獲得できなかったのち、オペレーションズリサーチの博士号を取得しました。現在は、ファイナンスサービスおよび保険業界のお客様が AWS で機械学習ソリューションを構築できるよう支援しています。休日には、読書と教育活動に従事しています。

Srinivasa Shaik は、ボストンを拠点とするAWSのソリューションアーキテクトです。企業顧客がクラウドへの移行を加速できるよう支援しています。彼はコンテナと機械学習技術に情熱を注いでいます。休日には、家族と過ごしたり、料理をしたり、旅行を楽しんでいます。

翻訳は機械学習プロトタイプソリューションアーキテクトの石橋直樹が担当しました。

Amazon Web Services ブログ