AWSをフル活用、リクルートのクラウドネイティブな機械学習プラットフォームの裏側：コンテナベースの機械学習基盤を大解剖（2）

リクルートのエンジニアが内製している機械学習基盤について詳しく解説していく本連載。第2回はクラウドネイティブな機械学習に活用しているAWSのサービスや設計について。

» 2022年08月01日 05時00分公開

この記事は会員限定です。会員登録（無料）すると全てご覧いただけます。

　前回はリクルートで機械学習プラットフォームを構成する「Crois」「Crafto」の2つを内製した理由と設計思想を紹介しました。第2回はジョブスケジューラーであるCroisの詳細な設計とCraftoでの活用を紹介します。

Croisでインフラ面の工夫が必要な部分

　Croisで工夫が必要な部分は2点あります。

高いスケーラビリティを持つこと
通常の使用において高いセキュリティを実現すること

　スケーラビリティの面では、巨大な1つのインスタンスを利用する大きなタスクから大量のインスタンスを並べて処理するタスクまで、縦にも横にも高いスケーラビリティが必要です。リクルートで稼働しているワークロードにおいても、Amazon EC2のx1インスタンスタイプを使うような単一で高vCPU、高メモリのものもあれば、c5インスタンスタイプのものを100台同時に動かすような並列度の高いものもあります。

　必要なリソースはビジネス要件次第で日々変化しますし、特定の時間帯に集中することも多々あります。使われていないときは最小限のインスタンスで稼働し、処理をする際は十分な処理能力を提供する、柔軟にスケールアップ、スケールアウトの両面をサポートできるアーキテクチャが必要です。このようなアーキテクチャを構築できるのはクラウドを使うことの利点ですし、醍醐味（だいごみ）でもあります。

　セキュリティの面では、データ漏えい対策としての暗号化はもちろん、チームを超えたアクセスができないような権限の仕組みが重要です。Croisはジョブスケジューラーという性質上、バッチの実行に必要なクラウドのシークレットキーやAPIキーなどを保存することが考えられます。チームを超えたアクセスができてしまうと、社員の異動や退職時に膨大な数のキーコンフィグの変更が必要となり、運用が回らなくなります。これを防ぐためにファイルや機密情報は暗号化した上で、特定のチームの権限でしか復号できない状態にする必要がありました。

　もう1つ重要なポイントは、Croisの開発者であっても内容へのアクセスが不可になっている必要があることです。もし開発者が機密情報を見られてしまうと同様の事態が発生します。開発者の異動や退職などの影響範囲は最小化されていなければなりません。

Croisのスケーラビリティを支える技術

　第1回の記事でCroisはクラウドインフラとしてAWS（Amazon Web Services）を採用し、ジョブ実行部分にはAWS Step Functions（以下、Step Functions）を利用していると述べました。Step FunctionsはAmazon State Language（ASL）で記述されるState Machineをフロー定義とし、その中で定義されたタスクを定義された通りに実行できるマネージドサービスです。Croisはユーザーが記述したワークフローをASLに変換し、それをState Machineとして登録したものを実行することでジョブ実行を実現しています。

　要件となっているスケーラビリティを実現するために、本体からジョブ実行に必要なフロー制御を切り離してStep Functionsに完全に移譲しています。Crois本体が関与してしまうと、ジョブの実行数が増えたときにCrois本体をスケールさせる必要が出てきてしまい運用工数が増えるためです。またCrois本体の障害が直接ジョブの失敗につながります。ジョブ実行のフロー制御をStep Functionsに移譲することは、スケーラビリティの確保だけでなく耐障害性を高めることも実現しているのです。

　CroisでどのようにStep Functionsを使ってジョブを実行しているのか紹介します。Crois上のワークフローとState Machineの対応を以下の図に示します。

　Crois上ではとてもシンプルなワークフローが、State Machine上では複雑なフローになっていることが分かるかと思います。State Machine上の各ステップでパラメーターのロード／ストア、パラメーターの注入処理やコンテナを実行するバックエンドシステムの選択、エラーコードによるリトライへの振り分けといった処理が行われるためです。こうした処理の多くはStep Functionsに組み込まれた機能では難しいため、Step FunctionsからAWS Lambda（以下、Lambda）の関数を呼ぶことで実現しています。

　Step FunctionsからCrois本体に何らかの情報を同期したいことがあります。例えば実行したタスクの終了ステータスをCrois本体に送信して、Web UIで終了を表示するなどです。しかし、Step Functions内からCrois本体と直接通信してしまうと、Step FunctionsはCrois本体の障害に関係なく実行され続けるという利点がなくなってしまいます。問題を解決するためにAmazon EventBridge（以下、EventBridge）とAWS Lambda、Amazon DynamoDBを組み合わせて情報の同期を行っています。

