def test_get_tasks_for_stepid(self):
mock_session = MockedSession()
gobworkflow.storage.storage.session = mock_session
mock_session.all = lambda: ['a', 'b']
result = get_tasks_for_stepid("someid")
self.assertEqual({"stepid": "someid"}, mock_session.filter_kwargs)
self.assertEqual(['a', 'b'], result)
def _create_tasks(self, jobid, stepid, process_id, tasks, key_prefix,
extra_msg, extra_header):
"""Create Task objects for the input list 'tasks'.
:param jobid:
:param stepid:
:param tasks:
:param key_prefix:
:param extra_msg:
:return:
"""
existing = get_tasks_for_stepid(stepid)
assert len(existing) == 0, f"Already have tasks for jobstep {stepid}"
for task in tasks:
task_def = {
'name': task['id'],
'dependencies': task['dependencies'],
'status': self.STATUS_NEW,
'jobid': jobid,
'stepid': stepid,
'key_prefix': key_prefix,
'extra_header': {
**extra_header,
},
'extra_msg': {
# Add global extra msg and extra_msg on task level
**extra_msg,
**task.get('extra_msg', {}),
},
'process_id': process_id,
}
task_save(task_def)
def _publish_complete(self, task):
"""Method is triggered when all tasks in a group have completed. Also triggered when tasks are stopped
because of failures. Handles final callback message to the user of the queue.
:param task:
:return:
"""
all_tasks = get_tasks_for_stepid(task.stepid)
warnings = [
warning for sublist in [t.summary['warnings'] for t in all_tasks]
for warning in sublist
]
errors = [
error for sublist in [t.summary['errors'] for t in all_tasks]
for error in sublist
]
msg = {
**task.extra_msg, 'header': {
'jobid': task.jobid,
'stepid': task.stepid,
**task.extra_header,
},
'summary': {
'warnings': warnings,
'errors': errors,
}
}
publish(WORKFLOW_EXCHANGE, task.key_prefix + '.' + TASK_COMPLETE, msg)
def _all_tasks_complete(self, stepid):
"""Returns whether all tasks for given stepid have STATUS_COMPLETED
:param stepid:
:return:
"""
tasks = get_tasks_for_stepid(stepid)
completed = [
task for task in tasks if task.status == self.STATUS_COMPLETED
]
return len(completed) == len(tasks)
def _abort_tasks(self, stepid):
"""Aborts all tasks belonging to stepid, as long as they are not queued or started yet.
:param stepid:
:return:
"""
all_tasks = get_tasks_for_stepid(stepid)
new = [task for task in all_tasks if task.status == self.STATUS_NEW]
for task in new:
if task_lock(task):
task_update({'id': task.id, 'status': self.STATUS_ABORTED})
task_unlock(task)
# Finish
self._publish_complete(all_tasks[0])
def _queue_free_tasks_for_jobstep(self, jobstep_id):
"""Queues the free tasks for jobstep.
:param stepid:
:return:
"""
tasks = get_tasks_for_stepid(jobstep_id)
completed = [
task.name for task in tasks if task.status == self.STATUS_COMPLETED
]
new = [task for task in tasks if task.status == self.STATUS_NEW]
for task in new:
if all([dep in completed for dep in task.dependencies]):
if task_lock(task):
if task_get(task.id).status == self.STATUS_NEW:
self._queue_task(task)
task_unlock(task)