aws-step-functions-data-science-sdk-python
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
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)
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!