Choice states cannot be chained

[piotrdrazikowski]

I'm currently adopting this package in place of own code to generate a state machine definition and one of our main patterns is to have a choice state that checks for a certain variable and kicks off previous step again if the condition was met, or move to the next step if not.
Example (failing) code:

import stepfunctions

first_job = stepfunctions.steps.Pass('First job')
second_job = stepfunctions.steps.Pass('Second job')

check = stepfunctions.steps.Choice('Check first job')
check.add_choice(
    rule=stepfunctions.steps.ChoiceRule.BooleanEquals(
        variable='$run_me_again',
        value=True
    ),
    next_step=first_job
)
check.default_choice(second_job) # This could be set automatically

chain = stepfunctions.steps.Chain([first_job, check, second_job])

This code currently gives State type `Choice` does not support method `next`.

I understand that the choice state was designed to branch the workflow and that terminates the chain. Please correct me if i'm wrong but in my opinion the presence of the next step can be treated as the default_choice for the choice state.

Here is a quick workaround that solves our issue:

class ChainableChoice(stepfunctions.steps.Choice):
    def next(self, next_step):
        self.default_choice(next_step)
        return next_step

I'd be more than happy to submit a proper PR, please let me know what do you think about it.

[vincentclaes]

i just had the same issue. it looks like that after a choice state you cannot add anymore steps to the Chain.

if you do

chain = stepfunctions.steps.Chain([first_job, check])

you step function should work

on top of that if you want to add downstream steps to "second_job" you need to define this codewise and not in the the Chain object. for example

second_job.next(third_job)

[ctippur]

I am having similar issue. @vincentclaes thanks for the pointer. This kind of took me out of the woods for the second iteration.
In my pipeline, I have several choice steps that needs to be in series.

second_job.next(third_job)
third_job.next(new_check)

At this time, I am getting the same error State type `Choice` does not support method `next`.

[ctippur]

looking at https://states-language.net/spec.html#choice-state
The values of the “And” and “Or” operators MUST be non-empty arrays of Choice Rules that MUST NOT contain “Next” fields; the “Next” field can only appear in a top-level Choice Rule.

Trying to understand this..

[shunjd]

Hi @ctippur, in the Choice state, there is no Next field, so we don't support chaining it.

[jennakwon06]

Why does it make sense for Choice state to not have Next field? It sounds like a natural thing for the State to have. I am wondering if I am conceptually missing something

[jennakwon06]

Let me explain our use case of needing Choice state to have the Next field.

Our use case is Machine learning workflow automation. So our customers are scientists that are submitting the list of Steps they would like to do in a JSON file. Each Step in the JSON can have "prior step" dependency, so essentially I am converting the JSON to a DAG of Steps. Each Step can be a "TrainingStep", "LambdaStep", or a Choice state with multiple Lambda Steps (execute-wait, etc).

When I have a DAG and try to convert it to a State Machine with Data Science SDK, I am creating these steps then programatically determining the "next" relationships (with DFS traversal using Chain/parallel). This means that I create a Choice state, then later determine what that Choice state's "next" field should be. But since Choice state does not support method "next", I am not able to do this.

I think supporting Choice state to have a "next" method is useful for more complex use cases like mine where we are programatically creating State Machines.

[wong-a]

By design, Choice states in Amazon States Language do not have a Next field. The purpose of the Choice state is to determine the next state transition based on the provided Choice Rules to enable conditional control flow. This differs from other non-terminal states where Next specifies an unconditional transition.

[nsankar]

@wong-a @shunjd This AWSs example post mentions about the next_step in the choice state. Not sure !? . Is this incorrect ?

https://github.com/aws/aws-step-functions-data-science-sdk-python/blob/08a0ed7f3aac667b5c704baf8016a8531a3509ed/doc/choicerules.rst

https://aws.amazon.com/blogs/machine-learning/building-machine-learning-workflows-with-amazon-sagemaker-processing-jobs-and-aws-step-functions/

[wong-a]

@nsankar Those links are correct. At the top-level, a Choice state has no Next field. The next state is determined by one or more ChoiceRules that specify the next transition based on a condition. You call choice.add_choice(rule, next_step) to add a rule to a Choice state. A Choice state may have a top-level Default state to transition to if none of the ChoiceRules match. This is done in the SDK by calling choice. default_choice(next_step).