　EventBridgeはAWS内のサービスの、例えば、Step Functionsの実行の終了やコンテナタスクのステータス変更といったイベントを検知すると、あらかじめ設定したルーティングルールに従ってデータを送信できるサービスです。その送信されたイベントデータを基に、Lambda関数が起動してDynamoDBに書き込むことでCrois本体へ情報を同期しています。ここでAmazon RDS（以下、RDS）のようなリレーショナルデータベース（RDB）を使わずにDynamoDBを使っている理由は、DBのスケーラビリティの差にあります。一度に大量のジョブが実行された場合には書き込みが大量に発生する可能性がありますが、RDSでは書き込み性能の急な変更は難しい印象です。DynamoDBはキャパシティーを事前に設定しなくともオンデマンドで対応してくれます。そのためDB性能がボトルネックになることなくStep Functionsのスケーラビリティを最大限に生かすことができるのです。

　Croisでユーザーから指定された処理を実際に実行するのはStep Functionsから起動されたコンテナです。Step Functionsのスケーラビリティを生かし、要件にあったような柔軟なスケーリングを実現するのは、マネージドのコンテナ実行サービスで実現できます。

　CroisではAWS BatchとAmazon ECS on AWS Fargate（以下、ECS Fargate）という2種類のコンテナ実行サービスを使い分けています。両者はともに定義したコンテナタスクを実行できるサービスですが、Croisにとって大きな差となるのは実行されるバックエンドの違いから最大処理性能や起動時間に違いがあることです。

　AWS Batchは要求リソースに対してAmazon EC2インスタンスを立ち上げてそのインスタンス上で処理を実行します。そのためCrois上のタスクが大量のCPU性能やメモリ、ディスクを消費するような場合であってもインスタンスの性能さえ確保できれば対応可能です。しかしインスタンスを立ち上げるという仕組みのためタスクが始まるまでに時間がかかるという問題があります。そうした性質があるためCroisではAWS Batchを機械学習の学習処理のような非常に重たいワークロードで時間がかかる処理に使用しています。

　これに対してECS FargateはAWS側であらかじめ用意されたインフラ上にコンテナを立ち上げる形になります。立ち上げるコンテナの性能はあらかじめ用意されたタイプから選ぶことになりAWS Batchほどの高スペックのものはありません。ただし既に準備されたインフラ上で起動するため起動速度が速いというメリットがあります。こうした性質からデータウェアハウス（DWH）へのクエリ実行などのコンテナ自体に処理性能が必要ないような処理に対して使われています。

　このようにCroisではStep Functionsのスケーラビリティを最大限生かせるように、自動的にスケールするマネージドサービスを組み合わせているのです。

Croisのセキュリティを支える技術

　次に、冒頭で述べたセキュリティの要件を満たすための構成を紹介します。Croisではプロジェクトという単位で権限を管理しており、案件やプロダクトの単位で切られることが多いです。

　プロジェクトにはそれぞれ固有のAWSのIAM（Identity and Access Management）ロールとKMS（Key Management Service）キーを発行しています。これがセキュリティの肝となる部分です。IAMロールはAWSサービスや、その上で動くアプリケーションに対してIAMポリシーと組み合わせることで他のAWSサービスへのアクセス権を管理することができる仕組みです。KMSキーは暗号化キーで、KMSキーには同じようにポリシーを指定することで暗号化や復号といった操作の権限を制限できます。

　CroisのプロジェクトにひもづくKMSキーはひもづくIAMロールでしか復号の操作ができないようになっており、このキーで暗号化されたデータに対して他のプロジェクトにひもづくIAMロールが復号を試みても復号できません。

　Croisはユーザーが作成したコンテナに対して実行時にエージェントと呼ぶ実行バイナリファイルを注入し、そのエージェントを介してユーザーの作成した処理を呼び出しています。エージェントを介して処理することでCroisの処理をコンテナ内の処理に介在させることが可能となり、セキュリティ面でも一役買っています。AWS BatchとECS Fargateでは注入方法が異なります。ECS Fargateを例に説明します。

　ECS Fargateでは複数のコンテナを同じグループとして立ち上げることができ、コンテナタスクの実行時にエージェントの入ったコンテナをセットで起動します。エージェントを共有ボリュームにコピーし、モジュールコンテナのコマンドを上書きしてエージェントを起動するようにしています。エージェントはCrois本体に必要な処理を行った後にユーザーの作成した本来の実行ファイルを実行します。

