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)