(Aws-Stepfunctions): (Issue with catch block in Custom State)

[nivedita104]

Describe the bug

I have a use case where I am using a custom step in one of our state machines and in the catch block, i want to add a new step that is not previously referenced in the state machine. When I tried deploying this, the statemachine throws an error that
"Invalid State Machine Definition: 'MISSING_TRANSITION_TARGET: Missing 'Next' target: testLambda at /States/CreateTrainingJobWithou
tItemsDataset/Catch[0]/Next, MISSING_TRANSITION_TARGET: Missing 'Next' target: testLambda at /States/CreateTrainingJobWithItemsDataset/Catch[0]/Next' (Service: AWSSte
pFunctions; Status Code: 400; Error Code: InvalidDefinition; Request ID: 77eb25f8-02b7-49b3-8cdd-b1ff1921ec5f; Proxy: null)" (RequestToken: 92e222c1-409c-cfb0-558a-15
f75529c9c7, HandlerErrorCode: InvalidRequest)"

Here, testLambda was the step I introduced in the catch block of the custom step.

I checked the generated template it was not generating the "testLambda" step. Hence i tried to use the .addCatch() function call that we use with the other steps and i got an error that the .addCatch handler is not supported for custom steps.

What is the solution for this? Can we not add new steps for catch block in the custom steps or am i missing something here?

Expected Behavior

The step mentioned in the Next step of the catch block of the custom step props should also get generated in the cfn template.

Current Behavior

"Invalid State Machine Definition: 'MISSING_TRANSITION_TARGET: Missing 'Next' target: testLambda at /States/CreateTrainingJobWithou
tItemsDataset/Catch[0]/Next, MISSING_TRANSITION_TARGET: Missing 'Next' target: testLambda at /States/CreateTrainingJobWithItemsDataset/Catch[0]/Next' (Service: AWSSte
pFunctions; Status Code: 400; Error Code: InvalidDefinition; Request ID: 77eb25f8-02b7-49b3-8cdd-b1ff1921ec5f; Proxy: null)" (RequestToken: 92e222c1-409c-cfb0-558a-15
f75529c9c7, HandlerErrorCode: InvalidRequest)"

Reproduction Steps

Default Custom state example in docs:

import * as dynamodb from '@aws-cdk/aws-dynamodb';

// create a table
const table = new dynamodb.Table(this, 'montable', {
  partitionKey: {
    name: 'id',
    type: dynamodb.AttributeType.STRING,
  },
});

const finalStatus = new sfn.Pass(this, 'final step');

// States language JSON to put an item into DynamoDB
// snippet generated from https://docs.aws.amazon.com/step-functions/latest/dg/tutorial-code-snippet.html#tutorial-code-snippet-1
const stateJson = {
  Type: 'Task',
  Resource: 'arn:aws:states:::dynamodb:putItem',
  Parameters: {
    TableName: table.tableName,
    Item: {
      id: {
        S: 'MyEntry',
      },
    },
  },
  ResultPath: null,
};

// custom state which represents a task to insert data into DynamoDB
const custom = new sfn.CustomState(this, 'my custom task', {
  stateJson,
});

const chain = sfn.Chain.start(custom)
  .next(finalStatus);

const sm = new sfn.StateMachine(this, 'StateMachine', {
  definition: chain,
  timeout: Duration.seconds(30),
});

// don't forget permissions. You need to assign them
table.grantWriteData(sm);

Now, we add a new pass state and catch block to the custom step :

const errorHandlerLambda = new Pass(
        this,
        "errorHandlerLambda"
    );

const stateJson = {
  Type: 'Task',
  Resource: 'arn:aws:states:::dynamodb:putItem',
  Parameters: {
    TableName: table.tableName,
    Item: {
      id: {
        S: 'MyEntry',
      },
    },
  },
Catch: [
          {
            ErrorEquals: [Errors.ALL],
            ResultPath: "$.JobDetails.errorInfo",
            Next: 'errorHandlerLambda'
          }
        ],
  ResultPath: null,
};

