AWS step functions task success callback using Boto3

Question

Overview

When using the start_execution method for an AWS Step Function with the SDK for Python (Boto3) I have added a 'time.sleep(6)' call to allow the step function execution to complete as a temporary fix.

If I don't add this the function execution fails as the state machine has not completed.

Code

def runStateMachine(dictInput):
    response = client.start_execution(
        stateMachineArn=arn,
        input=json.dumps(dictInput))
    time.sleep(6)
    executionArn = response["executionArn"]
    desc_exec_resp = client.describe_execution(
        executionArn=executionArn)
    return json.loads(desc_exec_resp["output"])

Question

Can I use some kind of callback or token to complete this without the sleep call?

If I understand correctly this SendTaskSuccess API call may be what I need, but I can't find any code examples.

Answer 1

I am curious how you can catch the status back in the same lambda invocation without waiting. Even if you are using a callback, it will end up in a totally separate invocation of lambda depending on the lambda concurrency settings.

API Call SendTaskSuccess used in Callbacks pattern where a workflow is paused until a task token is returned. The task will pause until it receives that task token back with a SendTaskSuccess or SendTaskFailure call.

I would break this handle execution output part into a separate lambda, save money and time.

For that, you can have a EventBridge (CloudWatch Events) for Step Functions execution status changes configured which can invoke your lambda about the full details. Below is an example of SUCCEEDED state of step function execution, and you can hook your lambda to it.

{
  "source": ["aws.states"],
  "detail-type": ["Step Functions Execution Status Change"],
  "detail": {
    "status": ["SUCCEEDED"],
    "stateMachineArn": ["arn:aws:states:eu-central-1:1234567890:stateMachine:Mapstate", ""]
  }
}

Callbacks are might not be good for this flow. Below is an example of send_task_success

Callback Pattern Example (Amazon SQS, Amazon SNS, Lambda)

import boto3

step_functions = boto3.client('stepfunctions')

def handler(event, context):
    # extract token from the event

    task_token = event['taskToken']
    step_functions.send_task_success(
        taskToken=task_token,
        output='json output of the task'
    )