class CentralPlannerTest(unittest.TestCase):
def setUp(self):
super(CentralPlannerTest, self).setUp()
conf = self.get_scheduler_config()
self.sch = CentralPlannerScheduler(**conf)
self.time = time.time
def get_scheduler_config(self):
return {
'retry_delay': 100,
'remove_delay': 1000,
'worker_disconnect_delay': 10,
'disable_persist': 10,
'disable_window': 10,
'disable_failures': 3,
'disable_hard_timeout': 60 * 60,
}
def tearDown(self):
super(CentralPlannerTest, self).tearDown()
if time.time != self.time:
time.time = self.time
def setTime(self, t):
time.time = lambda: t
def test_dep(self):
self.sch.add_task(worker=WORKER, task_id='B', deps=('A',))
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'B')
self.sch.add_task(worker=WORKER, task_id='B', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_failed_dep(self):
self.sch.add_task(worker=WORKER, task_id='B', deps=('A',))
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'B')
self.sch.add_task(worker=WORKER, task_id='B', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_broken_dep(self):
self.sch.add_task(worker=WORKER, task_id='B', deps=('A',))
self.sch.add_task(worker=WORKER, task_id='A', runnable=False)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'B')
self.sch.add_task(worker=WORKER, task_id='B', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_two_workers(self):
# Worker X wants to build A -> B
# Worker Y wants to build A -> C
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.sch.add_task(task_id='B', deps=('A',), worker='X')
self.sch.add_task(task_id='C', deps=('A',), worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None) # Worker Y is pending on A to be done
self.sch.add_task(worker='X', task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'C')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'B')
def test_retry(self):
# Try to build A but fails, will retry after 100s
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
for t in range(100):
self.setTime(t)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
self.sch.ping(worker=WORKER)
if t % 10 == 0:
self.sch.prune()
self.setTime(101)
self.sch.prune()
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
def test_disconnect_running(self):
# X and Y wants to run A.
# X starts but does not report back. Y does.
# After some timeout, Y will build it instead
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.sch.add_task(task_id='A', worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
for t in range(200):
self.setTime(t)
self.sch.ping(worker='Y')
if t % 10 == 0:
self.sch.prune()
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'A')
def test_remove_dep(self):
# X schedules A -> B, A is broken
# Y schedules C -> B: this should remove A as a dep of B
self.sch.add_task(task_id='A', worker='X', runnable=False)
self.sch.add_task(task_id='B', deps=('A',), worker='X')
# X can't build anything
self.assertEqual(self.sch.get_work(worker='X')['task_id'], None)
self.sch.add_task(task_id='B', deps=('C',), worker='Y') # should reset dependencies for A
self.sch.add_task(task_id='C', worker='Y', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_timeout(self):
# A bug that was earlier present when restarting the same flow
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.setTime(10000)
self.sch.add_task(task_id='A', worker='Y') # Will timeout X but not schedule A for removal
for i in range(2000):
self.setTime(10000 + i)
self.sch.ping(worker='Y')
self.sch.add_task(task_id='A', status=DONE, worker='Y') # This used to raise an exception since A was removed
def test_disallowed_state_changes(self):
# Test that we can not schedule an already running task
t = 'A'
self.sch.add_task(task_id=t, worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], t)
self.sch.add_task(task_id=t, worker='Y')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None)
def test_two_worker_info(self):
# Make sure the scheduler returns info that some other worker is running task A
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
r = self.sch.get_work(worker='Y')
self.assertEqual(r['task_id'], None) # Worker Y is pending on A to be done
s = r['running_tasks'][0]
self.assertEqual(s['task_id'], 'A')
self.assertEqual(s['worker'], 'X')
def test_assistant_get_work(self):
self.sch.add_task(worker='X', task_id='A')
self.sch.add_worker('Y', [])
self.assertEqual(self.sch.get_work(worker='Y', assistant=True)['task_id'], 'A')
# check that the scheduler recognizes tasks as running
running_tasks = self.sch.task_list('RUNNING', '')
self.assertEqual(len(running_tasks), 1)
self.assertEqual(list(running_tasks.keys()), ['A'])
self.assertEqual(running_tasks['A']['worker_running'], 'Y')
def test_assistant_get_work_external_task(self):
self.sch.add_task(worker='X', task_id='A', runnable=False)
self.assertTrue(self.sch.get_work(worker='Y', assistant=True)['task_id'] is None)
def test_task_fails_when_assistant_dies(self):
self.setTime(0)
self.sch.add_task(worker='X', task_id='A')
self.sch.add_worker('Y', [])
self.assertEqual(self.sch.get_work(worker='Y', assistant=True)['task_id'], 'A')
self.assertEqual(list(self.sch.task_list('RUNNING', '').keys()), ['A'])
# Y dies for 50 seconds, X stays alive
self.setTime(50)
self.sch.ping(worker='X')
self.assertEqual(list(self.sch.task_list('FAILED', '').keys()), ['A'])
def test_prune_with_live_assistant(self):
self.setTime(0)
self.sch.add_task(worker='X', task_id='A')
self.sch.get_work(worker='Y', assistant=True)
self.sch.add_task(worker='Y', task_id='A', status=DONE, assistant=True)
# worker X stops communicating, A should be marked for removal
self.setTime(600)
self.sch.ping(worker='Y')
self.sch.prune()
# A will now be pruned
self.setTime(2000)
self.sch.prune()
self.assertFalse(list(self.sch.task_list('', '')))
def test_fail_job_from_dead_worker_with_live_assistant(self):
self.setTime(0)
self.sch.add_task(worker='X', task_id='A')
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.sch.add_worker('Y', [('assistant', True)])
self.setTime(600)
self.sch.ping(worker='Y')
self.sch.prune()
self.assertEqual(['A'], list(self.sch.task_list('FAILED', '').keys()))
def test_assistant_request_runnable_task(self):
self.setTime(0)
self.sch.add_task(worker='X', task_id='A', runnable=True)
self.setTime(600)
self.sch.prune()
self.assertEqual('A', self.sch.get_work(worker='Y', assistant=True)['task_id'])
def test_assistant_request_external_task(self):
self.sch.add_task(worker='X', task_id='A', runnable=False)
self.assertIsNone(self.sch.get_work(worker='Y', assistant=True)['task_id'])
def test_prune_done_tasks(self, expected=None):
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
self.sch.add_task(worker=WORKER, task_id='B', deps=['A'], status=DONE)
self.sch.add_task(worker=WORKER, task_id='C', deps=['B'])
self.setTime(600)
self.sch.ping(worker='ASSISTANT')
self.sch.prune()
self.setTime(2000)
self.sch.ping(worker='ASSISTANT')
self.sch.prune()
self.assertEqual(set(expected or ()), set(self.sch.task_list('', '').keys()))
def test_keep_tasks_for_assistant(self):
self.sch.get_work(worker='ASSISTANT', assistant=True) # tell the scheduler this is an assistant
self.test_prune_done_tasks(['B', 'C'])
def test_keep_scheduler_disabled_tasks_for_assistant(self):
self.sch.get_work(worker='ASSISTANT', assistant=True) # tell the scheduler this is an assistant
# create a scheduler disabled task and a worker disabled task
for i in range(10):
self.sch.add_task(worker=WORKER, task_id='D', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='E', status=DISABLED)
# scheduler prunes the worker disabled task
self.assertEqual(set(['D', 'E']), set(self.sch.task_list(DISABLED, '')))
self.test_prune_done_tasks(['B', 'C', 'D'])
def test_keep_failed_tasks_for_assistant(self):
self.sch.get_work(worker='ASSISTANT', assistant=True) # tell the scheduler this is an assistant
self.sch.add_task(worker=WORKER, task_id='D', status=FAILED, deps='A')
self.test_prune_done_tasks(['A', 'B', 'C', 'D'])
def test_scheduler_resources_none_allow_one(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_resources_none_disallow_two(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.assertFalse(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_insufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=2)
self.assertFalse(self.sch.get_work(worker='X')['task_id'])
def test_scheduler_with_sufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=3)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_resources_used(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R1': 1})
self.sch.update_resources(R1=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_scheduler_overprovisioned_on_other_resource(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.sch.update_resources(R1=2)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R2': 2})
self.sch.update_resources(R1=1, R2=2)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_scheduler_with_priority_and_competing_resources(self):
self.sch.add_task(worker='X', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=10)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
self.sch.add_task(worker='Y', task_id='D', priority=0)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'D')
def test_do_not_lock_resources_when_not_ready(self):
""" Test to make sure that resources won't go unused waiting on workers """
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertEqual('C', self.sch.get_work(worker='Y')['task_id'])
def test_lock_resources_when_one_of_multiple_workers_is_ready(self):
self.sch.get_work(worker='X') # indicate to the scheduler that X is active
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 2)])
self.sch.add_worker('Y', [])
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_do_not_lock_resources_while_running_higher_priority(self):
""" Test to make sure that resources won't go unused waiting on workers """
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.assertEqual('C', self.sch.get_work(worker='Y')['task_id'])
def test_lock_resources_while_running_lower_priority(self):
""" Make sure resources will be made available while working on lower priority tasks """
self.sch.add_task(worker='X', task_id='A', priority=4)
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_lock_resources_for_second_worker(self):
self.sch.get_work(worker='Y') # indicate to the scheduler that Y is active
self.sch.add_task(worker='X', task_id='A', resources={'R': 1})
self.sch.add_task(worker='X', task_id='B', resources={'R': 1})
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=10)
self.sch.add_worker('X', {'workers': 2})
self.sch.add_worker('Y', {'workers': 1})
self.sch.update_resources(R=2)
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.assertFalse(self.sch.get_work(worker='X')['task_id'])
def test_can_work_on_lower_priority_while_waiting_for_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R': 1}, priority=0)
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.sch.add_task(worker='Y', task_id='B', resources={'R': 1}, priority=10)
self.sch.add_task(worker='Y', task_id='C', priority=0)
self.sch.update_resources(R=1)
self.assertEqual('C', self.sch.get_work(worker='Y')['task_id'])
def test_priority_update_with_pruning(self):
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.setTime(50) # after worker disconnects
self.sch.prune()
self.sch.add_task(task_id='B', deps=['A'], worker='X')
self.setTime(2000) # after remove for task A
self.sch.prune()
# Here task A that B depends on is missing
self.sch.add_task(worker=WORKER, task_id='C', deps=['B'], priority=100)
self.sch.add_task(worker=WORKER, task_id='B', deps=['A'])
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='D', priority=10)
self.check_task_order('ABCD')
def test_update_resources(self):
self.sch.add_task(worker=WORKER, task_id='A', deps=['B'])
self.sch.add_task(worker=WORKER, task_id='B', resources={'r': 2})
self.sch.update_resources(r=1)
# B requires too many resources, we can't schedule
self.check_task_order([])
self.sch.add_task(worker=WORKER, task_id='B', resources={'r': 1})
# now we have enough resources
self.check_task_order(['B', 'A'])
def test_hendle_multiple_resources(self):
self.sch.add_task(worker=WORKER, task_id='A', resources={'r1': 1, 'r2': 1})
self.sch.add_task(worker=WORKER, task_id='B', resources={'r1': 1, 'r2': 1})
self.sch.add_task(worker=WORKER, task_id='C', resources={'r1': 1})
self.sch.update_resources(r1=2, r2=1)
self.assertEqual('A', self.sch.get_work(worker=WORKER)['task_id'])
self.check_task_order('C')
def test_single_resource_lock(self):
self.sch.add_task(worker='X', task_id='A', resources={'r': 1})
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.sch.add_task(worker=WORKER, task_id='B', resources={'r': 2}, priority=10)
self.sch.add_task(worker=WORKER, task_id='C', resources={'r': 1})
self.sch.update_resources(r=2)
# Should wait for 2 units of r to be available for B before scheduling C
self.check_task_order([])
def test_no_lock_if_too_many_resources_required(self):
self.sch.add_task(worker=WORKER, task_id='A', resources={'r': 2}, priority=10)
self.sch.add_task(worker=WORKER, task_id='B', resources={'r': 1})
self.sch.update_resources(r=1)
self.check_task_order('B')
def test_multiple_resources_lock(self):
self.sch.get_work(worker='X') # indicate to the scheduler that X is active
self.sch.add_task(worker='X', task_id='A', resources={'r1': 1, 'r2': 1}, priority=10)
self.sch.add_task(worker=WORKER, task_id='B', resources={'r2': 1})
self.sch.add_task(worker=WORKER, task_id='C', resources={'r1': 1})
self.sch.update_resources(r1=1, r2=1)
# should preserve both resources for worker 'X'
self.check_task_order([])
def test_multiple_resources_no_lock(self):
self.sch.add_task(worker=WORKER, task_id='A', resources={'r1': 1}, priority=10)
self.sch.add_task(worker=WORKER, task_id='B', resources={'r1': 1, 'r2': 1}, priority=10)
self.sch.add_task(worker=WORKER, task_id='C', resources={'r2': 1})
self.sch.update_resources(r1=1, r2=2)
self.assertEqual('A', self.sch.get_work(worker=WORKER)['task_id'])
# C doesn't block B, so it can go first
self.check_task_order('C')
def test_allow_resource_use_while_scheduling(self):
self.sch.update_resources(r1=1)
self.sch.add_task(worker='SCHEDULING', task_id='A', resources={'r1': 1}, priority=10)
self.sch.add_task(worker=WORKER, task_id='B', resources={'r1': 1}, priority=1)
self.assertEqual('B', self.sch.get_work(worker=WORKER)['task_id'])
def test_stop_locking_resource_for_uninterested_worker(self):
self.setTime(0)
self.sch.update_resources(r1=1)
self.assertIsNone(self.sch.get_work(worker=WORKER)['task_id'])
self.sch.add_task(worker=WORKER, task_id='A', resources={'r1': 1}, priority=10)
self.sch.add_task(worker='LOW_PRIO', task_id='B', resources={'r1': 1}, priority=1)
self.assertIsNone(self.sch.get_work(worker='LOW_PRIO')['task_id'])
self.setTime(120)
self.assertEqual('B', self.sch.get_work(worker='LOW_PRIO')['task_id'])
def check_task_order(self, order):
for expected_id in order:
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], expected_id)
self.sch.add_task(worker=WORKER, task_id=expected_id, status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_priorities(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=10)
self.sch.add_task(worker=WORKER, task_id='B', priority=5)
self.sch.add_task(worker=WORKER, task_id='C', priority=15)
self.sch.add_task(worker=WORKER, task_id='D', priority=9)
self.check_task_order(['C', 'A', 'D', 'B'])
def test_priorities_default_and_negative(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=10)
self.sch.add_task(worker=WORKER, task_id='B')
self.sch.add_task(worker=WORKER, task_id='C', priority=15)
self.sch.add_task(worker=WORKER, task_id='D', priority=-20)
self.sch.add_task(worker=WORKER, task_id='E', priority=1)
self.check_task_order(['C', 'A', 'E', 'B', 'D'])
def test_priorities_and_dependencies(self):
self.sch.add_task(worker=WORKER, task_id='A', deps=['Z'], priority=10)
self.sch.add_task(worker=WORKER, task_id='B', priority=5)
self.sch.add_task(worker=WORKER, task_id='C', deps=['Z'], priority=3)
self.sch.add_task(worker=WORKER, task_id='D', priority=2)
self.sch.add_task(worker=WORKER, task_id='Z', priority=1)
self.check_task_order(['Z', 'A', 'B', 'C', 'D'])
def test_priority_update_dependency_after_scheduling(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=1)
self.sch.add_task(worker=WORKER, task_id='B', priority=5, deps=['A'])
self.sch.add_task(worker=WORKER, task_id='C', priority=10, deps=['B'])
self.sch.add_task(worker=WORKER, task_id='D', priority=6)
self.check_task_order(['A', 'B', 'C', 'D'])
def test_disable(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_disable_and_reenable(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.re_enable_task('A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
def test_disable_and_reenable_and_disable_again(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.re_enable_task('A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be still enabled
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled now
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_disable_and_done(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('DONE', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
def test_disable_by_worker(self):
self.sch.add_task(worker=WORKER, task_id='A', status=DISABLED)
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
def test_task_list_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
for c in 'ABCD':
sch.add_task(worker=WORKER, task_id=c)
self.assertEqual(set('ABCD'), set(sch.task_list('PENDING', '', False).keys()))
self.assertEqual({'num_tasks': 4}, sch.task_list('PENDING', ''))
def test_task_list_within_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=4)
for c in 'ABCD':
sch.add_task(worker=WORKER, task_id=c)
self.assertEqual(set('ABCD'), set(sch.task_list('PENDING', '').keys()))
def test_task_lists_some_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
for c in 'ABCD':
sch.add_task(worker=WORKER, task_id=c, status=DONE)
for c in 'EFG':
sch.add_task(worker=WORKER, task_id=c)
self.assertEqual(set('EFG'), set(sch.task_list('PENDING', '').keys()))
self.assertEqual({'num_tasks': 4}, sch.task_list('DONE', ''))
def test_task_list_filter_by_search(self):
self.sch.add_task(worker=WORKER, task_id='test_match_task')
self.sch.add_task(worker=WORKER, task_id='test_filter_task')
matches = self.sch.task_list('PENDING', '', search='match')
self.assertEqual(['test_match_task'], list(matches.keys()))
def test_task_list_filter_by_multiple_search_terms(self):
self.sch.add_task(worker=WORKER, task_id='abcd')
self.sch.add_task(worker=WORKER, task_id='abd')
self.sch.add_task(worker=WORKER, task_id='acd')
self.sch.add_task(worker=WORKER, task_id='ad')
self.sch.add_task(worker=WORKER, task_id='bc')
matches = self.sch.task_list('PENDING', '', search='b c')
self.assertEqual(set(['abcd', 'bc']), set(matches.keys()))
def test_search_results_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
sch.add_task(worker=WORKER, task_id='task_a')
sch.add_task(worker=WORKER, task_id='task_b')
sch.add_task(worker=WORKER, task_id='task_c')
sch.add_task(worker=WORKER, task_id='task_d')
self.assertEqual({'num_tasks': 4}, sch.task_list('PENDING', '', search='a'))
self.assertEqual(['task_a'], list(sch.task_list('PENDING', '', search='_a').keys()))
def test_priority_update_dependency_chain(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=10, deps=['B'])
self.sch.add_task(worker=WORKER, task_id='B', priority=5, deps=['C'])
self.sch.add_task(worker=WORKER, task_id='C', priority=1)
self.sch.add_task(worker=WORKER, task_id='D', priority=6)
self.check_task_order(['C', 'B', 'A', 'D'])
def test_priority_no_decrease_with_multiple_updates(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=1)
self.sch.add_task(worker=WORKER, task_id='B', priority=10, deps=['A'])
self.sch.add_task(worker=WORKER, task_id='C', priority=5, deps=['A'])
self.sch.add_task(worker=WORKER, task_id='D', priority=6)
self.check_task_order(['A', 'B', 'D', 'C'])
def test_unique_tasks(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='B')
self.sch.add_task(worker=WORKER, task_id='C')
self.sch.add_task(worker=WORKER + "_2", task_id='B')
response = self.sch.get_work(worker=WORKER)
self.assertEqual(3, response['n_pending_tasks'])
self.assertEqual(2, response['n_unique_pending'])
def test_pending_downstream_disable(self):
self.sch.add_task(worker=WORKER, task_id='A', status=DISABLED)
self.sch.add_task(worker=WORKER, task_id='B', deps=('A',))
self.sch.add_task(worker=WORKER, task_id='C', deps=('B',))
response = self.sch.get_work(worker=WORKER)
self.assertTrue(response['task_id'] is None)
self.assertEqual(0, response['n_pending_tasks'])
self.assertEqual(0, response['n_unique_pending'])
def test_pending_downstream_failure(self):
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='B', deps=('A',))
self.sch.add_task(worker=WORKER, task_id='C', deps=('B',))
response = self.sch.get_work(worker=WORKER)
self.assertTrue(response['task_id'] is None)
self.assertEqual(2, response['n_pending_tasks'])
self.assertEqual(2, response['n_unique_pending'])
def test_task_list_no_deps(self):
self.sch.add_task(worker=WORKER, task_id='B', deps=('A',))
self.sch.add_task(worker=WORKER, task_id='A')
task_list = self.sch.task_list('PENDING', '')
self.assertFalse('deps' in task_list['A'])
def test_task_first_failure_time(self):
self.sch.add_task(worker=WORKER, task_id='A')
test_task = self.sch._state.get_task('A')
self.assertIsNone(test_task.failures.first_failure_time)
time_before_failure = time.time()
test_task.add_failure()
time_after_failure = time.time()
self.assertLessEqual(time_before_failure,
test_task.failures.first_failure_time)
self.assertGreaterEqual(time_after_failure,
test_task.failures.first_failure_time)
def test_task_first_failure_time_remains_constant(self):
self.sch.add_task(worker=WORKER, task_id='A')
test_task = self.sch._state.get_task('A')
self.assertIsNone(test_task.failures.first_failure_time)
test_task.add_failure()
first_failure_time = test_task.failures.first_failure_time
test_task.add_failure()
self.assertEqual(first_failure_time, test_task.failures.first_failure_time)
def test_task_has_excessive_failures(self):
self.sch.add_task(worker=WORKER, task_id='A')
test_task = self.sch._state.get_task('A')
self.assertIsNone(test_task.failures.first_failure_time)
self.assertFalse(test_task.has_excessive_failures())
test_task.add_failure()
self.assertFalse(test_task.has_excessive_failures())
fake_failure_time = (test_task.failures.first_failure_time -
2 * 60 * 60)
test_task.failures.first_failure_time = fake_failure_time
self.assertTrue(test_task.has_excessive_failures())
def test_quadratic_behavior(self):
""" Test that get_work is not taking linear amount of time.
This is of course impossible to test, however, doing reasonable
assumptions about hardware. This time should finish in a timely
manner.
"""
# For 10000 it takes almost 1 second on my laptop. Prior to these
# changes it was being slow already at NUM_TASKS=300
NUM_TASKS = 10000
for i in range(NUM_TASKS):
self.sch.add_task(worker=str(i), task_id=str(i), resources={})
for i in range(NUM_TASKS):
self.assertEqual(self.sch.get_work(worker=str(i))['task_id'], str(i))
self.sch.add_task(worker=str(i), task_id=str(i), status=DONE)
def test_get_work_speed(self):
""" Test that get_work is fast for few workers and many DONEs.
In #986, @daveFNbuck reported that he got a slowdown.
"""
# This took almost 4 minutes without optimization.
# Now it takes 10 seconds on my machine.
NUM_PENDING = 1000
NUM_DONE = 200000
assert NUM_DONE >= NUM_PENDING
for i in range(NUM_PENDING):
self.sch.add_task(worker=WORKER, task_id=str(i), resources={})
for i in range(NUM_PENDING, NUM_DONE):
self.sch.add_task(worker=WORKER, task_id=str(i), status=DONE)
for i in range(NUM_PENDING):
res = int(self.sch.get_work(worker=WORKER)['task_id'])
self.assertTrue(0 <= res < NUM_PENDING)
self.sch.add_task(worker=WORKER, task_id=str(res), status=DONE)
class CentralPlannerTest(unittest.TestCase):
def setUp(self):
self.sch = CentralPlannerScheduler(retry_delay=100, remove_delay=1000, worker_disconnect_delay=10)
self.time = time.time
def tearDown(self):
if time.time != self.time:
time.time = self.time
def setTime(self, t):
time.time = lambda: t
def test_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)[1], 'A')
self.sch.add_task(WORKER, 'A', status=DONE)
self.assertEqual(self.sch.get_work(WORKER)[1], 'B')
self.sch.add_task(WORKER, 'B', status=DONE)
self.assertEqual(self.sch.get_work(WORKER), (0, None))
def test_failed_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)[1], 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.assertEqual(self.sch.get_work(WORKER)[1], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(WORKER, 'A', DONE)
self.assertEqual(self.sch.get_work(WORKER)[1], 'B')
self.sch.add_task(WORKER, 'B', DONE)
self.assertEqual(self.sch.get_work(WORKER), (0, None))
def test_broken_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A', runnable=False)
self.assertEqual(self.sch.get_work(WORKER)[1], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(WORKER, 'A', DONE)
self.assertEqual(self.sch.get_work(WORKER)[1], 'B')
self.sch.add_task(WORKER, 'B', DONE)
self.assertEqual(self.sch.get_work(WORKER), (0, None))
def test_two_workers(self):
# Worker X wants to build A -> B
# Worker Y wants to build A -> C
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.sch.add_task(task_id='B', deps=('A',), worker='X')
self.sch.add_task(task_id='C', deps=('A',), worker='Y')
self.assertEqual(self.sch.get_work(worker='X')[1], 'A')
self.assertEqual(self.sch.get_work(worker='Y')[1], None) # Worker Y is pending on A to be done
self.sch.add_task(worker='X', task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')[1], 'C')
self.assertEqual(self.sch.get_work(worker='X')[1], 'B')
def test_retry(self):
# Try to build A but fails, will retry after 100s
self.setTime(0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)[1], 'A')
self.sch.add_task(WORKER, 'A', FAILED)
for t in xrange(100):
self.setTime(t)
self.assertEqual(self.sch.get_work(WORKER)[1], None)
self.sch.ping(WORKER)
if t % 10 == 0:
self.sch.prune()
self.setTime(101)
self.sch.prune()
self.assertEqual(self.sch.get_work(WORKER)[1], 'A')
def test_disconnect_running(self):
# X and Y wants to run A.
# X starts but does not report back. Y does.
# After some timeout, Y will build it instead
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.sch.add_task(task_id='A', worker='Y')
self.assertEqual(self.sch.get_work(worker='X')[1], 'A')
for t in xrange(200):
self.setTime(t)
self.sch.ping(worker='Y')
if t % 10 == 0:
self.sch.prune()
self.assertEqual(self.sch.get_work(worker='Y')[1], 'A')
def test_remove_dep(self):
# X schedules A -> B, A is broken
# Y schedules C -> B: this should remove A as a dep of B
self.sch.add_task(task_id='A', worker='X', runnable=False)
self.sch.add_task(task_id='B', deps=('A',), worker='X')
# X can't build anything
self.assertEqual(self.sch.get_work(worker='X')[1], None)
self.sch.add_task(task_id='B', deps=('C',), worker='Y') # should reset dependencies for A
self.sch.add_task(task_id='C', worker='Y', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')[1], 'B')
def test_timeout(self):
# A bug that was earlier present when restarting the same flow
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.assertEqual(self.sch.get_work(worker='X')[1], 'A')
self.setTime(10000)
self.sch.add_task(task_id='A', worker='Y') # Will timeout X but not schedule A for removal
for i in xrange(2000):
self.setTime(10000 + i)
self.sch.ping(worker='Y')
self.sch.add_task(task_id='A', status=DONE, worker='Y') # This used to raise an exception since A was removed
def test_disallowed_state_changes(self):
# Test that we can not schedule an already running task
t = 'A'
self.sch.add_task(task_id=t, worker='X')
self.assertEqual(self.sch.get_work(worker='X')[1], t)
self.sch.add_task(task_id=t, worker='Y')
self.assertEqual(self.sch.get_work(worker='Y')[1], None)
class CentralPlannerTest(unittest.TestCase):
def setUp(self):
super(CentralPlannerTest, self).setUp()
conf = self.get_scheduler_config()
self.sch = CentralPlannerScheduler(**conf)
self.time = time.time
def get_scheduler_config(self):
return {
'retry_delay': 100,
'remove_delay': 1000,
'worker_disconnect_delay': 10,
'disable_persist': 10,
'disable_window': 10,
'disable_failures': 3,
'disable_hard_timeout': 60 * 60,
}
def tearDown(self):
super(CentralPlannerTest, self).tearDown()
if time.time != self.time:
time.time = self.time
def setTime(self, t):
time.time = lambda: t
def test_dep(self):
self.sch.add_task(worker=WORKER, task_id='B', deps=('A', ))
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'B')
self.sch.add_task(worker=WORKER, task_id='B', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_failed_dep(self):
self.sch.add_task(worker=WORKER, task_id='B', deps=('A', ))
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.assertEqual(
self.sch.get_work(worker=WORKER)['task_id'],
None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'B')
self.sch.add_task(worker=WORKER, task_id='B', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_broken_dep(self):
self.sch.add_task(worker=WORKER, task_id='B', deps=('A', ))
self.sch.add_task(worker=WORKER, task_id='A', runnable=False)
self.assertEqual(
self.sch.get_work(worker=WORKER)['task_id'],
None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'B')
self.sch.add_task(worker=WORKER, task_id='B', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_two_workers(self):
# Worker X wants to build A -> B
# Worker Y wants to build A -> C
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.sch.add_task(task_id='B', deps=('A', ), worker='X')
self.sch.add_task(task_id='C', deps=('A', ), worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'],
None) # Worker Y is pending on A to be done
self.sch.add_task(worker='X', task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'C')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'B')
def test_retry(self):
# Try to build A but fails, will retry after 100s
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
for t in range(100):
self.setTime(t)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
self.sch.ping(worker=WORKER)
if t % 10 == 0:
self.sch.prune()
self.setTime(101)
self.sch.prune()
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
def test_disconnect_running(self):
# X and Y wants to run A.
# X starts but does not report back. Y does.
# After some timeout, Y will build it instead
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.sch.add_task(task_id='A', worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
for t in range(200):
self.setTime(t)
self.sch.ping(worker='Y')
if t % 10 == 0:
self.sch.prune()
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'A')
def test_remove_dep(self):
# X schedules A -> B, A is broken
# Y schedules C -> B: this should remove A as a dep of B
self.sch.add_task(task_id='A', worker='X', runnable=False)
self.sch.add_task(task_id='B', deps=('A', ), worker='X')
# X can't build anything
self.assertEqual(self.sch.get_work(worker='X')['task_id'], None)
self.sch.add_task(task_id='B', deps=('C', ),
worker='Y') # should reset dependencies for A
self.sch.add_task(task_id='C', worker='Y', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_timeout(self):
# A bug that was earlier present when restarting the same flow
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.setTime(10000)
self.sch.add_task(
task_id='A',
worker='Y') # Will timeout X but not schedule A for removal
for i in range(2000):
self.setTime(10000 + i)
self.sch.ping(worker='Y')
self.sch.add_task(
task_id='A', status=DONE,
worker='Y') # This used to raise an exception since A was removed
def test_disallowed_state_changes(self):
# Test that we can not schedule an already running task
t = 'A'
self.sch.add_task(task_id=t, worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], t)
self.sch.add_task(task_id=t, worker='Y')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None)
def test_two_worker_info(self):
# Make sure the scheduler returns info that some other worker is running task A
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
r = self.sch.get_work(worker='Y')
self.assertEqual(r['task_id'],
None) # Worker Y is pending on A to be done
s = r['running_tasks'][0]
self.assertEqual(s['task_id'], 'A')
self.assertEqual(s['worker'], 'X')
def test_assistant_get_work(self):
self.sch.add_task(worker='X', task_id='A')
self.sch.add_worker('Y', [])
self.assertEqual(
self.sch.get_work(worker='Y', assistant=True)['task_id'], 'A')
# check that the scheduler recognizes tasks as running
running_tasks = self.sch.task_list('RUNNING', '')
self.assertEqual(len(running_tasks), 1)
self.assertEqual(list(running_tasks.keys()), ['A'])
self.assertEqual(running_tasks['A']['worker_running'], 'Y')
def test_assistant_get_work_external_task(self):
self.sch.add_task(worker='X', task_id='A', runnable=False)
self.assertTrue(
self.sch.get_work(worker='Y', assistant=True)['task_id'] is None)
def test_task_fails_when_assistant_dies(self):
self.setTime(0)
self.sch.add_task(worker='X', task_id='A')
self.sch.add_worker('Y', [])
self.assertEqual(
self.sch.get_work(worker='Y', assistant=True)['task_id'], 'A')
self.assertEqual(list(self.sch.task_list('RUNNING', '').keys()), ['A'])
# Y dies for 50 seconds, X stays alive
self.setTime(50)
self.sch.ping(worker='X')
self.assertEqual(list(self.sch.task_list('FAILED', '').keys()), ['A'])
def test_prune_with_live_assistant(self):
self.setTime(0)
self.sch.add_task(worker='X', task_id='A')
self.sch.get_work(worker='Y', assistant=True)
self.sch.add_task(worker='Y', task_id='A', status=DONE, assistant=True)
# worker X stops communicating, A should be marked for removal
self.setTime(600)
self.sch.ping(worker='Y')
self.sch.prune()
# A will now be pruned
self.setTime(2000)
self.sch.prune()
self.assertFalse(list(self.sch.task_list('', '')))
def test_fail_job_from_dead_worker_with_live_assistant(self):
self.setTime(0)
self.sch.add_task(worker='X', task_id='A')
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.sch.add_worker('Y', [('assistant', True)])
self.setTime(600)
self.sch.ping(worker='Y')
self.sch.prune()
self.assertEqual(['A'], list(self.sch.task_list('FAILED', '').keys()))
def test_prune_done_tasks(self, expected=None):
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
self.sch.add_task(worker=WORKER, task_id='B', deps=['A'], status=DONE)
self.sch.add_task(worker=WORKER, task_id='C', deps=['B'])
self.setTime(600)
self.sch.ping(worker='ASSISTANT')
self.sch.prune()
self.setTime(2000)
self.sch.ping(worker='ASSISTANT')
self.sch.prune()
self.assertEqual(set(expected or ()),
set(self.sch.task_list('', '').keys()))
def test_keep_tasks_for_assistant(self):
self.sch.get_work(
worker='ASSISTANT',
assistant=True) # tell the scheduler this is an assistant
self.test_prune_done_tasks(['B', 'C'])
def test_keep_scheduler_disabled_tasks_for_assistant(self):
self.sch.get_work(
worker='ASSISTANT',
assistant=True) # tell the scheduler this is an assistant
# create a scheduler disabled task and a worker disabled task
for i in range(10):
self.sch.add_task(worker=WORKER, task_id='D', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='E', status=DISABLED)
# scheduler prunes the worker disabled task
self.assertEqual(set(['D', 'E']), set(self.sch.task_list(DISABLED,
'')))
self.test_prune_done_tasks(['B', 'C', 'D'])
def test_keep_failed_tasks_for_assistant(self):
self.sch.get_work(
worker='ASSISTANT',
assistant=True) # tell the scheduler this is an assistant
self.sch.add_task(worker=WORKER, task_id='D', status=FAILED, deps='A')
self.test_prune_done_tasks(['A', 'B', 'C', 'D'])
def test_scheduler_resources_none_allow_one(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_resources_none_disallow_two(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.assertFalse(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_insufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=2)
self.assertFalse(self.sch.get_work(worker='X')['task_id'])
def test_scheduler_with_sufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=3)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_resources_used(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R1': 1})
self.sch.update_resources(R1=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_scheduler_overprovisioned_on_other_resource(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.sch.update_resources(R1=2)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R2': 2})
self.sch.update_resources(R1=1, R2=2)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_scheduler_with_priority_and_competing_resources(self):
self.sch.add_task(worker='X', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='X',
task_id='B',
resources={'R': 1},
priority=10)
self.sch.add_task(worker='Y',
task_id='C',
resources={'R': 1},
priority=1)
self.sch.update_resources(R=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
self.sch.add_task(worker='Y', task_id='D', priority=0)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'D')
def test_do_not_lock_resources_when_not_ready(self):
""" Test to make sure that resources won't go unused waiting on workers """
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X',
task_id='B',
resources={'R': 1},
priority=5)
self.sch.add_task(worker='Y',
task_id='C',
resources={'R': 1},
priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertEqual('C', self.sch.get_work(worker='Y')['task_id'])
def test_lock_resources_when_one_of_multiple_workers_is_ready(self):
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X',
task_id='B',
resources={'R': 1},
priority=5)
self.sch.add_task(worker='Y',
task_id='C',
resources={'R': 1},
priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 2)])
self.sch.add_worker('Y', [])
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_do_not_lock_resources_while_running_higher_priority(self):
""" Test to make sure that resources won't go unused waiting on workers """
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X',
task_id='B',
resources={'R': 1},
priority=5)
self.sch.add_task(worker='Y',
task_id='C',
resources={'R': 1},
priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.assertEqual('C', self.sch.get_work(worker='Y')['task_id'])
def test_lock_resources_while_running_lower_priority(self):
""" Make sure resources will be made available while working on lower priority tasks """
self.sch.add_task(worker='X', task_id='A', priority=4)
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.sch.add_task(worker='X',
task_id='B',
resources={'R': 1},
priority=5)
self.sch.add_task(worker='Y',
task_id='C',
resources={'R': 1},
priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_lock_resources_for_second_worker(self):
self.sch.add_task(worker='X', task_id='A', resources={'R': 1})
self.sch.add_task(worker='X', task_id='B', resources={'R': 1})
self.sch.add_task(worker='Y',
task_id='C',
resources={'R': 1},
priority=10)
self.sch.add_worker('X', {'workers': 2})
self.sch.add_worker('Y', {'workers': 1})
self.sch.update_resources(R=2)
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.assertFalse(self.sch.get_work(worker='X')['task_id'])
def test_can_work_on_lower_priority_while_waiting_for_resources(self):
self.sch.add_task(worker='X',
task_id='A',
resources={'R': 1},
priority=0)
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.sch.add_task(worker='Y',
task_id='B',
resources={'R': 1},
priority=10)
self.sch.add_task(worker='Y', task_id='C', priority=0)
self.sch.update_resources(R=1)
self.assertEqual('C', self.sch.get_work(worker='Y')['task_id'])
def test_priority_update_with_pruning(self):
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.setTime(50) # after worker disconnects
self.sch.prune()
self.sch.add_task(task_id='B', deps=['A'], worker='X')
self.setTime(2000) # after remove for task A
self.sch.prune()
# Here task A that B depends on is missing
self.sch.add_task(worker=WORKER, task_id='C', deps=['B'], priority=100)
self.sch.add_task(worker=WORKER, task_id='B', deps=['A'])
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='D', priority=10)
self.check_task_order('ABCD')
def test_update_resources(self):
self.sch.add_task(worker=WORKER, task_id='A', deps=['B'])
self.sch.add_task(worker=WORKER, task_id='B', resources={'r': 2})
self.sch.update_resources(r=1)
# B requires too many resources, we can't schedule
self.check_task_order([])
self.sch.add_task(worker=WORKER, task_id='B', resources={'r': 1})
# now we have enough resources
self.check_task_order(['B', 'A'])
def test_hendle_multiple_resources(self):
self.sch.add_task(worker=WORKER,
task_id='A',
resources={
'r1': 1,
'r2': 1
})
self.sch.add_task(worker=WORKER,
task_id='B',
resources={
'r1': 1,
'r2': 1
})
self.sch.add_task(worker=WORKER, task_id='C', resources={'r1': 1})
self.sch.update_resources(r1=2, r2=1)
self.assertEqual('A', self.sch.get_work(worker=WORKER)['task_id'])
self.check_task_order('C')
def test_single_resource_lock(self):
self.sch.add_task(worker='X', task_id='A', resources={'r': 1})
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.sch.add_task(worker=WORKER,
task_id='B',
resources={'r': 2},
priority=10)
self.sch.add_task(worker=WORKER, task_id='C', resources={'r': 1})
self.sch.update_resources(r=2)
# Should wait for 2 units of r to be available for B before scheduling C
self.check_task_order([])
def test_no_lock_if_too_many_resources_required(self):
self.sch.add_task(worker=WORKER,
task_id='A',
resources={'r': 2},
priority=10)
self.sch.add_task(worker=WORKER, task_id='B', resources={'r': 1})
self.sch.update_resources(r=1)
self.check_task_order('B')
def test_multiple_resources_lock(self):
self.sch.add_task(worker='X',
task_id='A',
resources={
'r1': 1,
'r2': 1
},
priority=10)
self.sch.add_task(worker=WORKER, task_id='B', resources={'r2': 1})
self.sch.add_task(worker=WORKER, task_id='C', resources={'r1': 1})
self.sch.update_resources(r1=1, r2=1)
# should preserve both resources for worker 'X'
self.check_task_order([])
def test_multiple_resources_no_lock(self):
self.sch.add_task(worker=WORKER,
task_id='A',
resources={'r1': 1},
priority=10)
self.sch.add_task(worker=WORKER,
task_id='B',
resources={
'r1': 1,
'r2': 1
},
priority=10)
self.sch.add_task(worker=WORKER, task_id='C', resources={'r2': 1})
self.sch.update_resources(r1=1, r2=2)
self.assertEqual('A', self.sch.get_work(worker=WORKER)['task_id'])
# C doesn't block B, so it can go first
self.check_task_order('C')
def check_task_order(self, order):
for expected_id in order:
self.assertEqual(
self.sch.get_work(worker=WORKER)['task_id'], expected_id)
self.sch.add_task(worker=WORKER, task_id=expected_id, status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_priorities(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=10)
self.sch.add_task(worker=WORKER, task_id='B', priority=5)
self.sch.add_task(worker=WORKER, task_id='C', priority=15)
self.sch.add_task(worker=WORKER, task_id='D', priority=9)
self.check_task_order(['C', 'A', 'D', 'B'])
def test_priorities_default_and_negative(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=10)
self.sch.add_task(worker=WORKER, task_id='B')
self.sch.add_task(worker=WORKER, task_id='C', priority=15)
self.sch.add_task(worker=WORKER, task_id='D', priority=-20)
self.sch.add_task(worker=WORKER, task_id='E', priority=1)
self.check_task_order(['C', 'A', 'E', 'B', 'D'])
def test_priorities_and_dependencies(self):
self.sch.add_task(worker=WORKER, task_id='A', deps=['Z'], priority=10)
self.sch.add_task(worker=WORKER, task_id='B', priority=5)
self.sch.add_task(worker=WORKER, task_id='C', deps=['Z'], priority=3)
self.sch.add_task(worker=WORKER, task_id='D', priority=2)
self.sch.add_task(worker=WORKER, task_id='Z', priority=1)
self.check_task_order(['Z', 'A', 'B', 'C', 'D'])
def test_priority_update_dependency_after_scheduling(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=1)
self.sch.add_task(worker=WORKER, task_id='B', priority=5, deps=['A'])
self.sch.add_task(worker=WORKER, task_id='C', priority=10, deps=['B'])
self.sch.add_task(worker=WORKER, task_id='D', priority=6)
self.check_task_order(['A', 'B', 'C', 'D'])
def test_disable(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_disable_and_reenable(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.re_enable_task('A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
def test_disable_and_reenable_and_disable_again(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.re_enable_task('A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be still enabled
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled now
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_disable_and_done(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('DONE', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
def test_disable_by_worker(self):
self.sch.add_task(worker=WORKER, task_id='A', status=DISABLED)
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
def test_task_list_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
for c in 'ABCD':
sch.add_task(worker=WORKER, task_id=c)
self.assertEqual(set('ABCD'),
set(sch.task_list('PENDING', '', False).keys()))
self.assertEqual({'num_tasks': 4}, sch.task_list('PENDING', ''))
def test_task_list_within_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=4)
for c in 'ABCD':
sch.add_task(worker=WORKER, task_id=c)
self.assertEqual(set('ABCD'), set(sch.task_list('PENDING', '').keys()))
def test_task_lists_some_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
for c in 'ABCD':
sch.add_task(worker=WORKER, task_id=c, status=DONE)
for c in 'EFG':
sch.add_task(worker=WORKER, task_id=c)
self.assertEqual(set('EFG'), set(sch.task_list('PENDING', '').keys()))
self.assertEqual({'num_tasks': 4}, sch.task_list('DONE', ''))
def test_task_list_filter_by_search(self):
self.sch.add_task(worker=WORKER, task_id='test_match_task')
self.sch.add_task(worker=WORKER, task_id='test_filter_task')
matches = self.sch.task_list('PENDING', '', search='match')
self.assertEqual(['test_match_task'], list(matches.keys()))
def test_task_list_filter_by_multiple_search_terms(self):
self.sch.add_task(worker=WORKER, task_id='abcd')
self.sch.add_task(worker=WORKER, task_id='abd')
self.sch.add_task(worker=WORKER, task_id='acd')
self.sch.add_task(worker=WORKER, task_id='ad')
self.sch.add_task(worker=WORKER, task_id='bc')
matches = self.sch.task_list('PENDING', '', search='b c')
self.assertEqual(set(['abcd', 'bc']), set(matches.keys()))
def test_search_results_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
sch.add_task(worker=WORKER, task_id='task_a')
sch.add_task(worker=WORKER, task_id='task_b')
sch.add_task(worker=WORKER, task_id='task_c')
sch.add_task(worker=WORKER, task_id='task_d')
self.assertEqual({'num_tasks': 4},
sch.task_list('PENDING', '', search='a'))
self.assertEqual(['task_a'],
list(
sch.task_list('PENDING', '', search='_a').keys()))
def test_priority_update_dependency_chain(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=10, deps=['B'])
self.sch.add_task(worker=WORKER, task_id='B', priority=5, deps=['C'])
self.sch.add_task(worker=WORKER, task_id='C', priority=1)
self.sch.add_task(worker=WORKER, task_id='D', priority=6)
self.check_task_order(['C', 'B', 'A', 'D'])
def test_priority_no_decrease_with_multiple_updates(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=1)
self.sch.add_task(worker=WORKER, task_id='B', priority=10, deps=['A'])
self.sch.add_task(worker=WORKER, task_id='C', priority=5, deps=['A'])
self.sch.add_task(worker=WORKER, task_id='D', priority=6)
self.check_task_order(['A', 'B', 'D', 'C'])
def test_unique_tasks(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='B')
self.sch.add_task(worker=WORKER, task_id='C')
self.sch.add_task(worker=WORKER + "_2", task_id='B')
response = self.sch.get_work(worker=WORKER)
self.assertEqual(3, response['n_pending_tasks'])
self.assertEqual(2, response['n_unique_pending'])
def test_pending_downstream_disable(self):
self.sch.add_task(worker=WORKER, task_id='A', status=DISABLED)
self.sch.add_task(worker=WORKER, task_id='B', deps=('A', ))
self.sch.add_task(worker=WORKER, task_id='C', deps=('B', ))
response = self.sch.get_work(worker=WORKER)
self.assertTrue(response['task_id'] is None)
self.assertEqual(0, response['n_pending_tasks'])
self.assertEqual(0, response['n_unique_pending'])
def test_pending_downstream_failure(self):
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='B', deps=('A', ))
self.sch.add_task(worker=WORKER, task_id='C', deps=('B', ))
response = self.sch.get_work(worker=WORKER)
self.assertTrue(response['task_id'] is None)
self.assertEqual(2, response['n_pending_tasks'])
self.assertEqual(2, response['n_unique_pending'])
def test_task_list_no_deps(self):
self.sch.add_task(worker=WORKER, task_id='B', deps=('A', ))
self.sch.add_task(worker=WORKER, task_id='A')
task_list = self.sch.task_list('PENDING', '')
self.assertFalse('deps' in task_list['A'])
def test_task_first_failure_time(self):
self.sch.add_task(worker=WORKER, task_id='A')
test_task = self.sch._state.get_task('A')
self.assertIsNone(test_task.failures.first_failure_time)
time_before_failure = time.time()
test_task.add_failure()
time_after_failure = time.time()
self.assertLessEqual(time_before_failure,
test_task.failures.first_failure_time)
self.assertGreaterEqual(time_after_failure,
test_task.failures.first_failure_time)
def test_task_first_failure_time_remains_constant(self):
self.sch.add_task(worker=WORKER, task_id='A')
test_task = self.sch._state.get_task('A')
self.assertIsNone(test_task.failures.first_failure_time)
test_task.add_failure()
first_failure_time = test_task.failures.first_failure_time
test_task.add_failure()
self.assertEqual(first_failure_time,
test_task.failures.first_failure_time)
def test_task_has_excessive_failures(self):
self.sch.add_task(worker=WORKER, task_id='A')
test_task = self.sch._state.get_task('A')
self.assertIsNone(test_task.failures.first_failure_time)
self.assertFalse(test_task.has_excessive_failures())
test_task.add_failure()
self.assertFalse(test_task.has_excessive_failures())
fake_failure_time = (test_task.failures.first_failure_time -
2 * 60 * 60)
test_task.failures.first_failure_time = fake_failure_time
self.assertTrue(test_task.has_excessive_failures())
def test_quadratic_behavior(self):
""" Test that get_work is not taking linear amount of time.
This is of course impossible to test, however, doing reasonable
assumptions about hardware. This time should finish in a timely
manner.
"""
# For 10000 it takes almost 1 second on my laptop. Prior to these
# changes it was being slow already at NUM_TASKS=300
NUM_TASKS = 10000
for i in range(NUM_TASKS):
self.sch.add_task(worker=str(i), task_id=str(i), resources={})
for i in range(NUM_TASKS):
self.assertEqual(
self.sch.get_work(worker=str(i))['task_id'], str(i))
self.sch.add_task(worker=str(i), task_id=str(i), status=DONE)
def test_get_work_speed(self):
""" Test that get_work is fast for few workers and many DONEs.
In #986, @daveFNbuck reported that he got a slowdown.
"""
# This took almost 4 minutes without optimization.
# Now it takes 10 seconds on my machine.
NUM_PENDING = 1000
NUM_DONE = 200000
assert NUM_DONE >= NUM_PENDING
for i in range(NUM_PENDING):
self.sch.add_task(worker=WORKER, task_id=str(i), resources={})
for i in range(NUM_PENDING, NUM_DONE):
self.sch.add_task(worker=WORKER, task_id=str(i), status=DONE)
for i in range(NUM_PENDING):
res = int(self.sch.get_work(worker=WORKER)['task_id'])
self.assertTrue(0 <= res < NUM_PENDING)
self.sch.add_task(worker=WORKER, task_id=str(res), status=DONE)
class WorkerExternalTaskTest(unittest.TestCase):
def setUp(self):
self.scheduler = CentralPlannerScheduler(retry_delay=0.01,
remove_delay=3,
worker_disconnect_delay=3,
disable_persist=3,
disable_window=5,
disable_failures=2)
def _assert_complete(self, tasks):
for t in tasks:
self.assert_(t.complete())
def _build(self, tasks):
w = luigi.worker.Worker(scheduler=self.scheduler, worker_processes=1)
for t in tasks:
w.add(t)
w.run()
w.stop()
def test_external_dependency_already_complete(self):
"""
Test that the test task completes when its dependency exists at the
start of the execution.
"""
tempdir = tempfile.mkdtemp(prefix='luigi-test-')
test_task = TestTask(tempdir=tempdir, complete_after=1)
luigi.build([test_task], local_scheduler=True)
assert os.path.exists(test_task.dep_path)
assert os.path.exists(test_task.output_path)
os.unlink(test_task.dep_path)
os.unlink(test_task.output_path)
os.rmdir(tempdir)
# complete() is called once per failure, twice per success
assert test_task.dependency.times_called == 2
@with_config({'core': {'retry-external-tasks': 'true',
'disable-num-failures': '4',
'max-reschedules': '4',
'worker-keep-alive': 'true',
'retry-delay': '0.01'}})
def test_external_dependency_completes_later(self):
"""
Test that an external dependency that is not `complete` when luigi is invoked, but \
becomes `complete` while the workflow is executing is re-evaluated and
allows dependencies to run.
"""
assert luigi.configuration.get_config().getboolean('core',
'retry-external-tasks',
False) is True
original_get_work = self.scheduler.get_work
def decorated_get_work(*args, **kwargs):
# need to call `prune()` to make the scheduler run the retry logic
self.scheduler.prune()
return original_get_work(*args, **kwargs)
self.scheduler.get_work = decorated_get_work
tempdir = tempfile.mkdtemp(prefix='luigi-test-')
with patch('random.randint', return_value=0.1):
test_task = TestTask(tempdir=tempdir, complete_after=3)
self._build([test_task])
assert os.path.exists(test_task.dep_path)
assert os.path.exists(test_task.output_path)
os.unlink(test_task.dep_path)
os.unlink(test_task.output_path)
os.rmdir(tempdir)
# complete() is called once per failure, twice per success
assert test_task.dependency.times_called == 4
class CentralPlannerTest(unittest.TestCase):
def setUp(self):
self.sch = CentralPlannerScheduler(retry_delay=100,
remove_delay=1000,
worker_disconnect_delay=10,
disable_persist=10,
disable_window=10,
disable_failures=3)
self.time = time.time
def tearDown(self):
if time.time != self.time:
time.time = self.time
def setTime(self, t):
time.time = lambda: t
def test_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_failed_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(WORKER, 'A', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_broken_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A', runnable=False)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(WORKER, 'A', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_two_workers(self):
# Worker X wants to build A -> B
# Worker Y wants to build A -> C
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.sch.add_task(task_id='B', deps=('A',), worker='X')
self.sch.add_task(task_id='C', deps=('A',), worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None) # Worker Y is pending on A to be done
self.sch.add_task(worker='X', task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'C')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'B')
def test_retry(self):
# Try to build A but fails, will retry after 100s
self.setTime(0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', FAILED)
for t in xrange(100):
self.setTime(t)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
self.sch.ping(WORKER)
if t % 10 == 0:
self.sch.prune()
self.setTime(101)
self.sch.prune()
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_disconnect_running(self):
# X and Y wants to run A.
# X starts but does not report back. Y does.
# After some timeout, Y will build it instead
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.sch.add_task(task_id='A', worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
for t in xrange(200):
self.setTime(t)
self.sch.ping(worker='Y')
if t % 10 == 0:
self.sch.prune()
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'A')
def test_remove_dep(self):
# X schedules A -> B, A is broken
# Y schedules C -> B: this should remove A as a dep of B
self.sch.add_task(task_id='A', worker='X', runnable=False)
self.sch.add_task(task_id='B', deps=('A',), worker='X')
# X can't build anything
self.assertEqual(self.sch.get_work(worker='X')['task_id'], None)
self.sch.add_task(task_id='B', deps=('C',), worker='Y') # should reset dependencies for A
self.sch.add_task(task_id='C', worker='Y', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_timeout(self):
# A bug that was earlier present when restarting the same flow
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.setTime(10000)
self.sch.add_task(task_id='A', worker='Y') # Will timeout X but not schedule A for removal
for i in xrange(2000):
self.setTime(10000 + i)
self.sch.ping(worker='Y')
self.sch.add_task(task_id='A', status=DONE, worker='Y') # This used to raise an exception since A was removed
def test_disallowed_state_changes(self):
# Test that we can not schedule an already running task
t = 'A'
self.sch.add_task(task_id=t, worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], t)
self.sch.add_task(task_id=t, worker='Y')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None)
def test_two_worker_info(self):
# Make sure the scheduler returns info that some other worker is running task A
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
r = self.sch.get_work(worker='Y')
self.assertEqual(r['task_id'], None) # Worker Y is pending on A to be done
s = r['running_tasks'][0]
self.assertEqual(s['task_id'], 'A')
self.assertEqual(s['worker'], 'X')
def test_scheduler_resources_none_allow_one(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_resources_none_disallow_two(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.assertFalse(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_insufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=2)
self.assertFalse(self.sch.get_work(worker='X')['task_id'])
def test_scheduler_with_sufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=3)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_resources_used(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R1': 1})
self.sch.update_resources(R1=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_scheduler_overprovisioned_on_other_resource(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.sch.update_resources(R1=2)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R2': 2})
self.sch.update_resources(R1=1, R2=2)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_scheduler_with_priority_and_competing_resources(self):
self.sch.add_task(worker='X', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=10)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
self.sch.add_task(worker='Y', task_id='D', priority=0)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'D')
def test_priority_update_with_pruning(self):
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.setTime(50) # after worker disconnects
self.sch.prune()
self.sch.add_task(task_id='B', deps=['A'], worker='X')
self.setTime(2000) # after remove for task A
self.sch.prune()
# Here task A that B depends on is missing
self.sch.add_task(WORKER, task_id='C', deps=['B'], priority=100)
self.sch.add_task(WORKER, task_id='B', deps=['A'])
self.sch.add_task(WORKER, task_id='A')
self.sch.add_task(WORKER, task_id='D', priority=10)
self.check_task_order('ABCD')
def test_update_resources(self):
self.sch.add_task(WORKER, task_id='A', deps=['B'])
self.sch.add_task(WORKER, task_id='B', resources={'r': 2})
self.sch.update_resources(r=1)
# B requires too many resources, we can't schedule
self.check_task_order([])
self.sch.add_task(WORKER, task_id='B', resources={'r': 1})
# now we have enough resources
self.check_task_order(['B', 'A'])
def test_hendle_multiple_resources(self):
self.sch.add_task(WORKER, task_id='A', resources={'r1': 1, 'r2': 1})
self.sch.add_task(WORKER, task_id='B', resources={'r1': 1, 'r2': 1})
self.sch.add_task(WORKER, task_id='C', resources={'r1': 1})
self.sch.update_resources(r1=2, r2=1)
self.assertEqual('A', self.sch.get_work(WORKER)['task_id'])
self.check_task_order('C')
def test_single_resource_lock(self):
self.sch.add_task(WORKER, task_id='A', resources={'r': 1}, priority=10)
self.sch.add_task(WORKER, task_id='B', resources={'r': 2}, priority=10)
self.sch.add_task(WORKER, task_id='C', resources={'r': 1})
self.sch.update_resources(r=2)
self.assertEqual('A', self.sch.get_work(WORKER)['task_id'])
# Should wait for 2 units of r to be available for B before scheduling C
self.check_task_order([])
def test_no_lock_if_too_many_resources_required(self):
self.sch.add_task(WORKER, task_id='A', resources={'r': 2}, priority=10)
self.sch.add_task(WORKER, task_id='B', resources={'r': 1})
self.sch.update_resources(r=1)
self.check_task_order('B')
def test_multiple_resources_lock(self):
self.sch.add_task(WORKER, task_id='A', resources={'r1': 1}, priority=10)
self.sch.add_task(WORKER, task_id='B', resources={'r1': 1, 'r2': 1}, priority=10)
self.sch.add_task(WORKER, task_id='C', resources={'r2': 1})
self.sch.update_resources(r1=1, r2=1)
self.assertEqual('A', self.sch.get_work(WORKER)['task_id'])
# Should wait for r1 to be available for B before using up r2
self.check_task_order([])
def test_multiple_resources_no_lock(self):
self.sch.add_task(WORKER, task_id='A', resources={'r1': 1}, priority=10)
self.sch.add_task(WORKER, task_id='B', resources={'r1': 1, 'r2': 1}, priority=10)
self.sch.add_task(WORKER, task_id='C', resources={'r2': 1})
self.sch.update_resources(r1=1, r2=2)
self.assertEqual('A', self.sch.get_work(WORKER)['task_id'])
# C doesn't block B, so it can go first
self.check_task_order('C')
def check_task_order(self, order):
for expected_id in order:
self.assertEqual(self.sch.get_work(WORKER)['task_id'], expected_id)
self.sch.add_task(WORKER, expected_id, status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_priorities(self):
self.sch.add_task(WORKER, 'A', priority=10)
self.sch.add_task(WORKER, 'B', priority=5)
self.sch.add_task(WORKER, 'C', priority=15)
self.sch.add_task(WORKER, 'D', priority=9)
self.check_task_order(['C', 'A', 'D', 'B'])
def test_priorities_default_and_negative(self):
self.sch.add_task(WORKER, 'A', priority=10)
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C', priority=15)
self.sch.add_task(WORKER, 'D', priority=-20)
self.sch.add_task(WORKER, 'E', priority=1)
self.check_task_order(['C', 'A', 'E', 'B', 'D'])
def test_priorities_and_dependencies(self):
self.sch.add_task(WORKER, 'A', deps=['Z'], priority=10)
self.sch.add_task(WORKER, 'B', priority=5)
self.sch.add_task(WORKER, 'C', deps=['Z'], priority=3)
self.sch.add_task(WORKER, 'D', priority=2)
self.sch.add_task(WORKER, 'Z', priority=1)
self.check_task_order(['Z', 'A', 'B', 'C', 'D'])
def test_priority_update_dependency_after_scheduling(self):
self.sch.add_task(WORKER, 'A', priority=1)
self.sch.add_task(WORKER, 'B', priority=5, deps=['A'])
self.sch.add_task(WORKER, 'C', priority=10, deps=['B'])
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['A', 'B', 'C', 'D'])
def test_disable(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_disable_and_reenable(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.re_enable_task('A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_disable_and_reenable_and_disable_again(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.re_enable_task('A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be still enabled
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled now
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_disable_and_done(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(WORKER, 'A', status=DONE)
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('DONE', '')), 1)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_disable_by_worker(self):
self.sch.add_task(WORKER, 'A', status=DISABLED)
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.sch.add_task(WORKER, 'A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_priority_update_dependency_chain(self):
self.sch.add_task(WORKER, 'A', priority=10, deps=['B'])
self.sch.add_task(WORKER, 'B', priority=5, deps=['C'])
self.sch.add_task(WORKER, 'C', priority=1)
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['C', 'B', 'A', 'D'])
def test_priority_no_decrease_with_multiple_updates(self):
self.sch.add_task(WORKER, 'A', priority=1)
self.sch.add_task(WORKER, 'B', priority=10, deps=['A'])
self.sch.add_task(WORKER, 'C', priority=5, deps=['A'])
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['A', 'B', 'D', 'C'])
def test_unique_tasks(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C')
self.sch.add_task(WORKER + "_2", 'B')
response = self.sch.get_work(WORKER)
self.assertEqual(3, response['n_pending_tasks'])
self.assertEqual(2, response['n_unique_pending'])
def test_prefer_more_dependents(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C', deps=['B'])
self.sch.add_task(WORKER, 'D', deps=['B'])
self.sch.add_task(WORKER, 'E', deps=['A'])
self.check_task_order('BACDE')
def test_prefer_readier_dependents(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C')
self.sch.add_task(WORKER, 'D')
self.sch.add_task(WORKER, 'F', deps=['A', 'B', 'C'])
self.sch.add_task(WORKER, 'G', deps=['A', 'B', 'C'])
self.sch.add_task(WORKER, 'E', deps=['D'])
self.check_task_order('DABCFGE')
def test_ignore_done_dependents(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C')
self.sch.add_task(WORKER, 'D', priority=1)
self.sch.add_task(WORKER, 'E', deps=['C', 'D'])
self.sch.add_task(WORKER, 'F', deps=['A', 'B'])
self.check_task_order('DCABEF')
class CentralPlannerTest(unittest.TestCase):
def setUp(self):
self.sch = CentralPlannerScheduler(retry_delay=100,
remove_delay=1000,
worker_disconnect_delay=10)
self.time = time.time
def tearDown(self):
if time.time != self.time:
time.time = self.time
def setTime(self, t):
time.time = lambda: t
def test_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A', ))
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_failed_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A', ))
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.assertEqual(
self.sch.get_work(WORKER)['task_id'],
None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(WORKER, 'A', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_broken_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A', ))
self.sch.add_task(WORKER, 'A', runnable=False)
self.assertEqual(
self.sch.get_work(WORKER)['task_id'],
None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(WORKER, 'A', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_two_workers(self):
# Worker X wants to build A -> B
# Worker Y wants to build A -> C
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.sch.add_task(task_id='B', deps=('A', ), worker='X')
self.sch.add_task(task_id='C', deps=('A', ), worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'],
None) # Worker Y is pending on A to be done
self.sch.add_task(worker='X', task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'C')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'B')
def test_retry(self):
# Try to build A but fails, will retry after 100s
self.setTime(0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', FAILED)
for t in xrange(100):
self.setTime(t)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
self.sch.ping(WORKER)
if t % 10 == 0:
self.sch.prune()
self.setTime(101)
self.sch.prune()
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_disconnect_running(self):
# X and Y wants to run A.
# X starts but does not report back. Y does.
# After some timeout, Y will build it instead
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.sch.add_task(task_id='A', worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
for t in xrange(200):
self.setTime(t)
self.sch.ping(worker='Y')
if t % 10 == 0:
self.sch.prune()
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'A')
def test_remove_dep(self):
# X schedules A -> B, A is broken
# Y schedules C -> B: this should remove A as a dep of B
self.sch.add_task(task_id='A', worker='X', runnable=False)
self.sch.add_task(task_id='B', deps=('A', ), worker='X')
# X can't build anything
self.assertEqual(self.sch.get_work(worker='X')['task_id'], None)
self.sch.add_task(task_id='B', deps=('C', ),
worker='Y') # should reset dependencies for A
self.sch.add_task(task_id='C', worker='Y', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_timeout(self):
# A bug that was earlier present when restarting the same flow
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.setTime(10000)
self.sch.add_task(
task_id='A',
worker='Y') # Will timeout X but not schedule A for removal
for i in xrange(2000):
self.setTime(10000 + i)
self.sch.ping(worker='Y')
self.sch.add_task(
task_id='A', status=DONE,
worker='Y') # This used to raise an exception since A was removed
def test_disallowed_state_changes(self):
# Test that we can not schedule an already running task
t = 'A'
self.sch.add_task(task_id=t, worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], t)
self.sch.add_task(task_id=t, worker='Y')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None)
def test_two_worker_info(self):
# Make sure the scheduler returns info that some other worker is running task A
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
r = self.sch.get_work(worker='Y')
self.assertEqual(r['task_id'],
None) # Worker Y is pending on A to be done
s = r['running_tasks'][0]
self.assertEqual(s['task_id'], 'A')
self.assertEqual(s['worker'], 'X')
def test_scheduler_resources_none_allow_one(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_resources_none_disallow_two(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.assertFalse(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_insufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=2)
self.assertFalse(self.sch.get_work(worker='X')['task_id'])
def test_scheduler_with_sufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=3)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_resources_used(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R1': 1})
self.sch.update_resources(R1=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_scheduler_overprovisioned_on_other_resource(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.sch.update_resources(R1=2)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R2': 2})
self.sch.update_resources(R1=1, R2=2)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_scheduler_with_priority_and_competing_resources(self):
self.sch.add_task(worker='X', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='X',
task_id='B',
resources={'R': 1},
priority=10)
self.sch.add_task(worker='Y',
task_id='C',
resources={'R': 1},
priority=1)
self.sch.update_resources(R=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
self.sch.add_task(worker='Y', task_id='D', priority=0)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'D')
def test_update_resources(self):
self.sch.add_task(WORKER, task_id='A', deps=['B'])
self.sch.add_task(WORKER, task_id='B', resources={'r': 2})
self.sch.update_resources(r=1)
# B requires too many resources, we can't schedule
self.check_task_order([])
self.sch.add_task(WORKER, task_id='B', resources={'r': 1})
# now we have enough resources
self.check_task_order(['B', 'A'])
def check_task_order(self, order):
for expected_id in order:
self.assertEqual(self.sch.get_work(WORKER)['task_id'], expected_id)
self.sch.add_task(WORKER, expected_id, status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_priorities(self):
self.sch.add_task(WORKER, 'A', priority=10)
self.sch.add_task(WORKER, 'B', priority=5)
self.sch.add_task(WORKER, 'C', priority=15)
self.sch.add_task(WORKER, 'D', priority=9)
self.check_task_order(['C', 'A', 'D', 'B'])
def test_priorities_default_and_negative(self):
self.sch.add_task(WORKER, 'A', priority=10)
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C', priority=15)
self.sch.add_task(WORKER, 'D', priority=-20)
self.sch.add_task(WORKER, 'E', priority=1)
self.check_task_order(['C', 'A', 'E', 'B', 'D'])
def test_priorities_and_dependencies(self):
self.sch.add_task(WORKER, 'A', deps=['Z'], priority=10)
self.sch.add_task(WORKER, 'B', priority=5)
self.sch.add_task(WORKER, 'C', deps=['Z'], priority=3)
self.sch.add_task(WORKER, 'D', priority=2)
self.sch.add_task(WORKER, 'Z', priority=1)
self.check_task_order(['Z', 'A', 'B', 'C', 'D'])
def test_priority_update_dependency_after_scheduling(self):
self.sch.add_task(WORKER, 'A', priority=1)
self.sch.add_task(WORKER, 'B', priority=5, deps=['A'])
self.sch.add_task(WORKER, 'C', priority=10, deps=['B'])
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['A', 'B', 'C', 'D'])
def test_priority_update_dependency_chain(self):
self.sch.add_task(WORKER, 'A', priority=10, deps=['B'])
self.sch.add_task(WORKER, 'B', priority=5, deps=['C'])
self.sch.add_task(WORKER, 'C', priority=1)
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['C', 'B', 'A', 'D'])
def test_priority_no_decrease_with_multiple_updates(self):
self.sch.add_task(WORKER, 'A', priority=1)
self.sch.add_task(WORKER, 'B', priority=10, deps=['A'])
self.sch.add_task(WORKER, 'C', priority=5, deps=['A'])
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['A', 'B', 'D', 'C'])
class CentralPlannerTest(unittest.TestCase):
def setUp(self):
self.sch = CentralPlannerScheduler(retry_delay=100,
remove_delay=1000,
worker_disconnect_delay=10,
disable_persist=10,
disable_window=10,
disable_failures=3)
self.time = time.time
def tearDown(self):
if time.time != self.time:
time.time = self.time
def setTime(self, t):
time.time = lambda: t
def test_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_failed_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(WORKER, 'A', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_broken_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A', runnable=False)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(WORKER, 'A', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_two_workers(self):
# Worker X wants to build A -> B
# Worker Y wants to build A -> C
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.sch.add_task(task_id='B', deps=('A',), worker='X')
self.sch.add_task(task_id='C', deps=('A',), worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None) # Worker Y is pending on A to be done
self.sch.add_task(worker='X', task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'C')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'B')
def test_retry(self):
# Try to build A but fails, will retry after 100s
self.setTime(0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', FAILED)
for t in range(100):
self.setTime(t)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
self.sch.ping(WORKER)
if t % 10 == 0:
self.sch.prune()
self.setTime(101)
self.sch.prune()
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_disconnect_running(self):
# X and Y wants to run A.
# X starts but does not report back. Y does.
# After some timeout, Y will build it instead
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.sch.add_task(task_id='A', worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
for t in range(200):
self.setTime(t)
self.sch.ping(worker='Y')
if t % 10 == 0:
self.sch.prune()
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'A')
def test_remove_dep(self):
# X schedules A -> B, A is broken
# Y schedules C -> B: this should remove A as a dep of B
self.sch.add_task(task_id='A', worker='X', runnable=False)
self.sch.add_task(task_id='B', deps=('A',), worker='X')
# X can't build anything
self.assertEqual(self.sch.get_work(worker='X')['task_id'], None)
self.sch.add_task(task_id='B', deps=('C',), worker='Y') # should reset dependencies for A
self.sch.add_task(task_id='C', worker='Y', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_timeout(self):
# A bug that was earlier present when restarting the same flow
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.setTime(10000)
self.sch.add_task(task_id='A', worker='Y') # Will timeout X but not schedule A for removal
for i in range(2000):
self.setTime(10000 + i)
self.sch.ping(worker='Y')
self.sch.add_task(task_id='A', status=DONE, worker='Y') # This used to raise an exception since A was removed
def test_disallowed_state_changes(self):
# Test that we can not schedule an already running task
t = 'A'
self.sch.add_task(task_id=t, worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], t)
self.sch.add_task(task_id=t, worker='Y')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None)
def test_two_worker_info(self):
# Make sure the scheduler returns info that some other worker is running task A
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
r = self.sch.get_work(worker='Y')
self.assertEqual(r['task_id'], None) # Worker Y is pending on A to be done
s = r['running_tasks'][0]
self.assertEqual(s['task_id'], 'A')
self.assertEqual(s['worker'], 'X')
def test_scheduler_resources_none_allow_one(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_resources_none_disallow_two(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.assertFalse(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_insufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=2)
self.assertFalse(self.sch.get_work(worker='X')['task_id'])
def test_scheduler_with_sufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=3)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_resources_used(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R1': 1})
self.sch.update_resources(R1=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_scheduler_overprovisioned_on_other_resource(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.sch.update_resources(R1=2)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R2': 2})
self.sch.update_resources(R1=1, R2=2)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_scheduler_with_priority_and_competing_resources(self):
self.sch.add_task(worker='X', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=10)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
self.sch.add_task(worker='Y', task_id='D', priority=0)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'D')
def test_do_not_lock_resources_when_not_ready(self):
""" Test to make sure that resources won't go unused waiting on workers """
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertEqual('C', self.sch.get_work(worker='Y')['task_id'])
def test_lock_resources_when_one_of_multiple_workers_is_ready(self):
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 2)])
self.sch.add_worker('Y', [])
self.assertFalse(self.sch.get_work('Y')['task_id'])
def test_do_not_lock_resources_while_running_higher_priority(self):
""" Test to make sure that resources won't go unused waiting on workers """
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertEqual('A', self.sch.get_work('X')['task_id'])
self.assertEqual('C', self.sch.get_work('Y')['task_id'])
def test_lock_resources_while_running_lower_priority(self):
""" Make sure resources will be made available while working on lower priority tasks """
self.sch.add_task(worker='X', task_id='A', priority=4)
self.assertEqual('A', self.sch.get_work('X')['task_id'])
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertFalse(self.sch.get_work('Y')['task_id'])
def test_lock_resources_for_second_worker(self):
self.sch.add_task(worker='X', task_id='A', resources={'R': 1})
self.sch.add_task(worker='X', task_id='B', resources={'R': 1})
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=10)
self.sch.add_worker('X', {'workers': 2})
self.sch.add_worker('Y', {'workers': 1})
self.sch.update_resources(R=2)
self.assertEqual('A', self.sch.get_work('X')['task_id'])
self.assertFalse(self.sch.get_work('X')['task_id'])
def test_can_work_on_lower_priority_while_waiting_for_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R': 1}, priority=0)
self.assertEqual('A', self.sch.get_work('X')['task_id'])
self.sch.add_task(worker='Y', task_id='B', resources={'R': 1}, priority=10)
self.sch.add_task(worker='Y', task_id='C', priority=0)
self.sch.update_resources(R=1)
self.assertEqual('C', self.sch.get_work('Y')['task_id'])
def test_priority_update_with_pruning(self):
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.setTime(50) # after worker disconnects
self.sch.prune()
self.sch.add_task(task_id='B', deps=['A'], worker='X')
self.setTime(2000) # after remove for task A
self.sch.prune()
# Here task A that B depends on is missing
self.sch.add_task(WORKER, task_id='C', deps=['B'], priority=100)
self.sch.add_task(WORKER, task_id='B', deps=['A'])
self.sch.add_task(WORKER, task_id='A')
self.sch.add_task(WORKER, task_id='D', priority=10)
self.check_task_order('ABCD')
def test_update_resources(self):
self.sch.add_task(WORKER, task_id='A', deps=['B'])
self.sch.add_task(WORKER, task_id='B', resources={'r': 2})
self.sch.update_resources(r=1)
# B requires too many resources, we can't schedule
self.check_task_order([])
self.sch.add_task(WORKER, task_id='B', resources={'r': 1})
# now we have enough resources
self.check_task_order(['B', 'A'])
def test_hendle_multiple_resources(self):
self.sch.add_task(WORKER, task_id='A', resources={'r1': 1, 'r2': 1})
self.sch.add_task(WORKER, task_id='B', resources={'r1': 1, 'r2': 1})
self.sch.add_task(WORKER, task_id='C', resources={'r1': 1})
self.sch.update_resources(r1=2, r2=1)
self.assertEqual('A', self.sch.get_work(WORKER)['task_id'])
self.check_task_order('C')
def test_single_resource_lock(self):
self.sch.add_task('X', task_id='A', resources={'r': 1})
self.assertEqual('A', self.sch.get_work('X')['task_id'])
self.sch.add_task(WORKER, task_id='B', resources={'r': 2}, priority=10)
self.sch.add_task(WORKER, task_id='C', resources={'r': 1})
self.sch.update_resources(r=2)
# Should wait for 2 units of r to be available for B before scheduling C
self.check_task_order([])
def test_no_lock_if_too_many_resources_required(self):
self.sch.add_task(WORKER, task_id='A', resources={'r': 2}, priority=10)
self.sch.add_task(WORKER, task_id='B', resources={'r': 1})
self.sch.update_resources(r=1)
self.check_task_order('B')
def test_multiple_resources_lock(self):
self.sch.add_task('X', task_id='A', resources={'r1': 1, 'r2': 1}, priority=10)
self.sch.add_task(WORKER, task_id='B', resources={'r2': 1})
self.sch.add_task(WORKER, task_id='C', resources={'r1': 1})
self.sch.update_resources(r1=1, r2=1)
# should preserve both resources for worker 'X'
self.check_task_order([])
def test_multiple_resources_no_lock(self):
self.sch.add_task(WORKER, task_id='A', resources={'r1': 1}, priority=10)
self.sch.add_task(WORKER, task_id='B', resources={'r1': 1, 'r2': 1}, priority=10)
self.sch.add_task(WORKER, task_id='C', resources={'r2': 1})
self.sch.update_resources(r1=1, r2=2)
self.assertEqual('A', self.sch.get_work(WORKER)['task_id'])
# C doesn't block B, so it can go first
self.check_task_order('C')
def check_task_order(self, order):
for expected_id in order:
self.assertEqual(self.sch.get_work(WORKER)['task_id'], expected_id)
self.sch.add_task(WORKER, expected_id, status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_priorities(self):
self.sch.add_task(WORKER, 'A', priority=10)
self.sch.add_task(WORKER, 'B', priority=5)
self.sch.add_task(WORKER, 'C', priority=15)
self.sch.add_task(WORKER, 'D', priority=9)
self.check_task_order(['C', 'A', 'D', 'B'])
def test_priorities_default_and_negative(self):
self.sch.add_task(WORKER, 'A', priority=10)
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C', priority=15)
self.sch.add_task(WORKER, 'D', priority=-20)
self.sch.add_task(WORKER, 'E', priority=1)
self.check_task_order(['C', 'A', 'E', 'B', 'D'])
def test_priorities_and_dependencies(self):
self.sch.add_task(WORKER, 'A', deps=['Z'], priority=10)
self.sch.add_task(WORKER, 'B', priority=5)
self.sch.add_task(WORKER, 'C', deps=['Z'], priority=3)
self.sch.add_task(WORKER, 'D', priority=2)
self.sch.add_task(WORKER, 'Z', priority=1)
self.check_task_order(['Z', 'A', 'B', 'C', 'D'])
def test_priority_update_dependency_after_scheduling(self):
self.sch.add_task(WORKER, 'A', priority=1)
self.sch.add_task(WORKER, 'B', priority=5, deps=['A'])
self.sch.add_task(WORKER, 'C', priority=10, deps=['B'])
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['A', 'B', 'C', 'D'])
def test_disable(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_disable_and_reenable(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.re_enable_task('A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_disable_and_reenable_and_disable_again(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.re_enable_task('A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be still enabled
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled now
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_disable_and_done(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(WORKER, 'A', status=DONE)
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('DONE', '')), 1)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_disable_by_worker(self):
self.sch.add_task(WORKER, 'A', status=DISABLED)
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.sch.add_task(WORKER, 'A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_task_list_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
for c in 'ABCD':
sch.add_task(WORKER, c)
self.assertEqual(set('ABCD'), set(sch.task_list('PENDING', '', False).keys()))
self.assertEqual({'num_tasks': 4}, sch.task_list('PENDING', ''))
def test_task_list_within_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=4)
for c in 'ABCD':
sch.add_task(WORKER, c)
self.assertEqual(set('ABCD'), set(sch.task_list('PENDING', '').keys()))
def test_task_lists_some_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
for c in 'ABCD':
sch.add_task(WORKER, c, 'DONE')
for c in 'EFG':
sch.add_task(WORKER, c)
self.assertEqual(set('EFG'), set(sch.task_list('PENDING', '').keys()))
self.assertEqual({'num_tasks': 4}, sch.task_list('DONE', ''))
def test_priority_update_dependency_chain(self):
self.sch.add_task(WORKER, 'A', priority=10, deps=['B'])
self.sch.add_task(WORKER, 'B', priority=5, deps=['C'])
self.sch.add_task(WORKER, 'C', priority=1)
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['C', 'B', 'A', 'D'])
def test_priority_no_decrease_with_multiple_updates(self):
self.sch.add_task(WORKER, 'A', priority=1)
self.sch.add_task(WORKER, 'B', priority=10, deps=['A'])
self.sch.add_task(WORKER, 'C', priority=5, deps=['A'])
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['A', 'B', 'D', 'C'])
def test_unique_tasks(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C')
self.sch.add_task(WORKER + "_2", 'B')
response = self.sch.get_work(WORKER)
self.assertEqual(3, response['n_pending_tasks'])
self.assertEqual(2, response['n_unique_pending'])
def test_prefer_more_dependents(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C', deps=['B'])
self.sch.add_task(WORKER, 'D', deps=['B'])
self.sch.add_task(WORKER, 'E', deps=['A'])
self.check_task_order('BACDE')
def test_prefer_readier_dependents(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C')
self.sch.add_task(WORKER, 'D')
self.sch.add_task(WORKER, 'F', deps=['A', 'B', 'C'])
self.sch.add_task(WORKER, 'G', deps=['A', 'B', 'C'])
self.sch.add_task(WORKER, 'E', deps=['D'])
self.check_task_order('DABCFGE')
def test_ignore_done_dependents(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C')
self.sch.add_task(WORKER, 'D', priority=1)
self.sch.add_task(WORKER, 'E', deps=['C', 'D'])
self.sch.add_task(WORKER, 'F', deps=['A', 'B'])
self.check_task_order('DCABEF')
class CentralPlannerTest(unittest.TestCase):
def setUp(self):
super(CentralPlannerTest, self).setUp()
conf = self.get_scheduler_config()
self.sch = CentralPlannerScheduler(**conf)
self.time = time.time
def get_scheduler_config(self):
return {
'retry_delay': 100,
'remove_delay': 1000,
'worker_disconnect_delay': 10,
'disable_persist': 10,
'disable_window': 10,
'disable_failures': 3,
}
def tearDown(self):
super(CentralPlannerTest, self).tearDown()
if time.time != self.time:
time.time = self.time
def setTime(self, t):
time.time = lambda: t
def test_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_failed_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(WORKER, 'A', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_broken_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A', runnable=False)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(WORKER, 'A', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_two_workers(self):
# Worker X wants to build A -> B
# Worker Y wants to build A -> C
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.sch.add_task(task_id='B', deps=('A',), worker='X')
self.sch.add_task(task_id='C', deps=('A',), worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None) # Worker Y is pending on A to be done
self.sch.add_task(worker='X', task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'C')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'B')
def test_retry(self):
# Try to build A but fails, will retry after 100s
self.setTime(0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', FAILED)
for t in range(100):
self.setTime(t)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
self.sch.ping(WORKER)
if t % 10 == 0:
self.sch.prune()
self.setTime(101)
self.sch.prune()
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_disconnect_running(self):
# X and Y wants to run A.
# X starts but does not report back. Y does.
# After some timeout, Y will build it instead
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.sch.add_task(task_id='A', worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
for t in range(200):
self.setTime(t)
self.sch.ping(worker='Y')
if t % 10 == 0:
self.sch.prune()
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'A')
def test_remove_dep(self):
# X schedules A -> B, A is broken
# Y schedules C -> B: this should remove A as a dep of B
self.sch.add_task(task_id='A', worker='X', runnable=False)
self.sch.add_task(task_id='B', deps=('A',), worker='X')
# X can't build anything
self.assertEqual(self.sch.get_work(worker='X')['task_id'], None)
self.sch.add_task(task_id='B', deps=('C',), worker='Y') # should reset dependencies for A
self.sch.add_task(task_id='C', worker='Y', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_timeout(self):
# A bug that was earlier present when restarting the same flow
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.setTime(10000)
self.sch.add_task(task_id='A', worker='Y') # Will timeout X but not schedule A for removal
for i in range(2000):
self.setTime(10000 + i)
self.sch.ping(worker='Y')
self.sch.add_task(task_id='A', status=DONE, worker='Y') # This used to raise an exception since A was removed
def test_disallowed_state_changes(self):
# Test that we can not schedule an already running task
t = 'A'
self.sch.add_task(task_id=t, worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], t)
self.sch.add_task(task_id=t, worker='Y')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None)
def test_two_worker_info(self):
# Make sure the scheduler returns info that some other worker is running task A
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
r = self.sch.get_work(worker='Y')
self.assertEqual(r['task_id'], None) # Worker Y is pending on A to be done
s = r['running_tasks'][0]
self.assertEqual(s['task_id'], 'A')
self.assertEqual(s['worker'], 'X')
def test_assistant_get_work(self):
self.sch.add_task(worker='X', task_id='A')
self.sch.add_worker('Y', [])
self.assertEqual(self.sch.get_work('Y', assistant=True)['task_id'], 'A')
# check that the scheduler recognizes tasks as running
running_tasks = self.sch.task_list('RUNNING', '')
self.assertEqual(len(running_tasks), 1)
self.assertEqual(list(running_tasks.keys()), ['A'])
self.assertEqual(running_tasks['A']['worker_running'], 'Y')
def test_assistant_get_work_external_task(self):
self.sch.add_task('X', task_id='A', runnable=False)
self.assertTrue(self.sch.get_work('Y', assistant=True)['task_id'] is None)
def test_task_fails_when_assistant_dies(self):
self.setTime(0)
self.sch.add_task(worker='X', task_id='A')
self.sch.add_worker('Y', [])
self.assertEqual(self.sch.get_work('Y', assistant=True)['task_id'], 'A')
self.assertEqual(list(self.sch.task_list('RUNNING', '').keys()), ['A'])
# Y dies for 50 seconds, X stays alive
self.setTime(50)
self.sch.ping('X')
self.assertEqual(list(self.sch.task_list('FAILED', '').keys()), ['A'])
def test_prune_with_live_assistant(self):
self.setTime(0)
self.sch.add_task(worker='X', task_id='A')
self.sch.get_work('Y', assistant=True)
self.sch.add_task(worker='Y', task_id='A', status=DONE, assistant=True)
# worker X stops communicating, A should be marked for removal
self.setTime(600)
self.sch.ping('Y')
self.sch.prune()
# A will now be pruned
self.setTime(2000)
self.sch.prune()
self.assertFalse(list(self.sch.task_list('', '')))
def test_prune_done_tasks(self, expected=None):
self.setTime(0)
self.sch.add_task(WORKER, task_id='A', status=DONE)
self.sch.add_task(WORKER, task_id='B', deps=['A'], status=DONE)
self.sch.add_task(WORKER, task_id='C', deps=['B'])
self.setTime(600)
self.sch.ping('ASSISTANT')
self.sch.prune()
self.setTime(2000)
self.sch.ping('ASSISTANT')
self.sch.prune()
self.assertEqual(set(expected or ()), set(self.sch.task_list('', '').keys()))
def test_keep_tasks_for_assistant(self):
self.sch.get_work('ASSISTANT', assistant=True) # tell the scheduler this is an assistant
self.test_prune_done_tasks(['B', 'C'])
def test_keep_scheduler_disabled_tasks_for_assistant(self):
self.sch.get_work('ASSISTANT', assistant=True) # tell the scheduler this is an assistant
# create a scheduler disabled task and a worker disabled task
for i in range(10):
self.sch.add_task(WORKER, 'D', status=FAILED)
self.sch.add_task(WORKER, 'E', status=DISABLED)
# scheduler prunes the worker disabled task
self.assertEqual(set(['D', 'E']), set(self.sch.task_list(DISABLED, '')))
self.test_prune_done_tasks(['B', 'C', 'D'])
def test_keep_failed_tasks_for_assistant(self):
self.sch.get_work('ASSISTANT', assistant=True) # tell the scheduler this is an assistant
self.sch.add_task(WORKER, 'D', status=FAILED, deps='A')
self.test_prune_done_tasks(['A', 'B', 'C', 'D'])
def test_scheduler_resources_none_allow_one(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_resources_none_disallow_two(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.assertFalse(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_insufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=2)
self.assertFalse(self.sch.get_work(worker='X')['task_id'])
def test_scheduler_with_sufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=3)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_resources_used(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R1': 1})
self.sch.update_resources(R1=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_scheduler_overprovisioned_on_other_resource(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.sch.update_resources(R1=2)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R2': 2})
self.sch.update_resources(R1=1, R2=2)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_scheduler_with_priority_and_competing_resources(self):
self.sch.add_task(worker='X', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=10)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
self.sch.add_task(worker='Y', task_id='D', priority=0)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'D')
def test_do_not_lock_resources_when_not_ready(self):
""" Test to make sure that resources won't go unused waiting on workers """
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertEqual('C', self.sch.get_work(worker='Y')['task_id'])
def test_lock_resources_when_one_of_multiple_workers_is_ready(self):
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 2)])
self.sch.add_worker('Y', [])
self.assertFalse(self.sch.get_work('Y')['task_id'])
def test_do_not_lock_resources_while_running_higher_priority(self):
""" Test to make sure that resources won't go unused waiting on workers """
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertEqual('A', self.sch.get_work('X')['task_id'])
self.assertEqual('C', self.sch.get_work('Y')['task_id'])
def test_lock_resources_while_running_lower_priority(self):
""" Make sure resources will be made available while working on lower priority tasks """
self.sch.add_task(worker='X', task_id='A', priority=4)
self.assertEqual('A', self.sch.get_work('X')['task_id'])
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertFalse(self.sch.get_work('Y')['task_id'])
def test_lock_resources_for_second_worker(self):
self.sch.add_task(worker='X', task_id='A', resources={'R': 1})
self.sch.add_task(worker='X', task_id='B', resources={'R': 1})
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=10)
self.sch.add_worker('X', {'workers': 2})
self.sch.add_worker('Y', {'workers': 1})
self.sch.update_resources(R=2)
self.assertEqual('A', self.sch.get_work('X')['task_id'])
self.assertFalse(self.sch.get_work('X')['task_id'])
def test_can_work_on_lower_priority_while_waiting_for_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R': 1}, priority=0)
self.assertEqual('A', self.sch.get_work('X')['task_id'])
self.sch.add_task(worker='Y', task_id='B', resources={'R': 1}, priority=10)
self.sch.add_task(worker='Y', task_id='C', priority=0)
self.sch.update_resources(R=1)
self.assertEqual('C', self.sch.get_work('Y')['task_id'])
def test_priority_update_with_pruning(self):
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.setTime(50) # after worker disconnects
self.sch.prune()
self.sch.add_task(task_id='B', deps=['A'], worker='X')
self.setTime(2000) # after remove for task A
self.sch.prune()
# Here task A that B depends on is missing
self.sch.add_task(WORKER, task_id='C', deps=['B'], priority=100)
self.sch.add_task(WORKER, task_id='B', deps=['A'])
self.sch.add_task(WORKER, task_id='A')
self.sch.add_task(WORKER, task_id='D', priority=10)
self.check_task_order('ABCD')
def test_update_resources(self):
self.sch.add_task(WORKER, task_id='A', deps=['B'])
self.sch.add_task(WORKER, task_id='B', resources={'r': 2})
self.sch.update_resources(r=1)
# B requires too many resources, we can't schedule
self.check_task_order([])
self.sch.add_task(WORKER, task_id='B', resources={'r': 1})
# now we have enough resources
self.check_task_order(['B', 'A'])
def test_hendle_multiple_resources(self):
self.sch.add_task(WORKER, task_id='A', resources={'r1': 1, 'r2': 1})
self.sch.add_task(WORKER, task_id='B', resources={'r1': 1, 'r2': 1})
self.sch.add_task(WORKER, task_id='C', resources={'r1': 1})
self.sch.update_resources(r1=2, r2=1)
self.assertEqual('A', self.sch.get_work(WORKER)['task_id'])
self.check_task_order('C')
def test_single_resource_lock(self):
self.sch.add_task('X', task_id='A', resources={'r': 1})
self.assertEqual('A', self.sch.get_work('X')['task_id'])
self.sch.add_task(WORKER, task_id='B', resources={'r': 2}, priority=10)
self.sch.add_task(WORKER, task_id='C', resources={'r': 1})
self.sch.update_resources(r=2)
# Should wait for 2 units of r to be available for B before scheduling C
self.check_task_order([])
def test_no_lock_if_too_many_resources_required(self):
self.sch.add_task(WORKER, task_id='A', resources={'r': 2}, priority=10)
self.sch.add_task(WORKER, task_id='B', resources={'r': 1})
self.sch.update_resources(r=1)
self.check_task_order('B')
def test_multiple_resources_lock(self):
self.sch.add_task('X', task_id='A', resources={'r1': 1, 'r2': 1}, priority=10)
self.sch.add_task(WORKER, task_id='B', resources={'r2': 1})
self.sch.add_task(WORKER, task_id='C', resources={'r1': 1})
self.sch.update_resources(r1=1, r2=1)
# should preserve both resources for worker 'X'
self.check_task_order([])
def test_multiple_resources_no_lock(self):
self.sch.add_task(WORKER, task_id='A', resources={'r1': 1}, priority=10)
self.sch.add_task(WORKER, task_id='B', resources={'r1': 1, 'r2': 1}, priority=10)
self.sch.add_task(WORKER, task_id='C', resources={'r2': 1})
self.sch.update_resources(r1=1, r2=2)
self.assertEqual('A', self.sch.get_work(WORKER)['task_id'])
# C doesn't block B, so it can go first
self.check_task_order('C')
def check_task_order(self, order):
for expected_id in order:
self.assertEqual(self.sch.get_work(WORKER)['task_id'], expected_id)
self.sch.add_task(WORKER, expected_id, status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_priorities(self):
self.sch.add_task(WORKER, 'A', priority=10)
self.sch.add_task(WORKER, 'B', priority=5)
self.sch.add_task(WORKER, 'C', priority=15)
self.sch.add_task(WORKER, 'D', priority=9)
self.check_task_order(['C', 'A', 'D', 'B'])
def test_priorities_default_and_negative(self):
self.sch.add_task(WORKER, 'A', priority=10)
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C', priority=15)
self.sch.add_task(WORKER, 'D', priority=-20)
self.sch.add_task(WORKER, 'E', priority=1)
self.check_task_order(['C', 'A', 'E', 'B', 'D'])
def test_priorities_and_dependencies(self):
self.sch.add_task(WORKER, 'A', deps=['Z'], priority=10)
self.sch.add_task(WORKER, 'B', priority=5)
self.sch.add_task(WORKER, 'C', deps=['Z'], priority=3)
self.sch.add_task(WORKER, 'D', priority=2)
self.sch.add_task(WORKER, 'Z', priority=1)
self.check_task_order(['Z', 'A', 'B', 'C', 'D'])
def test_priority_update_dependency_after_scheduling(self):
self.sch.add_task(WORKER, 'A', priority=1)
self.sch.add_task(WORKER, 'B', priority=5, deps=['A'])
self.sch.add_task(WORKER, 'C', priority=10, deps=['B'])
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['A', 'B', 'C', 'D'])
def test_disable(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_disable_and_reenable(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.re_enable_task('A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_disable_and_reenable_and_disable_again(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.re_enable_task('A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be still enabled
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled now
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_disable_and_done(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(WORKER, 'A', status=DONE)
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('DONE', '')), 1)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_disable_by_worker(self):
self.sch.add_task(WORKER, 'A', status=DISABLED)
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.sch.add_task(WORKER, 'A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_task_list_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
for c in 'ABCD':
sch.add_task(WORKER, c)
self.assertEqual(set('ABCD'), set(sch.task_list('PENDING', '', False).keys()))
self.assertEqual({'num_tasks': 4}, sch.task_list('PENDING', ''))
def test_task_list_within_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=4)
for c in 'ABCD':
sch.add_task(WORKER, c)
self.assertEqual(set('ABCD'), set(sch.task_list('PENDING', '').keys()))
def test_task_lists_some_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
for c in 'ABCD':
sch.add_task(WORKER, c, 'DONE')
for c in 'EFG':
sch.add_task(WORKER, c)
self.assertEqual(set('EFG'), set(sch.task_list('PENDING', '').keys()))
self.assertEqual({'num_tasks': 4}, sch.task_list('DONE', ''))
def test_priority_update_dependency_chain(self):
self.sch.add_task(WORKER, 'A', priority=10, deps=['B'])
self.sch.add_task(WORKER, 'B', priority=5, deps=['C'])
self.sch.add_task(WORKER, 'C', priority=1)
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['C', 'B', 'A', 'D'])
def test_priority_no_decrease_with_multiple_updates(self):
self.sch.add_task(WORKER, 'A', priority=1)
self.sch.add_task(WORKER, 'B', priority=10, deps=['A'])
self.sch.add_task(WORKER, 'C', priority=5, deps=['A'])
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['A', 'B', 'D', 'C'])
def test_unique_tasks(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C')
self.sch.add_task(WORKER + "_2", 'B')
response = self.sch.get_work(WORKER)
self.assertEqual(3, response['n_pending_tasks'])
self.assertEqual(2, response['n_unique_pending'])
def test_pending_downstream_disable(self):
self.sch.add_task(WORKER, 'A', status=DISABLED)
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'C', deps=('B',))
response = self.sch.get_work(WORKER)
self.assertTrue(response['task_id'] is None)
self.assertEqual(0, response['n_pending_tasks'])
self.assertEqual(0, response['n_unique_pending'])
def test_pending_downstream_failure(self):
self.sch.add_task(WORKER, 'A', status=FAILED)
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'C', deps=('B',))
response = self.sch.get_work(WORKER)
self.assertTrue(response['task_id'] is None)
self.assertEqual(2, response['n_pending_tasks'])
self.assertEqual(2, response['n_unique_pending'])
def test_prefer_more_dependents(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C', deps=['B'])
self.sch.add_task(WORKER, 'D', deps=['B'])
self.sch.add_task(WORKER, 'E', deps=['A'])
self.check_task_order('BACDE')
def test_prefer_readier_dependents(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C')
self.sch.add_task(WORKER, 'D')
self.sch.add_task(WORKER, 'F', deps=['A', 'B', 'C'])
self.sch.add_task(WORKER, 'G', deps=['A', 'B', 'C'])
self.sch.add_task(WORKER, 'E', deps=['D'])
self.check_task_order('DABCFGE')
def test_ignore_done_dependents(self):
self.sch.add_task(WORKER, 'A')
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C')
self.sch.add_task(WORKER, 'D', priority=1)
self.sch.add_task(WORKER, 'E', deps=['C', 'D'])
self.sch.add_task(WORKER, 'F', deps=['A', 'B'])
self.check_task_order('DCABEF')
class WorkerExternalTaskTest(unittest.TestCase):
def setUp(self):
self.scheduler = CentralPlannerScheduler(retry_delay=0.01,
remove_delay=3,
worker_disconnect_delay=3,
disable_persist=3,
disable_window=5,
disable_failures=2)
def _assert_complete(self, tasks):
for t in tasks:
self.assert_(t.complete())
def _build(self, tasks):
w = luigi.worker.Worker(scheduler=self.scheduler, worker_processes=1)
for t in tasks:
w.add(t)
w.run()
w.stop()
def test_external_dependency_already_complete(self):
"""
Test that the test task completes when its dependency exists at the
start of the execution.
"""
tempdir = tempfile.mkdtemp(prefix='luigi-test-')
test_task = TestTask(tempdir=tempdir, complete_after=1)
luigi.build([test_task], local_scheduler=True)
assert os.path.exists(test_task.dep_path)
assert os.path.exists(test_task.output_path)
os.unlink(test_task.dep_path)
os.unlink(test_task.output_path)
os.rmdir(tempdir)
# complete() is called once per failure, twice per success
assert test_task.dependency.times_called == 2
@with_config({
'core': {
'retry-external-tasks': 'true',
'disable-num-failures': '4',
'max-reschedules': '4',
'worker-keep-alive': 'true',
'retry-delay': '0.01'
}
})
def test_external_dependency_completes_later(self):
"""
Test that an external dependency that is not `complete` when luigi is invoked, but \
becomes `complete` while the workflow is executing is re-evaluated and
allows dependencies to run.
"""
assert luigi.configuration.get_config().getboolean(
'core', 'retry-external-tasks', False) is True
original_get_work = self.scheduler.get_work
def decorated_get_work(*args, **kwargs):
# need to call `prune()` to make the scheduler run the retry logic
self.scheduler.prune()
return original_get_work(*args, **kwargs)
self.scheduler.get_work = decorated_get_work
tempdir = tempfile.mkdtemp(prefix='luigi-test-')
with patch('random.randint', return_value=0.1):
test_task = TestTask(tempdir=tempdir, complete_after=3)
self._build([test_task])
assert os.path.exists(test_task.dep_path)
assert os.path.exists(test_task.output_path)
os.unlink(test_task.dep_path)
os.unlink(test_task.output_path)
os.rmdir(tempdir)
# complete() is called once per failure, twice per success
assert test_task.dependency.times_called == 4
class CentralPlannerTest(unittest.TestCase):
def setUp(self):
super(CentralPlannerTest, self).setUp()
conf = self.get_scheduler_config()
self.sch = CentralPlannerScheduler(**conf)
self.time = time.time
def get_scheduler_config(self):
return {
"retry_delay": 100,
"remove_delay": 1000,
"worker_disconnect_delay": 10,
"disable_persist": 10,
"disable_window": 10,
"disable_failures": 3,
"disable_hard_timeout": 60 * 60,
}
def tearDown(self):
super(CentralPlannerTest, self).tearDown()
if time.time != self.time:
time.time = self.time
def setTime(self, t):
time.time = lambda: t
def test_dep(self):
self.sch.add_task(worker=WORKER, task_id="B", deps=("A",))
self.sch.add_task(worker=WORKER, task_id="A")
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], "A")
self.sch.add_task(worker=WORKER, task_id="A", status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], "B")
self.sch.add_task(worker=WORKER, task_id="B", status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], None)
def test_failed_dep(self):
self.sch.add_task(worker=WORKER, task_id="B", deps=("A",))
self.sch.add_task(worker=WORKER, task_id="A")
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], "A")
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
self.assertEqual(
self.sch.get_work(worker=WORKER)["task_id"], None
) # can still wait and retry: TODO: do we want this?
self.sch.add_task(worker=WORKER, task_id="A", status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], "B")
self.sch.add_task(worker=WORKER, task_id="B", status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], None)
def test_broken_dep(self):
self.sch.add_task(worker=WORKER, task_id="B", deps=("A",))
self.sch.add_task(worker=WORKER, task_id="A", runnable=False)
self.assertEqual(
self.sch.get_work(worker=WORKER)["task_id"], None
) # can still wait and retry: TODO: do we want this?
self.sch.add_task(worker=WORKER, task_id="A", status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], "B")
self.sch.add_task(worker=WORKER, task_id="B", status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], None)
def test_two_workers(self):
# Worker X wants to build A -> B
# Worker Y wants to build A -> C
self.sch.add_task(worker="X", task_id="A")
self.sch.add_task(worker="Y", task_id="A")
self.sch.add_task(task_id="B", deps=("A",), worker="X")
self.sch.add_task(task_id="C", deps=("A",), worker="Y")
self.assertEqual(self.sch.get_work(worker="X")["task_id"], "A")
self.assertEqual(self.sch.get_work(worker="Y")["task_id"], None) # Worker Y is pending on A to be done
self.sch.add_task(worker="X", task_id="A", status=DONE)
self.assertEqual(self.sch.get_work(worker="Y")["task_id"], "C")
self.assertEqual(self.sch.get_work(worker="X")["task_id"], "B")
def test_retry(self):
# Try to build A but fails, will retry after 100s
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id="A")
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], "A")
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
for t in range(100):
self.setTime(t)
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], None)
self.sch.ping(worker=WORKER)
if t % 10 == 0:
self.sch.prune()
self.setTime(101)
self.sch.prune()
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], "A")
def test_resend_task(self):
self.sch.add_task(worker=WORKER, task_id="A")
self.sch.add_task(worker=WORKER, task_id="B")
for _ in range(10):
self.assertEqual("A", self.sch.get_work(worker=WORKER, current_tasks=[])["task_id"])
self.assertEqual("B", self.sch.get_work(worker=WORKER, current_tasks=["A"])["task_id"])
def test_resend_multiple_tasks(self):
self.sch.add_task(worker=WORKER, task_id="A")
self.sch.add_task(worker=WORKER, task_id="B")
self.sch.add_task(worker=WORKER, task_id="C")
# get A and B running
self.assertEqual("A", self.sch.get_work(worker=WORKER)["task_id"])
self.assertEqual("B", self.sch.get_work(worker=WORKER)["task_id"])
for _ in range(10):
self.assertEqual("A", self.sch.get_work(worker=WORKER, current_tasks=[])["task_id"])
self.assertEqual("A", self.sch.get_work(worker=WORKER, current_tasks=["B"])["task_id"])
self.assertEqual("B", self.sch.get_work(worker=WORKER, current_tasks=["A"])["task_id"])
self.assertEqual("C", self.sch.get_work(worker=WORKER, current_tasks=["A", "B"])["task_id"])
def test_disconnect_running(self):
# X and Y wants to run A.
# X starts but does not report back. Y does.
# After some timeout, Y will build it instead
self.setTime(0)
self.sch.add_task(task_id="A", worker="X")
self.sch.add_task(task_id="A", worker="Y")
self.assertEqual(self.sch.get_work(worker="X")["task_id"], "A")
for t in range(200):
self.setTime(t)
self.sch.ping(worker="Y")
if t % 10 == 0:
self.sch.prune()
self.assertEqual(self.sch.get_work(worker="Y")["task_id"], "A")
def test_do_not_overwrite_tracking_url_while_running(self):
self.sch.add_task(task_id="A", worker="X", status="RUNNING", tracking_url="trackme")
self.assertEqual("trackme", self.sch.task_list("RUNNING", "")["A"]["tracking_url"])
# not wiped out by another working scheduling as pending
self.sch.add_task(task_id="A", worker="Y", status="PENDING")
self.assertEqual("trackme", self.sch.task_list("RUNNING", "")["A"]["tracking_url"])
def test_do_update_tracking_url_while_running(self):
self.sch.add_task(task_id="A", worker="X", status="RUNNING", tracking_url="trackme")
self.assertEqual("trackme", self.sch.task_list("RUNNING", "")["A"]["tracking_url"])
self.sch.add_task(task_id="A", worker="X", status="RUNNING", tracking_url="stage_2")
self.assertEqual("stage_2", self.sch.task_list("RUNNING", "")["A"]["tracking_url"])
def test_keep_tracking_url_on_done_and_fail(self):
for status in ("DONE", "FAILED"):
self.sch.add_task(task_id="A", worker="X", status="RUNNING", tracking_url="trackme")
self.assertEqual("trackme", self.sch.task_list("RUNNING", "")["A"]["tracking_url"])
self.sch.add_task(task_id="A", worker="X", status=status)
self.assertEqual("trackme", self.sch.task_list(status, "")["A"]["tracking_url"])
def test_drop_tracking_url_when_rescheduled_while_not_running(self):
for status in ("DONE", "FAILED", "PENDING"):
self.sch.add_task(task_id="A", worker="X", status=status, tracking_url="trackme")
self.assertEqual("trackme", self.sch.task_list(status, "")["A"]["tracking_url"])
self.sch.add_task(task_id="A", worker="Y", status="PENDING")
self.assertIsNone(self.sch.task_list("PENDING", "")["A"]["tracking_url"])
def test_reset_tracking_url_on_new_run(self):
self.sch.add_task(task_id="A", worker="X", status="PENDING", tracking_url="trackme")
self.assertEqual("trackme", self.sch.task_list("PENDING", "")["A"]["tracking_url"])
self.sch.add_task(task_id="A", worker="Y", status="RUNNING")
self.assertIsNone(self.sch.task_list("RUNNING", "")["A"]["tracking_url"])
def test_remove_dep(self):
# X schedules A -> B, A is broken
# Y schedules C -> B: this should remove A as a dep of B
self.sch.add_task(task_id="A", worker="X", runnable=False)
self.sch.add_task(task_id="B", deps=("A",), worker="X")
# X can't build anything
self.assertEqual(self.sch.get_work(worker="X")["task_id"], None)
self.sch.add_task(task_id="B", deps=("C",), worker="Y") # should reset dependencies for A
self.sch.add_task(task_id="C", worker="Y", status=DONE)
self.assertEqual(self.sch.get_work(worker="Y")["task_id"], "B")
def test_timeout(self):
# A bug that was earlier present when restarting the same flow
self.setTime(0)
self.sch.add_task(task_id="A", worker="X")
self.assertEqual(self.sch.get_work(worker="X")["task_id"], "A")
self.setTime(10000)
self.sch.add_task(task_id="A", worker="Y") # Will timeout X but not schedule A for removal
for i in range(2000):
self.setTime(10000 + i)
self.sch.ping(worker="Y")
self.sch.add_task(task_id="A", status=DONE, worker="Y") # This used to raise an exception since A was removed
def test_disallowed_state_changes(self):
# Test that we can not schedule an already running task
t = "A"
self.sch.add_task(task_id=t, worker="X")
self.assertEqual(self.sch.get_work(worker="X")["task_id"], t)
self.sch.add_task(task_id=t, worker="Y")
self.assertEqual(self.sch.get_work(worker="Y")["task_id"], None)
def test_two_worker_info(self):
# Make sure the scheduler returns info that some other worker is running task A
self.sch.add_task(worker="X", task_id="A")
self.sch.add_task(worker="Y", task_id="A")
self.assertEqual(self.sch.get_work(worker="X")["task_id"], "A")
r = self.sch.get_work(worker="Y")
self.assertEqual(r["task_id"], None) # Worker Y is pending on A to be done
s = r["running_tasks"][0]
self.assertEqual(s["task_id"], "A")
self.assertEqual(s["worker"], "X")
def test_assistant_get_work(self):
self.sch.add_task(worker="X", task_id="A")
self.sch.add_worker("Y", [])
self.assertEqual(self.sch.get_work(worker="Y", assistant=True)["task_id"], "A")
# check that the scheduler recognizes tasks as running
running_tasks = self.sch.task_list("RUNNING", "")
self.assertEqual(len(running_tasks), 1)
self.assertEqual(list(running_tasks.keys()), ["A"])
self.assertEqual(running_tasks["A"]["worker_running"], "Y")
def test_assistant_get_work_external_task(self):
self.sch.add_task(worker="X", task_id="A", runnable=False)
self.assertTrue(self.sch.get_work(worker="Y", assistant=True)["task_id"] is None)
def test_task_fails_when_assistant_dies(self):
self.setTime(0)
self.sch.add_task(worker="X", task_id="A")
self.sch.add_worker("Y", [])
self.assertEqual(self.sch.get_work(worker="Y", assistant=True)["task_id"], "A")
self.assertEqual(list(self.sch.task_list("RUNNING", "").keys()), ["A"])
# Y dies for 50 seconds, X stays alive
self.setTime(50)
self.sch.ping(worker="X")
self.assertEqual(list(self.sch.task_list("FAILED", "").keys()), ["A"])
def test_prune_with_live_assistant(self):
self.setTime(0)
self.sch.add_task(worker="X", task_id="A")
self.sch.get_work(worker="Y", assistant=True)
self.sch.add_task(worker="Y", task_id="A", status=DONE, assistant=True)
# worker X stops communicating, A should be marked for removal
self.setTime(600)
self.sch.ping(worker="Y")
self.sch.prune()
# A will now be pruned
self.setTime(2000)
self.sch.prune()
self.assertFalse(list(self.sch.task_list("", "")))
def test_fail_job_from_dead_worker_with_live_assistant(self):
self.setTime(0)
self.sch.add_task(worker="X", task_id="A")
self.assertEqual("A", self.sch.get_work(worker="X")["task_id"])
self.sch.add_worker("Y", [("assistant", True)])
self.setTime(600)
self.sch.ping(worker="Y")
self.sch.prune()
self.assertEqual(["A"], list(self.sch.task_list("FAILED", "").keys()))
def test_assistant_request_runnable_task(self):
self.setTime(0)
self.sch.add_task(worker="X", task_id="A", runnable=True)
self.setTime(600)
self.sch.prune()
self.assertEqual("A", self.sch.get_work(worker="Y", assistant=True)["task_id"])
def test_assistant_request_external_task(self):
self.sch.add_task(worker="X", task_id="A", runnable=False)
self.assertIsNone(self.sch.get_work(worker="Y", assistant=True)["task_id"])
def test_prune_done_tasks(self, expected=None):
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id="A", status=DONE)
self.sch.add_task(worker=WORKER, task_id="B", deps=["A"], status=DONE)
self.sch.add_task(worker=WORKER, task_id="C", deps=["B"])
self.setTime(600)
self.sch.ping(worker="ASSISTANT")
self.sch.prune()
self.setTime(2000)
self.sch.ping(worker="ASSISTANT")
self.sch.prune()
self.assertEqual(set(expected or ()), set(self.sch.task_list("", "").keys()))
def test_keep_tasks_for_assistant(self):
self.sch.get_work(worker="ASSISTANT", assistant=True) # tell the scheduler this is an assistant
self.test_prune_done_tasks(["B", "C"])
def test_keep_scheduler_disabled_tasks_for_assistant(self):
self.sch.get_work(worker="ASSISTANT", assistant=True) # tell the scheduler this is an assistant
# create a scheduler disabled task and a worker disabled task
for i in range(10):
self.sch.add_task(worker=WORKER, task_id="D", status=FAILED)
self.sch.add_task(worker=WORKER, task_id="E", status=DISABLED)
# scheduler prunes the worker disabled task
self.assertEqual(set(["D", "E"]), set(self.sch.task_list(DISABLED, "")))
self.test_prune_done_tasks(["B", "C", "D"])
def test_keep_failed_tasks_for_assistant(self):
self.sch.get_work(worker="ASSISTANT", assistant=True) # tell the scheduler this is an assistant
self.sch.add_task(worker=WORKER, task_id="D", status=FAILED, deps="A")
self.test_prune_done_tasks(["A", "B", "C", "D"])
def test_scheduler_resources_none_allow_one(self):
self.sch.add_task(worker="X", task_id="A", resources={"R1": 1})
self.assertEqual(self.sch.get_work(worker="X")["task_id"], "A")
def test_scheduler_resources_none_disallow_two(self):
self.sch.add_task(worker="X", task_id="A", resources={"R1": 2})
self.assertFalse(self.sch.get_work(worker="X")["task_id"], "A")
def test_scheduler_with_insufficient_resources(self):
self.sch.add_task(worker="X", task_id="A", resources={"R1": 3})
self.sch.update_resources(R1=2)
self.assertFalse(self.sch.get_work(worker="X")["task_id"])
def test_scheduler_with_sufficient_resources(self):
self.sch.add_task(worker="X", task_id="A", resources={"R1": 3})
self.sch.update_resources(R1=3)
self.assertEqual(self.sch.get_work(worker="X")["task_id"], "A")
def test_scheduler_with_resources_used(self):
self.sch.add_task(worker="X", task_id="A", resources={"R1": 1})
self.assertEqual(self.sch.get_work(worker="X")["task_id"], "A")
self.sch.add_task(worker="Y", task_id="B", resources={"R1": 1})
self.sch.update_resources(R1=1)
self.assertFalse(self.sch.get_work(worker="Y")["task_id"])
def test_scheduler_overprovisioned_on_other_resource(self):
self.sch.add_task(worker="X", task_id="A", resources={"R1": 2})
self.sch.update_resources(R1=2)
self.assertEqual(self.sch.get_work(worker="X")["task_id"], "A")
self.sch.add_task(worker="Y", task_id="B", resources={"R2": 2})
self.sch.update_resources(R1=1, R2=2)
self.assertEqual(self.sch.get_work(worker="Y")["task_id"], "B")
def test_scheduler_with_priority_and_competing_resources(self):
self.sch.add_task(worker="X", task_id="A")
self.assertEqual(self.sch.get_work(worker="X")["task_id"], "A")
self.sch.add_task(worker="X", task_id="B", resources={"R": 1}, priority=10)
self.sch.add_task(worker="Y", task_id="C", resources={"R": 1}, priority=1)
self.sch.update_resources(R=1)
self.assertFalse(self.sch.get_work(worker="Y")["task_id"])
self.sch.add_task(worker="Y", task_id="D", priority=0)
self.assertEqual(self.sch.get_work(worker="Y")["task_id"], "D")
def test_do_not_lock_resources_when_not_ready(self):
""" Test to make sure that resources won't go unused waiting on workers """
self.sch.add_task(worker="X", task_id="A", priority=10)
self.sch.add_task(worker="X", task_id="B", resources={"R": 1}, priority=5)
self.sch.add_task(worker="Y", task_id="C", resources={"R": 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker("X", [("workers", 1)])
self.assertEqual("C", self.sch.get_work(worker="Y")["task_id"])
def test_lock_resources_when_one_of_multiple_workers_is_ready(self):
self.sch.get_work(worker="X") # indicate to the scheduler that X is active
self.sch.add_task(worker="X", task_id="A", priority=10)
self.sch.add_task(worker="X", task_id="B", resources={"R": 1}, priority=5)
self.sch.add_task(worker="Y", task_id="C", resources={"R": 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker("X", [("workers", 2)])
self.sch.add_worker("Y", [])
self.assertFalse(self.sch.get_work(worker="Y")["task_id"])
def test_do_not_lock_resources_while_running_higher_priority(self):
""" Test to make sure that resources won't go unused waiting on workers """
self.sch.add_task(worker="X", task_id="A", priority=10)
self.sch.add_task(worker="X", task_id="B", resources={"R": 1}, priority=5)
self.sch.add_task(worker="Y", task_id="C", resources={"R": 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker("X", [("workers", 1)])
self.assertEqual("A", self.sch.get_work(worker="X")["task_id"])
self.assertEqual("C", self.sch.get_work(worker="Y")["task_id"])
def test_lock_resources_while_running_lower_priority(self):
""" Make sure resources will be made available while working on lower priority tasks """
self.sch.add_task(worker="X", task_id="A", priority=4)
self.assertEqual("A", self.sch.get_work(worker="X")["task_id"])
self.sch.add_task(worker="X", task_id="B", resources={"R": 1}, priority=5)
self.sch.add_task(worker="Y", task_id="C", resources={"R": 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker("X", [("workers", 1)])
self.assertFalse(self.sch.get_work(worker="Y")["task_id"])
def test_lock_resources_for_second_worker(self):
self.sch.get_work(worker="Y") # indicate to the scheduler that Y is active
self.sch.add_task(worker="X", task_id="A", resources={"R": 1})
self.sch.add_task(worker="X", task_id="B", resources={"R": 1})
self.sch.add_task(worker="Y", task_id="C", resources={"R": 1}, priority=10)
self.sch.add_worker("X", {"workers": 2})
self.sch.add_worker("Y", {"workers": 1})
self.sch.update_resources(R=2)
self.assertEqual("A", self.sch.get_work(worker="X")["task_id"])
self.assertFalse(self.sch.get_work(worker="X")["task_id"])
def test_can_work_on_lower_priority_while_waiting_for_resources(self):
self.sch.add_task(worker="X", task_id="A", resources={"R": 1}, priority=0)
self.assertEqual("A", self.sch.get_work(worker="X")["task_id"])
self.sch.add_task(worker="Y", task_id="B", resources={"R": 1}, priority=10)
self.sch.add_task(worker="Y", task_id="C", priority=0)
self.sch.update_resources(R=1)
self.assertEqual("C", self.sch.get_work(worker="Y")["task_id"])
def test_priority_update_with_pruning(self):
self.setTime(0)
self.sch.add_task(task_id="A", worker="X")
self.setTime(50) # after worker disconnects
self.sch.prune()
self.sch.add_task(task_id="B", deps=["A"], worker="X")
self.setTime(2000) # after remove for task A
self.sch.prune()
# Here task A that B depends on is missing
self.sch.add_task(worker=WORKER, task_id="C", deps=["B"], priority=100)
self.sch.add_task(worker=WORKER, task_id="B", deps=["A"])
self.sch.add_task(worker=WORKER, task_id="A")
self.sch.add_task(worker=WORKER, task_id="D", priority=10)
self.check_task_order("ABCD")
def test_update_resources(self):
self.sch.add_task(worker=WORKER, task_id="A", deps=["B"])
self.sch.add_task(worker=WORKER, task_id="B", resources={"r": 2})
self.sch.update_resources(r=1)
# B requires too many resources, we can't schedule
self.check_task_order([])
self.sch.add_task(worker=WORKER, task_id="B", resources={"r": 1})
# now we have enough resources
self.check_task_order(["B", "A"])
def test_hendle_multiple_resources(self):
self.sch.add_task(worker=WORKER, task_id="A", resources={"r1": 1, "r2": 1})
self.sch.add_task(worker=WORKER, task_id="B", resources={"r1": 1, "r2": 1})
self.sch.add_task(worker=WORKER, task_id="C", resources={"r1": 1})
self.sch.update_resources(r1=2, r2=1)
self.assertEqual("A", self.sch.get_work(worker=WORKER)["task_id"])
self.check_task_order("C")
def test_single_resource_lock(self):
self.sch.add_task(worker="X", task_id="A", resources={"r": 1})
self.assertEqual("A", self.sch.get_work(worker="X")["task_id"])
self.sch.add_task(worker=WORKER, task_id="B", resources={"r": 2}, priority=10)
self.sch.add_task(worker=WORKER, task_id="C", resources={"r": 1})
self.sch.update_resources(r=2)
# Should wait for 2 units of r to be available for B before scheduling C
self.check_task_order([])
def test_no_lock_if_too_many_resources_required(self):
self.sch.add_task(worker=WORKER, task_id="A", resources={"r": 2}, priority=10)
self.sch.add_task(worker=WORKER, task_id="B", resources={"r": 1})
self.sch.update_resources(r=1)
self.check_task_order("B")
def test_multiple_resources_lock(self):
self.sch.get_work(worker="X") # indicate to the scheduler that X is active
self.sch.add_task(worker="X", task_id="A", resources={"r1": 1, "r2": 1}, priority=10)
self.sch.add_task(worker=WORKER, task_id="B", resources={"r2": 1})
self.sch.add_task(worker=WORKER, task_id="C", resources={"r1": 1})
self.sch.update_resources(r1=1, r2=1)
# should preserve both resources for worker 'X'
self.check_task_order([])
def test_multiple_resources_no_lock(self):
self.sch.add_task(worker=WORKER, task_id="A", resources={"r1": 1}, priority=10)
self.sch.add_task(worker=WORKER, task_id="B", resources={"r1": 1, "r2": 1}, priority=10)
self.sch.add_task(worker=WORKER, task_id="C", resources={"r2": 1})
self.sch.update_resources(r1=1, r2=2)
self.assertEqual("A", self.sch.get_work(worker=WORKER)["task_id"])
# C doesn't block B, so it can go first
self.check_task_order("C")
def test_allow_resource_use_while_scheduling(self):
self.sch.update_resources(r1=1)
self.sch.add_task(worker="SCHEDULING", task_id="A", resources={"r1": 1}, priority=10)
self.sch.add_task(worker=WORKER, task_id="B", resources={"r1": 1}, priority=1)
self.assertEqual("B", self.sch.get_work(worker=WORKER)["task_id"])
def test_stop_locking_resource_for_uninterested_worker(self):
self.setTime(0)
self.sch.update_resources(r1=1)
self.assertIsNone(self.sch.get_work(worker=WORKER)["task_id"])
self.sch.add_task(worker=WORKER, task_id="A", resources={"r1": 1}, priority=10)
self.sch.add_task(worker="LOW_PRIO", task_id="B", resources={"r1": 1}, priority=1)
self.assertIsNone(self.sch.get_work(worker="LOW_PRIO")["task_id"])
self.setTime(120)
self.assertEqual("B", self.sch.get_work(worker="LOW_PRIO")["task_id"])
def check_task_order(self, order):
for expected_id in order:
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], expected_id)
self.sch.add_task(worker=WORKER, task_id=expected_id, status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], None)
def test_priorities(self):
self.sch.add_task(worker=WORKER, task_id="A", priority=10)
self.sch.add_task(worker=WORKER, task_id="B", priority=5)
self.sch.add_task(worker=WORKER, task_id="C", priority=15)
self.sch.add_task(worker=WORKER, task_id="D", priority=9)
self.check_task_order(["C", "A", "D", "B"])
def test_priorities_default_and_negative(self):
self.sch.add_task(worker=WORKER, task_id="A", priority=10)
self.sch.add_task(worker=WORKER, task_id="B")
self.sch.add_task(worker=WORKER, task_id="C", priority=15)
self.sch.add_task(worker=WORKER, task_id="D", priority=-20)
self.sch.add_task(worker=WORKER, task_id="E", priority=1)
self.check_task_order(["C", "A", "E", "B", "D"])
def test_priorities_and_dependencies(self):
self.sch.add_task(worker=WORKER, task_id="A", deps=["Z"], priority=10)
self.sch.add_task(worker=WORKER, task_id="B", priority=5)
self.sch.add_task(worker=WORKER, task_id="C", deps=["Z"], priority=3)
self.sch.add_task(worker=WORKER, task_id="D", priority=2)
self.sch.add_task(worker=WORKER, task_id="Z", priority=1)
self.check_task_order(["Z", "A", "B", "C", "D"])
def test_priority_update_dependency_after_scheduling(self):
self.sch.add_task(worker=WORKER, task_id="A", priority=1)
self.sch.add_task(worker=WORKER, task_id="B", priority=5, deps=["A"])
self.sch.add_task(worker=WORKER, task_id="C", priority=10, deps=["B"])
self.sch.add_task(worker=WORKER, task_id="D", priority=6)
self.check_task_order(["A", "B", "C", "D"])
def test_disable(self):
self.sch.add_task(worker=WORKER, task_id="A")
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list("DISABLED", "")), 1)
self.assertEqual(len(self.sch.task_list("FAILED", "")), 0)
self.sch.add_task(worker=WORKER, task_id="A")
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], None)
def test_disable_and_reenable(self):
self.sch.add_task(worker=WORKER, task_id="A")
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list("DISABLED", "")), 1)
self.assertEqual(len(self.sch.task_list("FAILED", "")), 0)
self.sch.re_enable_task("A")
# should be enabled at this point
self.assertEqual(len(self.sch.task_list("DISABLED", "")), 0)
self.assertEqual(len(self.sch.task_list("FAILED", "")), 1)
self.sch.add_task(worker=WORKER, task_id="A")
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], "A")
def test_disable_and_reenable_and_disable_again(self):
self.sch.add_task(worker=WORKER, task_id="A")
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list("DISABLED", "")), 1)
self.assertEqual(len(self.sch.task_list("FAILED", "")), 0)
self.sch.re_enable_task("A")
# should be enabled at this point
self.assertEqual(len(self.sch.task_list("DISABLED", "")), 0)
self.assertEqual(len(self.sch.task_list("FAILED", "")), 1)
self.sch.add_task(worker=WORKER, task_id="A")
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], "A")
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
# should be still enabled
self.assertEqual(len(self.sch.task_list("DISABLED", "")), 0)
self.assertEqual(len(self.sch.task_list("FAILED", "")), 1)
self.sch.add_task(worker=WORKER, task_id="A")
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], "A")
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
# should be disabled now
self.assertEqual(len(self.sch.task_list("DISABLED", "")), 1)
self.assertEqual(len(self.sch.task_list("FAILED", "")), 0)
self.sch.add_task(worker=WORKER, task_id="A")
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], None)
def test_disable_and_done(self):
self.sch.add_task(worker=WORKER, task_id="A")
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list("DISABLED", "")), 1)
self.assertEqual(len(self.sch.task_list("FAILED", "")), 0)
self.sch.add_task(worker=WORKER, task_id="A", status=DONE)
# should be enabled at this point
self.assertEqual(len(self.sch.task_list("DISABLED", "")), 0)
self.assertEqual(len(self.sch.task_list("DONE", "")), 1)
self.sch.add_task(worker=WORKER, task_id="A")
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], "A")
def test_automatic_re_enable(self):
self.sch = CentralPlannerScheduler(disable_failures=2, disable_persist=100)
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
# should be disabled now
self.assertEqual(DISABLED, self.sch.task_list("", "")["A"]["status"])
# re-enables after 100 seconds
self.setTime(101)
self.assertEqual(FAILED, self.sch.task_list("", "")["A"]["status"])
def test_automatic_re_enable_with_one_failure_allowed(self):
self.sch = CentralPlannerScheduler(disable_failures=1, disable_persist=100)
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
# should be disabled now
self.assertEqual(DISABLED, self.sch.task_list("", "")["A"]["status"])
# re-enables after 100 seconds
self.setTime(101)
self.assertEqual(FAILED, self.sch.task_list("", "")["A"]["status"])
def test_no_automatic_re_enable_after_manual_disable(self):
self.sch = CentralPlannerScheduler(disable_persist=100)
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id="A", status=DISABLED)
# should be disabled now
self.assertEqual(DISABLED, self.sch.task_list("", "")["A"]["status"])
# should not re-enable after 100 seconds
self.setTime(101)
self.assertEqual(DISABLED, self.sch.task_list("", "")["A"]["status"])
def test_no_automatic_re_enable_after_auto_then_manual_disable(self):
self.sch = CentralPlannerScheduler(disable_failures=2, disable_persist=100)
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
# should be disabled now
self.assertEqual(DISABLED, self.sch.task_list("", "")["A"]["status"])
# should remain disabled once set
self.sch.add_task(worker=WORKER, task_id="A", status=DISABLED)
self.assertEqual(DISABLED, self.sch.task_list("", "")["A"]["status"])
# should not re-enable after 100 seconds
self.setTime(101)
self.assertEqual(DISABLED, self.sch.task_list("", "")["A"]["status"])
def test_disable_by_worker(self):
self.sch.add_task(worker=WORKER, task_id="A", status=DISABLED)
self.assertEqual(len(self.sch.task_list("DISABLED", "")), 1)
self.sch.add_task(worker=WORKER, task_id="A")
# should be enabled at this point
self.assertEqual(len(self.sch.task_list("DISABLED", "")), 0)
self.sch.add_task(worker=WORKER, task_id="A")
self.assertEqual(self.sch.get_work(worker=WORKER)["task_id"], "A")
def test_task_list_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
for c in "ABCD":
sch.add_task(worker=WORKER, task_id=c)
self.assertEqual(set("ABCD"), set(sch.task_list("PENDING", "", False).keys()))
self.assertEqual({"num_tasks": 4}, sch.task_list("PENDING", ""))
def test_task_list_within_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=4)
for c in "ABCD":
sch.add_task(worker=WORKER, task_id=c)
self.assertEqual(set("ABCD"), set(sch.task_list("PENDING", "").keys()))
def test_task_lists_some_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
for c in "ABCD":
sch.add_task(worker=WORKER, task_id=c, status=DONE)
for c in "EFG":
sch.add_task(worker=WORKER, task_id=c)
self.assertEqual(set("EFG"), set(sch.task_list("PENDING", "").keys()))
self.assertEqual({"num_tasks": 4}, sch.task_list("DONE", ""))
def test_task_list_filter_by_search(self):
self.sch.add_task(worker=WORKER, task_id="test_match_task")
self.sch.add_task(worker=WORKER, task_id="test_filter_task")
matches = self.sch.task_list("PENDING", "", search="match")
self.assertEqual(["test_match_task"], list(matches.keys()))
def test_task_list_filter_by_multiple_search_terms(self):
self.sch.add_task(worker=WORKER, task_id="abcd")
self.sch.add_task(worker=WORKER, task_id="abd")
self.sch.add_task(worker=WORKER, task_id="acd")
self.sch.add_task(worker=WORKER, task_id="ad")
self.sch.add_task(worker=WORKER, task_id="bc")
matches = self.sch.task_list("PENDING", "", search="b c")
self.assertEqual(set(["abcd", "bc"]), set(matches.keys()))
def test_search_results_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
sch.add_task(worker=WORKER, task_id="task_a")
sch.add_task(worker=WORKER, task_id="task_b")
sch.add_task(worker=WORKER, task_id="task_c")
sch.add_task(worker=WORKER, task_id="task_d")
self.assertEqual({"num_tasks": 4}, sch.task_list("PENDING", "", search="a"))
self.assertEqual(["task_a"], list(sch.task_list("PENDING", "", search="_a").keys()))
def test_priority_update_dependency_chain(self):
self.sch.add_task(worker=WORKER, task_id="A", priority=10, deps=["B"])
self.sch.add_task(worker=WORKER, task_id="B", priority=5, deps=["C"])
self.sch.add_task(worker=WORKER, task_id="C", priority=1)
self.sch.add_task(worker=WORKER, task_id="D", priority=6)
self.check_task_order(["C", "B", "A", "D"])
def test_priority_no_decrease_with_multiple_updates(self):
self.sch.add_task(worker=WORKER, task_id="A", priority=1)
self.sch.add_task(worker=WORKER, task_id="B", priority=10, deps=["A"])
self.sch.add_task(worker=WORKER, task_id="C", priority=5, deps=["A"])
self.sch.add_task(worker=WORKER, task_id="D", priority=6)
self.check_task_order(["A", "B", "D", "C"])
def test_unique_tasks(self):
self.sch.add_task(worker=WORKER, task_id="A")
self.sch.add_task(worker=WORKER, task_id="B")
self.sch.add_task(worker=WORKER, task_id="C")
self.sch.add_task(worker=WORKER + "_2", task_id="B")
response = self.sch.get_work(worker=WORKER)
self.assertEqual(3, response["n_pending_tasks"])
self.assertEqual(2, response["n_unique_pending"])
def test_pending_downstream_disable(self):
self.sch.add_task(worker=WORKER, task_id="A", status=DISABLED)
self.sch.add_task(worker=WORKER, task_id="B", deps=("A",))
self.sch.add_task(worker=WORKER, task_id="C", deps=("B",))
response = self.sch.get_work(worker=WORKER)
self.assertTrue(response["task_id"] is None)
self.assertEqual(0, response["n_pending_tasks"])
self.assertEqual(0, response["n_unique_pending"])
def test_pending_downstream_failure(self):
self.sch.add_task(worker=WORKER, task_id="A", status=FAILED)
self.sch.add_task(worker=WORKER, task_id="B", deps=("A",))
self.sch.add_task(worker=WORKER, task_id="C", deps=("B",))
response = self.sch.get_work(worker=WORKER)
self.assertTrue(response["task_id"] is None)
self.assertEqual(2, response["n_pending_tasks"])
self.assertEqual(2, response["n_unique_pending"])
def test_task_list_no_deps(self):
self.sch.add_task(worker=WORKER, task_id="B", deps=("A",))
self.sch.add_task(worker=WORKER, task_id="A")
task_list = self.sch.task_list("PENDING", "")
self.assertFalse("deps" in task_list["A"])
def test_task_first_failure_time(self):
self.sch.add_task(worker=WORKER, task_id="A")
test_task = self.sch._state.get_task("A")
self.assertIsNone(test_task.failures.first_failure_time)
time_before_failure = time.time()
test_task.add_failure()
time_after_failure = time.time()
self.assertLessEqual(time_before_failure, test_task.failures.first_failure_time)
self.assertGreaterEqual(time_after_failure, test_task.failures.first_failure_time)
def test_task_first_failure_time_remains_constant(self):
self.sch.add_task(worker=WORKER, task_id="A")
test_task = self.sch._state.get_task("A")
self.assertIsNone(test_task.failures.first_failure_time)
test_task.add_failure()
first_failure_time = test_task.failures.first_failure_time
test_task.add_failure()
self.assertEqual(first_failure_time, test_task.failures.first_failure_time)
def test_task_has_excessive_failures(self):
self.sch.add_task(worker=WORKER, task_id="A")
test_task = self.sch._state.get_task("A")
self.assertIsNone(test_task.failures.first_failure_time)
self.assertFalse(test_task.has_excessive_failures())
test_task.add_failure()
self.assertFalse(test_task.has_excessive_failures())
fake_failure_time = test_task.failures.first_failure_time - 2 * 60 * 60
test_task.failures.first_failure_time = fake_failure_time
self.assertTrue(test_task.has_excessive_failures())
def test_quadratic_behavior(self):
""" Test that get_work is not taking linear amount of time.
This is of course impossible to test, however, doing reasonable
assumptions about hardware. This time should finish in a timely
manner.
"""
# For 10000 it takes almost 1 second on my laptop. Prior to these
# changes it was being slow already at NUM_TASKS=300
NUM_TASKS = 10000
for i in range(NUM_TASKS):
self.sch.add_task(worker=str(i), task_id=str(i), resources={})
for i in range(NUM_TASKS):
self.assertEqual(self.sch.get_work(worker=str(i))["task_id"], str(i))
self.sch.add_task(worker=str(i), task_id=str(i), status=DONE)
def test_get_work_speed(self):
""" Test that get_work is fast for few workers and many DONEs.
In #986, @daveFNbuck reported that he got a slowdown.
"""
# This took almost 4 minutes without optimization.
# Now it takes 10 seconds on my machine.
NUM_PENDING = 1000
NUM_DONE = 200000
assert NUM_DONE >= NUM_PENDING
for i in range(NUM_PENDING):
self.sch.add_task(worker=WORKER, task_id=str(i), resources={})
for i in range(NUM_PENDING, NUM_DONE):
self.sch.add_task(worker=WORKER, task_id=str(i), status=DONE)
for i in range(NUM_PENDING):
res = int(self.sch.get_work(worker=WORKER)["task_id"])
self.assertTrue(0 <= res < NUM_PENDING)
self.sch.add_task(worker=WORKER, task_id=str(res), status=DONE)
class CentralPlannerTest(unittest.TestCase):
def setUp(self):
self.sch = CentralPlannerScheduler(retry_delay=100, remove_delay=1000, worker_disconnect_delay=10)
self.time = time.time
def tearDown(self):
if time.time != self.time:
time.time = self.time
def setTime(self, t):
time.time = lambda: t
def test_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_failed_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', status=FAILED)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(WORKER, 'A', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_broken_dep(self):
self.sch.add_task(WORKER, 'B', deps=('A',))
self.sch.add_task(WORKER, 'A', runnable=False)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(WORKER, 'A', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'B')
self.sch.add_task(WORKER, 'B', DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_two_workers(self):
# Worker X wants to build A -> B
# Worker Y wants to build A -> C
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.sch.add_task(task_id='B', deps=('A',), worker='X')
self.sch.add_task(task_id='C', deps=('A',), worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None) # Worker Y is pending on A to be done
self.sch.add_task(worker='X', task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'C')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'B')
def test_retry(self):
# Try to build A but fails, will retry after 100s
self.setTime(0)
self.sch.add_task(WORKER, 'A')
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
self.sch.add_task(WORKER, 'A', FAILED)
for t in xrange(100):
self.setTime(t)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
self.sch.ping(WORKER)
if t % 10 == 0:
self.sch.prune()
self.setTime(101)
self.sch.prune()
self.assertEqual(self.sch.get_work(WORKER)['task_id'], 'A')
def test_disconnect_running(self):
# X and Y wants to run A.
# X starts but does not report back. Y does.
# After some timeout, Y will build it instead
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.sch.add_task(task_id='A', worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
for t in xrange(200):
self.setTime(t)
self.sch.ping(worker='Y')
if t % 10 == 0:
self.sch.prune()
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'A')
def test_remove_dep(self):
# X schedules A -> B, A is broken
# Y schedules C -> B: this should remove A as a dep of B
self.sch.add_task(task_id='A', worker='X', runnable=False)
self.sch.add_task(task_id='B', deps=('A',), worker='X')
# X can't build anything
self.assertEqual(self.sch.get_work(worker='X')['task_id'], None)
self.sch.add_task(task_id='B', deps=('C',), worker='Y') # should reset dependencies for A
self.sch.add_task(task_id='C', worker='Y', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_timeout(self):
# A bug that was earlier present when restarting the same flow
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.setTime(10000)
self.sch.add_task(task_id='A', worker='Y') # Will timeout X but not schedule A for removal
for i in xrange(2000):
self.setTime(10000 + i)
self.sch.ping(worker='Y')
self.sch.add_task(task_id='A', status=DONE, worker='Y') # This used to raise an exception since A was removed
def test_disallowed_state_changes(self):
# Test that we can not schedule an already running task
t = 'A'
self.sch.add_task(task_id=t, worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], t)
self.sch.add_task(task_id=t, worker='Y')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None)
def test_two_worker_info(self):
# Make sure the scheduler returns info that some other worker is running task A
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
r = self.sch.get_work(worker='Y')
self.assertEqual(r['task_id'], None) # Worker Y is pending on A to be done
s = r['running_tasks'][0]
self.assertEqual(s['task_id'], 'A')
self.assertEqual(s['worker'], 'X')
def test_scheduler_resources_none_allow_one(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_resources_none_disallow_two(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.assertFalse(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_insufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=2)
self.assertFalse(self.sch.get_work(worker='X')['task_id'])
def test_scheduler_with_sufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=3)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_resources_used(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R1': 1})
self.sch.update_resources(R1=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_scheduler_overprovisioned_on_other_resource(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.sch.update_resources(R1=2)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R2': 2})
self.sch.update_resources(R1=1, R2=2)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_scheduler_with_priority_and_competing_resources(self):
self.sch.add_task(worker='X', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=10)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
self.sch.add_task(worker='Y', task_id='D', priority=0)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'D')
def test_update_resources(self):
self.sch.add_task(WORKER, task_id='A', deps=['B'])
self.sch.add_task(WORKER, task_id='B', resources={'r': 2})
self.sch.update_resources(r=1)
# B requires too many resources, we can't schedule
self.check_task_order([])
self.sch.add_task(WORKER, task_id='B', resources={'r': 1})
# now we have enough resources
self.check_task_order(['B', 'A'])
def check_task_order(self, order):
for expected_id in order:
self.assertEqual(self.sch.get_work(WORKER)['task_id'], expected_id)
self.sch.add_task(WORKER, expected_id, status=DONE)
self.assertEqual(self.sch.get_work(WORKER)['task_id'], None)
def test_priorities(self):
self.sch.add_task(WORKER, 'A', priority=10)
self.sch.add_task(WORKER, 'B', priority=5)
self.sch.add_task(WORKER, 'C', priority=15)
self.sch.add_task(WORKER, 'D', priority=9)
self.check_task_order(['C', 'A', 'D', 'B'])
def test_priorities_default_and_negative(self):
self.sch.add_task(WORKER, 'A', priority=10)
self.sch.add_task(WORKER, 'B')
self.sch.add_task(WORKER, 'C', priority=15)
self.sch.add_task(WORKER, 'D', priority=-20)
self.sch.add_task(WORKER, 'E', priority=1)
self.check_task_order(['C', 'A', 'E', 'B', 'D'])
def test_priorities_and_dependencies(self):
self.sch.add_task(WORKER, 'A', deps=['Z'], priority=10)
self.sch.add_task(WORKER, 'B', priority=5)
self.sch.add_task(WORKER, 'C', deps=['Z'], priority=3)
self.sch.add_task(WORKER, 'D', priority=2)
self.sch.add_task(WORKER, 'Z', priority=1)
self.check_task_order(['Z', 'A', 'B', 'C', 'D'])
def test_priority_update_dependency_after_scheduling(self):
self.sch.add_task(WORKER, 'A', priority=1)
self.sch.add_task(WORKER, 'B', priority=5, deps=['A'])
self.sch.add_task(WORKER, 'C', priority=10, deps=['B'])
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['A', 'B', 'C', 'D'])
def test_priority_update_dependency_chain(self):
self.sch.add_task(WORKER, 'A', priority=10, deps=['B'])
self.sch.add_task(WORKER, 'B', priority=5, deps=['C'])
self.sch.add_task(WORKER, 'C', priority=1)
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['C', 'B', 'A', 'D'])
def test_priority_no_decrease_with_multiple_updates(self):
self.sch.add_task(WORKER, 'A', priority=1)
self.sch.add_task(WORKER, 'B', priority=10, deps=['A'])
self.sch.add_task(WORKER, 'C', priority=5, deps=['A'])
self.sch.add_task(WORKER, 'D', priority=6)
self.check_task_order(['A', 'B', 'D', 'C'])
class CentralPlannerTest(unittest.TestCase):
def setUp(self):
super(CentralPlannerTest, self).setUp()
conf = self.get_scheduler_config()
self.sch = CentralPlannerScheduler(**conf)
self.time = time.time
def get_scheduler_config(self):
return {
'retry_delay': 100,
'remove_delay': 1000,
'worker_disconnect_delay': 10,
'disable_persist': 10,
'disable_window': 10,
'disable_failures': 3,
'disable_hard_timeout': 60 * 60,
}
def tearDown(self):
super(CentralPlannerTest, self).tearDown()
if time.time != self.time:
time.time = self.time
def setTime(self, t):
time.time = lambda: t
def test_dep(self):
self.sch.add_task(worker=WORKER, task_id='B', deps=('A',))
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'B')
self.sch.add_task(worker=WORKER, task_id='B', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_failed_dep(self):
self.sch.add_task(worker=WORKER, task_id='B', deps=('A',))
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'B')
self.sch.add_task(worker=WORKER, task_id='B', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_broken_dep(self):
self.sch.add_task(worker=WORKER, task_id='B', deps=('A',))
self.sch.add_task(worker=WORKER, task_id='A', runnable=False)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None) # can still wait and retry: TODO: do we want this?
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'B')
self.sch.add_task(worker=WORKER, task_id='B', status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_two_workers(self):
# Worker X wants to build A -> B
# Worker Y wants to build A -> C
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.sch.add_task(task_id='B', deps=('A',), worker='X')
self.sch.add_task(task_id='C', deps=('A',), worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None) # Worker Y is pending on A to be done
self.sch.add_task(worker='X', task_id='A', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'C')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'B')
def test_retry(self):
# Try to build A but fails, will retry after 100s
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
for t in range(100):
self.setTime(t)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
self.sch.ping(worker=WORKER)
if t % 10 == 0:
self.sch.prune()
self.setTime(101)
self.sch.prune()
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
def test_resend_task(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='B')
for _ in range(10):
self.assertEqual('A', self.sch.get_work(worker=WORKER, current_tasks=[])['task_id'])
self.assertEqual('B', self.sch.get_work(worker=WORKER, current_tasks=['A'])['task_id'])
def test_resend_multiple_tasks(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='B')
self.sch.add_task(worker=WORKER, task_id='C')
# get A and B running
self.assertEqual('A', self.sch.get_work(worker=WORKER)['task_id'])
self.assertEqual('B', self.sch.get_work(worker=WORKER)['task_id'])
for _ in range(10):
self.assertEqual('A', self.sch.get_work(worker=WORKER, current_tasks=[])['task_id'])
self.assertEqual('A', self.sch.get_work(worker=WORKER, current_tasks=['B'])['task_id'])
self.assertEqual('B', self.sch.get_work(worker=WORKER, current_tasks=['A'])['task_id'])
self.assertEqual('C', self.sch.get_work(worker=WORKER, current_tasks=['A', 'B'])['task_id'])
def test_disconnect_running(self):
# X and Y wants to run A.
# X starts but does not report back. Y does.
# After some timeout, Y will build it instead
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.sch.add_task(task_id='A', worker='Y')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
for t in range(200):
self.setTime(t)
self.sch.ping(worker='Y')
if t % 10 == 0:
self.sch.prune()
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'A')
def test_do_not_overwrite_tracking_url_while_running(self):
self.sch.add_task(task_id='A', worker='X', status='RUNNING', tracking_url='trackme')
self.assertEqual('trackme', self.sch.task_list('RUNNING', '')['A']['tracking_url'])
# not wiped out by another working scheduling as pending
self.sch.add_task(task_id='A', worker='Y', status='PENDING')
self.assertEqual('trackme', self.sch.task_list('RUNNING', '')['A']['tracking_url'])
def test_do_update_tracking_url_while_running(self):
self.sch.add_task(task_id='A', worker='X', status='RUNNING', tracking_url='trackme')
self.assertEqual('trackme', self.sch.task_list('RUNNING', '')['A']['tracking_url'])
self.sch.add_task(task_id='A', worker='X', status='RUNNING', tracking_url='stage_2')
self.assertEqual('stage_2', self.sch.task_list('RUNNING', '')['A']['tracking_url'])
def test_keep_tracking_url_on_done_and_fail(self):
for status in ('DONE', 'FAILED'):
self.sch.add_task(task_id='A', worker='X', status='RUNNING', tracking_url='trackme')
self.assertEqual('trackme', self.sch.task_list('RUNNING', '')['A']['tracking_url'])
self.sch.add_task(task_id='A', worker='X', status=status)
self.assertEqual('trackme', self.sch.task_list(status, '')['A']['tracking_url'])
def test_drop_tracking_url_when_rescheduled_while_not_running(self):
for status in ('DONE', 'FAILED', 'PENDING'):
self.sch.add_task(task_id='A', worker='X', status=status, tracking_url='trackme')
self.assertEqual('trackme', self.sch.task_list(status, '')['A']['tracking_url'])
self.sch.add_task(task_id='A', worker='Y', status='PENDING')
self.assertIsNone(self.sch.task_list('PENDING', '')['A']['tracking_url'])
def test_reset_tracking_url_on_new_run(self):
self.sch.add_task(task_id='A', worker='X', status='PENDING', tracking_url='trackme')
self.assertEqual('trackme', self.sch.task_list('PENDING', '')['A']['tracking_url'])
self.sch.add_task(task_id='A', worker='Y', status='RUNNING')
self.assertIsNone(self.sch.task_list('RUNNING', '')['A']['tracking_url'])
def test_remove_dep(self):
# X schedules A -> B, A is broken
# Y schedules C -> B: this should remove A as a dep of B
self.sch.add_task(task_id='A', worker='X', runnable=False)
self.sch.add_task(task_id='B', deps=('A',), worker='X')
# X can't build anything
self.assertEqual(self.sch.get_work(worker='X')['task_id'], None)
self.sch.add_task(task_id='B', deps=('C',), worker='Y') # should reset dependencies for A
self.sch.add_task(task_id='C', worker='Y', status=DONE)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_start_time(self):
self.setTime(100)
self.sch.add_task(worker=WORKER, task_id='A')
self.setTime(200)
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
self.assertEqual(100, self.sch.task_list(DONE, '')['A']['start_time'])
def test_last_updated_does_not_change_with_same_status_update(self):
for t, status in ((100, PENDING), (300, DONE), (500, DISABLED)):
self.setTime(t)
self.sch.add_task(worker=WORKER, task_id='A', status=status)
self.assertEqual(t, self.sch.task_list(status, '')['A']['last_updated'])
self.setTime(t + 100)
self.sch.add_task(worker=WORKER, task_id='A', status=status)
self.assertEqual(t, self.sch.task_list(status, '')['A']['last_updated'])
def test_last_updated_shows_running_start(self):
self.setTime(100)
self.sch.add_task(worker=WORKER, task_id='A', status=PENDING)
self.assertEqual(100, self.sch.task_list(PENDING, '')['A']['last_updated'])
self.setTime(200)
self.assertEqual('A', self.sch.get_work(worker=WORKER)['task_id'])
self.assertEqual(200, self.sch.task_list('RUNNING', '')['A']['last_updated'])
self.setTime(300)
self.sch.add_task(worker=WORKER, task_id='A', status=PENDING)
self.assertEqual(200, self.sch.task_list('RUNNING', '')['A']['last_updated'])
def test_last_updated_with_failure_and_recovery(self):
self.setTime(100)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual('A', self.sch.get_work(worker=WORKER)['task_id'])
self.setTime(200)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.assertEqual(200, self.sch.task_list(FAILED, '')['A']['last_updated'])
self.setTime(1000)
self.sch.prune()
self.assertEqual(1000, self.sch.task_list(PENDING, '')['A']['last_updated'])
def test_timeout(self):
# A bug that was earlier present when restarting the same flow
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.setTime(10000)
self.sch.add_task(task_id='A', worker='Y') # Will timeout X but not schedule A for removal
for i in range(2000):
self.setTime(10000 + i)
self.sch.ping(worker='Y')
self.sch.add_task(task_id='A', status=DONE, worker='Y') # This used to raise an exception since A was removed
def test_disallowed_state_changes(self):
# Test that we can not schedule an already running task
t = 'A'
self.sch.add_task(task_id=t, worker='X')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], t)
self.sch.add_task(task_id=t, worker='Y')
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], None)
def test_two_worker_info(self):
# Make sure the scheduler returns info that some other worker is running task A
self.sch.add_task(worker='X', task_id='A')
self.sch.add_task(worker='Y', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
r = self.sch.get_work(worker='Y')
self.assertEqual(r['task_id'], None) # Worker Y is pending on A to be done
s = r['running_tasks'][0]
self.assertEqual(s['task_id'], 'A')
self.assertEqual(s['worker'], 'X')
def test_assistant_get_work(self):
self.sch.add_task(worker='X', task_id='A')
self.sch.add_worker('Y', [])
self.assertEqual(self.sch.get_work(worker='Y', assistant=True)['task_id'], 'A')
# check that the scheduler recognizes tasks as running
running_tasks = self.sch.task_list('RUNNING', '')
self.assertEqual(len(running_tasks), 1)
self.assertEqual(list(running_tasks.keys()), ['A'])
self.assertEqual(running_tasks['A']['worker_running'], 'Y')
def test_assistant_get_work_external_task(self):
self.sch.add_task(worker='X', task_id='A', runnable=False)
self.assertTrue(self.sch.get_work(worker='Y', assistant=True)['task_id'] is None)
def test_task_fails_when_assistant_dies(self):
self.setTime(0)
self.sch.add_task(worker='X', task_id='A')
self.sch.add_worker('Y', [])
self.assertEqual(self.sch.get_work(worker='Y', assistant=True)['task_id'], 'A')
self.assertEqual(list(self.sch.task_list('RUNNING', '').keys()), ['A'])
# Y dies for 50 seconds, X stays alive
self.setTime(50)
self.sch.ping(worker='X')
self.assertEqual(list(self.sch.task_list('FAILED', '').keys()), ['A'])
def test_prune_with_live_assistant(self):
self.setTime(0)
self.sch.add_task(worker='X', task_id='A')
self.sch.get_work(worker='Y', assistant=True)
self.sch.add_task(worker='Y', task_id='A', status=DONE, assistant=True)
# worker X stops communicating, A should be marked for removal
self.setTime(600)
self.sch.ping(worker='Y')
self.sch.prune()
# A will now be pruned
self.setTime(2000)
self.sch.prune()
self.assertFalse(list(self.sch.task_list('', '')))
def test_re_enable_failed_task_assistant(self):
self.setTime(0)
self.sch.add_worker('X', [('assistant', True)])
self.sch.add_task(worker='X', task_id='A', status=FAILED, assistant=True)
# should be failed now
self.assertEqual(FAILED, self.sch.task_list('', '')['A']['status'])
# resets to PENDING after 100 seconds
self.setTime(101)
self.sch.ping(worker='X') # worker still alive
self.assertEqual('PENDING', self.sch.task_list('', '')['A']['status'])
def test_fail_job_from_dead_worker_with_live_assistant(self):
self.setTime(0)
self.sch.add_task(worker='X', task_id='A')
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.sch.add_worker('Y', [('assistant', True)])
self.setTime(600)
self.sch.ping(worker='Y')
self.sch.prune()
self.assertEqual(['A'], list(self.sch.task_list('FAILED', '').keys()))
def test_assistant_request_runnable_task(self):
self.setTime(0)
self.sch.add_task(worker='X', task_id='A', runnable=True)
self.setTime(600)
self.sch.prune()
self.assertEqual('A', self.sch.get_work(worker='Y', assistant=True)['task_id'])
def test_assistant_request_external_task(self):
self.sch.add_task(worker='X', task_id='A', runnable=False)
self.assertIsNone(self.sch.get_work(worker='Y', assistant=True)['task_id'])
def test_prune_done_tasks(self, expected=None):
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
self.sch.add_task(worker=WORKER, task_id='B', deps=['A'], status=DONE)
self.sch.add_task(worker=WORKER, task_id='C', deps=['B'])
self.setTime(600)
self.sch.ping(worker='ASSISTANT')
self.sch.prune()
self.setTime(2000)
self.sch.ping(worker='ASSISTANT')
self.sch.prune()
self.assertEqual(set(expected or ()), set(self.sch.task_list('', '').keys()))
def test_keep_tasks_for_assistant(self):
self.sch.get_work(worker='ASSISTANT', assistant=True) # tell the scheduler this is an assistant
self.test_prune_done_tasks(['B', 'C'])
def test_keep_scheduler_disabled_tasks_for_assistant(self):
self.sch.get_work(worker='ASSISTANT', assistant=True) # tell the scheduler this is an assistant
# create a scheduler disabled task and a worker disabled task
for i in range(10):
self.sch.add_task(worker=WORKER, task_id='D', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='E', status=DISABLED)
# scheduler prunes the worker disabled task
self.assertEqual(set(['D', 'E']), set(self.sch.task_list(DISABLED, '')))
self.test_prune_done_tasks(['B', 'C', 'D'])
def test_keep_failed_tasks_for_assistant(self):
self.sch.get_work(worker='ASSISTANT', assistant=True) # tell the scheduler this is an assistant
self.sch.add_task(worker=WORKER, task_id='D', status=FAILED, deps='A')
self.test_prune_done_tasks(['A', 'B', 'C', 'D'])
def test_scheduler_resources_none_allow_one(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_resources_none_disallow_two(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.assertFalse(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_insufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=2)
self.assertFalse(self.sch.get_work(worker='X')['task_id'])
def test_scheduler_with_sufficient_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 3})
self.sch.update_resources(R1=3)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
def test_scheduler_with_resources_used(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 1})
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R1': 1})
self.sch.update_resources(R1=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_scheduler_overprovisioned_on_other_resource(self):
self.sch.add_task(worker='X', task_id='A', resources={'R1': 2})
self.sch.update_resources(R1=2)
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='Y', task_id='B', resources={'R2': 2})
self.sch.update_resources(R1=1, R2=2)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'B')
def test_scheduler_with_priority_and_competing_resources(self):
self.sch.add_task(worker='X', task_id='A')
self.assertEqual(self.sch.get_work(worker='X')['task_id'], 'A')
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=10)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
self.sch.add_task(worker='Y', task_id='D', priority=0)
self.assertEqual(self.sch.get_work(worker='Y')['task_id'], 'D')
def test_do_not_lock_resources_when_not_ready(self):
""" Test to make sure that resources won't go unused waiting on workers """
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertEqual('C', self.sch.get_work(worker='Y')['task_id'])
def test_lock_resources_when_one_of_multiple_workers_is_ready(self):
self.sch.get_work(worker='X') # indicate to the scheduler that X is active
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 2)])
self.sch.add_worker('Y', [])
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_do_not_lock_resources_while_running_higher_priority(self):
""" Test to make sure that resources won't go unused waiting on workers """
self.sch.add_task(worker='X', task_id='A', priority=10)
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.assertEqual('C', self.sch.get_work(worker='Y')['task_id'])
def test_lock_resources_while_running_lower_priority(self):
""" Make sure resources will be made available while working on lower priority tasks """
self.sch.add_task(worker='X', task_id='A', priority=4)
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.sch.add_task(worker='X', task_id='B', resources={'R': 1}, priority=5)
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=1)
self.sch.update_resources(R=1)
self.sch.add_worker('X', [('workers', 1)])
self.assertFalse(self.sch.get_work(worker='Y')['task_id'])
def test_lock_resources_for_second_worker(self):
self.sch.get_work(worker='Y') # indicate to the scheduler that Y is active
self.sch.add_task(worker='X', task_id='A', resources={'R': 1})
self.sch.add_task(worker='X', task_id='B', resources={'R': 1})
self.sch.add_task(worker='Y', task_id='C', resources={'R': 1}, priority=10)
self.sch.add_worker('X', {'workers': 2})
self.sch.add_worker('Y', {'workers': 1})
self.sch.update_resources(R=2)
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.assertFalse(self.sch.get_work(worker='X')['task_id'])
def test_can_work_on_lower_priority_while_waiting_for_resources(self):
self.sch.add_task(worker='X', task_id='A', resources={'R': 1}, priority=0)
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.sch.add_task(worker='Y', task_id='B', resources={'R': 1}, priority=10)
self.sch.add_task(worker='Y', task_id='C', priority=0)
self.sch.update_resources(R=1)
self.assertEqual('C', self.sch.get_work(worker='Y')['task_id'])
def test_priority_update_with_pruning(self):
self.setTime(0)
self.sch.add_task(task_id='A', worker='X')
self.setTime(50) # after worker disconnects
self.sch.prune()
self.sch.add_task(task_id='B', deps=['A'], worker='X')
self.setTime(2000) # after remove for task A
self.sch.prune()
# Here task A that B depends on is missing
self.sch.add_task(worker=WORKER, task_id='C', deps=['B'], priority=100)
self.sch.add_task(worker=WORKER, task_id='B', deps=['A'])
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='D', priority=10)
self.check_task_order('ABCD')
def test_update_resources(self):
self.sch.add_task(worker=WORKER, task_id='A', deps=['B'])
self.sch.add_task(worker=WORKER, task_id='B', resources={'r': 2})
self.sch.update_resources(r=1)
# B requires too many resources, we can't schedule
self.check_task_order([])
self.sch.add_task(worker=WORKER, task_id='B', resources={'r': 1})
# now we have enough resources
self.check_task_order(['B', 'A'])
def test_hendle_multiple_resources(self):
self.sch.add_task(worker=WORKER, task_id='A', resources={'r1': 1, 'r2': 1})
self.sch.add_task(worker=WORKER, task_id='B', resources={'r1': 1, 'r2': 1})
self.sch.add_task(worker=WORKER, task_id='C', resources={'r1': 1})
self.sch.update_resources(r1=2, r2=1)
self.assertEqual('A', self.sch.get_work(worker=WORKER)['task_id'])
self.check_task_order('C')
def test_single_resource_lock(self):
self.sch.add_task(worker='X', task_id='A', resources={'r': 1})
self.assertEqual('A', self.sch.get_work(worker='X')['task_id'])
self.sch.add_task(worker=WORKER, task_id='B', resources={'r': 2}, priority=10)
self.sch.add_task(worker=WORKER, task_id='C', resources={'r': 1})
self.sch.update_resources(r=2)
# Should wait for 2 units of r to be available for B before scheduling C
self.check_task_order([])
def test_no_lock_if_too_many_resources_required(self):
self.sch.add_task(worker=WORKER, task_id='A', resources={'r': 2}, priority=10)
self.sch.add_task(worker=WORKER, task_id='B', resources={'r': 1})
self.sch.update_resources(r=1)
self.check_task_order('B')
def test_multiple_resources_lock(self):
self.sch.get_work(worker='X') # indicate to the scheduler that X is active
self.sch.add_task(worker='X', task_id='A', resources={'r1': 1, 'r2': 1}, priority=10)
self.sch.add_task(worker=WORKER, task_id='B', resources={'r2': 1})
self.sch.add_task(worker=WORKER, task_id='C', resources={'r1': 1})
self.sch.update_resources(r1=1, r2=1)
# should preserve both resources for worker 'X'
self.check_task_order([])
def test_multiple_resources_no_lock(self):
self.sch.add_task(worker=WORKER, task_id='A', resources={'r1': 1}, priority=10)
self.sch.add_task(worker=WORKER, task_id='B', resources={'r1': 1, 'r2': 1}, priority=10)
self.sch.add_task(worker=WORKER, task_id='C', resources={'r2': 1})
self.sch.update_resources(r1=1, r2=2)
self.assertEqual('A', self.sch.get_work(worker=WORKER)['task_id'])
# C doesn't block B, so it can go first
self.check_task_order('C')
def test_allow_resource_use_while_scheduling(self):
self.sch.update_resources(r1=1)
self.sch.add_task(worker='SCHEDULING', task_id='A', resources={'r1': 1}, priority=10)
self.sch.add_task(worker=WORKER, task_id='B', resources={'r1': 1}, priority=1)
self.assertEqual('B', self.sch.get_work(worker=WORKER)['task_id'])
def test_stop_locking_resource_for_uninterested_worker(self):
self.setTime(0)
self.sch.update_resources(r1=1)
self.assertIsNone(self.sch.get_work(worker=WORKER)['task_id'])
self.sch.add_task(worker=WORKER, task_id='A', resources={'r1': 1}, priority=10)
self.sch.add_task(worker='LOW_PRIO', task_id='B', resources={'r1': 1}, priority=1)
self.assertIsNone(self.sch.get_work(worker='LOW_PRIO')['task_id'])
self.setTime(120)
self.assertEqual('B', self.sch.get_work(worker='LOW_PRIO')['task_id'])
def check_task_order(self, order):
for expected_id in order:
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], expected_id)
self.sch.add_task(worker=WORKER, task_id=expected_id, status=DONE)
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_priorities(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=10)
self.sch.add_task(worker=WORKER, task_id='B', priority=5)
self.sch.add_task(worker=WORKER, task_id='C', priority=15)
self.sch.add_task(worker=WORKER, task_id='D', priority=9)
self.check_task_order(['C', 'A', 'D', 'B'])
def test_priorities_default_and_negative(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=10)
self.sch.add_task(worker=WORKER, task_id='B')
self.sch.add_task(worker=WORKER, task_id='C', priority=15)
self.sch.add_task(worker=WORKER, task_id='D', priority=-20)
self.sch.add_task(worker=WORKER, task_id='E', priority=1)
self.check_task_order(['C', 'A', 'E', 'B', 'D'])
def test_priorities_and_dependencies(self):
self.sch.add_task(worker=WORKER, task_id='A', deps=['Z'], priority=10)
self.sch.add_task(worker=WORKER, task_id='B', priority=5)
self.sch.add_task(worker=WORKER, task_id='C', deps=['Z'], priority=3)
self.sch.add_task(worker=WORKER, task_id='D', priority=2)
self.sch.add_task(worker=WORKER, task_id='Z', priority=1)
self.check_task_order(['Z', 'A', 'B', 'C', 'D'])
def test_priority_update_dependency_after_scheduling(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=1)
self.sch.add_task(worker=WORKER, task_id='B', priority=5, deps=['A'])
self.sch.add_task(worker=WORKER, task_id='C', priority=10, deps=['B'])
self.sch.add_task(worker=WORKER, task_id='D', priority=6)
self.check_task_order(['A', 'B', 'C', 'D'])
def test_disable(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_disable_and_reenable(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.re_enable_task('A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
def test_disable_and_reenable_and_disable_again(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.re_enable_task('A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be still enabled
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled now
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], None)
def test_disable_and_done(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.assertEqual(len(self.sch.task_list('FAILED', '')), 0)
self.sch.add_task(worker=WORKER, task_id='A', status=DONE)
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.assertEqual(len(self.sch.task_list('DONE', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
def test_automatic_re_enable(self):
self.sch = CentralPlannerScheduler(disable_failures=2, disable_persist=100)
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled now
self.assertEqual(DISABLED, self.sch.task_list('', '')['A']['status'])
# re-enables after 100 seconds
self.setTime(101)
self.assertEqual(FAILED, self.sch.task_list('', '')['A']['status'])
def test_automatic_re_enable_with_one_failure_allowed(self):
self.sch = CentralPlannerScheduler(disable_failures=1, disable_persist=100)
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled now
self.assertEqual(DISABLED, self.sch.task_list('', '')['A']['status'])
# re-enables after 100 seconds
self.setTime(101)
self.assertEqual(FAILED, self.sch.task_list('', '')['A']['status'])
def test_no_automatic_re_enable_after_manual_disable(self):
self.sch = CentralPlannerScheduler(disable_persist=100)
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id='A', status=DISABLED)
# should be disabled now
self.assertEqual(DISABLED, self.sch.task_list('', '')['A']['status'])
# should not re-enable after 100 seconds
self.setTime(101)
self.assertEqual(DISABLED, self.sch.task_list('', '')['A']['status'])
def test_no_automatic_re_enable_after_auto_then_manual_disable(self):
self.sch = CentralPlannerScheduler(disable_failures=2, disable_persist=100)
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# should be disabled now
self.assertEqual(DISABLED, self.sch.task_list('', '')['A']['status'])
# should remain disabled once set
self.sch.add_task(worker=WORKER, task_id='A', status=DISABLED)
self.assertEqual(DISABLED, self.sch.task_list('', '')['A']['status'])
# should not re-enable after 100 seconds
self.setTime(101)
self.assertEqual(DISABLED, self.sch.task_list('', '')['A']['status'])
def test_disable_by_worker(self):
self.sch.add_task(worker=WORKER, task_id='A', status=DISABLED)
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 1)
self.sch.add_task(worker=WORKER, task_id='A')
# should be enabled at this point
self.assertEqual(len(self.sch.task_list('DISABLED', '')), 0)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual(self.sch.get_work(worker=WORKER)['task_id'], 'A')
def test_disable_worker(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.disable_worker(worker=WORKER)
work = self.sch.get_work(worker=WORKER)
self.assertEqual(0, work['n_unique_pending'])
self.assertEqual(0, work['n_pending_tasks'])
self.assertIsNone(work['task_id'])
def test_disable_worker_leaves_jobs_running(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.get_work(worker=WORKER)
self.sch.disable_worker(worker=WORKER)
self.assertEqual(['A'], list(self.sch.task_list('RUNNING', '').keys()))
self.assertEqual(['A'], list(self.sch.worker_list()[0]['running'].keys()))
def test_disable_worker_cannot_pick_up_failed_jobs(self):
self.setTime(0)
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.get_work(worker=WORKER)
self.sch.disable_worker(worker=WORKER)
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
# increase time and prune to make the job pending again
self.setTime(1000)
self.sch.ping(worker=WORKER)
self.sch.prune()
# we won't try the job again
self.assertIsNone(self.sch.get_work(worker=WORKER)['task_id'])
# not even if other stuff is pending, changing the pending tasks code path
self.sch.add_task(worker='other_worker', task_id='B')
self.assertIsNone(self.sch.get_work(worker=WORKER)['task_id'])
def test_disable_worker_cannot_continue_scheduling(self):
self.sch.disable_worker(worker=WORKER)
self.sch.add_task(worker=WORKER, task_id='A')
self.assertIsNone(self.sch.get_work(worker=WORKER)['task_id'])
def test_disable_worker_can_finish_task(self, new_status=DONE, new_deps=[]):
self.sch.add_task(worker=WORKER, task_id='A')
self.assertEqual('A', self.sch.get_work(worker=WORKER)['task_id'])
self.sch.disable_worker(worker=WORKER)
self.assertEqual(['A'], list(self.sch.task_list('RUNNING', '').keys()))
for dep in new_deps:
self.sch.add_task(worker=WORKER, task_id=dep, status='PENDING')
self.sch.add_task(worker=WORKER, task_id='A', status=new_status, new_deps=new_deps)
self.assertFalse(self.sch.task_list('RUNNING', '').keys())
self.assertEqual(['A'], list(self.sch.task_list(new_status, '').keys()))
self.assertIsNone(self.sch.get_work(worker=WORKER)['task_id'])
for task in self.sch.task_list('', '').values():
self.assertFalse(task['workers'])
def test_disable_worker_can_fail_task(self):
self.test_disable_worker_can_finish_task(new_status=FAILED)
def test_disable_worker_stays_disabled_on_new_deps(self):
self.test_disable_worker_can_finish_task(new_status='PENDING', new_deps=['B', 'C'])
def test_prune_worker(self):
self.setTime(1)
self.sch.add_worker(worker=WORKER, info={})
self.setTime(10000)
self.sch.prune()
self.setTime(20000)
self.sch.prune()
self.assertFalse(self.sch.worker_list())
def test_task_list_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
for c in 'ABCD':
sch.add_task(worker=WORKER, task_id=c)
self.assertEqual(set('ABCD'), set(sch.task_list('PENDING', '', False).keys()))
self.assertEqual({'num_tasks': 4}, sch.task_list('PENDING', ''))
def test_task_list_within_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=4)
for c in 'ABCD':
sch.add_task(worker=WORKER, task_id=c)
self.assertEqual(set('ABCD'), set(sch.task_list('PENDING', '').keys()))
def test_task_lists_some_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
for c in 'ABCD':
sch.add_task(worker=WORKER, task_id=c, status=DONE)
for c in 'EFG':
sch.add_task(worker=WORKER, task_id=c)
self.assertEqual(set('EFG'), set(sch.task_list('PENDING', '').keys()))
self.assertEqual({'num_tasks': 4}, sch.task_list('DONE', ''))
def add_task(self, family, **params):
task_id = str(hash((family, str(params)))) # use an unhelpful task id
self.sch.add_task(worker=WORKER, family=family, params=params, task_id=task_id)
return task_id
def search_pending(self, term, expected_keys):
actual_keys = set(self.sch.task_list('PENDING', '', search=term).keys())
self.assertEqual(expected_keys, actual_keys)
def test_task_list_filter_by_search_family_name(self):
task1 = self.add_task('MySpecialTask')
task2 = self.add_task('OtherSpecialTask')
self.search_pending('Special', {task1, task2})
self.search_pending('Task', {task1, task2})
self.search_pending('My', {task1})
self.search_pending('Other', {task2})
def test_task_list_filter_by_search_long_family_name(self):
task = self.add_task('TaskClassWithAVeryLongNameAndDistinctEndingUUDDLRLRAB')
self.search_pending('UUDDLRLRAB', {task})
def test_task_list_filter_by_param_name(self):
task1 = self.add_task('ClassA', day='2016-02-01')
task2 = self.add_task('ClassB', hour='2016-02-01T12')
self.search_pending('day', {task1})
self.search_pending('hour', {task2})
def test_task_list_filter_by_long_param_name(self):
task = self.add_task('ClassA', a_very_long_param_name_ending_with_uuddlrlrab='2016-02-01')
self.search_pending('uuddlrlrab', {task})
def test_task_list_filter_by_param_value(self):
task1 = self.add_task('ClassA', day='2016-02-01')
task2 = self.add_task('ClassB', hour='2016-02-01T12')
self.search_pending('2016-02-01', {task1, task2})
self.search_pending('T12', {task2})
def test_task_list_filter_by_long_param_value(self):
task = self.add_task('ClassA', param='a_very_long_param_value_ending_with_uuddlrlrab')
self.search_pending('uuddlrlrab', {task})
def test_task_list_filter_by_param_name_value_pair(self):
task = self.add_task('ClassA', param='value')
self.search_pending('param=value', {task})
def test_task_list_does_not_filter_by_task_id(self):
task = self.add_task('Class')
self.search_pending(task, set())
def test_task_list_filter_by_multiple_search_terms(self):
expected = self.add_task('ClassA', day='2016-02-01', num='5')
self.add_task('ClassA', day='2016-03-01', num='5')
self.add_task('ClassB', day='2016-02-01', num='5')
self.add_task('ClassA', day='2016-02-01', val='5')
self.search_pending('ClassA 2016-02-01 num', {expected})
def test_search_results_beyond_limit(self):
sch = CentralPlannerScheduler(max_shown_tasks=3)
for i in range(4):
sch.add_task(worker=WORKER, family='Test', params={'p': str(i)}, task_id='Test_%i' % i)
self.assertEqual({'num_tasks': 4}, sch.task_list('PENDING', '', search='Test'))
self.assertEqual(['Test_0'], list(sch.task_list('PENDING', '', search='0').keys()))
def test_priority_update_dependency_chain(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=10, deps=['B'])
self.sch.add_task(worker=WORKER, task_id='B', priority=5, deps=['C'])
self.sch.add_task(worker=WORKER, task_id='C', priority=1)
self.sch.add_task(worker=WORKER, task_id='D', priority=6)
self.check_task_order(['C', 'B', 'A', 'D'])
def test_priority_no_decrease_with_multiple_updates(self):
self.sch.add_task(worker=WORKER, task_id='A', priority=1)
self.sch.add_task(worker=WORKER, task_id='B', priority=10, deps=['A'])
self.sch.add_task(worker=WORKER, task_id='C', priority=5, deps=['A'])
self.sch.add_task(worker=WORKER, task_id='D', priority=6)
self.check_task_order(['A', 'B', 'D', 'C'])
def test_unique_tasks(self):
self.sch.add_task(worker=WORKER, task_id='A')
self.sch.add_task(worker=WORKER, task_id='B')
self.sch.add_task(worker=WORKER, task_id='C')
self.sch.add_task(worker=WORKER + "_2", task_id='B')
response = self.sch.get_work(worker=WORKER)
self.assertEqual(3, response['n_pending_tasks'])
self.assertEqual(2, response['n_unique_pending'])
def test_pending_downstream_disable(self):
self.sch.add_task(worker=WORKER, task_id='A', status=DISABLED)
self.sch.add_task(worker=WORKER, task_id='B', deps=('A',))
self.sch.add_task(worker=WORKER, task_id='C', deps=('B',))
response = self.sch.get_work(worker=WORKER)
self.assertTrue(response['task_id'] is None)
self.assertEqual(0, response['n_pending_tasks'])
self.assertEqual(0, response['n_unique_pending'])
def test_pending_downstream_failure(self):
self.sch.add_task(worker=WORKER, task_id='A', status=FAILED)
self.sch.add_task(worker=WORKER, task_id='B', deps=('A',))
self.sch.add_task(worker=WORKER, task_id='C', deps=('B',))
response = self.sch.get_work(worker=WORKER)
self.assertTrue(response['task_id'] is None)
self.assertEqual(2, response['n_pending_tasks'])
self.assertEqual(2, response['n_unique_pending'])
def test_task_list_no_deps(self):
self.sch.add_task(worker=WORKER, task_id='B', deps=('A',))
self.sch.add_task(worker=WORKER, task_id='A')
task_list = self.sch.task_list('PENDING', '')
self.assertFalse('deps' in task_list['A'])
def test_task_first_failure_time(self):
self.sch.add_task(worker=WORKER, task_id='A')
test_task = self.sch._state.get_task('A')
self.assertIsNone(test_task.failures.first_failure_time)
time_before_failure = time.time()
test_task.add_failure()
time_after_failure = time.time()
self.assertLessEqual(time_before_failure,
test_task.failures.first_failure_time)
self.assertGreaterEqual(time_after_failure,
test_task.failures.first_failure_time)
def test_task_first_failure_time_remains_constant(self):
self.sch.add_task(worker=WORKER, task_id='A')
test_task = self.sch._state.get_task('A')
self.assertIsNone(test_task.failures.first_failure_time)
test_task.add_failure()
first_failure_time = test_task.failures.first_failure_time
test_task.add_failure()
self.assertEqual(first_failure_time, test_task.failures.first_failure_time)
def test_task_has_excessive_failures(self):
self.sch.add_task(worker=WORKER, task_id='A')
test_task = self.sch._state.get_task('A')
self.assertIsNone(test_task.failures.first_failure_time)
self.assertFalse(test_task.has_excessive_failures())
test_task.add_failure()
self.assertFalse(test_task.has_excessive_failures())
fake_failure_time = (test_task.failures.first_failure_time -
2 * 60 * 60)
test_task.failures.first_failure_time = fake_failure_time
self.assertTrue(test_task.has_excessive_failures())
def test_quadratic_behavior(self):
""" Test that get_work is not taking linear amount of time.
This is of course impossible to test, however, doing reasonable
assumptions about hardware. This time should finish in a timely
manner.
"""
# For 10000 it takes almost 1 second on my laptop. Prior to these
# changes it was being slow already at NUM_TASKS=300
NUM_TASKS = 10000
for i in range(NUM_TASKS):
self.sch.add_task(worker=str(i), task_id=str(i), resources={})
for i in range(NUM_TASKS):
self.assertEqual(self.sch.get_work(worker=str(i))['task_id'], str(i))
self.sch.add_task(worker=str(i), task_id=str(i), status=DONE)
def test_get_work_speed(self):
""" Test that get_work is fast for few workers and many DONEs.
In #986, @daveFNbuck reported that he got a slowdown.
"""
# This took almost 4 minutes without optimization.
# Now it takes 10 seconds on my machine.
NUM_PENDING = 1000
NUM_DONE = 200000
assert NUM_DONE >= NUM_PENDING
for i in range(NUM_PENDING):
self.sch.add_task(worker=WORKER, task_id=str(i), resources={})
for i in range(NUM_PENDING, NUM_DONE):
self.sch.add_task(worker=WORKER, task_id=str(i), status=DONE)
for i in range(NUM_PENDING):
res = int(self.sch.get_work(worker=WORKER)['task_id'])
self.assertTrue(0 <= res < NUM_PENDING)
self.sch.add_task(worker=WORKER, task_id=str(res), status=DONE)