Step Functions から Lambda を呼び Jupyter Notebook の実行まで繋げてみる

こんにちは。AS 部の兼安です。

先日中学の同級生と久しぶりに話しました。その時、初めて触った PC は中学の理科の先生が貸してくれた FM-7 だったことを思い出しました。 当時でも相当古かった機種なので、誰も知らないですよね・・・。雑誌に載っているBASICプログラムを昼休憩に写経する。そんな中学1年生でした。

本題に行きます。 本記事はこちらの記事の続きです。

blog.serverworks.co.jp

Step Functions から Lambda を呼び Jupyter Notebook の実行まで繋げてみます。

本記事の目的

本記事の目的は、Step Functions から Lambda を呼び出し、終了条件を満たすまで待機させるための構成を作る流れを掴んでもらうことです。 Step Functionsの個々のパラメータについて詳しくは触れていません。

AWS Step Functions とは

aws.amazon.com

AWS Step Functions は、AWS のマネージドサービスの一つで、サーバーレスのワークフローを可視化し、構築および実行するためのサービスです。 Lambda などを順次実行させたり、分岐させたりすることができます。

似たサービスとしてAmazon SageMaker PipelinesやAWS Batch が存在しますが、本記事ではAWS Step Functionsに焦点を当てています。

AWS Step Functions から ノートブックを実行する Lambda を呼び出す

AWS Step Functions では、Lambda 関数をタスクとしてワークフローに組み込むことができます。

Workflow Studioで見たLambda 関数呼び出し部分

Lambda 関数は、Step Functions のステートマシン定義内で、Lambda の ARN を指定することで呼び出すことができます。

呼び出し先の Lambda は、前回の記事をご覧ください。

{
  "Comment": "A description of my state machine",
  "StartAt": "state-machine-start",
  "States": {

    // 途中省略

    "notebook-execution": {
      "Comment": "Image processing notebook execution",
      "Next": "wait",
      "Parameters": {
        "FunctionName": "ノートブックを実行するLambdaのARN",
        "Payload.$": "$"
      },
      "Resource": "arn:aws:states:::lambda:invoke",
      "ResultPath": "$.notebook_result",
      "Retry": [
        {
          "BackoffRate": 2,
          "ErrorEquals": [
            "Lambda.ServiceException",
            "Lambda.AWSLambdaException",
            "Lambda.SdkClientException",
            "Lambda.TooManyRequestsException"
          ],
          "IntervalSeconds": 2,
          "MaxAttempts": 6
        }
      ],
      "Type": "Task"
    },

AWS Step Functions が Lambda を呼び出すには、以下のポリシーが必要です。

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "InvokeLambdaFunction",
      "Effect": "Allow",
      "Action": ["lambda:InvokeFunction"],
      "Resource": ["LambdaのARN"]
    }
  ]
}

このポリシーをアタッチしたロールを作成し、Step Functions のステートマシン定義で指定します。

ステートマシンのアクセス許可

Lambda にパラメータを渡す

Step Functions は、Lambda 関数にパラメータを渡すことができます。 パラメータは、Lambda 関数の入力イベントとして渡され、Lambda のイベントハンドラ関数の引数：event になります。

Step Functions は起動時に JSON でパラメータを指定することができるので、それを Lambda 関数に渡します。 今回の例では実行するノートブック名を渡すことにします。

Step Functionsの入力イベント

Lambda 関数の呼び出しタスクのこの部分で、呼び出し先の Lambda の指定と渡すパラメータの指定を行っています。 この時点の"$"には StepFunctions の入力イベントが入っていて、これを丸ごとを渡すよう指定しています。

"Parameters": {
  "FunctionName": "LambdaのARN",
  "Payload.$": "$" # StepFunctionsの入力イベントを丸ごと渡す
},

Lambda の返却値を保存する

Lambda 関数の実行が完了すると、返却値を受けとることができます。 返却値は次のタスクに渡すことができます。 これを使って、複数の Lambda 関数を連鎖的に実行するワークフローを作成することが可能です。

返却値は専用の名前をつけて保存するようにしました。 これは、Lambda の返却値で、元々持っていたパラメータが上書きされるのを防ぐためです。

ポイントはLambda 関数の呼び出しタスクのこの部分です。 この記述で、このタスク以降で当該タスクの返却値を$.notebook_resultという名前でアクセス可能になります。

"ResultPath": "$.notebook_result",

ノートブックの実行完了を待つ待機イベントを作る

本記事の Step Functions から起動する Lambda は、ノートブックを非同期で実行するようにしています。 従って、Lambda が終了してもノートブックは完了していないので、別途ノートブックの実行成否を確認する必要があります。 これを実現するために、Step Functions の待機イベントとループを使います。

Step Functionsを用いた待機イベント

ポイントは「status-retrieval」と「status-decision」と名付けた 2 つのタスクです。 「wait」はその名の通り一定時間待機するためのものです。 「status-retrieval-completed」は特に何もしていません。

「status-retrieval」は、ノートブックの実行状態をチェックする Lambda を呼び出すためのものです。 ノートブックの実行状態をチェックする Lambda は前回の記事をご覧ください。

「status-retrieval」のタスクは、ノートブックを実行成否を確認する Lambda を呼び出します。 同時に、InputPath で、前のタスクの返却値を渡します。 前のタスクの返却値には、ノートブックを実行しているジョブのキーが入っているので、これを渡すのが目的です。

ノートブックの実行成否を確認する Lambda は、実行成否を返すので、その返却値を ResultPath で指定した名前で保存します。

"status-retrieval": {
    "Comment": "Image processing status retrieval",
    # 前のタスクの返却値を渡す
    "InputPath": "$.notebook_result.Payload",
    "Next": "status-retrieval-completed",
    "Parameters": {
    "FunctionName": "ノートブックの実行成否を確認するLambdaのARN",
    "Payload.$": "$"
    },
    "Resource": "arn:aws:states:::lambda:invoke",
    # 返却値を別名で保存する
    "ResultPath": "$.status_check_result",

「status-decision」のタスクは、ノートブックの実行成否を確認する Lambda の返却値をチェックして、ノートブックの実行が完了しているかどうかを判断します。 返却値が成功でなければ、「wait」に戻り待機します。

"status-decision": {
    # 返却値がSUCCEEDEDならば次へ進む
    "Choices": [
    {
        "Next": "state-machine-completed",
        "StringEquals": "SUCCEEDED",
        "Variable": "$.status_check_result.Payload"
    }
    ],
    # デフォルトの遷移先、これによりSUCCEEDED以外の場合はwaitに戻る
    "Default": "wait",
    "Type": "Choice"
},

以上で、Step Functions から Lambda を呼び Jupyter Notebook の実行まで繋げることができました。

本運用では、ノートブックは一つとは限らないと思います。 その場合は、今回の例を複数繋げることで、ノートブックの実行と待機を繋いでいくことができます。

補足

今回使った Step Functions のステートマシン定義の全文はこちらです。

兼安 聡(執筆記事の一覧)

アプリケーションサービス部 DS1課所属
AWS12冠。
広島在住で今日も明日も修行中です。