This will throw the error I mentioned above.
I cannot add the errorHandlerLambda to the step function definition as it will only be used as the next step in case the custom step fails.

Possible Solution

No response

Additional Information/Context

No response

CDK CLI Version

"aws-cdk-lib": "^2.60.0", CDKBuild = 4.x;

Framework Version

No response

Node.js Version

12.x

OS

Linux

Language

Typescript

Language Version

No response

Other information

No response

[peterwoodworth]

Could you post the full example where you add the errorHandlerLambda to the definition?

Also, why not use DynamoPutItem?

[nivedita104]

I just gave you an example to recreate the issue. The custom step example is the one that is present in the doc. I just added the catch block to it.
PFB the state machine file that I used for testing:

import { DeploymentStack, DeploymentStackProps } from "@amzn/pipelines";
import { App, Duration } from "aws-cdk-lib";
import {
  Choice,
  Condition,
  CustomState,
  JsonPath,
  LogLevel,
  Pass,
  Result,
  StateMachine,
  TaskInput,
} from "aws-cdk-lib/aws-stepfunctions";
import { LogGroup } from "aws-cdk-lib/aws-logs";
import { Constants } from "../config/constants";
import { Function } from "aws-cdk-lib/aws-lambda";
import {
  EvaluateExpression,
  LambdaInvoke,
  S3DataDistributionType,
  S3DataType,
} from "aws-cdk-lib/aws-stepfunctions-tasks";
import { mote_iam as moteiam } from "@amzn/motecdk";
import {
  getSageMakerPassRolePolicyStatement,
  getSageMakerTrainingJobPolicyStatement,
  getStepFunctionsEventRulePolicyStatement,
} from "../util/iamUtil";
import {
  BACKOFF_RATE,
  MAX_RETRY_NUMBER,
  lambdaRetryProps,
  sagemakerRetryableExceptions, sagemakerExceptionsToCatch,
} from "../util/retryPropsUtil";
import { AwsRegion } from "@amzn/awsstarkcommoninfrastructurecdk";

export interface ModelTrainingWorkflowProps extends DeploymentStackProps {
  readonly disambiguator: string;
  readonly parentWorkflowName: string;
  readonly populateSagemakerTrainingJobLambda: Function;
  readonly sagemakerRole: moteiam.SecureRole;
  readonly region: AwsRegion;
  readonly accountId: string;
  readonly logGroup: LogGroup;
}

export class ModelTrainingWorkflow extends DeploymentStack {
  readonly stateMachine: StateMachine;

