Fargateでタスクを定時実行してSageMaker推論エンドポイントを更新するcfnテンプレートファイルを作成した

概要

当エントリは『クラスメソッド Amazon SageMaker Advent Calendar 2019』1日目のエントリです。

機械学習のワークフローも色々と実現方法があると感じているのですが、ささっとシンプルにやるのなら「ECSタスク定義(Fargate)をCloudWatchEventsで定時実行する」というのが一番いいのかなと思い、今回は「前処理→SageMakerでのトレーニングジョブ→推論エンドポイントの更新、をFargateで定時実行する」cfnテンプレートファイルを作成してみました。

シンプルですが、「Cronで実行タイミングを柔軟に指定できる」、「Lambdaと違って実行時間制限に怯えなくていい」、「Fargateはサーバレスなのでサーバの管理をしなくていい」等の嬉しいポイントがあるので、簡単なバッチJOBくらいだったらこれで十分に対応できます。

目次

1.やること

SageMakerのチュートリアルの一つである「DeepAR-Electricity.ipynb」を参考に、下記のリソースを展開するcfnテンプレートファイルを作成しました。

  • 指定したタイミングでトリガーを発火する「CloudWatchEvents」
  • 「ソースデータの前処理→モデルの学習→推論エンドポイントの更新」を実行する「ECSタスク定義(Fargate)」
  • 「ECSタスク定義(Fargate)」に必要なVPC
  • 「CloudWatchEventsRules」、「ECSタスク定義(Fargate)」に必要なIAMロール

ちなみに、今回はSageMakerのトレーニングJOBが終了したらそのまま推論エンドポイントを更新するのですが、「SageMakerの推論エンドポイントの更新はダウンタイム無しで実現できる」という点は重要です。

すでに実稼働環境にデプロイされているモデルを停止中の状態にすることなく、エンドポイントを変更することができます。たとえば、新しいモデルバリアントを追加したり、既存のモデルバリアントの ML コンピューティングインスタンス設定を更新したり、モデルバリアント間のトラフィックの分散を変更することができます。

参照:Amazon SageMaker ホスティングサービスでモデルをデプロイする

また、2つ以上のインスタンスをデプロイすると異なるAZに起動してくれたりと、推論エンドポイントは運用する上では地味に嬉しいことを裏側で配慮してくれます。

各本番稼働用バリアントについて、デプロイする ML コンピューティングインスタンスの数を指定します。2 つ以上のインスタンスを指定すると、Amazon SageMaker はそれらを複数のアベイラビリティーゾーンで起動します。これにより、継続的な可用性が保証されます。

参照:Amazon SageMaker ホスティングサービスでモデルをデプロイする

ちょっと話が逸れてしまいましたが、話を実装の方に戻します。

2.技術要素のポイント

今回使う要素は下記の3点です。

CoudWatchEvents

ECSタスク(Fargate)を定時実行するために利用しています。
Cron式でどのタイミングで実行するか、を指定できるのでとても便利です。
Lambdaの起動に使っている方も多いのではないでしょうか?

You can create rules that self-trigger on an automated schedule in CloudWatch Events using cron or rate expressions. All scheduled events use UTC time zone and the minimum precision for schedules is 1 minute.

参照:Schedule Expressions for Rules

ECS

処理自体はコンテナで実行するので、ECSを使っています。
ローカルで開発したコンテナイメージをECRリポジトリにPushして、ECSタスク定義でイメージや利用するリソース量を指定して「タスク定義を実行」することで今回やりたいことを実現しています。

タスク定義の中で起動モードをFargateとして実行することで、EC2のようなサーバーを管理する必要がなくなり、実行した分だけの課金で済むのでありがたいです。

AWS Fargate は、サーバーやクラスターの管理の必要なしにコンテナを実行するための、Amazon ECS に対応したコンピューティングエンジンです。AWS Fargate を使用すると、コンテナを実行するために仮想マシンのクラスターをプロビジョニング、設定、スケールする必要がありません。これにより、サーバータイプの選択、クラスターをスケールするタイミングの決定、クラスターのパッキングの最適化を行う必要がなくなります。AWS Fargate により、サーバーやクラスターの操作や検討が不要になります。

参照:AWS Fargate

機械学習に使えそうなコンテナイメージもいくつかAWSが用意してくれているので、もしよかったらこちらも見てみてください。

SageMaker

AWS上で機械学習をするにあたってのプラットフォームです。
自前のアルゴリズムを使えるのはもちろんとして、機械学習の開発者用インフラ準備ハイパーパラメータのチューニングいくつかのアルゴリズムを手軽に使える機能や、モデルのデプロイアノテーションの仕組みマーケットプレイスに公開されているアルゴリズムやモデルの利用等、幅広いサービスが提供されています。

