時間のかかる推論をSQSとワーカーでどうにかするインフラ構築

f:id:tawamura:20191218120558p:plain

この記事はBASE Advent Calendar 2019の22日目の記事です。

devblog.thebase.in

こんにちは、Data Strategyチームのid:tawamuraです。BASEには今年の8月に入社し、今月で5ヶ月目になります。

DSチームでは、ネットショップ作成サービス「BASE」のデータを集計し、機械学習など様々な利活用を行なっています。主に挙げられるのはおすすめ商品のレコメンデーション、商品のカテゴリ・属性推定、不正利用の検知などです。
弊社では、DSチームが作成したこれらの学習モデルを、DS側で立てたAPIサーバを介して利用してもらう形を取り入れています。ALBを通したAPIプロキシサーバが、各推論サーバ(ECSインスタンスやLambda)にリクエストを投げるようにしています。

基本的にDSチームの立てたAPIのレスポンスは即時で返されています。しかし、今回私が作成した推論APIは、実行にかなりの時間がかかってしまうものでした。リクエストで受け取る単一のIDについての結果を返すために、推論で必要な素性を都度RDSから収集しているのが直接の原因です。しかもリクエスト内容によって素性となる対象データ量が異なるため、実行時間が数秒〜数十秒まで変動してしまいます。
幸いこのAPIは元々リアルタイム性が求められるものではなかったので、ひとまず実行時間を受け入れつつ、この重い処理をどうにかして扱えるように色々と試行錯誤した、という話を紹介できればと思います。 ちなみに、この話は本番環境への導入前の開発環境での話となります。

初期案. リクエストに対して計算し、そのままレスポンスを返す

f:id:tawamura:20191218114428j:plain

(CDNやVPC、セキュリティグループなどは省略しています)

まず最初に、APIプロキシサーバを介してリクエストを投げそのまま計算結果レスポンスを返す、従来通りの構成を検討しました。他のDS APIも多くはこの形を採用していました(今はLambdaに移行しつつあります)。最初の段階では実行に時間がかかるものがあるとわかっていなかったため、この形で試しています。 推論サーバにはECSインスタンスを使用しており、aiohttpを使用してリクエストに対し推論結果を返すWebサーバを立てています。

この場合、実行時間が数十秒かかる際に、APIプロキシサーバから事前に設定されたタイムアウトエラーが返されます。これ自体は正しい挙動です。レスポンスが返ってくるまで数十秒も待つという仕様は、リクエストを投げるサービス側にとっても良くありません。

案1. 推論部分を分離、SNS→SQS→Lambdaで推論

f:id:tawamura:20191218114527j:plain

そのため、このAPIが扱っている「リクエストを受け取る、レスポンスを返す」「リクエストについての素性データ抽出、推論」という機能を別々のAPIとして分離することにしました。 また、この頃チーム内でAPIサーバとしてECSではなくLambdaを使用するようにしていったこともあり、Lambdaでのリクエスト処理に移行することにしました。

まず、サービス側から叩くリクエストAPIとしてのLambdaの挙動は、リクエストについての計算済み結果があれば即時返す(今回は保存場所にDynamoDBを使用)、無い場合は「計算開始」または「計算中」のステータスを即時返すというものに限定します。これにより、タイムアウトなくレスポンスを返すという部分についてはクリアできます。
そして、「計算開始」または「計算中」となった場合は、同時に計算イベントを投げ、別で計算させるようにします。今回この計算イベントの投げ先にSNSを指定しました。

SNSに投げられたイベントは、サブスクライブしているSQSを経て、メッセージとして「リクエストについての素性データ抽出、推論」を行う計算APIとしてのLambdaに到達し計算されます。 ここでの計算結果がDynamoDBに格納され、次回以降リクエストAPIを叩かれた場合に、即時に計算結果が返却されることになります。さらに、計算が終了したことがわかるように、計算終了時に終了イベントをSNSに投げています。ここでは、DynamoDBで計算が完了したレコードに注視し、イベントを投げていますが、計算終了時に直接SNSにイベントを投げても同じかと思います。

計算APIのタイムアウトについてですが、Lambdaはデフォルトでタイムアウトの時間を最大値900秒に設定することができます。そこまで時間がかかることはほぼあり得ないと思われるので、今回のユースケースの場合、全てのリクエストについて理論上計算可能となります。今回はとりあえず120秒としました。 その間、SQSに溜められた計算はLambdaの処理を待つ形になります。 エラーログとして、10回以上処理が失敗(10回以上受信)したものはSQSのデッドレターキューに流すようにします。

