Skip to content

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

Open
@a13zen

Description

@a13zen

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)

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions