AWS - Step functions, use execution input within a TuningStep

user_1115237 0
Viewed 799

I've written a simple AWS step functions workflow with a single step:

from stepfunctions.inputs import ExecutionInput
from stepfunctions.steps import Chain, TuningStep
from stepfunctions.workflow import Workflow
import train_utils


def main():
    workflow_execution_role = 'arn:aws:iam::MY ARN'
    execution_input = ExecutionInput(schema={
        'app_id': str
    })
    estimator = train_utils.get_estimator()
    tuner = train_utils.get_tuner(estimator)

    tuning_step = TuningStep(state_id="HP Tuning", tuner=tuner, data={
        'train': f's3://my-bucket/{execution_input["app_id"]}/data/'},
                             wait_for_completion=True,
                             job_name='HP-Tuning')

    workflow_definition = Chain([
        tuning_step
    ])

    workflow = Workflow(
        name='HP-Tuning',
        definition=workflow_definition,
        role=workflow_execution_role,
        execution_input=execution_input
    )
    workflow.create()


if __name__ == '__main__':
    main()

My goal is to have the train input pulled from the execution JSON provided at runtime. When I execute the workflow (from the step functions console), providing the JSON {"app_id": "My App ID"} the tuning step does not get the right data, instead it gets a to_string representation of the stepfunctions.inputs.placeholders.ExecutionInput. Furthermore when looking at the generated ASL I can see that the execution input was rendered as a string:

... 
"DataSource": {
   "S3DataSource": {
   "S3DataType": "S3Prefix",
   "S3Uri": "s3://my-bucket/<stepfunctions.inputs.placeholders.ExecutionInput object at 0x12261f7d0>/data/",
    "S3DataDistributionType": "FullyReplicated"
    }
},
...

What am I doing wrong?

Update: As mentioned by @yoodan the SDK is probably behind, so I'll have to edit the definition before calling create. I can see there is a way to review the definition before calling create, but can I modify the graph definition? how?

hyperparameters aws-step-functions state-machine machine-learning python
3 Answers

The python SDK for step functions generates corresponding code, we need a string concatenation / format built into the Amazon States Language to accomplish what you desire.

Recently in August 2020, Amazon States Language introduced built-in functions such as string format into it's language spec. https://states-language.net/#appendix-b

Unfortunately, the python SDK is not up to date and does not support the new changes.

As a work around, maybe manually modify the definition before calling workflow create?

user_143367590

your definition of step function is good and it looks like it should work.
it is unclear from your code example where the execution actually happens (you stated directly from console) which can lead to several options that may cause the problem, and i believe thats the source pf the issue.
please provide more information regarding that.

are you somehow exporting your created Workflow object so it can be executed?
it seems that the something is missing... as a sanity check, append the following to your main function:

workflow.execute(inputs={"app_id": "My App ID"})

and check logs again

user_142421110

It seems like what you're looking for is not supported by the SDK. https://github.com/aws/aws-step-functions-data-science-sdk-python/issues/79

You can however, change the definition before creating the state machine. Here is a function that will iterate the definition dict and replace a place holder surrounded with {{PH}} with the right intrinsic functions syntax, for example :

s3://my-bucket/{{app_id}}/data/

def get_updated_definition(data):
    if isinstance(data, dict):
        for k, v in data.copy().items():
            if isinstance(v, dict):  # For DICT
                data[k] = get_updated_definition(v)
            elif isinstance(v, list):  # For LIST
                data[k] = [get_updated_definition(i) for i in v]
            elif isinstance(v, str) and re.search(r'{{([a-z_]+)}}', v):  # Update Key-Value
                # data.pop(k)
                # OR
                del data[k]
                keys = re.findall(r'{{([a-z_]+)}}', v)
                data[f"{k}.$"] = f"States.Format('{re.sub(r'{{[a-z_]+}}', '{}', v)}',{','.join(['$.'+k for k in keys])})"

    return data

usage:

workflow_definition = get_updated_definition(workflow.definition.to_dict())
create_state_machine(json.dumps(workflow_definition)) #<-- implement using boto3 (https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/stepfunctions.html#SFN.Client.create_state_machine)
user_11160680
Related