  constructor(scope: App, name: string, props: ModelTrainingWorkflowProps) {
    super(scope, name, props);

    const setRetryCount = new Pass(this, "settingRetryCount to 0",
        {
          result: Result.fromNumber(-1),
          resultPath: "$.Input.RetryCount",
        }
        );

    const addOneToRetryCount = new EvaluateExpression(this, "addOneToRetryCount",
        {
          expression: '$.Input.RetryCount + 1',
          resultPath: '$.Input.RetryCount',
        }
    );
    const populateSagemakerTrainingJobAttributes = new LambdaInvoke(
      this,
      "PopulateSagemakerTrainingJobAttributes",
      {
        lambdaFunction: props.populateSagemakerTrainingJobLambda,
        payload: TaskInput.fromObject({
          "sagemakerJobName.$": "$$.Execution.Name",
          "audienceGeneratorDetails.$": "$.Input.Payload",
          "s3Path.$": "$.Input.Account.s3Path",
          "accountId.$": "$.Input.Account.accountId",
          "region.$": "$.Input.Account.region",
        }),
        resultPath: "$.JobDetails.Training",
      }
    ).addRetry(lambdaRetryProps);

    const modelTrainingWithoutItemsTask = new CustomState(
      this,
      "CreateTrainingJobWithoutItemsDataset",
      ModelTrainingWorkflow.getTrainingJobCustomStateProps(
        props.sagemakerRole.roleArn,
        false
      )
    );
    const testLambda = new Pass(
        this,
        "testLambda",
    );
    const modelTrainingWithItemsTask = new CustomState(
      this,
      "CreateTrainingJobWithItemsDataset",
      ModelTrainingWorkflow.getTrainingJobCustomStateProps(
        props.sagemakerRole.roleArn,
        true
      )
    );

    const itemsDataSetChoice = new Choice(this, "ItemsDatasetChoice")
      .when(
        Condition.booleanEquals(
          "$.JobDetails.Training.Payload.itemDataSetAvailable",
          true
        ),
        modelTrainingWithItemsTask
      )
      .otherwise(modelTrainingWithoutItemsTask);

    const definition = setRetryCount
        .next(addOneToRetryCount)
        .next(populateSagemakerTrainingJobAttributes)
        .next(itemsDataSetChoice);

    this.stateMachine = new StateMachine(
      this,
      Constants.MODEL_TRAINING_WORKFLOW,
      {
        definition,
        stateMachineName: `${props.parentWorkflowName}-${Constants.MODEL_TRAINING_WORKFLOW}-${props.disambiguator}`,
        timeout: Duration.days(7),
        logs: {
          destination: props.logGroup,
          level: LogLevel.ALL,
        },
      }
    );
    this.stateMachine.addToRolePolicy(
      getSageMakerPassRolePolicyStatement(props.sagemakerRole.roleArn)
    );
    this.stateMachine.addToRolePolicy(
      getSageMakerTrainingJobPolicyStatement(props.region, props.accountId)
    );
    this.stateMachine.addToRolePolicy(
      getStepFunctionsEventRulePolicyStatement(
        props.region,
        props.accountId,
        "SageMakerTrainingJobs"
      )
    );
  }

  private static getTrainingJobCustomStateProps(
    roleArn: string,
    itemsChannel: boolean
  ) {
    const inputDataConfig = [
      {
        ChannelName: "train",
        DataSource: {
          S3DataSource: {
            S3DataDistributionType: S3DataDistributionType.FULLY_REPLICATED,
            S3DataType: S3DataType.S3_PREFIX,
            "S3Uri.$": "$.JobDetails.Training.Payload.trainingS3Channel",
          },
        },
      },
      {
        ChannelName: "valid",
        DataSource: {
          S3DataSource: {
            S3DataDistributionType: S3DataDistributionType.FULLY_REPLICATED,
            S3DataType: S3DataType.S3_PREFIX,
            "S3Uri.$": "$.JobDetails.Training.Payload.validationS3Channel",
          },
        },
      },
      {
        ChannelName: "test",
        DataSource: {
          S3DataSource: {
            S3DataDistributionType: S3DataDistributionType.FULLY_REPLICATED,
            S3DataType: S3DataType.S3_PREFIX,
            "S3Uri.$": "$.JobDetails.Training.Payload.validationS3Channel",
          },
        },
      },
    ];
    if (itemsChannel) {
      inputDataConfig.push({
        ChannelName: "items",
        DataSource: {
          S3DataSource: {
            S3DataDistributionType: S3DataDistributionType.FULLY_REPLICATED,
            S3DataType: S3DataType.S3_PREFIX,
            "S3Uri.$": "$.JobDetails.Training.Payload.items3Channel",
          },
        },
      });
    }
    return {
      stateJson: {
        Resource: "arn:aws:states:::sagemaker:createTrainingJob.sync",
        Parameters: {
          "TrainingJobName.$": "$.JobDetails.Training.Payload.trainingJobName",
          RoleArn: roleArn,
          AlgorithmSpecification: {
            "TrainingInputMode.$":
              "$.JobDetails.Training.Payload.trainingInputMode",
            "TrainingImage.$": "$.JobDetails.Training.Payload.trainingImage",
            MetricDefinitions: Constants.METRICS_DEFINITION_SEQUENCE_MODEL,
          },
          InputDataConfig: inputDataConfig,
          OutputDataConfig: {
            "S3OutputPath.$": "$.JobDetails.Training.Payload.outputS3Location",
            "KmsKeyId.$": "$.JobDetails.Training.Payload.artifactsKmsKeyArn",
          },
          ResourceConfig: {
            "InstanceCount.$":
              "$.JobDetails.Training.Payload.resourceConfig.instanceCount",
            "InstanceType.$":
              "$.JobDetails.Training.Payload.resourceConfig.instanceType",
            "VolumeSizeInGB.$":
              "$.JobDetails.Training.Payload.resourceConfig.volumeSize",
            "VolumeKmsKeyId.$":
              "$.JobDetails.Training.Payload.artifactsKmsKeyArn",
          },
          StoppingCondition: {
            "MaxRuntimeInSeconds.$": "$.JobDetails.Training.Payload.maxRuntime",
          },
          "HyperParameters.$": "$.JobDetails.Training.Payload.hyperparameters",
        },
        Type: "Task",
        Retry: [
          {
            ErrorEquals: sagemakerRetryableExceptions,
            MaxAttempts: MAX_RETRY_NUMBER,
            BackoffRate: BACKOFF_RATE,
            IntervalSeconds: 30,
          },
        ],
        Catch: [
          {
            ErrorEquals: sagemakerExceptionsToCatch,
            ResultPath: "$.JobDetails.Training.errorInfo",
            Next: 'testLambda'
          }
        ]
      },
    };
  }
}

