Fargateタスクで機械学習モデルの訓練・推論をしたのでポイントを整理してみた

2022.08.02

データアナリティクス事業本部の鈴木です。

Fargate起動タイプを使用したタスクでXGBoostを動かしたいことがありました。FargateタスクはLambdaより長い時間実行できるのでバッチ処理を動かす際に便利です。いざやってみるとプライベートサブネットで動かす場合には、関連サービスとの連携のための設定が少し難しく感じたので、ポイントをまとめてみました。

Fargateについて

Amazon ECSの、コンテナをサーバレスで実行できる機能です。

データ分析基盤や機械学習システムでは、15分以上を超える可能性があるバッチ処理は頻繁に登場するので、その要件を実現できる心強い機能です。AWS Fargate を使用して、イベント駆動型およびスケジュールされたワークロードを大規模に実行のドキュメントでも

このパターンは次のビジネスユースケースに役立ちます。

・ランタイム(15 分の制限)またはメモリに関する制限のために AWS Lambda では実行できない、スケジュールされたイベント駆動型のワークロードをビジネスで実行する必要がある場合

のようにFargateのユースケースが紹介されています。

一方で、コンテナレジストリとの通信が必須なため、プライベートサブネットで実行するようなケースでは、通信経路や権限面で要件に気をつける必要はあります。以前からFargateを使って機械学習のモデルを動かしてみたいなと思っていたので、この要件について実際に動かしてみてこの記事でまとめてみました。

なお、現在はFargateでGPUを使えないので、GPUを使わず動作させるモデルやユースケースを想定しています。

エフェメラルストレージは最大で200GiBまで設定できるため、ある程度大きなデータを扱う際でも安心です。

構成

以下の構成でFargateでXGBoostのサンプルコードが実行できるかを確認してみました。

検証した構成の図

環境構築と動かしたコードについては、以下のGitHubレポジトリに格納しておきました。

動かしてみた

まず、S3バケットに学習データ、XGBoostのハイパーパラメータを記載したJSONファイル、推論対象のデータを配置しておきました。 データを配置するバケットはタスク定義の環境変数に合ったものとしておきました。

最初にS3に配置したオブジェクト

学習データ、推論対象のデータは実際はAthenaの検索結果などをデータマートの位置付けにあるキーに配置しますが、今回は簡単のためタスク定義で事前に設定してあるキーに配置しておきました。

また、ECRに学習および推論用のDockerイメージをPushしておきました。

DockerイメージのPush

次に学習用のクラスターの画面のタスクタブから、新しいタスクの実行を押してタスクの実行に必要な情報を入力していきます。

クラスターからの実行

タスクの実行セクションで以下のように入力しました。ほかはデフォルトとしました。

  • 起動タイプ: FARGATE
  • オペレーティングシステムファミリー: Linux
  • タスク定義およびリビジョン: CloudFormationで作成したもの

タスクの実行セクション

VPCとセキュリティグループのセクションでは、タスクを起動するVPC・サブネット・セキュリティグループを指定しました。

VPCとセキュリティグループのセクション

入力としては以上で、右下のタスクの実行ボタンを押すとタスクが実行されました。

クラスターの画面に戻ると必要なタスクのステータスRUNNINGであるタスクができていることが分かります。タスクをクリックすると詳細を確認できます。

Runningのタスク

詳細タブでは、プラットフォームのバージョンやどのタスク定義・タスクロールを使っているかなどが確認できました。

タスクの詳細

また、ログ設定セクションからは作成されたログストリームに直接飛ぶこともできるので、エラーの内容も確認しやすかったです。

ログの設定

タスクが正常に終了すると、期待通りS3に学習済みモデルを出力できました。

学習用コードの終了

推論用コンテナも同様に実行し、推論結果がアップロードされることを確認しました。

推論用コードの終了

ネットワーク要件のポイント

Fargateを使ってコンテナを実行する場合のネットワーク要件は、以下に記載があります。こちらの記載を参考にしつつ、動作確認をしていくと良いです。

実行環境ごとに、必要だったリソースについて記載します。

プライベートサブネットで実行する場合

必要なエンドポイントは以下のドキュメントを参考に作成していくことになります。

私の検証環境では以下のエンドポイントを作成しました。

  • ECRからのイメージのプル
    • com.amazonaws.リージョンコード.ecr.dkr
    • com.amazonaws.リージョンコード.ecr.api
    • com.amazonaws.リージョンコード.s3
  • CloudWatch Logsへのログの送信
    • com.amazonaws.リージョンコード.logs
  • S3からのデータやモデルのやりとり
    • com.amazonaws.リージョンコード.s3(実際には、上記ECR用と同じもの)

また、インターフェース型のエンドポイントには、Fargateタスクからアクセスできるようなセキュリティグループを付与しておきました。今回の環境だと、タスクに付けるセキュリティグループと同じものをエンドポイントにもつけておき、同じセキュリティグループ間では全ての通信を許可するようにしておきました。

Fargateタスクをプライベートサブネットで実行した際にNAT無しで通信を実現するのに必要なエンドポイントは、以下の記事に詳しくまとまっているのでご参考ください。

パブリックサブネットで実行する場合

基本的にはパブリックIPの自動割り当てENABLEDになっており、タスクがパブリックIPアドレスを受け取るようにしていれば実行することが可能でした。