　エージェントとIAMロール、KMSキーが連携してコンテナの成果物である出力ファイルをどう暗号化しているのか見ていきましょう。エージェントはユーザーの作成した処理が終わると成果物のファイルをオブジェクトストレージであるAmazon S3（以下、S3）へアップロードします。

　このときにサーバ側の暗号化オプションを利用して暗号化キーとしてCroisプロジェクトに割り当てられたKMSキーを指定することでファイルが暗号化された状態でS3へ保存されます。後続のコンテナタスクでS3に保存されたファイルを利用したい場合はエージェントがS3からファイルをダウンロードします。このコンテナタスクにはアップロードに使われたものと同一のIAMロールが割り当てられており、KMSキーのポリシーで許可された復号が行われます。そのため暗号化されたファイルを復号した状態でダウンロードできます。これが他のCroisプロジェクトのコンテナタスクになると、割り当てられたIAMロールが異なるため復号できず、ファイルもダウンロードできないという流れです。

　DBのパスワードなど秘密情報は事前にユーザーがCroisのAPIを使ってCrois本体に登録します。Croisは送信された情報をDBには保存せずAWS Secrets Managerに送信します。AWS Secrets Managerは秘密情報を管理、保存できるマネージドサービスです。保存時にCroisプロジェクトに割り当てられたKMSキーを指定することで暗号化して保存されます。

　先ほどと同様にコンテナタスクの開始時にエージェントは必要な情報をAWS Secrets Managerからダウンロードします。KMSキーのポリシーで保存したプロジェクトと同一のプロジェクト内のコンテナタスクであれば復号された状態で情報を取得できますが、違うプロジェクトであれば取得できません。Croisの開発者であってもKMSポリシーに指定されたIAMロールとは一致しないため復号できない状態で保存されています。

　このようにCroisではクラウドサービスがする暗号化キーとそのポリシーをうまく利用することで強固なセキュリティを実現し、社内で横断して利用されても問題のない状態にしています。

Craftoの構成

　それではCroisを活用したCraftoのシステム構成を紹介します。

　機械学習アプリケーションのCraftoは、ワークフローエンジンとしてCroisを活用しているプロダクトで、データセットの登録やモデル構築などの機能をWebアプリから実行することができます。それぞれの機能は、基本的にCroisの存在を意識することなくボタン一つで実行できます。内部の構成はAPIを介してCroisワークフローを利用するものとなっています。

　機械学習システムにおけるワークフローでは、コンピュータリソースの柔軟性は重要です。大小さまざまなデータセット、複雑さの異なるモデル、並列での実行など、Croisワークフローのスケーラビリティの高さは、機械学習システムの多くのシーンで役立ちます。

　Craftoでは、機械学習機能のほとんどをCroisワークフローで実行しているため、モデル構築が大量に実行されるような場合でも、リソースの枯渇でサービスが停止するような障害が発生したことはありません。データセットのサイズに合わせて自動でコンピュータリソースを設定するため、パフォーマンス効率も考慮されています。

　Craftoの機械学習機能は、データセット登録や予測機能などがありますが、中心となる機能はモデル構築です。モデル構築では、データセットのダウンロード、データの分割、学習、評価、モデルのデプロイなどの複数ステップを実行するワークフローを作成します。このとき、CraftoからはCroisのワークフロー作成およびジョブ実行のAPIを利用します。ジョブ実行時には、それぞれのステップに対応したパラメーターを渡すことで、さまざまな設定が可能です。例えば、データ分割のステップでは分割手法や（訓練データと検証データの）分割割合を指定でき、学習のステップでは最適化指標の設定ができます。これらのパラメーターを基にCroisワークフローの実行が完了すると、ハイパーパラメーターの探索結果や各種スコア、構築済みモデル、特徴量重要度などの評価結果がS3に出力されます。Crafto側では、各ステップの終了時にこれらの結果を取得し、順次画面上に表示していきます。

　機械学習機能をCroisワークフローとして切り出す構成は、コンピュータリソースの柔軟性があるという利点を挙げましたが、それ以外にもWebアプリエンジニアとML（機械学習）エンジニアでの開発範囲が明確になることで、機能の分業ができることも大きなメリットです。それぞれの機能は独立して開発可能なため、各機能間の依存性は小さく、テストも容易です。開発効率や保守性という観点でも優れた構成であると考えています。

　次回の記事では、Croisを活用していて起きた障害や問題を例に挙げながら、クラウドネイティブなアプリケーションで起こりうるトラブルとその対応策、意識すべき点をご紹介します。またCroisおよびCraftoを社内でさらに利用してもらうために目指している方向性も紹介します。