AWS Step Functions Data Science SDK for Amazon SageMaker を使用してモデルの再トレーニングとデプロイを自動化する

機械学習 (ML) が企業のコアビジネスの大きな部分を占めるようになると、モデルの作成からデプロイまでの時間を短縮することに重点が置かれるようになりました。2019 年 11 月、AWS は AWS Step Functions Data Science SDK for Amazon SageMaker をリリースしました。これは、開発者が Python で Step Functions ベースの機械学習ワークフローを作成できるオープンソース SDK です。SDK を使用して、モデルの開発に使用するものと同じツールを使用して、再利用可能なモデルデプロイワークフローを作成できるようになりました。このソリューションの完全なノートブックは、GitHub リポジトリの「automate_model_retraining_workflow」フォルダにあります。

この記事では、Data Science SDK の機能を、スケジュールされたモデルの再トレーニングとデプロイの一般的なユースケースで示します。また、ML モデルをトレーニングするためのサーバーレスワークフローを作成し、検証データセットに対してモデルのパフォーマンスを確認します。さらに、モデルの精度が設定されたしきい値を超えた場合に、モデルを運用環境にデプロイします。最後に、この記事では、定期的なスケジュールに基づいてワークフローをトリガーする方法を示します。

次の図は、AWS Step Functions を利用した上記のサーバーレスワークフローを示しています。

この記事では、以下の AWS のサービスを使用します。

AWS Step Functions により、複数の AWS のサービスをサーバーレスワークフローに統合できます。1 つのステップの出力が次のステップへの入力として機能するワークフローを設計および実行し、エラー処理をワークフローに埋め込むことができます。
Amazon SageMaker は、さまざまなタイプの ML モデルを構築、トレーニング、デプロイするためのツールを開発者とデータサイエンティストに提供する完全マネージド型サービスです。
AWS Glue は、抽出、変換、ロード (ETL) を行える完全マネージド型サービスです。AWS Glue をサポート対象のデータストアに対して使うと、コードを生成して抽出し、ターゲットデータストアにロードします。AWS Glue は分散 Apache Spark 環境で実行されます。これにより、インフラストラクチャを管理せずに Spark を利用できます。
AWS Lambda は、サーバーをプロビジョニングまたは管理することなく、コードを実行できるコンピューティングサービスです。Lambda はトリガーされたときにのみコードを実行し、1 日あたり数回のリクエストから毎秒数千回ものリクエストまで自動的にスケーリングします。
Amazon EventBridge は、さまざまな SaaS アプリケーション、AWS のサービス、およびアプリケーションのデータを簡単に接続できるサーバーレスイベントバスです。

SDK の概要

SDK は、AWS Step Functions を使用する新しい方法を提供します。Step Functions は、一連の個別のステップで構成されるステートマシンです。各ステップは、作業の実行、選択、並列実行の開始、またはタイムアウトの管理を行うことができます。個々のステップを開発し、Step Functions を使用して、ワークフロー全体のトリガー、調整、および状態を処理できます。Data Science SDK の前に、JSON ベースの Amazon States Language を使用して Step Functions を定義する必要がありました。SDK では、Python コードを使用して Step Functions を簡単に作成、実行、視覚化できます。

この記事では、Step Functions ステップの作成、パラメータの操作、サービス固有の機能の統合、およびこれらのステップをリンクしてワークフローを作成および視覚化する方法など、SDK の概要を説明します。記事を通してコード例をいくつか用意する一方、プロセス全体の詳細な Amazon SageMaker ノートブックも作成しました。詳細については、GitHub リポジトリをご覧ください。

ステップ、パラメータ、動的ワークフロー

Step Functions 内で、各ステップはその出力を次のステップに渡します。この出力を次のステップで使用して、動的なワークフローを作成できます。Step Functions を実行する度に入力パラメータを渡すこともできます。パラメータにより、一般的なワークフローを保つことができ、他のプロジェクトをサポートできます。

SDK を使用して、ワークフローに必要な入力パラメータを定義するには、次のコードを参照してください。

execution_input = ExecutionInput(schema={
    'TrainingJobName': str,
    'GlueJobName': str,
    'ModelName': str,
    'EndpointName': str,
    'LambdaFunctionName': str
})

組み込みのサービス統合

Data Science SDK は複数の AWS のサービスと統合します。統合により、API 呼び出しを記述することなく、サポートされているサービスを直接制御できます。この記事では、AWS Glue、Lambda、および Amazon SageMaker の統合を使用します。詳細については、「AWS Step Functions のサービス統合」を参照してください。

