def query_extra_data(self, perf_index, num_cores=1, node_id=None, node_name=None):
"""Returns data for the next subtask.
:param int perf_index:
:param int num_cores:
:param str | None node_id:
:param str | None node_name:
:rtype: ComputeTaskDef"""
# create new subtask_id
subtask_id = uuid.uuid4().get_hex()
with self._lock:
# check if a task has been assigned to this node
if node_id in self.assigned_nodes:
return self.ExtraData(should_wait=True)
# assign a task
self.assigned_nodes[node_id] = subtask_id
self.assigned_subtasks[subtask_id] = node_id
# create subtask-specific data, 4 bits go for one char (hex digit)
data = random.getrandbits(self.task_params.subtask_data_size * 4)
self.subtask_ids.append(subtask_id)
self.subtask_data[subtask_id] = '%x' % data
subtask_def = ComputeTaskDef()
subtask_def.task_id = self.task_id
subtask_def.subtask_id = subtask_id
subtask_def.src_code = self.src_code
subtask_def.task_owner = self.header.task_owner
subtask_def.environment = self.header.environment
subtask_def.return_address = self.header.task_owner_address
subtask_def.return_port = self.header.task_owner_port
subtask_def.deadline = timeout_to_deadline(5 * 60)
subtask_def.extra_data = {
'data_file': self.shared_data_file,
'subtask_data': self.subtask_data[subtask_id],
'difficulty': self.task_params.difficulty,
'result_size': self.task_params.result_size,
'result_file': 'result.' + subtask_id[0:6]
}
return self.ExtraData(ctd=subtask_def)
def test_react_to_task_to_compute(self):
conn = Mock()
ts = TaskSession(conn)
ts.key_id = "KEY_ID"
ts.task_manager = Mock()
ts.task_computer = Mock()
ts.task_server = Mock()
ts.task_server.get_subtask_ttl.return_value = 31313
env = Mock()
env.docker_images = [DockerImage("dockerix/xii", tag="323")]
env.allow_custom_main_program_file = False
env.get_source_code.return_value = None
ts.task_server.get_environment_by_id.return_value = env
def __reset_mocks():
ts.task_manager.reset_mock()
ts.task_computer.reset_mock()
conn.reset_mock()
# msg.ctd is None -> failure
msg = MessageTaskToCompute()
with self.assertLogs(logger, level="WARNING"):
ts._react_to_task_to_compute(msg)
ts.task_manager.comp_task_keeper.receive_subtask.assert_not_called()
ts.task_computer.session_closed.assert_called_with()
assert conn.close.called
# No source code in the local environment -> failure
__reset_mocks()
ctd = ComputeTaskDef()
ctd.key_id = "KEY_ID"
ctd.subtask_id = "SUBTASKID"
ctd.task_owner = Node()
ctd.task_owner.key = "KEY_ID"
ctd.return_address = "10.10.10.10"
ctd.return_port = 1112
ctd.docker_images = [DockerImage("dockerix/xiii", tag="323")]
msg = MessageTaskToCompute(ctd)
ts._react_to_task_to_compute(msg)
ts.task_manager.comp_task_keeper.receive_subtask.assert_not_called()
ts.task_computer.session_closed.assert_called_with()
assert conn.close.called
# Source code from local environment -> proper execution
__reset_mocks()
env.get_source_code.return_value = "print 'Hello world'"
ts._react_to_task_to_compute(msg)
ts.task_manager.comp_task_keeper.receive_subtask.assert_called_with(
ctd)
ts.task_computer.session_closed.assert_not_called()
ts.task_server.add_task_session.assert_called_with("SUBTASKID", ts)
ts.task_computer.task_given.assert_called_with(ctd)
conn.close.assert_not_called()
# Wrong key id -> failure
__reset_mocks()
ctd.key_id = "KEY_ID2"
ts._react_to_task_to_compute(MessageTaskToCompute(ctd))
ts.task_manager.comp_task_keeper.receive_subtask.assert_not_called()
ts.task_computer.session_closed.assert_called_with()
assert conn.close.called
# Wrong task owner key id -> failure
__reset_mocks()
ctd.key_id = "KEY_ID"
ctd.task_owner.key = "KEY_ID2"
ts._react_to_task_to_compute(MessageTaskToCompute(ctd))
ts.task_manager.comp_task_keeper.receive_subtask.assert_not_called()
ts.task_computer.session_closed.assert_called_with()
assert conn.close.called
# Wrong return port -> failure
__reset_mocks()
ctd.task_owner.key = "KEY_ID"
ctd.return_port = 0
ts._react_to_task_to_compute(MessageTaskToCompute(ctd))
ts.task_manager.comp_task_keeper.receive_subtask.assert_not_called()
ts.task_computer.session_closed.assert_called_with()
assert conn.close.called
# Proper port and key -> proper execution
__reset_mocks()
ctd.task_owner.key = "KEY_ID"
ctd.return_port = 1319
ts._react_to_task_to_compute(MessageTaskToCompute(ctd))
conn.close.assert_not_called()
# Allow custom code / no code in ComputeTaskDef -> failure
__reset_mocks()
env.allow_custom_main_program_file = True
ctd.src_code = ""
ts._react_to_task_to_compute(MessageTaskToCompute(ctd))
ts.task_manager.comp_task_keeper.receive_subtask.assert_not_called()
ts.task_computer.session_closed.assert_called_with()
assert conn.close.called
# Allow custom code / code in ComputerTaskDef -> proper execution
__reset_mocks()
ctd.src_code = "print 'Hello world!'"
ts._react_to_task_to_compute(MessageTaskToCompute(ctd))
ts.task_computer.session_closed.assert_not_called()
ts.task_server.add_task_session.assert_called_with("SUBTASKID", ts)
ts.task_computer.task_given.assert_called_with(ctd)
conn.close.assert_not_called()
# No environment available -> failure
__reset_mocks()
ts.task_server.get_environment_by_id.return_value = None
ts._react_to_task_to_compute(MessageTaskToCompute(ctd))
assert ts.err_msg.startswith("Wrong environment")
ts.task_manager.comp_task_keeper.receive_subtask.assert_not_called()
ts.task_computer.session_closed.assert_called_with()
assert conn.close.called
# Envrionment is Docker environment but with different images -> failure
__reset_mocks()
ts.task_server.get_environment_by_id.return_value = \
DockerEnvironment([DockerImage("dockerix/xii", tag="323"),
DockerImage("dockerix/xiii", tag="325"),
DockerImage("dockerix/xiii")])
ts._react_to_task_to_compute(MessageTaskToCompute(ctd))
assert ts.err_msg.startswith("Wrong docker images")
ts.task_manager.comp_task_keeper.receive_subtask.assert_not_called()
ts.task_computer.session_closed.assert_called_with()
assert conn.close.called
# Envrionment is Docker environment with proper images, but no srouce code -> failure
__reset_mocks()
de = DockerEnvironment([
DockerImage("dockerix/xii", tag="323"),
DockerImage("dockerix/xiii", tag="325"),
DockerImage("dockerix/xiii", tag="323")
])
ts.task_server.get_environment_by_id.return_value = de
ts._react_to_task_to_compute(MessageTaskToCompute(ctd))
assert ts.err_msg.startswith("No source code")
ts.task_manager.comp_task_keeper.receive_subtask.assert_not_called()
ts.task_computer.session_closed.assert_called_with()
assert conn.close.called
# Proper Docker environment with source code
__reset_mocks()
file_name = os.path.join(self.path, "main_program_file")
with open(file_name, 'w') as f:
f.write("Hello world!")
de.main_program_file = file_name
ts._react_to_task_to_compute(MessageTaskToCompute(ctd))
ts.task_server.add_task_session.assert_called_with("SUBTASKID", ts)
ts.task_computer.task_given.assert_called_with(ctd)
conn.close.assert_not_called()
def test_computation(self):
task_server = mock.MagicMock()
task_server.get_task_computer_root.return_value = self.path
task_server.config_desc = config_desc()
tc = TaskComputer("ABC", task_server, use_docker_machine_manager=False)
ctd = ComputeTaskDef()
ctd.task_id = "xyz"
ctd.subtask_id = "xxyyzz"
ctd.return_address = "10.10.10.10"
ctd.return_port = 10203
ctd.key_id = "key"
ctd.task_owner = "owner"
ctd.src_code = "cnt=0\nfor i in range(10000):\n\tcnt += 1\noutput={'data': cnt, 'result_type': 0}"
ctd.extra_data = {}
ctd.short_description = "add cnt"
ctd.deadline = timeout_to_deadline(10)
self.assertEqual(len(tc.assigned_subtasks), 0)
tc.task_given(ctd)
self.assertEqual(tc.assigned_subtasks["xxyyzz"], ctd)
self.assertLessEqual(tc.assigned_subtasks["xxyyzz"].deadline,
timeout_to_deadline(10))
self.assertEqual(tc.task_to_subtask_mapping["xyz"], "xxyyzz")
tc.task_server.request_resource.assert_called_with(
"xyz", tc.resource_manager.get_resource_header("xyz"),
"10.10.10.10", 10203, "key", "owner")
assert tc.task_resource_collected("xyz")
tc.task_server.unpack_delta.assert_called_with(
tc.dir_manager.get_task_resource_dir("xyz"), None, "xyz")
assert len(tc.current_computations) == 0
assert tc.assigned_subtasks.get("xxyyzz") is None
task_server.send_task_failed.assert_called_with(
"xxyyzz", "xyz", "Host direct task not supported", "10.10.10.10",
10203, "key", "owner", "ABC")
tc.support_direct_computation = True
tc.task_given(ctd)
assert tc.task_resource_collected("xyz")
assert not tc.waiting_for_task
assert len(tc.current_computations) == 1
self.__wait_for_tasks(tc)
prev_task_failed_count = task_server.send_task_failed.call_count
self.assertFalse(tc.counting_task)
self.assertEqual(len(tc.current_computations), 0)
self.assertIsNone(tc.assigned_subtasks.get("xxyyzz"))
assert task_server.send_task_failed.call_count == prev_task_failed_count
self.assertTrue(task_server.send_results.called)
args = task_server.send_results.call_args[0]
self.assertEqual(args[0], "xxyyzz")
self.assertEqual(args[1], "xyz")
self.assertEqual(args[2]["data"], 10000)
self.assertGreater(args[3], 0)
self.assertLess(args[3], 10)
self.assertEqual(args[4], "10.10.10.10")
self.assertEqual(args[5], 10203)
self.assertEqual(args[6], "key")
self.assertEqual(args[7], "owner")
self.assertEqual(args[8], "ABC")
ctd.subtask_id = "aabbcc"
ctd.src_code = "raise Exception('some exception')"
ctd.deadline = timeout_to_deadline(5)
tc.task_given(ctd)
self.assertEqual(tc.assigned_subtasks["aabbcc"], ctd)
self.assertLessEqual(tc.assigned_subtasks["aabbcc"].deadline,
timeout_to_deadline(5))
self.assertEqual(tc.task_to_subtask_mapping["xyz"], "aabbcc")
tc.task_server.request_resource.assert_called_with(
"xyz", tc.resource_manager.get_resource_header("xyz"),
"10.10.10.10", 10203, "key", "owner")
self.assertTrue(tc.task_resource_collected("xyz"))
self.__wait_for_tasks(tc)
self.assertFalse(tc.counting_task)
self.assertEqual(len(tc.current_computations), 0)
self.assertIsNone(tc.assigned_subtasks.get("aabbcc"))
task_server.send_task_failed.assert_called_with(
"aabbcc", "xyz", 'some exception', "10.10.10.10", 10203, "key",
"owner", "ABC")
ctd.subtask_id = "aabbcc2"
ctd.src_code = "print 'Hello world'"
ctd.timeout = timeout_to_deadline(5)
tc.task_given(ctd)
self.assertTrue(tc.task_resource_collected("xyz"))
self.__wait_for_tasks(tc)
task_server.send_task_failed.assert_called_with(
"aabbcc2", "xyz", "Wrong result format", "10.10.10.10", 10203,
"key", "owner", "ABC")
ctd.subtask_id = "xxyyzz2"
ctd.timeout = timeout_to_deadline(1)
tc.task_given(ctd)
self.assertTrue(tc.task_resource_collected("xyz"))
tt = tc.current_computations[0]
tc.task_computed(tc.current_computations[0])
self.assertEqual(len(tc.current_computations), 0)
task_server.send_task_failed.assert_called_with(
"xxyyzz2", "xyz", "Wrong result format", "10.10.10.10", 10203,
"key", "owner", "ABC")
tt.end_comp()
time.sleep(0.5)
if tt.is_alive():
tt.join(timeout=5)