def test_accept_node_tasks(self):
"""Tests successfully calling accept_node_tasks()"""
node = MagicMock()
node.hostname = 'host_1'
node.id = 1
health_task = HealthTask('1234', 'agent_1')
pull_task = PullTask('1234', 'agent_1')
node.is_ready_for_new_job = MagicMock()
node.is_ready_for_new_job.return_value = True
node.is_ready_for_next_job_task = MagicMock()
node.is_ready_for_next_job_task.return_value = True
node.get_next_tasks = MagicMock()
node.get_next_tasks.return_value = [health_task, pull_task]
node_task_resources = NodeResources()
node_task_resources.add(health_task.get_resources())
node_task_resources.add(pull_task.get_resources())
offered_resources = NodeResources([Cpus(100.0), Mem(5000.0)])
expected_remaining_resources = NodeResources()
expected_remaining_resources.add(offered_resources)
expected_remaining_resources.subtract(node_task_resources)
task_resources = NodeResources()
watermark_resources = NodeResources([Cpus(100.0), Mem(5000.0)])
resource_set = ResourceSet(offered_resources, task_resources,
watermark_resources)
scheduling_node = SchedulingNode('agent_1', node, [], [], resource_set)
waiting_tasks = []
had_waiting_task = scheduling_node.accept_node_tasks(
now(), waiting_tasks)
self.assertFalse(had_waiting_task)
self.assertEqual(len(scheduling_node.allocated_tasks), 2)
self.assertTrue(
scheduling_node.allocated_resources.is_equal(node_task_resources))
self.assertTrue(
scheduling_node._remaining_resources.is_equal(
expected_remaining_resources))
self.assertListEqual(waiting_tasks, [])
def test_add_allocated_offers_remove_all_tasks(self):
"""Tests calling add_allocated_offers() when there are not enough resources for the job exes or node tasks"""
node = MagicMock()
node.hostname = 'host_1'
node.id = 1
health_task = HealthTask('1234', 'agent_1')
pull_task = PullTask('1234', 'agent_1')
node.is_ready_for_new_job = MagicMock()
node.is_ready_for_new_job.return_value = True
node.is_ready_for_next_job_task = MagicMock()
node.is_ready_for_next_job_task.return_value = True
node.get_next_tasks = MagicMock()
node.get_next_tasks.return_value = [health_task, pull_task]
offered_resources = NodeResources([Cpus(100.0), Mem(500.0)])
watermark_resources = NodeResources([Cpus(100.0), Mem(500.0)])
resource_set = ResourceSet(offered_resources, NodeResources(),
watermark_resources)
scheduling_node = SchedulingNode('agent_1', node, [], [], resource_set)
running_job_exe_1 = job_test_utils.create_running_job_exe(
agent_id=self.agent_id,
resources=NodeResources([Cpus(1.0), Mem(10.0)]))
running_job_exe_2 = job_test_utils.create_running_job_exe(
agent_id=self.agent_id,
resources=NodeResources([Cpus(2.0), Mem(20.0)]))
node_task_resources = NodeResources()
node_task_resources.add(health_task.get_resources())
node_task_resources.add(pull_task.get_resources())
all_required_resources = NodeResources()
all_required_resources.add(node_task_resources)
all_required_resources.add(
running_job_exe_1.next_task().get_resources())
all_required_resources.add(
running_job_exe_2.next_task().get_resources())
expected_remaining_resources = NodeResources()
expected_remaining_resources.add(offered_resources)
expected_remaining_resources.subtract(node_task_resources)
# Set up node with node tasks and job exes (there would never be queued job exes since they would be scheduled
# before add_allocated_offers() was called
scheduling_node.accept_node_tasks(now(), [])
scheduling_node.accept_job_exe_next_task(running_job_exe_1, [])
scheduling_node.accept_job_exe_next_task(running_job_exe_2, [])
self.assertEqual(len(scheduling_node.allocated_tasks), 2)
self.assertEqual(len(scheduling_node._allocated_running_job_exes), 2)
self.assertEqual(len(scheduling_node._allocated_queued_job_exes), 0)
self.assertTrue(
scheduling_node.allocated_resources.is_equal(
all_required_resources))
# Set up offers (not enough for job exes or node tasks)
offer_1 = ResourceOffer('offer_1', 'agent_1', '1234',
NodeResources([Cpus(0.1),
Mem(600.0)]), now(), None)
scheduling_node.add_allocated_offers([offer_1])
self.assertListEqual(scheduling_node.allocated_offers, [offer_1])
# All allocated tasks and job exes should be gone
self.assertEqual(len(scheduling_node.allocated_tasks), 0)
self.assertEqual(len(scheduling_node._allocated_running_job_exes), 0)
self.assertEqual(len(scheduling_node._allocated_queued_job_exes), 0)
self.assertTrue(
scheduling_node.allocated_resources.is_equal(NodeResources()))
self.assertTrue(
scheduling_node._remaining_resources.is_equal(offered_resources))