AWS Step Functionsでバッチジョブワークフローの実行基盤を構築する

<2023年12月19日追記>
本記事中に「①失敗したジョブから再実行が出来ない」という課題の記載がありますが、現時点で解消されています。詳しくは、「AWS Step Functions 待望の『リドライブ』機能がGAされました | TECH | NRI Digital」の記事をご参照ください。

こんにちは、NRIデジタルの島です。
AWSパブリッククラウド上でバッチジョブのワークフローを構築する際の選択肢はあまり多くありません。
弊社社内においても「JP1 ※1」や「Senju ※2」等のジョブ運用管理ツールを使ってAWSサービスと連携している事例が多いと思っております。

AWSサービスでは「Amazon Managed Workflows for Apache Airflow (MWAA)」や「AWS Step Functions（以下Step Functions」、「AWS Glue Workflow」がその候補（※3）になると思いますが、今回は「Step Functions」にて実運用に耐えられるバッチジョブのワークフローが構築できるのかを検証してみました。

※1　日立製作所が開発・発売している統合システム運用管理ツール
▶統合システム運用管理JP1：日立

※2 弊社親会社のNRIが開発・発売している運用管理ツール
▶システム運用管理ならNRIのSenju Family

※3 AWSの各ワークフローサービスの特徴

Step Functionsとは

アプリケーションの一連のワークフローを構築・管理するマネージドなサーバレスワークフローサービスです。

Step Functionsを利用するとマイクロサービス等の疎結合化が必要な分散アプリケーションやバッチアプリケーションのジョブフローにおいて、AWS Lambda（以下Lambda）やAWS Batch（以下Batch）等様々なAWSサービスと連携したワークフローを容易に構築できます。記述はJSONベースの「ASL（Amazon States Language）」という言語で行いますが、「Workflow Studio」というツールを使用することで、ドラッグ&ドロップでグラフィカルにワークフローを構築することもできます。また、エラー発生時の処理やリトライ機能も備えており、ワークフローエンジンに必要な機能を一通り備えております。
▶AWS Step Functionsとは

Step Functionsの構成要素は下記2つのみです。

State Machine（ステートマシン）

ワークフロー全体を表す単位で、下記Statesの集合体

States（ステート）

処理実行単位で、入力に基づいて下記のアクションを実行して出力を他のStates（ステート）に渡す

• Task（ステートマシンで何らかの作業をする）
• Choice（実行の選択肢間で選択する）
• Fail or Success（失敗または成功で実行を停止する）
• Pass（入力を単純に出力に渡す、または一部の固定データを出力する）
• Wait（一定時間または指定された時刻/日付まで遅延を提供する）
• Parallel（実行の並列ブランチを開始する）
• Map（ステップを動的に反復する）

※各ステートは独立性が高くなる為、個々にスケーリングや拡張が可能です

ユースケースとしては、下記公式ページの通り「大規模データの処理基盤」「マイクロサービスのオーケストレーション」など色々と考えられます。
▶ユースケース – AWS Step Functions

今年のAWS DevDayではコネヒト（株）様より「家族ノート」アプリにて「Step FunctionsでECS/Fargate,Lambda,EventBridgeを組み合わせたワークフロー」を構築した事例が紹介されました。
▶AWS Dev Day Online Japan 2021に登壇してきました！-コネヒト開発者ブログ

なお、Step Functionsと連携可能なAWSサービスはこれまで17サービスでしたが、先日9月の発表で200を超えるサービスと連携可能となりました。
▶AWS Step Functions が AWS SDK 統合で 200 を超える AWS のサービスのサポートを追加

このアップデートにより、何らかのAWSサービスと連携する為にLambdaを経由する必要がなくなり、より使いやすくなっています。

Step Functionsの課題

Step FunctionsはAWSクラウドで様々なワークフローを構築する際の選択肢として最有力候補になってくると思います。
前項にてユースケースを共有致しましたが、筆者はその中でもバッチジョブフローの実行基盤として利用するケースが今後増えてくるのでは？と考えております。
しかしながら、仮にバッチジョブフローの実行基盤として利用しようとした場合、実運用する際に課題になると思われる事項もございます。

①失敗したジョブから再実行が出来ない

Step Functionsは、JP1やSenjuのように途中のジョブ（ステート）から再実行することが仕様上出来ません。

その為、再実行したい場合頭から流し直しする必要があります。
バッチの処理は数時間かかる場合もある為、運用上この仕様は致命的かと思います。

②並列実行時に一部のジョブが失敗した場合、その他のジョブもキャンセルされてしまう

また、バッチのジョブワークフローを組む場合、一部のジョブを並列実行して処理したい場面もあります。
Step Functionsは並列実行もサポートしております。
しかしながら、並列実行時に一部のジョブが異常終了すると、実行完了していない他のジョブがキャンセルされてしまう仕様となっています。

