Amazon Web Services ブログ

CyberCRX が AWS Step Functions 分散 Map を使用して ML の処理時間を 8 日間から 56 分に短縮した方法

2022年 12 月、Sébastien Stormacq が AWS Step Functions の分散 Map ステートの公開について書きました。クラウドで大規模な並列ワークロードをオーケストレーションできる新機能についての記事でした。CyberGRX という会社のデータシステムエンジニアである Charles Burton 氏はその記事でこのことを知り、ワークフローをリファクタリングすることにより、機械学習 (ML) 処理ジョブの所要時間を 8 日から 56 分へと短縮したのです。以前、ジョブの実行時には、エンジニアが常時監視する必要がありましたが、現在は 1 時間足らずで実行でき、かつてのようなサポートも必要ありません。さらに、AWS Step Functions 分散 Map を使用した新しい実装は、当初よりもコストが下がりました。

CyberGRX がこのソリューションで達成したことは、サーバーレステクノロジーが内包していることの好例です。エンジニアやデータサイエンティストが本来の重要な業務に集中できるよう、差別化されていない面倒な作業をできるだけクラウドに割り当てられるようになったのです。ここでは、CyberGRX の大規模かつ増大中の機械学習の知見を使用して、サードパーティのサイバーリスクを評価するという、CyberGRX の主要なサービスのために、モデルとプロセスを継続的に改善していくことを意味します。

ビジネス上の課題とは?
CyberGRX は、サードパーティのサイバーリスク (TPCRM) データを利用者に提供しています。サードパーティ企業がリスク評価アンケートにどのように回答するかを、高い信頼性をもって予測します。そのためには、プラットフォーム上のすべての企業で予測モデルを実行する必要があります。現在、同社が所有する予測データは、225,000 社分を超えています。新しい会社が増えたり、ある会社のデータが変化したりするたびに、データセット全体を処理して予測モデルを作り直します。時間が経つにつれて、CyberGRX のデータサイエンティストはモデルを改善するか、あるいは、新しい機能を追加しますが、その場合にもモデルの再生成が必要になります。

225,000 社の企業を対象に、できるだけ少ない実務リソースで、この仕事をタイムリーに実行することが課題です。ジョブは会社ごとに一連の操作を実行し、会社の計算自体は、他の会社とは別に実行されます。つまり、すべての企業を同時に処理できることが理想的です。しかし、このような大規模な並列化の実装には、解決しにくい問題が伴います。

1 つ目のイテレーション
そのことを念頭に置いて、同社は Kubernetes と Argo Workflows を使用してパイプラインとなる 1 つ目のイテレーションを構築しました。Argo Workflows は、Kubernetes 上で並列ジョブをオーケストレーションするためのオープンソースのコンテナネイティブなワークフローエンジンです。これらはインフラストラクチャですでに使用されていたため、同社にとって使い慣れたツールでした。

ところが、プラットフォーム上の全企業を対象にジョブを実行しようとするとすぐに、システムが効率的に処理できる能力の限界にぶつかりました。このソリューションは集中管理型のコントローラー、Argo Workflows に依存していたため、堅牢ではなく、この間にコントローラーは最大容量までスケールされました。当時、同社が担っていたのは 15 万社に過ぎませんでした。また、すべての企業でジョブを実行するのに約 8 日かかりました。その間にシステムがクラッシュすれば、再起動する必要もありました。作業には多大な労力が必要で、常に現場のエンジニアが作業の監視とトラブルシューティングを行う必要がありました。

転機となったのは、Charles が 2022 年の初めにアナリティクスチームに加わったときでした。彼にとっての初仕事の 1 つは、当時担っていた約 17 万社に対しフルモデルの実行を行うことでした。モデルの実行は 1 週間続き、日曜日の午前 2 時に終了しました。その時、彼はシステムの改良が不可欠と実感しました。

2 つ目のイテレーション
前回モデルを再実行したときの苦痛を頭の中で思い浮かべながら、Charles はワークフローを書き換える方法をじっくり考えました。彼が最初に考えたのは AWS LambdaSQS を使用することでしたが、そのソリューションにはオーケストレーターが必要であることに気付きました。そのため、プロセスの自動化、マイクロサービスのオーケストレーション、データと ML パイプラインの作成に役立つだけでなく、必要に応じてスケールできるサーバーレスサービスである Step Functions を選択しました。

