aws / aws-step-functions-data-science-sdk-python

Step Functions Data Science SDK for building machine learning (ML) workflows and pipelines on AWS
Apache License 2.0
285 stars 87 forks source link

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

Open a13zen opened 2 years ago

a13zen commented 2 years ago

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)
ca-nguyen commented 2 years ago

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!