はじめに

API Gateway + Lambda で画像分類 AI を動かそうとしたところ、API Gateway のタイムアウト制限に引っかかってしまいました。
Step Functions を使うことで回避することができたので、その方法について書いていきます。

実行環境

macOS Big Sur 11.3.1
SAM CLI 1.27.2
Docker 20.10.7

やりたいこと

API Gateway と Lambda の間に Step Functions を挟んで、API Gateway から Step Functionsを呼び出します。

Step Functions の StartExecution API を呼び出すと、　Lambda のステートマシンでタスクを実行し、そのタスク名をレスポンスとして API Gateway に返します。
API Gateway にレスポンスが返ってきているためタイムアウトになりません。

また、Step Functions の DescribeExecution API　を呼び出すと、タスクの状況と処理結果を一緒に返してくれます。

構成はこちらの記事を参考にしました。

手順

Lambdaの作成

こちらの記事を参考にLambdaを作成しました。

自身のIAMユーザーに以下のポリシーを付与しておいてください。

AWSLambda_FullAccess
AmazonAPIGatewayAdministrator
AmazonAPIGatewayPushToCloudWatchLogs
AmazonEC2ContainerRegistryFullAccess
IAMFullAccess
AWSCloudFormationFullAccess
AWSStepFunctionsFullAccess

Lambda ステートマシンを作成

Lambda の作成ができたら、Step Functions で Lambda ステートマシンを作成します。

Step Functions のコンソールで「ステートマシンの作成」をクリックし、「コードをワークフローで記述」、タイプは「標準」を選択してください。

ステートマシンの定義を下記の内容に変更してください。Resource には、プロジェクト作成時に作られた Lambda 関数の ARN を指定します。

{
  "StartAt": "CallLambda",
  "States": {
    "CallLambda": {
      "Type": "Task",
      "Resource": <Lambda関数のARN>,
      "End": true
    }
  }
}

詳細を以下のように指定します。ステートマシン名は自由に変更してください。

これで Step Functions の準備は終わりです。

API Gateway 作成

次に API Gateway を作成していきます。

ここからはこちらも参考にしています。

IAM ロールの作成

まず API Gateway 用の IAM ロールを作成します。

IAM コンソールから「ロール」＞「ロールの作成」を選択してください。

「AWS サービス」から API Gateway を選択して「次のステップ：アクセス権限」へ進みます。

「Attached アクセス権限ポリシー」と「タグの追加(オプション)」では何もせずに次のステップへ進んで大丈夫です。

任意のロール名を入力し「ロールの作成」をしてください。今回はAPIGatewayToStepFunctionsとします。

IAM ロールにポリシーをアタッチする

「ロール」ページで先ほど作成したロールを選択します。

下記のような「ロール ARN」が表示されるのでメモをしておいてください。

arn:aws:iam::123456789012:role/APIGatewayToStepFunctions

「アクセス権限」タブの「ポリシーをアタッチします」を選択します。

AWSStepFunctionsFullAccessにチェックをし、「ポリシーのアタッチ」をします。

Step Functions と連携する API を作成

API Gateway を作成します。API タイプは「REST API」で構築してください。
API名を入力してAPIを作成します。

新しいリソースを作成します。「アクション」から「リソースの作成」を選択してください。

任意のリソース名を入力し「API Gateway CORS を有効にする」にチェックを入れて「リソースの作成」をします。

作成したリソースに POST メソッドを作成します。「アクション」から「メソッドの作成」を選択してください。
リストの中から POST を選択し、チェックマークをオンにします。下図を参考に設定してください。実行ロールには先ほどメモしたロール ARN を使用してください。

作成したリソースの下にもう 1 つリソースと POST メソッドを作成します。下図を参考に設定してください。

CORS の有効化を 2 つのリソースそれぞれに対して行います。作成したリソースを選択した状態で「アクション」から「CORS の有効化」を行ってください。

Access-Control-Allow-Headers を'*'に変更し,「CORS を有効にして既存の CORS ヘッダーを置換」します。

API のデプロイを行います。「アクション」から「API のデプロイ」を選択してください。「デプロイされるステージ」を「新しいステージ」、「ステージ名」に任意のステージ名を入力してデプロイをしてください。

これで Step Functions と API Gateway の準備は完了です。

動作確認

では実際に動作を確認してみましょう。
curl を用いて POST リクエストを送ります。URL は自分で作成したリソースの URL に変更してください。

curl -X POST -d '{"input": "{}","stateMachineArn": "<Step Functions の ARN>"}' https://xxx.execute-api.ap-northeast-1.amazonaws.com/prod/execution

このような結果が返ってくれば大丈夫です。

{
"executionArn":"arn:aws:states:ap-northeast-1:<アカウントID>:execution:test-function:xxxxx",
"startDate":1.63046260246E9
}%

状態を返してくれるリソースの方にもリクエストを送ってみましょう。先ほど返ってきたexecutionArnを渡します。

curl -X POST -d '{"executionArn":"arn:aws:states:ap-northeast-1:<アカウントID>:execution:test-function:xxxxx"}' https://xxx.execute-api.ap-northeast-1.amazonaws.com/prod/execution/status

実行途中ならstatusとしてRUNNINGが返ってきます。

{
"executionArn":"arn:aws:states:ap-northeast-1:<アカウントID>:execution:test-function:xxxxx",
"input":"{}",
"inputDetails":{"__type":"com.amazonaws.swf.base.model#CloudWatchEventsExecutionDataDetails","included":true},
"name":"xxxxx",
"startDate":1.63046260246E9,
"stateMachineArn":"arn:aws:states:ap-northeast-1:<アカウントID>:stateMachine:test-function",
"status":"RUNNING",
"stopDate":1.630462645972E9,
"traceHeader":"Root=1-612ee28a-044d88e25fb4be157a0e5046;Sampled=1"
}%

処理が終了するとoutputと、statusとしてSUCCEEDEDが返ってきます。
outputには Lambda の処理結果が入っています。

{
"executionArn":"arn:aws:states:ap-northeast-1:<アカウントID>:execution:test-function:xxxxx",
"input":"{}",
"inputDetails":{"__type":"com.amazonaws.swf.base.model#CloudWatchEventsExecutionDataDetails","included":true},
"name":"xxxxx",
"output":"",
"outputDetails":{"__type":"com.amazonaws.swf.base.model#CloudWatchEventsExecutionDataDetails","included":true},
"startDate":1.630464429504E9,
"stateMachineArn":"arn:aws:states:ap-northeast-1:<アカウントID>:stateMachine:test-function",
"status":"SUCCEEDED",
"stopDate":1.630464457648E9,
"traceHeader":"Root=1-612ee9ad-1a6f0b711a2d53ca6d079d6b;Sampled=1"
}%