こちらは処理要件にもよりますが、正常に処理しているジョブは実行を継続してほしいという要件のほうが多いのではないでしょうか。
このような仕様となっているのは、前項のユースケースで紹介した通り、Step Functionsがバッチのワークフローに特化したサービスではない為と推測します。

次項から、上記の課題を踏まえ、バッチの実運用でも利用できるような方式を検討、検証した結果を共有したいと思います。

実行ステータス管理

まず「①失敗したジョブから再実行が出来ない」についてです。
Step Functionsが任意のステートから実行出来ない仕様である以上、AWS側に対応してもらえない限り「異常終了したジョブ（ステート）から再実行する」ことは出来ません。
（今後のアップデートで可能になるかもしれませんが・・・）
その為、頭から流し直すことになりますが、その際各ジョブ実行前に、実行ステータスに応じて実行判断する方式が考えられます。

イメージは下記のようになります。
※本検証ではステータス管理に「Amazon DynamoDB（以下DynamoDB）」を使用しています。

では実際にシンプルなフローで試してみます。
「GenerateRandomIntValue」というがバッチ処理のアプリケーション（ランダム値を返すだけの単純なLambda関数）をコールしているジョブ（ステート）です。
（アプリケーションの処理の内容に意味はありません）

2番目の「GenerateRandomIntValue2」を異常終了させてみます。

今回は正常に実行完了したジョブだけ、DynamoDBに登録しています。
※id列はバッチの「一連のシーケンス実行単位を表すID」です

頭から再実行します。
先ほど成功した1番目のジョブの実行はスキップされ、2番目のジョブのみ実行されます。
結果、再実行したいジョブから実行できるようになりました。

誰もが考えつく方式かと思いますが、結局これが一番シンプルな方式かと存じます。

並列実行

次に「②並列実行時に一部のジョブが失敗した場合、その他のジョブもキャンセルされてしまう」について考えてみます。
前項の実行ステータス管理を取り込んだフローは下記のようになりますが、前述した通り一部のジョブが異常終了すると、実行完了していない他のジョブがキャンセルされてしまいます。

これを解消する手段として、下記公式ページにも公開されている「エラー制御機能」を取り入れてみたいと思います。
▶Step Functions エラー処理-AWS Step Functions

ジョブフローは下記のようになります。

エラー処理の概要です。

このようにすることで、並列に並べた一部のジョブでエラーが発生しても、下記の通り他のジョブがキャンセルすることなく実行可能となり、ジョブフロー全体としては「失敗」とすることが可能です。

複雑性の排除

ここまでで前述した2つの課題を解決することが出来ましたが、1点問題があります。
今回のようなシンプルなジョブフローなら良いのですが、並列処理が入れ子になるなどより複雑なジョブフローになると、フロー全体がごちゃごちゃして可読性が悪く、結果メンテナンス性も悪くなることが懸念されます。

Step Functionsは下記公式ページにもある通り、ステートマシンから別のステートマシンを呼び出すことが出来ます。
▶統合サービスとしてAWS Step Functions実行を管理する-AWS Step Functions

これを利用して、各ジョブを論理的な単位で別のステートマシンとして独立させ、親のステートマシンから呼び出す方式にすることで、シンプルなジョブフローが構成可能です。

下記のようなイメージです。

こうすることにより、親のステートマシンを参照すればジョブフロー全体の流れがシンプルにわかり、かつ独立した各ジョブは前述のエラー制御のような複雑なことをしなくても他ジョブに干渉することがなくなる（他ジョブがキャンセルされることもなくなる）のでは？という考え方です。

実際の親のステートマシンは下記のようにシンプルなものになります。

ただ、子ステートマシンを同期で呼び出す方式にすると、残念ながらステートマシンを分割しても一部の並列ジョブがエラーとなった場合、結局はキャンセルされてしまうようです。

この事象は、下記に公開されている「EventBridgeを利用したWait-For-Callback方式」を採用し、非同期呼び出しをすることで解消できます。
▶AWS Step Functions と Amazon EventBridge のサービス統合のご紹介

結果、最終系の全体イメージは下記のようになります。

実際に上記イメージでフローを構成し動作させてみたところ、想定通り子ステートマシンはキャンセルされず実行されることが出来ました。

※コールバック処理が必要な分、若干子ステートマシンが複雑になります
※親ステートマシン上「キャンセル」されているように見えますが、実際は該当の子ステートマシンは最後まで実行されます

さいごに

かなり試行錯誤しましたが、方式を工夫することでStep Functions上でも実運用に耐えうるバッチジョブワークフローが構築できることがわかりました。Step Functionsは本稿前半でお伝えした通り、200以上のAWSサービスと連携可能で、方式設計のパターンもたくさんあると思います。本稿で紹介したやり方以外にもより良い方式があるかもしれませんので、今後更に深掘りしていければと思います。

以上