Amazon SageMaker は、すべての開発者とデータサイエンティストに機械学習モデルの構築、トレーニング、デプロイ手段を提供します。Amazon SageMaker は、機械学習のワークフロー全体をカバーする完全マネージド型サービスです。

参照:SageMaker

今回は「モデルの学習」、「モデルのデプロイ」の機能を利用しています。

3.やってみた

今回のテンプレートファイルでデプロイされるリソースの内容はこちらの通りです。

処理の流れとしては、下記の通りとなります。

  • 1.CLoudWatchEventsで指定した時間にECSタスク定義(Fargate)を実行する
  • 2.Fargateが起動して、ECRリポジトリからコンテナイメージを取得する
  • 3.Fargateにて処理が実行される
  • 4.この処理では「S3からデータを取得して前処理した結果をS3に出力」→「SageMakerでモデルを学習」→「指定した名前の推論エンドポイントが存在する場合はアップデート、無い場合は新規作成」するといった順番に処理を実行する

もう少しイメージを掴むために、実際にJOBが動いているところのコンソール画面の画像を載せておきます。

今回のcfnテンプレートファイルでは、CloudWatchEventsで指定したルールを無効にしているので、実行タイミングを有効化すると動きます。

「CloudWatchEventsRule」で実行できると、下記のように「ECSクラスター」の画面から確認できます。
「開始元」を見ると、「CloudWatchEventsRule」から実行されていることがわかりますね。

前処理が終わると、SageMakerのトレーニングジョブが実行されます。

初回実行時は、推論エンドポイントを新規作成するようにしているので、エンドポイントを新規作成しています。

2回目以降の実行時(同名の推論エンドポイントが存在する時)は推論エンドポイントをUPDATEするようにしています。

シンプルな作りですが、ちょっと何かしらを試しに定時実行するバッチにしたい時にはこのくらいでも十分じゃないでしょうか?

4.cfnテンプレートファイルについて

cfnテンプレートの中身も少しだけ説明します。

4-1.「CloudWatchEventsRules」で「ECSタスク定義」を定時実行(cronベース)

CloudWatchEventsRulesでスケジュールを指定し、ECSのタスク定義を実行します。
(下記はコンソール画面上からルールを定義するイメージ)

今回のスクリプトではcfnテンプレートファイルの中にIAMロールやスケジュールが指定されています。

ポイントとしては下記の2点です。

  • パブリックサブネットで実行する場合は、「AssignPublicIp」はENABLEDにする
  • CloudWatchEventsRulesが実行主体なので、IAMロールの信頼関係を指定する
  ECSEventRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub ${StackName}-ECSEventRole
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - events.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AmazonEC2ContainerServiceEventsRole
  TaskScheduleEvents:
    Type: AWS::Events::Rule
    Properties:
      Description: "my Fargate Schedule"
      ScheduleExpression: cron(40 15 ? * SAT *)
      Name: !Sub ${StackName}-schedule
      State: DISABLED
      Targets:
        - Arn: !GetAtt 
            - EcsCluster
            - Arn
          RoleArn: !GetAtt 
            - ECSEventRole
            - Arn
          Id: ID123
          EcsParameters:
            TaskCount: 1
            TaskDefinitionArn: !Ref EcsTaskDefinition
            LaunchType: FARGATE
            NetworkConfiguration:
              AwsVpcConfiguration:
                AssignPublicIp: ENABLED
                Subnets:
                  - !Ref FrontendSubnet
                SecurityGroups:
                  - !Ref VPCDefaultSecurityGroup

4-2.Fargateのタスク定義を実行する(run task)

続いて、「ECSタスク定義」をFargateで実行します。
Fargateを実行するためのVPCやIAMロール、ECR等を作成する必要があります。
IAMロールには、SageMakerのアクセス権限を忘れずに付与しましょう。

  EcsTaskRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub ${StackName}-EcsTaskRole
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - ecs-tasks.amazonaws.com
                - sagemaker.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess

5.まとめ

「とりあえず実行するだけなら今回のような構成が一番シンプルそうだしcfnテンプレートまでできたら便利だな」と感じたのでcfnテンプレートファイルを作成しました。
コンテナイメージもAWSが提供しているDeep Learning Containers イメージを利用したので、準備がとても楽でした。

最近、コンテナ周りのアップデートが多いので、アップデートに追随してより良い使い方を追求していきたいですね。

6.参考