モデルの再トレーニングを行うには、まず最新のデータを取得する必要があります。生データを充実させる一方、ML モデルでサポートされているファイルの種類と場所に保存する必要もあります。AWS Glue は、ほとんどのデータストアに接続し、Python のカスタムスクリプトをサポートし、サーバーの管理を必要としません。AWS Glue を使用して、本番データストアからデータを読み取り、変換されたデータを Amazon S3 に書き込むことにより、ワークフローを開始します。

Data Science SDK では、AWS Glue ジョブをワークフローに簡単に追加できます。AWS Glue ジョブ自体は、データソースの場所、ETL の Python コード、および使用するファイルの宛先を指定します。SDK に必要なのは、GlueStartJobRunStep のパラメータとして使う AWS Glue ジョブの名前です。詳細については、YouTube で「AWS Glue ETL の使用開始」をご覧ください。

入力パラメータを使用して、実行時に AWS Glue ジョブを選択できるようにします。

etl_step = steps.GlueStartJobRunStep(
    'Extract, Transform, Load',
    parameters={"JobName": execution_input['GlueJobName']}
)

入力データを抽出して保存したら、SDK の TrainingStep を使用してモデルをトレーニングします。Amazon SageMaker は基礎となる計算リソースを処理しますが、お客様はトレーニング用のアルゴリズム、ハイパーパラメータ、およびデータソースを指定する必要があります。次のコードを参照してください。

training_step = steps.TrainingStep(
    'Model Training',
    estimator=xgb,
    data={
      'train': sagemaker.s3_input(train_data, content_type='csv'),
      'validation': sagemaker.s3_input(validation_data, content_type='csv')},
    job_name=execution_input['TrainingJobName']
)

前のコードの推定子 xgb は、XGBoost アルゴリズムとそのハイパーパラメータをカプセル化します。推定子の定義方法の詳細については、GitHub リポジトリをご覧ください。

Step Functions ワークフローは、トレーニングが完了するまでトレーニングステップに残ります。その後、トレーニング結果を取得して、新しいモデルの精度に基づいてワークフローを分岐できるようにする必要があります。Step Functions LambdaStep を使用して Lambda を呼び出し、Amazon SageMaker トレーニングジョブをクエリして結果を返す単純な Python 関数を実行します。SDK で Lambda ステートを追加するには、関数名とペイロードを指定します。この記事では、JSON パスを使用して Lambda 関数ペイロードの TrainingJobName を選択し、クエリするトレーニングジョブを認識します。次のコードを参照してください。

lambda_step = steps.compute.LambdaStep(
    'Query Training Results',
    parameters={"FunctionName": execution_input['LambdaFunctionName'],
        'Payload':{"TrainingJobName.$": "$.TrainingJobName"}
    }
)

トレーニング後にモデルをデプロイするには、SDK の ModelStep と EndpointConfigStep を使用して、トレーニングアーティファクトからモデルオブジェクトとデプロイ設定を作成する必要があります。次のコードを参照してください。

model_step = steps.ModelStep(
    'Save Model',
    model=training_step.get_expected_model(),
    model_name=execution_input['ModelName'],
    result_path='$.ModelStepResults'
)

endpoint_config_step = steps.EndpointConfigStep(
    "Create Model Endpoint Config",
    endpoint_config_name=execution_input['ModelName'],
    model_name=execution_input['ModelName'],
    initial_instance_count=1,
    instance_type='ml.m4.xlarge'
)

最後に、ワークフローは EndpointStep を使用して、新しいモデルをマネージド API エンドポイントとしてデプロイできます。「update」パラメータにより、新しいエンドポイントを作成するのではなく、既存の Amazon SageMaker エンドポイントを更新します。次のコードを参照してください。

endpoint_step = steps.EndpointStep(
    'Update Model Endpoint',
    endpoint_name=execution_input['EndpointName'],
    endpoint_config_name=execution_input['ModelName'],
    update=True
)

制御フローとリンク状態

Step Functions SDK の Choice ステートは、前のステップの出力に基づいた分岐ロジックをサポートしています。このステートを追加することにより、動的で複雑なワークフローを作成できます。

この記事では、Amazon SageMaker トレーニングステップの結果に基づいて分岐するステップを作成します。次のコードを参照してください。

check_accuracy_step = steps.states.Choice(
    'Accuracy > 90%'
)

ステップに分岐および分岐ロジックを追加します。Choice ステートは、複数のデータ型と複合ブール式をサポートします。ただし、この記事では、2 つの数値を比較します。最初は 0.90 の固定しきい値、2 番目は TrainingStep からの検証データセットのモデル精度です。トレーニング結果には、モデルのエラーが表示されます。これは、 (＃間違ったケース)/(＃すべてのケース) として計算します。その結果、測定誤差が 10% (.10) 未満の場合、モデルの精度は 90% を超えます。