The custom step in our state machine is used for starting a sagemaker training job. Though we started of by using the "SageMakerCreateTrainingJob", this did not support kms keys and hence we had to switch to a custom step and now, we want to add error handling to it and are facing difficulties with the catch block.

[vishak28]

@nivedita104 were you able to find a work-around for this?

[studna]

I had the same issue and I switched from using CustomState to tasks.CallAwsService. More complex flows I have to write as raw ASL and register it as:

new CfnStateMachine(this, 'MagicWorkflow', {
  roleArn,
  definition, // json definition that you can either write manually or export from AWS Workflow Studio
})

I would love to know if there's a better way.

[AaronWoodrow]

We're having exactly the same issue. We have a state machine the we'd like to define using chained CustomStates similar to the original example above but Catch is directing to another CustomState that it cannot find.

[object-Object]

I also ran into this issue, and found this workaround (Python):

custom_state._add_choice(sfn.Condition.is_null("$"), handler_state)

This is hacky, but it works because State.addChoice() basically has the same functionality as State._addCatch(), which is what subclasses call to add a catch handler to the state graph. CustomState doesn't use the private choices list (or catches) in its toStateJson() implementation, so this doesn't actually add a choice object to the output.

I think the TypeScript equivalent would be this:

customState.addChoice(sfn.Condition.is_null("$"), handlerState)

You might also be able to call _addCatch() directly, or create a subclass which implements addCatch() in the same way as the existing constructs. We just can't do that from Python because _addCatch() is marked as @internal.

[mschwieb]

For TypeScript I got around this with this hack:

const customStateExample = new CustomState(...);
const stateThatDependsOnCustomStateExample = ...;

// HACK: because CDK doesn't support addCatch for CustomState yet: https://github.com/aws/aws-cdk/issues/25798
// @ts-ignore
customStateExample.stateJson = {
  // @ts-ignore
  ...customStateExample.stateJson,
  Catch: [
    {
      ErrorEquals: ['States.ALL'],
      ResultPath: '$.retryConfig',
      Next: stateThatDependsOnCustomStateExample.id,
    },
  ],
};
 

[object-Object]

@mschwieb Does that actually add stateThatDependsOnCustomStateExample to the state machine definition if it's only used as the error handler for customStateExample (ie. no non-custom state has stateThatDependsOnCustomStateExample as a .next() state)? That's what this issue is about.

[mschwieb]

