def wait_for_state(self, goal_state='SUCCEEDED', failed_state='ABORTED'):
"""
Waits for the workflow to reach a particular state
:param goal_state: The state to reach
:param failed_state: The failed state of the update
"""
state = ''
attempts = 0
start = time.time()
log.info('%s waiting for state %s', self.workflow_id, goal_state)
state_transition_failure = False
while attempts < self.config.max_retry_attempts:
try:
request = update_svc.GetUpdateRequest(
updateId=peloton.UpdateID(value=self.workflow_id), )
resp = self.client.update_svc.GetUpdate(
request,
metadata=self.client.jobmgr_metadata,
timeout=self.config.rpc_timeout_sec,
)
update_info = resp.updateInfo
new_state = update.State.Name(update_info.status.state)
if state != new_state:
log.info('%s transitioned to state %s', self.workflow_id,
new_state)
state = new_state
if state == goal_state:
break
# If we assert here, we will log the exception,
# and continue with the finally block. Set a flag
# here to indicate failure and then break the loop
# in the finally block
if state == failed_state:
state_transition_failure = True
except Exception as e:
log.warn(e)
finally:
if state_transition_failure:
break
time.sleep(self.config.sleep_time_sec)
attempts += 1
if state_transition_failure:
log.info('goal_state:%s current_state:%s attempts: %s', goal_state,
state, str(attempts))
assert False
if attempts == self.config.max_retry_attempts:
log.info('%s max attempts reached to wait for goal state',
self.workflow_id)
log.info('goal_state:%s current_state:%s', goal_state, state)
assert False
end = time.time()
elapsed = end - start
log.info('%s state transition took %s seconds', self.workflow_id,
elapsed)
assert state == goal_state
def get_state(self):
"""
get the current state of workflow
"""
request = update_svc.GetUpdateRequest(
updateId=peloton.UpdateID(value=self.workflow_id), )
resp = self.client.update_svc.GetUpdate(
request,
metadata=self.client.jobmgr_metadata,
timeout=self.config.rpc_timeout_sec,
)
return resp.updateInfo.status.state
def abort(self):
"""
aborts the given workflow
"""
request = update_svc.AbortUpdateRequest(
updateId=peloton.UpdateID(value=self.workflow_id), )
resp = self.client.update_svc.AbortUpdate(
request,
metadata=self.client.jobmgr_metadata,
timeout=self.config.rpc_timeout_sec,
)
return resp