設定として、計算APIの起動Lambda数は2、可視性タイムアウトはLambdaのタイムアウトに合わせて120秒にしました(可視性タイムアウトは、メッセージ受信時にSQSの該当メッセージを見えなくし、別のワーカーによって重複して呼ばれることを避けるためのものです)。 わかりやすさのため計算API Lambdaに「×2」と表記していますが、実際は他のECSなども複数サーバーが動いています。

多くの計算イベントが意図せずデッドレターキューへ😖

f:id:tawamura:20191218114836j:plain

これでうまくいくと思っていたのですが、実際に数百程度の計算イベントを投げてみると半分くらいのイベントがデッドレターキューに入ってしまっていることがわかりました。このデッドレターキューに流れてくるものの想定は、計算APIでの処理が10回失敗してしまったものだと思っていたので、ここがあまりにも多いのはちょっとおかしいです。一部の計算イベントについて時間がかかることで、一部のイベントがLambdaを占有してしまうことは想像がつくのですが、それでもSQSで処理を待つことで最終的には全ての計算が行われることを期待していました。

調査の結果、計算APIのLambdaでログをとっていたものの、そこでは一切エラーログは出ていないことがわかりました。つまり、Lambdaで一度も処理されることなく、デッドレターキューに入っているものがほとんどだということです。
おそらく勘違いをしていたのですが、SQSから受け取って計算APIのLambdaでの処理が失敗される度に、受信回数が1回増えるのではなく、「SQSから受け取ったが処理できるLambdaがいないのでSQSに返却された場合も受信回数が1回増える」なのではないか、ということがわかりました。 似たような事例が別の方からも挙げられていました。
medium.com

対策として考えられるのは、Lambdaの同時実行数を増やす、可視性タイムアウトを伸ばしてみる、などかと思います。 ただ、大きな前提問題としてこの計算APIはRDSへの接続を行なって動作しています。Lambdaの同時実行数を増やす場合、RDSへの接続数も比例して増えてしまいます。この結果、RDBMSの最大同時接続数を超えてしまう場合、他のDS APIのパフォーマンスに影響を及ぼすため、できれば避けたいものでした。可視性タイムアウトも根本的な解決にはなりそうにないとのことで、検討を見送りました。

実際に、弊社でサポートしていただいているAWSの方に意見を伺った際に、Lambdaの同時実行数が少ない場合に同様の現象が起きるということが確認されたようです。

案2. 推論部分を分離、SNS→SQS→ECS workerで推論🎉

f:id:tawamura:20191218114939j:plain

AWSの方との話をした際に、別途こちらで考えてみた構成案を話させていただいており、現状はその案をとっていただくと良さそう、という回答をいただきました。 SQS→Lambdaに繋ぐのではなく、SQSのメッセージをECSインスタンスが取りに行って同期的に処理するというものです。

案1の問題は、LambdaがSQSにメッセージを取りに行く回数が実行時間に対して多く、処理する前にメッセージの受信回数が上限を超え、デッドレターキューに流れてしまうという点でした。では、メッセージの受信動作をこちらで制御し、計算中はメッセージを新たに取りに行かないようにしよう、というのがこの案のモチベーションです。
ECSのインスタンスとして2つのworkerを立てて、それぞれが同じSQSからメッセージを取って、順次計算し結果を格納する、格納し終わったらまたメッセージがあるか見に行く、という動作をし続けます。こちらも各メッセージを処理する際のタイムアウトの時間は十分長くとるようにします。
この方法により、時間は多少かかりつつも問題なく全リクエストをさばけるようになりました。

・・というか、この案はどうやらLambdaがSQSをソースで扱う前の、従来のSQSのポーリングの仕組みになるようです。 振り返ってみると、APIがそれぞれRDSへ接続をしており、実行時間もかかる、というのがそもそもLambdaとの相性としてあまりよくなかったかなと思います。

まとめ

今回は時間のかかる推論APIをどうにかして処理させるために試行錯誤した話を紹介しました。
もちろん、推論の速度をあげるべきと言う指摘はもっともです。ただ、今回の場合は即時に結果が返る必要のない条件下での推論でもあったため、速度の追求よりも導入の実現性にフォーカスして対応を進めておりました。

AWS上のアーキテクチャを活用して上記のようなシステムを構築しましたが、自分自身AWSに触れるのはこの会社にきてからが初めてでした。ですが、チームのメンバーやテックリードの方々の手厚いサポートにより、ここまで理解と実装を進めることができました。

ちなみに私が作成したAPIはとある「不正検知API」です。
BASEではユーザのみなさまが安心してサービスをお使いいただける環境を提供するため、多方面からの協力も得ながら不正利用の検知と対応に尽力しております。
これからもネットショップ作成サービス「BASE」をどうぞよろしくお願いいたします。

明日はProduct Managementの山田さんとBASE BANKの清水さんです。お楽しみに。