パブリックIP設定

それでも動かない場合は、パブリックサブネットにインターネットへのルートが何かしらの理由でないか、Fargateからの通信が遮断されていることが考えられそうです。

IAMロールのポイント

タスクを実行するために、2種のロールを作成して適切にタスク定義に設定しておく必要があります。

ロール 役割 アタッチするポリシー例
タスク実行ロール ECSコンテナエージェントで使用する AmazonECSTaskExecutionRolePolicy
タスクロール コンテナで使用する コンテナでしたいアクションを許可したポリシー

IAMロールついても、以下の記事に詳しくまとまっているのでご参考ください。

タスク実行ロールの注意点としては、AmazonECSTaskExecutionRolePolicyにはCloudWatchのログストリームを作成しログをPUTする権限はありますが、ロググループを作る権限はないため、ロググループをあらかじめ作成しておくか、タスク実行ロールに追加でlogs:CreateLogGroupを付与する必要があります。ロググループがない場合はタスクが失敗してしまうため注意が必要です。

タスクロールは、我々がコンテナ内の機械学習関連のコードに許可したい権限を付与するところです。S3などからデータを取得したり、学習済みのモデルをアップロードしたりするための権限を付与しておきましょう。

何かエラーが出てしまった場合は、新しいタスクの実行をコンソールから入れていくような場合は、詳細オプションセクションにアタッチされるIAMロールが記載されているので、意図したものになっているか確認するとトラブル解決がしやすかったです。私の場合は、CloudFormationからロール含めて作成していたので、どちらがどちらか分からなくなり、入れ替わって動かなくなっていたりしました。

詳細オプションセクション

タスクを別のリソースから実行する場合のポイント

データ分析基盤・機械学習システムでは、多くのケースでタスクはワークフローエンジンがインストールされた別のコンピューティングリソースなどから実行されると思います。この場合、コンピューティングリソース側にタスクロールとタスク実行ロールのPassRoleの権限とecs:RunTaskを付与しておく必要があります。

例えばEC2上にワークフローエンジンがいる場合は、EC2にアタッチしたIAMロールに以下のような権限をつけておくことになります。

{
    "Effect": "Allow",
    "Action": "ecs:RunTask",
    "Resource": [
        "*"
    ]
},
{
    "Action": [
        "iam:PassRole"
    ],
    "Effect": "Allow",
    "Resource": [
        "タスク実行ロールのARN",
        "タスクロールのARN"
    ]
}

また、多くの場合、Fargateタスクを起動する場所とワークフローエンジンの乗ったリソースがある場所は同じプライベートサブネットだと思うので、タスクの起動用にエンドポイントを準備していて、ワークフローエンジンがインターネットに出る経路がない場合は、追加でcom.amazonaws.リージョンコード.ecsを該当サブネットに作成しておくようにしましょう。

例として、ワークフローエンジンから実行する場合の図だと以下のようにするとboto3で新規タスクを実行できました。赤字が追加のリソースです。

ワークフローエンジンから実行する場合

パフォーマンスチューニングについて

Fargateではタスクに割り当てるリソース(CPUの値・メモリの値など)を変更することができます。

検証では、まずはだいたいこれくらい必要だろうという値を指定してタスクを起動し、メモリが溢れてしまうような現象が起きれば徐々に大きくして、十分そうな大きさに上げてみました。

このとき重要なのは、CPUの値とメモリの値に許容される組み合わせがあることです。

例えば、記事執筆時点ではCPUの値が256の場合、メモリの値は512 MiB1 GB2 GBから選ぶことになります。またオペレーティングシステムはLinuxである必要があります。

ほかの組み合わせについては以下のドキュメントに対応表があるのでご確認ください。

タスク定義パラメータ タスクサイズ - Amazon ECS

実験管理について

学習時の実験管理については、Fargateの場合はあくまでもコンテナを実行してくれるだけなので、必要であれば自分で作る必要があります。

環境をAWSで揃える条件だと、選択肢はおおよそ以下です。

  • S3かCloudWatch Logsに出力するようにする
  • MLflow環境を構築して連携する
  • TensorBoard環境を構築して連携する

始めのうちはとりあえずS3かCloudWatch Logsに出力するのでも問題ないかもしれませんが、より便利なツールを使って可視化や比較をしたくなると思います。

MLflowの場合は以下のブログ記事にてFargateでMLflowサーバーを立てる構成が紹介されていました。

またTensorBoardの場合は以下のブログ記事に構成の例が紹介されていました。

本格的にしたい場合はMLflow環境を構築した方がいいと思いますが、RDSなどの準備も必要になるため、取り急ぎはTensorBoardの方がお手軽かもしれません。

特にXGBoostなどのモデルでTensorBoardで可視化できる形式でメトリックを出力する例は、以下の記事でご紹介しました。

最後に

今回はFargateタスクで機械学習モデル、特にXGBoostを動かす例と、そのために必要な各種要件をご紹介しました。いろいろと注意する点がありましたが、ここまで分かれば、かなり楽に導入できるのではないかと思います。

Fargateタスクを使って手元のモデルをAWSで動かしてみたい方は是非チャレンジしてみてください!

参考資料

文中に記載していない参考資料を記載します。