詳細については、Choice Rules を参照してください。

次の比較ルールを追加します。

threshold_rule = steps.choice_rule.ChoiceRule.NumericLessThan(variable=lambda_step.output()['Payload']['trainingMetrics'][0]['Value'], value=.10)

check_accuracy_step.add_choice(rule=threshold_rule, next_step=endpoint_config_step)
check_accuracy_step.default_choice(next_step=fail_step)

選択ルールは、ルールが正常に合格した場合のワークフローの次のステップを指定します。ここまでで、ステップを作成しましたが、それらをリンクして実行順序を作成していません。SDK を使用して、2 つの異なる方法でステップをリンクできます。まず、next() メソッドを使用して、個々のステップの次のステップを指定できます。次のコードを参照してください。

endpoint_config_step.next(endpoint_step)

Chain() メソッドを使用して、複数のステップを一度にリンクすることもできます。次のコードを参照してください。

workflow_definition = steps.Chain([
    etl_step,
    training_step,
    model_step,
    lambda_step,
    check_accuracy_step
])

ワークフローの作成

すべてのステップを定義して順序付けした後、次のコードを使用して Step Functions 自体を作成します。

workflow = Workflow(
    name='MyInferenceRoutine_{}'.format(id),
    definition=workflow_definition,
    role=workflow_execution_role,
    execution_input=execution_input
)

workflow.create()

ワークフローを作成すると、workflow.render_graph() は、以下に示す Step Functions コンソールに表示されるものと同様のワークフローの図を返します。

これで、新しいデプロイパイプラインを実行する準備ができました。execute() メソッドにより SDK を使用してモデルを手動で実行するか、このタスクを自動化できます。

EventBridge トリガーを使用したワークフローのスケジューリング

EventBridge トリガーを使用して、ワークフローのスケジュールを立てることができます。この記事では、EventBridge 内でルールを作成して、設定されたスケジュールでターゲット Step Functions を呼び出す方法を示します。詳細については、AWS Resource の「イベントでトリガーする EventBridge ルールを作成する」を参照してください。

次の手順を実行します。

AWS マネジメントコンソールの Services で、[Amazon EventBridge] を選択します。
[Rules] を選択します。
[Create rule] を選択します。
Name and description の Name で、ルールの名前を入力します。この記事では、automate-model-retraining-trigger という名前を入力します。
オプションのステップとして、Description には、ステップの説明を入力します。
Define pattern では、[Schedule] を選択します。
Fix rate every では、[1 Hours] を選択します。
Select event bus で、[AWS default event bus] を選択します。
[Enable the rule on the selected event bus] を選択します。
Select targets の Target で、[Step Functions state machine] を選択します。
State machine では、マシンを選択します。
[Configure input] を選択してから、[Constant (JSON text)] をクリックします。
入力パラメータを JSON テキストとしてワークフローへ入力します。
[Create a new role for this specific resource] を選択します。
ロールの名前を入力します。既存のロールがある場合は、代わりに [Use existing role] を選択します。
[Create] をクリックします。

まとめ

この記事では、AWS Step Functions Data Science SDK for Amazon SageMaker の概要を説明しました。Python を使用して再利用可能なデプロイモデルワークフローを作成する方法を示しました。ワークフローには、データを抽出および変換する AWS Glue ジョブ、新しいデータで ML モデルをトレーニングするトレーニングステップ、トレーニング結果をクエリする Lambda ステップ、モデルアーティファクトを作成するモデルステップ、デプロイパラメータを定義するエンドポイント設定ステップ、および更新されたモデルを既存のエンドポイントにデプロイするエンドポイントステップが含まれていました。この記事では、EventBridge を使用して、指定されたスケジュールに従ってワークフローを自動的にトリガーする方法の概要も説明しました。

SDK に関連する他の技術ドキュメントとノートブックの例については、AWS Step Functions Data Science SDK for Amazon SageMaker 発表ページをご覧ください。

ご質問またはご提案などについては、以下からコメントを残してください。

著者について

Sean Wilkinson は、サーバーレスと機械学習を専門とする AWS のソリューションアーキテクトです。

Julia Soscia は、ニューヨーク市を拠点とするアマゾンウェブサービスのソリューションアーキテクトです。彼女が力を入れているのは、お客様が AWS クラウドプラットフォーム上で Well-Architected 環境を作成できるようにすることです。彼女は経験豊富なデータアナリストで、分析と機械学習に重点を置いています。

Amazon Web Services ブログ