yes, stateThatDependsOnCustomStateExample is added to the SFN definition as it is a state reached when there is an error thrown. It is similar to the example described by OP.

[object-Object]

Here's a reasonably minimal set of test cases to show the issue. Both tests with addNext=false are failing as of CDK 2.93.0.

test/cdk-ts.test.ts

import * as cdk from 'aws-cdk-lib';
import { Template, Match, Capture } from 'aws-cdk-lib/assertions';
import { AppendStack, InlineStack } from '../lib/cdk-ts-stack';

describe.each([InlineStack, AppendStack])("%p", (clazz) => {
    it.each([true, false])("adds ErrorHandler to the definition when addNext=%p", (addNext) => {
        // synth the app
        const app = new cdk.App();
        const stack = new clazz(app, "Stack", addNext);
        const template = Template.fromStack(stack);
    
        // capture the states definition
        const statesCapture = new Capture();
        
        template.hasResourceProperties("AWS::StepFunctions::StateMachine", {
            DefinitionString: Match.serializedJson({
                StartAt: "CustomState",
                States: statesCapture,
            })
        })
    
        // assert ErrorHandler was added to the definition
        expect(statesCapture.asObject()).toHaveProperty("ErrorHandler");
    })
})

lib/cdk-ts-stack.ts

import * as cdk from 'aws-cdk-lib';
import { Construct } from 'constructs';
import * as sfn from 'aws-cdk-lib/aws-stepfunctions';

// add the catcher in the state definition
export class InlineStack extends cdk.Stack {
  constructor(scope: Construct, id: string, addNext: boolean) {
    super(scope, id);

    let errorHandler = new sfn.Pass(this, "ErrorHandler");

    // this could be any Task, Map, or Parallel which can't be represented with existing constructs
    let customState = new sfn.CustomState(this, "CustomState", {
      stateJson: {
        Type: "Parallel",
        Branches: [],
        Catch: {
          ErrorEquals: ["States.ALL"],
          Next: errorHandler.id,
        },
      }
    })
    
    new sfn.StateMachine(this, "StateMachine", {
      definitionBody: sfn.DefinitionBody.fromChainable(
        // for the test, conditionally add ErrorHandler as the next state after CustomState
        addNext ? customState.next(errorHandler) : customState
      )
    })
  }
}

// add the catcher after defining the state by modifying private properties
export class AppendStack extends cdk.Stack {
  constructor(scope: Construct, id: string, addNext: boolean) {
    super(scope, id);

    let errorHandler = new sfn.Pass(this, "ErrorHandler");

    // this could be any Task, Map, or Parallel which can't be represented with existing constructs
    let customState = new sfn.CustomState(this, "CustomState", {
      stateJson: {
        Type: "Parallel",
        Branches: [],
      }
    })

    // @ts-ignore
    customState.stateJson = {
      // @ts-ignore
      ...customState.stateJson,
      Catch: {
        ErrorEquals: ["States.ALL"],
        Next: errorHandler.id,
      }
    }
    
    new sfn.StateMachine(this, "StateMachine", {
      definitionBody: sfn.DefinitionBody.fromChainable(
        // for the test, conditionally add ErrorHandler as the next state after CustomState
        addNext ? customState.next(errorHandler) : customState
      )
    })
  }
}

[Function InlineStack]
    √ adds ErrorHandler to the definition when addNext=true (68 ms)
    × adds ErrorHandler to the definition when addNext=false (45 ms)
  [Function AppendStack]
    √ adds ErrorHandler to the definition when addNext=true (20 ms)
    × adds ErrorHandler to the definition when addNext=false (19 ms)

  ● [Function InlineStack] › adds ErrorHandler to the definition when addNext=false

    expect(received).toHaveProperty(path)

    Expected path: "ErrorHandler"
    Received path: []

    Received value: {"CustomState": {"Branches": [], "Catch": {"ErrorEquals": ["States.ALL"], "Next": "ErrorHandler"}, "End": true, "Type": "Parallel"}}

  ● [Function AppendStack] › adds ErrorHandler to the definition when addNext=false

    expect(received).toHaveProperty(path)

    Expected path: "ErrorHandler"
    Received path: []

    Received value: {"CustomState": {"Branches": [], "Catch": {"ErrorEquals": ["States.ALL"], "Next": "ErrorHandler"}, "End": true, "Type": "Parallel"}}