Charles は、約 2 週間で Step Functions が動作する新しいバージョンのワークフローを完成させました。彼が取った最初のステップは、Lambda のコンテナイメージパッケージ形式を使用して既存の Docker イメージを Lambda で実行するように調整することでした。コンテナはすでに彼のデータ処理タスクに対応していたので、今回の更新は簡単でした。ジョブを開始した時点で必要な機能がすべて用意されていることを確認するために、Lambda による同時実行のプロビジョニングをスケジュールしました。また、Lambda がこの最大数の同時実行を一度に処理できるように、予約済みの同時実行も設定しました。同時に実行される多数の関数に対応するために、アカウントあたりの Lambda の同時実行クォータを引き上げました

そして、ステップが並行して実行されることを確認するために、彼は Step Functions と Map ステートを使用しました。Map ステートにより、Charles はデータセット内のアイテムごとに一連のワークフローステップを実行することができました。イテレーションは並行して実行されます。Step Functions の Map ステートは 40 件の同時実行が可能で、CyberGRX ではさらに並列化が必要であったため、複数のステートマシンを並行して起動するソリューションを開発しました。こうすることで、全社で高速にイテレーション処理を行うことができました。この複雑なソリューションを作成するには、システムの同時実行のヒューリスティックを処理し、入力データを複数のステートマシンに分割するプリプロセッサが必要でした。

この 2 つ目のイテレーションはすでに、1 つ目のイテレーションよりも優れていました。問題なく実行を完了でき、20 万社を超える企業を 90 分間で処理できるようになっていました。しかし、プリプロセッサはシステムの非常に複雑な部分であり、並列化の量で Lambda および Step Functions API の限界に達していました。

AWS Step Functions を使った 2 つ目のイテレーション

3 つ目、最終イテレーション
その後、AWS re:Invent 2022 の期間中、AWS は Step Functions 用の分散 Map を発表しました。これは、大規模な並列ワークロードを調整する Step Functions を記述できる新しいタイプの Map ステートです。この新機能を使用すると、Amazon Simple Storage Service (Amazon S3) に保存されている何百万個ものオブジェクトを簡単にイテレーション処理でき、分散 Map では最大 10,000 件の並列サブワークフローを起動してデータを処理できます。

Charles は、ニュースブログで 10,000 件の並行ワークフロー実行に関する記事を読んだとき、すぐにこの新しいステートを試してみようと考えました。2 週間後、Charles はワークフローの新しいイテレーションを構築しました。

分散 Map ステートでは入力が複数のプロセッサに分割され、さまざまな実行の同時処理が処理されるので、Charles は複雑なプリプロセッサコードを排除できました。

新しいプロセスはこれまでになくシンプルでした。今では、ジョブを実行したいときはいつでも、入力データを含むファイルを Amazon S3 にアップロードするだけです。このアクションは、分散 Map を持つステートマシンを対象とする Amazon EventBridge ルールをトリガーします。次に、ステートマシンはそのファイルを入力として実行し、その結果を Amazon Simple Notification Service (Amazon SNS) トピックに公開します。

AWS Step Functions を使った最終イテレーション

その影響とは?
3 つ目のイテレーションが完了してから数週間後、プラットフォーム上の 227,000 社すべてでジョブを実行する必要がありました。ジョブが終わったとき、Charles のチームは驚きました。プロセス全体にわずか 56 分しかかからなかったからです。このジョブでは 56 分間の間に 570 億回以上の計算が実行されたと推定されています。

分散 Map ステートの処理

以下の画像は、ワークフローが実行中の Lambda 関数 1 つの同時実行数を示す Amazon CloudWatch のグラフです。この間、約 1 万個の関数が並行して実行されています。

Lambda の同時実行 CloudWatch グラフ

ジョブの実行時間を簡略化および短縮することは、CyberGRX とデータサイエンスチームに多くの可能性をもたらします。データサイエンティストの一人がモデルに加えた改善点をテストするためにジョブを実行したいと思った瞬間から、そのメリットを実感できました。データサイエンティストたちはエンジニアの助けを必要とせずに、独自にこれを運用することができました。

また、予測モデル自体が CyberGRX の主要サービスの 1 つであるため、予測分析を日々改善することができ、同社の製品はより競争力のあるものへと変貌しました。

AWS Step Functions の使用方法の詳細については、こちらをご覧ください。

また、この新機能のテストや詳細を確認するには、Serverless Land の Serverless Workflows Collection (サーバーレスワーフクローコレクション) も参照してください。

– Marcia

原文はこちらです。