aws-step-functions-data-science-sdk-python icon indicating copy to clipboard operation
aws-step-functions-data-science-sdk-python copied to clipboard

Feature Request: Add Workflow.create_or_update() to update workflow if it exists or create it if it doesn't

Open a13zen opened this issue 3 years ago • 1 comments

When running notebooks in an automated fashion many times the StepFunction created by the workflow.create() method already exists. The same for when running the workflow.update() method when defining the Workflow for the first time.

We could extend the Workflow class with a create_or_update method that would check if the stepfunction exists, if it does, attach to that otherwise, update it.

# Grab account/region from Boto/Sagemaker Session

class CustomWorkflow(Workflow):
    def __init__(self, name, definition, role, update_sleep_seconds=8, **kwargs):
        super(CustomWorkflow, self).__init__(name, definition, role, **kwargs)
        self.update_sleep_seconds = update_sleep_seconds
        self.expected_arn = f"arn:aws:states:{region}:{account_id}:stateMachine:{self.name}"
        
    def create_or_update(self):
        try:
            target_wf = Workflow.attach(self.expected_arn)
            print('State Machine Exists. Attaching and updating')
            target_wf.update(definition=self.definition, role=self.role)
            # Sleep to allow infrastructure to update before executing a new job
            time.sleep(self.update_sleep_seconds)
        except boto3.client('stepfunctions').exceptions.StateMachineDoesNotExist:
            print('State Machine Not found, creating')
            self.expected_arn = self.create()
        return Workflow.attach(self.expected_arn)

a13zen avatar Sep 16 '21 17:09 a13zen

Hi @a13zen, Having a function that would do the check for you and call the create/update function accordingly would surely be useful and save you he hassle to having to edit the notebook between runs. Thank you for bringing this to our attention!

ca-nguyen avatar Sep 30 '21 18:09 ca-nguyen