def setUp(self): self.zombie_db_client = mock_client_factory.create() self.mock_db_client = self.zombie_db_client.mock # Create the basic image sources, systems, and benchmarks. self.image_collections = [ make_image_collection(self.mock_db_client).identifier for _ in range(2) ] self.systems = [ self.mock_db_client.system_collection.insert_one( mock_core.MockSystem().serialize()).inserted_id for _ in range(2) ] self.benchmarks = [ self.mock_db_client.benchmarks_collection.insert_one( mock_core.MockBenchmark().serialize()).inserted_id for _ in range(2) ] # Create run system tasks and trial results self.run_system_tasks = {} self.trial_results = [] for image_collection_id in self.image_collections: for system_id in self.systems: trial_result_id = self.mock_db_client.trials_collection.insert_one( core.trial_result.TrialResult( system_id, True, core.sequence_type.ImageSequenceType.NON_SEQUENTIAL, {}).serialize()).inserted_id self.trial_results.append(trial_result_id) self.run_system_tasks[ trial_result_id] = self.mock_db_client.tasks_collection.insert_one( run_system_task.RunSystemTask( system_id, image_collection_id, result=trial_result_id, state=batch_analysis.task.JobState.DONE).serialize( )).inserted_id self.benchmark_trial_tasks = {} self.benchmark_results = {} for trial_result_id in self.trial_results: self.benchmark_trial_tasks[trial_result_id] = [] self.benchmark_results[trial_result_id] = [] for benchmark_id in self.benchmarks: result_id = self.mock_db_client.results_collection.insert_one( core.benchmark.BenchmarkResult( benchmark_id, trial_result_id, True).serialize()).inserted_id self.benchmark_results[trial_result_id].append(result_id) self.benchmark_trial_tasks[trial_result_id].append( self.mock_db_client.tasks_collection.insert_one( benchmark_trial_task.BenchmarkTrialTask( trial_result_id, benchmark_id, result=result_id, state=batch_analysis.task.JobState.DONE).serialize( )).inserted_id)
def get_run_system_task(self, system_id, image_source_id): """ Get a task to run a system. Most of the parameters are resources requirements passed to the job system. :param system_id: The id of the vision system to test :param image_source_id: The id of the image source to test with :return: A RunSystemTask """ if system_id not in self._run_system_tasks: self._run_system_tasks[system_id] = {} if image_source_id not in self._run_system_tasks[system_id]: self._run_system_tasks[system_id][ image_source_id] = rst.RunSystemTask( system_id=system_id, image_source_id=image_source_id) return self._run_system_tasks[system_id][image_source_id]
def test_do_task_checks_run_system_task_is_unique(self): mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) system_id = bson.ObjectId() image_source_id = bson.ObjectId() task = run_system_task.RunSystemTask(system_id, image_source_id) subject.do_task(task) self.assertTrue(mock_collection.find.called) query = mock_collection.find.call_args[0][0] self.assertIn('system_id', query) self.assertEqual(system_id, query['system_id']) self.assertIn('image_source_id', query) self.assertEqual(image_source_id, query['image_source_id'])
def test_do_task_saves_new_task(self): # Mockthe method chain on the pymongo cursor mock_cursor = mock.MagicMock() mock_cursor.limit.return_value = mock_cursor mock_cursor.count.return_value = 0 mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_collection.find.return_value = mock_cursor mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) system_id = bson.ObjectId() image_source_id = bson.ObjectId() task = run_system_task.RunSystemTask(system_id, image_source_id) subject.do_task(task) self.assertTrue(mock_collection.insert.called) s_task = task.serialize() del s_task[ '_id'] # This gets set after the insert call, clear it again self.assertEqual(s_task, mock_collection.insert.call_args[0][0])
def get_run_system_task(self, system_id, image_source_id, repeat=0, num_cpus=1, num_gpus=0, memory_requirements='3GB', expected_duration='1:00:00'): """ Get a task to run a system. Most of the parameters are resources requirements passed to the job system. :param system_id: The id of the vision system to test :param image_source_id: The id of the image source to test with :param repeat: The repeat of this trial, so we can run the same system more than once. :param num_cpus: The number of CPUs required for the job. Default 1. :param num_gpus: The number of GPUs required for the job. Default 0. :param memory_requirements: The memory required for this job. Default 3 GB. :param expected_duration: The expected time this job will take. Default 1 hour. :return: A RunSystemTask """ existing = self._collection.find_one({ 'system_id': system_id, 'image_source_id': image_source_id, 'repeat': repeat }) if existing is not None: return self._db_client.deserialize_entity(existing) else: return run_system_task.RunSystemTask( system_id=system_id, image_source_id=image_source_id, repeat=repeat, num_cpus=num_cpus, num_gpus=num_gpus, memory_requirements=memory_requirements, expected_duration=expected_duration)
def setUp(self): self.zombie_db_client = mock_client_factory.create() self.mock_db_client = self.zombie_db_client.mock # Create the basic image sources, systems, and benchmarks. self.image_collections = [ make_image_collection(self.mock_db_client).identifier for _ in range(2) ] self.systems = [ self.mock_db_client.system_collection.insert_one( mock_core.MockSystem().serialize()).inserted_id for _ in range(2) ] self.benchmarks = [ self.mock_db_client.benchmarks_collection.insert_one( mock_core.MockBenchmark().serialize()).inserted_id for _ in range(2) ] # Add generate dataset tasks for all the image collections. self.generate_dataset_tasks = { image_collection_id: self.mock_db_client.tasks_collection.insert_one( generate_dataset_task.GenerateDatasetTask( bson.ObjectId(), bson.ObjectId(), {}, result=image_collection_id, state=batch_analysis.task.JobState.DONE).serialize() ).inserted_id for image_collection_id in self.image_collections } # Add controllers that follow the image sources self.controllers = { image_collection_id: self.mock_db_client.image_source_collection.insert_one( traj_follow_controller.TrajectoryFollowController( trajectory={}, trajectory_source=image_collection_id, sequence_type=core.sequence_type.ImageSequenceType. NON_SEQUENTIAL).serialize()).inserted_id for image_collection_id in self.image_collections } # Create run system tasks and trial results self.run_system_tasks = {} self.trial_results = {} for image_collection_id in self.image_collections: self.run_system_tasks[image_collection_id] = [] self.trial_results[image_collection_id] = [] for system_id in self.systems: trial_result_id = self.mock_db_client.trials_collection.insert_one( core.trial_result.TrialResult( system_id, True, core.sequence_type.ImageSequenceType.NON_SEQUENTIAL, {}).serialize()).inserted_id self.trial_results[image_collection_id].append(trial_result_id) self.run_system_tasks[image_collection_id].append( self.mock_db_client.tasks_collection.insert_one( run_system_task.RunSystemTask( system_id, image_collection_id, result=trial_result_id, state=batch_analysis.task.JobState.DONE).serialize( )).inserted_id)