Consider the following if/else code block that executes a different function by comparing the input string.

def make_a_choice(input):
    if input == "A":
        a()
    elif input == "B":
        b()
    else:
        c() # Fall back to c() if none match

To model something similar in the Step Functions SDK as a Choice state would look something like this:

# Some other states:
step_a = Pass(state_id="A")
step_b = Pass(state_id="B")
step_c = Pass(state_id="C")

choice_state = Choice("Make a Choice")

# If input is "A", go to step "A" next
choice_state.add_choice(
    rule=ChoiceRule.StringEquals(variable="$", value="A"),
    next_step=step_a # this could also be a Chain instead of a single state
)

# Else, if input is "B", go to step "B" next
choice_state.add_choice(
    rule=ChoiceRule.StringEquals(variable="$", value="B"),
    next_step=step_b
)

# Else, go to step "C" next
choice_state.default_choice(step_c)

I created a gist demonstrating how to use Choice states in the SDK with several examples so you can see what the code looks like and how they can be chained to other states in the workflow: https://gist.github.com/wong-a/e0ef01e4cb717ee277ed90d8eda8dbb7. The embedded images of the graph's don't render in gist, but will if you open it in Jupyter and they are attached. Hopefully this clears up how to use Choice states.

Based on the feedback here and some internal conversations with the Step Functions team, the current behaviour isn't very obvious. We are considering making ergonomic improvements to make Choice states chainable, where would set the top-level Default field instead of Next.

---

Relevant docs:

• https://aws-step-functions-data-science-sdk.readthedocs.io/en/stable/states.html#stepfunctions.steps.states.Choice
• https://aws-step-functions-data-science-sdk.readthedocs.io/en/stable/choicerules.html?highlight=choice%20rule
• https://docs.aws.amazon.com/step-functions/latest/dg/amazon-states-language-choice-state.html
• https://states-language.net/spec.html#transition
• https://states-language.net/spec.html#choice-state

[nsankar]

@wong-a Thank you for enlightening and providing the examples. It really helps.

I have a couple of related questions.

(1) Is there a way to get the actual elapsed time from a wait state and use it in a choice state ? something like as follows ? Kindly let me know how this can be done if feasible.

 wait_state = Wait(
        state_id='Wait120',
        seconds=120,
        )
//Possible to Call a lambda state when the elapsed time is 60 secs in the wait state? something like this ?
choice_wait_state1 = Choice(
    state_id="60secs time elapsed?"
)
choice_wait_state.add_choice(ChoiceRule.NumericEquals(wait_state.output()["seconds"], 60),  my_lambda_state)

(2) Which are the states including the common SFN states for which state output () is applicable or can be used? This is not in the document unless we see an example, for instance, as an example output () is applicable to the processing step which can be then used in a choice state as cited below:

processing_step = ProcessingStep(
    "SageMaker processing step",
    processor=script_processor,
    job_name=execution_input["PreprocessingJobName"], ## placeholder per schema above. Later in workflow execute() below pass the uuid of jobname
    inputs=inputs,
    outputs=outputs,
    # container_arguments=["--train-test-split-ratio", "0.2"],
    container_entrypoint=["python3", "/opt/ml/processing/input/code/myprogram.py"],
)

choice_state1.add_choice(ChoiceRule.StringEquals(processing_step**.output()**["ProcessingJobStatus"], "Completed"), lambda_state1)

[wong-a]

@nsankar Sorry for the delayed reply but I'm glad that helped.

(1) : If you want to start something after a certain time, you should use a Wait state. The state after a Wait state is only entered once the time condition is met (120 seconds in your example).

It is not possible to compare elapsed time using a ChoiceRule unless you calculate that using a Task to compute it.

(2): Can you create a separate issue? This one is for Choice state chaining. It seems like there is some documentation to improve around Placeholder usage but I don't quite understand the problem you are coming across.

[anatoly-scherbakov]

Thank you for the detailed description of the problem. In my particular case, I created a very ad hoc solution by overriding Choice.next() method and adding Succeed() to the end of the Chain(). Described that in this article in a bit more detail. But I did not arrive at a more abstract/generalized approach.

[wong-a]

Released in v2.1.0: https://github.com/aws/aws-step-functions-data-science-sdk-python/releases/tag/v2.1.0