Menu
AWS Step Functions
Developer Guide

Creating a Job Status Poller

In this tutorial, you create an AWS Step Functions state machine that uses an AWS Lambda function to implement a Wait state loop that checks on an AWS Batch job. A Wait state is a state type that waits for a trigger to perform a single unit of work.

Step 1: Create an AWS Batch Job Definition

Important

Ensure that your AWS Batch job is under the same AWS account as your state machine.

  1. Follow the instructions to create a job definition in the AWS Batch User Guide.

  2. Note the following:

    • Job definition name (for example, my-job-definition)

    • Job queue name (for example, my-job-queue)

Step 2: Creating an IAM Role for Lambda

Both Lambda and Step Functions can execute code and access AWS resources (for example, data stored in Amazon S3 buckets). To maintain security, you must grant Lambda and Step Functions access to these resources.

In the Getting Started tutorial, this is done automatically for Step Functions: An IAM role is created when you create the state machine. However, Lambda requires you to assign an IAM role when you create a Lambda function in the same way Step Functions requires you to assign an IAM role when you create a state machine.

To create a role for Lambda

  1. Log in to the IAM console and choose Roles, Create role.

  2. On the Select type of trusted entity page, under AWS service, select Lambda from the list and then choose Next: Permissions.

    Note

    The role is automatically provided with a trust relationship that allows Lambda to use the role.

  3. On the Attach permissions policy page, type Batch into the filter, choose AWSBatchFullAccess, and then choose Next: Review.

  4. On the Review page, type BatchToStepFunctions for Role Name and then choose Create role.

The IAM role appears in the list of roles.

Step 3: Creating a Lambda Function

The job submission Lambda function receives input: the AWS Batch job definition name, compute environment name, and job queue name. The job status Lambda function retrieves the status of the first.

To create the job submission Lambda function

Important

Ensure that your Lambda function is under the same AWS account as your state machine.

  1. Log in to the Lambda console and choose Create a function.

  2. In the Blueprints section, type batch-submit into the filter, and then choose the batch-submit-job-python27 blueprint.

  3. In the Basic information section, configure your Lambda function:

    1. For Name, type SubmitJob.

    2. For Role, select Choose an existing role.

    3. For Existing role, select the Lambda role that you created earlier.

      Note

      If the IAM role that you created doesn't appear in the list, the role might still need a few minutes to propagate to Lambda.

  4. The following code is displayed in the Lambda function code pane:

    from __future__ import print_function
    
    import json
    import boto3
    
    print('Loading function')
    
    batch = boto3.client('batch')
    
    
    def lambda_handler(event, context):
        # Log the received event
        print("Received event: " + json.dumps(event, indent=2))
        # Get parameters for the SubmitJob call
        # http://docs.aws.amazon.com/batch/latest/APIReference/API_SubmitJob.html
        jobName = event['jobName']
        jobQueue = event['jobQueue']
        jobDefinition = event['jobDefinition']
        # containerOverrides and parameters are optional
        if event.get('containerOverrides'):
            containerOverrides = event['containerOverrides']
        else:
            containerOverrides = {}
        if event.get('parameters'):
            parameters = event['parameters']
        else:
            parameters = {}
    
        try:
            # Submit a Batch Job
            response = batch.submit_job(jobQueue=jobQueue, jobName=jobName, jobDefinition=jobDefinition,
                                        containerOverrides=containerOverrides, parameters=parameters)
            # Log response from AWS Batch
            print("Response: " + json.dumps(response, indent=2))
            # Return the jobId
            jobId = response['jobId']
            return {
                'jobId': jobId
            }
        except Exception as e:
            print(e)
            message = 'Error submitting Batch Job'
            print(message)
            raise Exception(message)

    The lambda_handler method assembles the job status poller using the event field of the input data. This input data is provided by the jobName, jobDefinition, and jobQueue objects passed into your function.

    Note

    jobName corresponds to the name of the AWS Batch job to be submitted.

    You will add input data for your function later, when you start a new execution.

  5. Choose Create function.

    When your Lambda function is created, note its ARN in the upper-right corner of the page. For example:

    arn:aws:lambda:us-east-1:123456789012:function:SubmitJob

To create the job status Lambda function

Important

Ensure that your Lambda function is under the same AWS account as your state machine.

  1. On the Functions page, choose Create function.

  2. In the Blueprints section, type batch-get into the filter, and then choose the batch-get-job-python27 blueprint.

  3. In the Basic information section, configure your Lambda function:

    1. For Name, type CheckJob.

    2. For Role, select Choose an existing role.

    3. For Existing role, select the Lambda role that you created earlier.

      Note

      If the IAM role that you created doesn't appear in the list, the role might still need a few minutes to propagate to Lambda.

  4. The following code is displayed in the Lambda function code pane:

    from __future__ import print_function
    
    import json
    import boto3
    
    print('Loading function')
    
    batch = boto3.client('batch')
    
    
    def lambda_handler(event, context):
        # Log the received event
        print("Received event: " + json.dumps(event, indent=2))
        # Get jobId from the event
        jobId = event['jobId']
    
        try:
            # Call DescribeJobs
            response = batch.describe_jobs(jobs=[jobId])
            # Log response from AWS Batch
            print("Response: " + json.dumps(response, indent=2))
            # Return the jobtatus
            jobStatus = response['jobs'][0]['status']
            return jobStatus
        except Exception as e:
            print(e)
            message = 'Error getting Batch Job status'
            print(message)
            raise Exception(message)

    The lambda_handler method assembles the job status check using the event field of the input data. This input data is provided by the jobId object passed into your function.

  5. Choose Create function.

    When your Lambda function is created, note its ARN in the upper-right corner of the page. For example:

    arn:aws:lambda:us-east-1:123456789012:function:CheckJob