[joonb14]

I tried above solutions, but still did not solve my problem.
So I just used StepFunctionsStartExecution to execute my CustomState and applied addCatch to the task. This is messy work-around but it worked for me.
P.S I created 2 files for each child state machine and parent state machine to prevent circular dependency.

• ChildStateMachine

...
export class ChildStateMachine extends Construct {
  readonly stateMachine : StateMachine;
  constructor(scope: Construct, id: string, props:  ... ) {
    super(scope, id);

    // My CustomState
    const customState = new sfn.CustomState( ... )
    
    // Define a state machine
    this.stateMachine = new sfn.StateMachine(this, "ChildStateMachine", {
        definitionBody: sfn.DefinitionBody.fromChainable(customState),
        role: stateMachineRole,
        timeout: Duration.hours(1),
    });
  }
}

• ParentStateMachine

...
export class ParentStateMachine extends Construct {
  readonly stateMachine : StateMachine;
  constructor(scope: Construct, id: string, props:  ... ) {
    super(scope, id);

    // Include the state machine in a Task state with callback pattern
    const task = new sfntasks.StepFunctionsStartExecution(this, 'ChildTask', {
        stateMachine: childStateMachine,
        integrationPattern: sfn.IntegrationPattern.RUN_JOB,
        taskTimeout: sfn.Timeout.duration(Duration.hours(1)),
    });
    
    // add Catch
    // errorHandlingLambdaInvoke is my sfntasks.LambdaInvoke which has sfn.Fail state in next
    // Use your error in props errors.
    task.addCatch(errorHandlingLambdaInvoke, {
        errors: ['States.TaskFailed', ... ],
        resultPath: '$.errorInfo',
    });

    // Define a second state machine with the Task state above
    this.stateMachine = new sfn.StateMachine(this, 'ParentStateMachine', {
        definition: task,
        role: stateMachineRole,
        timeout: Duration.hours(1),
    });
  }
}

• In certain stack

...
    const child = new ChildStateMachine( ... );
    const parent = new ParentStateMachine(this, "ParentStateMachine", { child.stateMachine, ... };
...

[mfekadu]

This worked for me

Suppose there's a chain able failurePath

const failurePath: IChainable = ...;

And there exists a customState which already has its own Catch

const customState = new CustomState(
    scope,
    "MyCustomState",
    {
        stateJson: {
            ...
            "Catch": [ 
                ...,
                "Next": failurePath.startState.stateId
            ],
            ...
        }
    }
);

After all that, our failure path might not find its way into the final state machine definition body unless some other state calls the addCatch method.

So we can just take advantage of the core _addCatch which every State object inherits, per class methods inside of state.js

// Sometimes the failure path is not added to the state machine definition, unless _addCatch is called
// https://github.com/aws/aws-cdk/issues/25798#issuecomment-1749916284
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
// @ts-ignore
customState._addCatch(failurePath.startState) // intentionally omit the CatchProps because my CustomState already sets it up

Still hacky, sadly... but arguably close to the valid solution.

AWS folks, please just implement the addCatch method on CustomState. That's all we need here.

[ocratravis]

I've run into the exact same issue.

Thanks to everyone who has suggested workarounds, but I hate resorting to hacks. Would love to see more robust support in CDK.

[github-actions]

⚠️COMMENT VISIBILITY WARNING⚠️

Comments on closed issues are hard for our team to see.
If you need more assistance, please either tag a team member or open a new issue that references this one.
If you wish to keep having a conversation with other community members under this issue feel free to do so.