AI事業本部の協業リテールメディアdivでソフトウェアエンジニアをしている 中澤 といいます。直近では、プロダクト開発以外にAI 事業本部の新卒研修の運営を行なったりもしていました。
私が所属しているチームで最近、定期バッチを行うワークフロー管理ツールを AWS Step Functionsへ移行したので、移行の背景や得た知見を記事として公開します。
移行前の構成
私たちのチームでは、ワークフロー管理ツールを AWS Step Functions に置き換える前には、Prefect を使っていました。 Prefect に関しては、弊社ブログの別記事があるので、Prefect について知りたい方はそちらも参考にしてみてください。
Prefect を利用している時の構成では、Prefect 側でワークフローのスケジュール管理やワークフロー内のタスク実行を Prefect、実際のワークフローのタスク実行自体は、Goで書かれたアプリケーションをAmazon ECS の ECS Task を利用し実行していました。
AWS Step Functions へ置き換えた背景
Prefect自体は、ワークフローの状態を Python で管理できることから Data Scientist や ML エンジニアがメインでバッチを管理する場合には、開発負荷が減ることやAirflowを用いてワークフロー管理をしているチームでの乗り換え先としては有用だと思います。しかしながら、現状の私たちのチームで Prefectを利用し続けることは以下の点でチームに対して運用負荷がかかっていました。
- Prefect を使用する上で 稀に発生する Python 由来のバグに対応できるメンバーが限られること
- 監視ツールとの統合が現状では不十分であり、監視負荷がかかること
- コスト面の改善の余地があること
それぞれもう少し詳細に書きます。
1. Prefect を使用する上で 稀に発生する Python 由来のバグに対応できるメンバーが限られること
Prefect を使う上では、agent をホストする必要があります。その際に、利用しているバージョン等によっては、Python由来のライブラリ依存関係によりエラー（もしくはバグ）が生じることがありました。
その際には、以下のような対応を行なっていました。
- PrefectのIsuueを確認し、対応バージョンを確認する
- ライブラリのバージョンアップをする
アップデートの際には Python の深い知識が求められるわけではないですが、 開発チームのメインスタックが Go であることから頭の切り替えが必要だったため、認知負荷が発生していました。
2. 監視ツールとの統合が現状では不十分であり、監視負荷がかかること
私たちのチームでは、監視ツールとして、Datadog を利用していますが、PrefectとDatadogの統合機能がまだありません。そのため、Prefect で実行しているバッチや Prefect の agent の ECS Task からのログを Datadog に転送し、それをキャッチして監視を行っていました。そのため、アプリケーション側のログ設定漏れが発生すると バッチ失敗の検知が遅れることがありました。また、ワークフローの状態は Prefect の管理画面、ログは Datadog、 のようにアラートが発生した場合に見るべき箇所が複数あることも監視負荷の増加に繋がっていました。
3. コスト面の改善の余地があること
弊社では、自社開発の認証基盤を利用しており、AWSやDatadogにはそちらからログインを行なっています。そのため、運用時には、複数のログイン先を利用し、それぞれの環境にログインする必要がありました。PrefectにもSAMLを利用したSSOの機構があるため、ID管理の一本化やログイン作業の簡略化のためSSO導入を検討しました。SSOを利用するためには、Prefectの有料プランを利用する必要があり、そのコストは、私たちが運用しているバッチのワークフローに対してはコストが高くなってしまうことがわかりました。
移行判断
前述した課題があり、移行先を探していたところ、以下の理由で AWS Step Functions を移行先として選択しました。
- ワークフロー管理のためにECS Task等のコンピューティングリソースを起動する必要がないこと
- バッチ管理がPython、インフラ管理がTerraformと分かれていたのが、Terraformで統一できること
- ワークフロー失敗時のステップ単位での再実行サポートが入ったこと
- 他のAWS リソースとの統合、Datadogとの統合が簡単であること
これらの理由から、ワークフローの管理ツールを Prefect から AWS Step Functionsに置き換えました。
移行後の構成
AWS Step Functionsへの移行後は以下のような構成になりました。ワークフローのスケジュール管理は Amazon EventBridge Scheduler、ワークフローの管理は AWS Step Functionsで行っています。ワークフロー内のタスクは引き続きECS Task で実行していて、日付処理が必要な場合にはAWS Lambdaを実行するステップ、重複実行制御のためのステップが存在しています。
※ AWS Step Functions では、日付の処理を始めとした動的な値をワークフロー内で利用したい場合にAWS Lambda を利用する必要がありました。
移行計画から完全移行まで
ワークフローの移行に際して、まず以下の観点から最初に移行するワークフローを2つ選び、検証・移行を行いました。
- 実行頻度が低く(日次）、リカバリが簡単でビジネス影響が少ないワークフローであること
- 複数のタスクを実行し、AWS Step Functionsの ParallelやMapを使うワークフローであること
- 動的な日付処理が必要なワークフローであること
私たちのチームでは、上記の観点で選定したワークフローが実行できることで AWS Step Functions がユースケースを満たしていて、完全に移行が可能であることが確かめられました。
二つのワークフローを移行した後は、Step Functionsを使ったワークフローの設定方法をチームメンバーに手を動かしながら理解してもらえるために巻き込みながら順次移行し、検証から完全移行まで約１ヶ月で移行しました。
この後は、移行した際に発生した課題とその解決策についていくつか紹介します。
移行の際に発生した課題と解決策
重複制御をどうやって制御するか
私たちのチームが開発しているシステムでは、ワークフローの実行が重複で発生した場合や実行中に同じワークフローが実行された場合に、同時実行数を制限する要件があります。また、Amazon EventBridge Scheduler自体の実行保証が at-least-once であるため、ワークフロー自体に冪等性があっても、コストがかかる実行を行う場合には、Amazon EventBridge Scheduler以外の機能で同時実行数を制御する必要がありました。
同時実行数制御はAWS Step Functions 自体の機能を利用して、制御できます。詳しくは別のブログ記事で紹介しているので、気になる方はそちらをみてもらえると嬉しいです。
新規ワークフロー構築時の Terraform 記述量をどの程度減らせるか
最初の ワークフロー検証時に、Terraform で記載してみると思いのほか記述量が増えることに気がつきました。運用負荷の中にはもちろん 開発時の負荷も入ると思うため、できるだけ管理するべき Terraform のコードは減らしたかったです。そのため、State Machineの定義を行う JSON ファイルと ワークフローの実行タイミング以外は、書かないように工夫を行いました。
一部のコードを紹介します。
ここでは、単一のワークフローの設定を行うLocal Valueを利用して全体を定義しています。これを行うことで、どのワークフローがどのタイミングで実行されるかがコードから簡単にわかります。
ここで定義したものをもとに、for each を利用し、Amazon EventBridge Scheduler や AWS Step Functionsの定義等の設定を行っています。IAM Roleや State Machineの定義自体は JSON で行っており、ワークフローで設定している key をもとにファイルから取得します。これを行うことで、副次的ですが、強制的に命名のルールを担保できます。
コンテナのコマンドに他のステップで出力した値をどうやって利用するか
前述の通り、私たちのチームでは、ワークフロー内のタスクは Go で書かれたアプリケーションを ECS Task で実行しています。そのため、タスクごとに ECS Task の Container Override を利用して、コマンドやコマンドのオプションを定義しています。
State Machine内で定義している Container Override を一部抜粋すると以下のようにしています。
しかしながら、AWS Lambda 等を利用した際に、パスを通じて出力した値を利用しようとし、以下のように記載した場合には設定できません。
そのため、組み込み関数のStates.FormatとStates.StringSplit を組み合わせて、配列を作成するようにしています。
まとめ
本記事では、ワークフロー管理ツールを Prefect から AWS Step Functions置き換えた内容とその途中で発生した課題と解決策について記載しました。 実際に移行してみて開発チームの運用負荷が下がっていることを実感しています。
AWS Step Functions は昨年入った、外部APIへのリクエストやステップごとの再起動といった複数のアップデートでかなり改善された印象があります。加えて、他のAWS サービスとの統合も多いので、AWS 内で構築するシステムでワークフローを構築する際のファーストチョイスになりうるのではないかと思います。
私たちのチームでは、リテールメディアにおける広告プロダクト開発に取り組んでいます。
興味のある方はカジュアル面談のページや私のX(@kensuke__19)のDMから気軽にお声がけください。