Step 4: Testing the Job Submission Lambda Function

Test your job submission Lambda function to see it in operation.

To test the job submission Lambda function

  1. On the SubmitJob page, choose Test.

  2. On the Configure test event dialog box, type SubmitJob for Event name.

  3. Replace the example data with the following, using the job definition name and job queue name of the AWS Batch job definition that you created earlier in the jobDefinition and jobQueue fields:

    Note

    jobName corresponds to the name of the AWS Batch job to be submitted.

    Copy
    { "jobName": "my-job", "jobDefinition": "my-job-definition", "jobQueue": "my-job-queue" }

    The "jobName", "jobDefinition", and "jobQueue" entries correspond to the fields in your Lambda function, completing the notification. You will use the same input data when running the function as a Step Functions task.

  4. Choose Create.

  5. On the SubmitJob page, Test your Lambda function using the new data.

    The results of the test are displayed at the top of the page.

Step 5: Creating a State Machine

Use the Step Functions console to create a state machine with a Task state. Add a reference to your Lambda function in the Task state. The Lambda function is invoked when an execution of the state machine reaches the Task state.

To create the state machine

  1. Log in to the Step Functions console, and then choose Get Started.

  2. On the Create a state machine page, Name your state machine, for example JobStatusPoller.

    Note

    State machine names must be 1-80 characters in length, must be unique for your account and region, and must not contain any of the following:

    • Whitespace

    • Whitespace characters (? *)

    • Bracket characters (< > { } [ ])

    • Special characters (: ; , \ | ^ ~ $ # % & ` ")

    • Control characters (\\u0000 - \\u001f or \\u007f - \\u009f).

  3. Choose the Job status poller blueprint.

  4. In the Code pane, add the ARNs of the two Lambda functions that you created earlier to the corresponding Resource fields, for example:

    Copy
    { "Comment": "A state machine that submits a Job to AWS Batch and monitors the Job until it completes.", "StartAt": "Submit Job", "States": { "Submit Job": { "Type": "Task", "Resource": "arn:aws:lambda:us-east-1:123456789012:function:SubmitJob", "ResultPath": "$.guid", "Next": "Wait X Seconds" }, "Wait X Seconds": { "Type": "Wait", "SecondsPath": "$.wait_time", "Next": "Get Job Status" }, "Get Job Status": { "Type": "Task", "Resource": "arn:aws:lambda:us-east-1:123456789012:function:CheckJob", "Next": "Job Complete?", "InputPath": "$.guid", "ResultPath": "$.status" }, "Job Complete?": { "Type": "Choice", "Choices": [ { "Variable": "$.status", "StringEquals": "FAILED", "Next": "Job Failed" }, { "Variable": "$.status", "StringEquals": "SUCCEEDED", "Next": "Get Final Job Status" } ], "Default": "Wait X Seconds" }, "Job Failed": { "Type": "Fail", "Cause": "AWS Batch Job Failed", "Error": "DescribeJob returned FAILED" }, "Get Final Job Status": { "Type": "Task", "Resource": "arn:aws:lambda:us-east-1:123456789012:function:CheckJob", "InputPath": "$.guid", "End": true } } }

    This is a description of your state machine using the Amazon States Language. It defines a Task state named Submit Job, a Wait state named Wait X Seconds, and a Task state named Get Job Status. For more information, see State Machine Structure.

  5. Use the graph in the Visual Workflow pane to check that your Amazon States Language code describes your state machine correctly.

    If you don't see the graph, choose 
       refresh
    in the Visual Workflow pane.

  6. Choose Create State Machine.

    On the IAM role for your state machine executions dialog box, Step Functions creates and selects an IAM role automatically.

    Note

    If you delete the IAM role that Step Functions creates, Step Functions can't recreate it later. Similarly, if you modify the role (for example, by removing Step Functions from the principals in the IAM policy), Step Functions can't restore its original settings later. For more information about creating an IAM role manually, see Creating IAM Roles for AWS Step Functions.

  7. Choose OK.

    The state machine is created and an acknowledgement page is displayed.

Step 6: Starting a New Execution

After you create your state machine, you can start an execution.

To start a new execution

  1. On the JobStatusPoller page, choose New execution.

    The New execution page is displayed.

  2. (Optional) To help identify your execution, you can enter an ID for it. To specify the ID, use the Enter your execution id here text box. If you don't enter an ID, Step Functions generates a unique ID automatically.

  3. In the execution input area, replace the example data with the following:

    Copy
    { "jobName": "my-job", "jobDefinition": "my-job-definition", "jobQueue": "my-job-queue", "wait_time": 5 }

    Note

    wait_time instructs the Wait state to loop every five seconds.

  4. Choose Start Execution.

    A new execution of your state machine starts, and a new page showing your running execution is displayed.

  5. (Optional) In the Execution Details section, choose Info to view the Execution Status and the Started and Closed timestamps.

  6. To view the changing status of your AWS Batch job and the looping results of your execution, choose Output.