Invoke AWS Lambda with AWS Step Function

First published on: Medium


In the previous tutorial, A Choice-based Step Function, we have seen the flow control of a Step Function based on choices. In this tutorial, we will try to invoke an AWS Lambda with a Step Function.


AWS Lambda Function

Below I’ll give a very brief description of what is Lambda Function! I’ll not explain the details of a Lambda function because this tutorial is not intended for this purpose!

AWS Lambda is an awesome component of AWS. If you’re planning to go Serverless, this is a really cool thing.

Lambda function is basically a container-based system where you can write your codes and the codes get executed on demand. What I mean is your codes get executed only when the Lambda function is invoked.

What do you do with your server? You put your codes there and the server always runs. If you expose those codes through an API or something else then invoking those API can execute the code. But the code is always there because the server is always running. So you have to bear the cost for every second whether you execute your codes or not.

With a Lambda function, the scenario is a bit different. Your code is contained in a container with all the configuration variables but the container is not always running. The container only runs when it’s invoked. The rest of the time, simply put, the container sleeps. So, when you invoke the Lambda function, the container wakes up and starts executing your codes. So you bear the cost only when the container is awake and executing your codes. When the container is in sleep, you are not bearing any costs! This is one of the many reasons why Lambda Function is so cool! There are also so many other important features of Lambda Function that depict why it’s cool and important but that’s another story for another time!!!


Invoke AWS Lambda Function with AWS Step Function

If you are designing an event-driven based serverless system this design can help you in many ways. You may fall into various scenarios where you want to use a Step Function to invoke your Lambda Function. In the following, we will look into this design. So let’s get started.

First, let’s create a Lambda Function.


Create a Lambda Function

  1. Go to the Lambda console and click on the Create function button.
  2. Choose Author from scratch option.
  3. Provide a function name. For this example, the name of our function will be step-function-invoking-lambda-function.
  4. From the runtime section, you can pick whatever runtime you want. I’ll pick Python3.8.
  5. Choose a role. You can either create a new role or choose from an existing role. If you don’t have any role, you can create a new role from AWS policy templates. I’ll choose from the existing role.
  6. At the bottom of the page, click the Create Function button to create the function.
The UI of the Lambda Function console. The UI may look different based on the time you’re viewing this.
The UI of the Lambda Function console. The UI may look different based on the time you’re viewing this.

In the above picture, you can see the UI of the Lambda Function console. There’s an inline code editor where you can write code.

Inline editor to write your code
Inline editor to write your code

For this tutorial, our Lambda Function will contain a sorted array. Our Step function will invoke the Lambda function with an integer. Our Lambda function will check if the integer exists in the array and will return True if so. Otherwise, it will return False.

Here’s the code:

def lambda_handler(event, context):
    print(event)
    target_value = int(event['target_value'])
    
    print("Lambda is Invoked with the Integer:: ", target_value)
    
    nums = [-1,0,3,5,9,12]
    
    found, index = search(nums, target_value)
    
    if found:
        statusCode = 200
        msg = 'Target value {0} found in position {1}'.format(target_value, index)
    
    else:
        statusCode = 400
        msg = 'Target value {0} was not found in the List!'.format(target_value)
    
    return {
        'statusCode': statusCode,
        'body': msg
    }
    
    
def search(nums, target_value):
    head, tail = 0, len(nums)-1
    
    while head <= tail:
        mid = (head + tail) // 2
        
        if nums[mid] == target_value:
            return (True,  mid)
        
        elif nums[mid] > target_value:
            tail = mid - 1
            
        else:
            head = mid + 1
            
    return (False, -1)

Now, let’s create a Step Function.


Create a Step Function

In the 3rd tutorial of this series, I’ve written how to create a Step Function. You can follow this 3rd part, Create a Step Function to create a Step Function.

The definition of our Step function is the following:

{
  "StartAt": "InvokeLambdaState",
  "States": {
    "InvokeLambdaState": {
      "Type": "Task",
      "Resource": "arn:aws:lambda:us-east-2:XXXX:function:step-function-invoking-lambda-function",
      "Parameters": {
        "target_value": 100
      },
      "Next": "CheckStatusFromLambda"
    },
    "CheckStatusFromLambda": {
      "Type": "Choice",
      "Choices": [
        {
          "Variable": "$.statusCode",
          "NumericEquals": 200,           
          "Next": "SuccessState"
        },
        {
          "Variable": "$.statusCode",
          "NumericEquals": 400,           
          "Next": "FailedState"
        }
      ]
    },
    "SuccessState": {
      "Type": "Pass",
      "End": true
    },
    "FailedState": {
      "Type": "Fail"
    }
  }
}

If you are confused with this JSON, you can follow the second and the third part of the series Amazon State Language part-1 and Amazon State Language part-2. In those tutorial, I’ve explained what is Amazon State Language and all the keywords you can see in the above definition.

In the above definition, inside the InvokeLambdaState, the resource

"Resource": "arn:aws:lambda:us-east-2:XXXXXXX:function:step-function-invoking-lambda-function"

is the ARN of your Lambda function. You can find this from image 1.

The Step Function flow will look like this:

The flow of the Step Function
The flow of the Step Function

The flow of the Step FunctionWe will invoke the Lambda function with parameter

"target_value": 9

The Lambda function will return statusCode either 200 or 400. We have a Choice state CheckStatusFromLambda where we check the statusCode. If it is 200 we go to the SuccessState. Otherwise, we go to the FailedState.

If you don’t know what is Choice State you can check out the 5th tutorial A Choice-based Step Function of this series.

Now if you execute the Step Function, you’ll see the below output:

We found the value in our list and the Lambda function returned success.
We found the value in our list and the Lambda function returned success.

Remember, our list looked like this:

nums = [-1,0,3,5,9,12]

Now, if we invoke the Lambda function with parameter

"target_value": 100

we will see the below output:

Value 100 doesn't exist in our list. So the Lambda function returned 400
Value 100 doesn't exist in our list. So the Lambda function returned 400

So this is the mechanism to invoke a Lambda Function from a Step Function.

This is a naive example, I understand. But the intention of this tutorial is to show how do you invoke a Lambda function from a Step function and how do you propagate messages between them. Of course, there are so many cool things you can do with this design.


Use Cases:

Of course, there are no fixed use cases but I’ll try to mention one or two.

In an event-driven based serverless system, you can use a Step function as the backbone, and based on the different state you can control the flow. Now in these flow, you can publish messages by invoking a Lambda. You can have multiple Lambda functions. All those can be assigned different tasks. And which will get executed can be controlled by your step function based on the event and flow of the Step function.

How to do so? Well, that’s a different story for some other time! :D

Until then, enjoy your life!