def test_add_less_gpu(self): node_id = 2 gpu_count = 3 GPUManager.define_node_gpus(node_id, gpu_count) self.assertEqual(GPUManager.get_gpu_count_for_node(node_id), 3) gpu_count = 1 GPUManager.define_node_gpus(node_id, gpu_count) self.assertEqual(GPUManager.get_gpu_count_for_node(node_id), 3)
def test_calls_where_node_has_no_gpus(self): node_id = 7 job_id = 10 gpu_count = 2 required_gpus = 2 GPUManager.define_node_gpus(node_id, gpu_count) node_id = 8 self.assertFalse( GPUManager.reserve_gpus_for_job(node_id, required_gpus)) self.assertFalse( GPUManager.assign_gpus_for_job(node_id, job_id, required_gpus)) nvidia_label = GPUManager.get_nvidia_docker_label(node_id, job_id) self.assertEqual(nvidia_label, "")
def test_reserve_gpu(self): node_id = 4 gpu_count = 2 required_gpus = 2 GPUManager.define_node_gpus(node_id, gpu_count) self.assertTrue(GPUManager.reserve_gpus_for_job( node_id, required_gpus)) job_id = 11 self.assertFalse( GPUManager.reserve_gpus_for_job(node_id, required_gpus)) gpu_count = 4 GPUManager.define_node_gpus(node_id, gpu_count) self.assertTrue(GPUManager.reserve_gpus_for_job( node_id, required_gpus))
def __init__(self, agent_id, node, tasks, running_job_exes, resource_set): """Constructor :param agent_id: The agent ID :type agent_id: string :param node: The node :type node: :class:`scheduler.node.node_class.Node` :param tasks: The current tasks running on the node :type tasks: :func:`list` :param running_job_exes: The current job executions running on the node :type running_job_exes: :func:`list` :param resource_set: The set of resources for the node :type resource_set: :class:`scheduler.resources.agent.ResourceSet` """ self.agent_id = agent_id # Set agent ID separately from node since it can change during scheduling self.hostname = node.hostname self.node_id = node.id self.is_ready_for_new_job = node.is_ready_for_new_job( ) # Cache this for consistency self.is_ready_for_next_job_task = node.is_ready_for_next_job_task( ) # Cache this for consistency self.is_ready_for_system_task = node.is_ready_for_system_task( ) # Cache this for consistency self.allocated_offers = [] self.allocated_resources = NodeResources() self.allocated_tasks = [ ] # Tasks that have been allocated resources from this node self._node = node self._allocated_queued_job_exes = [ ] # New queued job executions that have been allocated resources self._allocated_running_job_exes = [ ] # Running job executions that have been allocated resources self._running_job_exes = running_job_exes self._running_tasks = tasks self._offered_resources = NodeResources( ) # The amount of resources that were originally offered self._offered_resources.add(resource_set.offered_resources) self._remaining_resources = NodeResources() self._remaining_resources.add(self._offered_resources) self._task_resources = resource_set.task_resources self._watermark_resources = resource_set.watermark_resources if int(resource_set.offered_resources.gpus) > 0: GPUManager.define_node_gpus( self.node_id, int(resource_set.offered_resources.gpus))
def test_get_nvidia_label(self): node_id = 6 job_id = 10 gpu_count = 2 required_gpus = 2 GPUManager.define_node_gpus(node_id, gpu_count) GPUManager.reserve_gpus_for_job(node_id, required_gpus) GPUManager.assign_gpus_for_job(node_id, job_id, required_gpus) nvidia_label = GPUManager.get_nvidia_docker_label(node_id, job_id) self.assertEqual(nvidia_label, "0,1") gpu_count = 4 job_id = 11 GPUManager.define_node_gpus(node_id, gpu_count) GPUManager.reserve_gpus_for_job(node_id, required_gpus) GPUManager.assign_gpus_for_job(node_id, job_id, required_gpus) nvidia_label = GPUManager.get_nvidia_docker_label(node_id, job_id) self.assertEqual(nvidia_label, "2,3")
def test_assign_gpus(self): node_id = 5 job_id = 10 gpu_count = 2 required_gpus = 2 GPUManager.define_node_gpus(node_id, gpu_count) GPUManager.reserve_gpus_for_job(node_id, required_gpus) self.assertTrue( GPUManager.assign_gpus_for_job(node_id, job_id, required_gpus)) job_id = 11 self.assertFalse( GPUManager.reserve_gpus_for_job( node_id, required_gpus)) # shouldnt have enough GPUs gpu_count = 4 GPUManager.define_node_gpus(node_id, gpu_count) GPUManager.reserve_gpus_for_job(node_id, required_gpus) self.assertTrue( GPUManager.assign_gpus_for_job(node_id, job_id, required_gpus))
def test_release_gpu(self): node_id = 7 job_id = 10 gpu_count = 2 required_gpus = 2 GPUManager.define_node_gpus(node_id, gpu_count) GPUManager.reserve_gpus_for_job(node_id, required_gpus) self.assertTrue( GPUManager.assign_gpus_for_job(node_id, job_id, required_gpus)) job_id = 11 self.assertFalse( GPUManager.reserve_gpus_for_job( node_id, required_gpus)) # shouldnt have enough GPUs GPUManager.release_gpus(node_id, 10) self.assertTrue(GPUManager.reserve_gpus_for_job( node_id, required_gpus)) #gpus should be avail again self.assertTrue( GPUManager.assign_gpus_for_job( node_id, job_id, required_gpus)) #gpus should be avail again nvidia_label = GPUManager.get_nvidia_docker_label(node_id, job_id) self.assertEqual(nvidia_label, "0,1")
def test_add_additional_GPU(self): node_id = 3 gpu_count = 4 GPUManager.define_node_gpus(node_id, gpu_count) self.assertEqual(GPUManager.get_gpu_count_for_node(node_id), 4)
def test_add_new_node_gpus(self): node_id = 1 gpu_count = 3 GPUManager.define_node_gpus(node_id, gpu_count) self.assertEqual(GPUManager.get_gpu_count_for_node(node_id), gpu_count)