def test__filter_and_sort_workers(self):
bundle = Mock(spec=RunBundle, metadata=Mock(spec=MetadataSpec))
bundle.dependencies = []
bundle.metadata.request_queue = None
bundle_resources = RunResources(cpus=1,
gpus=0,
docker_image='',
time=100,
memory=1000,
disk=1000,
network=False)
# gpu worker should be last in the filtered and sorted list
sorted_workers_list = self.bundle_manager._filter_and_sort_workers(
self.get_sample_workers_list(), bundle, bundle_resources)
self.assertEqual(len(sorted_workers_list), 3)
self.assertEqual(sorted_workers_list[0]['gpus'], 0)
self.assertEqual(sorted_workers_list[1]['gpus'], 0)
self.assertEqual(sorted_workers_list[-1]['gpus'], 1)
# gpu worker should be the only worker in the filtered and sorted list
bundle_resources.gpus = 1
sorted_workers_list = self.bundle_manager._filter_and_sort_workers(
self.get_sample_workers_list(), bundle, bundle_resources)
self.assertEqual(len(sorted_workers_list), 1)
self.assertEqual(sorted_workers_list[0]['gpus'], 1)
def _compute_bundle_resources(self, bundle, user_info=None):
return RunResources(
cpus=self._compute_request_cpus(bundle),
gpus=self._compute_request_gpus(bundle),
docker_image=self._get_docker_image(bundle),
# _compute_request_time contains database queries that may reduce efficiency
time=self._compute_request_time(bundle, user_info),
memory=self._compute_request_memory(bundle),
# _compute_request_disk contains database queries that may reduce efficiency
disk=self._compute_request_disk(bundle, user_info),
network=bundle.metadata.request_network,
)
def setUp(self):
self.codalab_manager = Mock(CodaLabManager)
self.codalab_manager.config = {
'workers': {
'default_cpu_image': 'codalab/default-cpu:latest',
'default_gpu_image': 'codalab/default-gpu:latest',
}
}
self.bundle_manager = BundleManager(self.codalab_manager)
self.bundle = Mock(spec=RunBundle, metadata=Mock(spec=MetadataSpec))
self.bundle.metadata.request_queue = None
self.bundle_resources = RunResources(
cpus=0, gpus=0, docker_image='', time=100, memory=1000, disk=1000, network=False
)
self.bundle.dependencies = []
dep = namedtuple('dep', ['parent_uuid', 'parent_path'])
for i in range(6):
self.bundle.dependencies.append(dep(parent_uuid=i, parent_path=''))
self.workers_list = self.get_sample_workers_list()