def setUp(self):
self.communicator_mock = MagicMock()
self.missing_data = [
{
"connector_name": "CONNECTOR",
"url": "URL",
"data_id": "1",
"from_storage_location_id": 1,
"to_storage_location_id": 2,
"to_connector": "local",
}
]
self.DOWNLOAD_STARTED_LOCK = Message.command(
"download_started",
{
ExecutorProtocol.STORAGE_LOCATION_ID: 2,
ExecutorProtocol.DOWNLOAD_STARTED_LOCK: True,
},
)
self.DOWNLOAD_STARTED_NO_LOCK = Message.command(
"download_started",
{
ExecutorProtocol.STORAGE_LOCATION_ID: 2,
ExecutorProtocol.DOWNLOAD_STARTED_LOCK: False,
},
)
self.MISSING_DATA = Message.command(ExecutorProtocol.MISSING_DATA_LOCATIONS, "")
self.MISSING_DATA_RESPONSE = Response(
ResponseStatus.OK.value, self.missing_data.copy()
)
return super().setUp()
async def _get_response(self, peer_identity: PeerIdentity,
message: Message) -> Response:
"""Handle command received over 0MQ."""
# Logging is handled separately. The worker might have non-numeric peer identifier.
if message.command_name == "log":
try:
return self.handle_log(message)
except:
logger.exception("Error in handle_log method.")
return message.respond("Error writing logs.",
ResponseStatus.ERROR)
if peer_identity not in self.peers:
try:
peer = Processor(
peer_identity,
message.sequence_number,
self,
)
self.peers[peer_identity] = peer
except:
logger.exception("Error creating Processor instance.")
return message.respond("Error creating Processor instance",
ResponseStatus.ERROR)
peer = self.peers[peer_identity]
try:
response = await database_sync_to_async(peer.process_command
)(message)
except:
logger.exception("Error in process_command method.")
return message.respond("Error processing command",
ResponseStatus.ERROR)
else:
return response
def test_transfer_downloadmulti(self):
self.missing_data = [
{
"connector_name": "CONNECTOR",
"url": "URL",
"data_id": "1",
"from_storage_location_id": 1,
"to_storage_location_id": 2,
},
{
"connector_name": "CONNECTOR",
"url": "URL",
"data_id": "1",
"from_storage_location_id": 1,
"to_storage_location_id": 2,
},
]
self.MISSING_DATA_RESPONSE = Response(
ResponseStatus.OK.value, self.missing_data.copy()
)
commands = [
1,
(self.MISSING_DATA, self.MISSING_DATA_RESPONSE),
(Message.command("update_status", "PP"), self.RESULT_OK),
(self.DOWNLOAD_STARTED_LOCK, self.DOWNLOAD_IN_PROGRESS),
(self.DOWNLOAD_STARTED_LOCK, self.DOWNLOAD_STARTED),
(self.DOWNLOAD_STARTED_NO_LOCK, self.DOWNLOAD_FINISHED),
]
download_command = MagicMock(side_effect=lambda a, b: coroutine(True))
send_command = MagicMock(side_effect=partial(send, commands))
result = self._test_workflow(send_command, download_command, commands, True)
self.assertTrue(result)
download_command.assert_called_once()
def test_handle_get_referenced_files(self):
obj = Message.command(ExecutorProtocol.GET_REFERENCED_FILES, "")
storage_location = StorageLocation.objects.create(
file_storage=self.file_storage,
connector_name="local",
status=StorageLocation.STATUS_DONE,
url=str(self.file_storage.id),
)
path = Path(storage_location.get_path(filename="output.txt"))
path.parent.mkdir(exist_ok=True, parents=True)
path.touch()
data = Data.objects.get(id=1)
data.process.output_schema = [{
"name": "output_file",
"type": "basic:file:"
}]
data.process.save()
data.output = {"output_file": {"file": "output.txt"}}
data.save()
response = self.processor.handle_get_referenced_files(
obj, self.manager)
expected = Response(
ResponseStatus.OK.value,
[
"jsonout.txt",
"stdout.txt",
"output.txt",
],
)
self.assertEqual(response, expected)
def test_handle_missing_data_locations(self):
obj = Message.command(ExecutorProtocol.MISSING_DATA_LOCATIONS, "")
parent = Data.objects.get(id=2)
child = Data.objects.get(id=1)
DataDependency.objects.create(parent=parent,
child=child,
kind=DataDependency.KIND_IO)
storage_location = StorageLocation.objects.create(
file_storage=parent.location,
connector_name="not_local",
status=StorageLocation.STATUS_DONE,
url="url",
)
response = self.processor.handle_missing_data_locations(
obj, self.manager)
self.assertEqual(StorageLocation.all_objects.count(), 3)
created = StorageLocation.all_objects.last()
expected = Response(
ResponseStatus.OK.value,
{
"url": {
"data_id": parent.id,
"from_connector": "not_local",
"from_storage_location_id": storage_location.id,
"to_storage_location_id": created.id,
"to_connector": "local",
}
},
)
self.assertEqual(response, expected)
def handle_get_referenced_files(self, message: Message,
manager: "Processor") -> Response[List]:
"""Get a list of files referenced by the data object.
The list also depends on the Data object itself, more specifically on
its output field. So the method is not idempotent.
"""
return message.respond_ok(referenced_files(manager.data))
def handle_missing_data_locations(
self, message: Message, manager: "Processor"
) -> Response[Union[str, List[Dict[str, Any]]]]:
"""Handle an incoming request to get missing data locations."""
missing_data = []
dependencies = (Data.objects.filter(
children_dependency__child=manager.data,
children_dependency__kind=DataDependency.KIND_IO,
).exclude(location__isnull=True).exclude(
pk=manager.data.id).distinct())
for parent in dependencies:
file_storage = parent.location
if not file_storage.has_storage_location(STORAGE_LOCAL_CONNECTOR):
from_location = file_storage.default_storage_location
if from_location is None:
manager._log_exception(
"No storage location exists (handle_get_missing_data_locations).",
extra={"file_storage_id": file_storage.id},
)
return message.respond_error("No storage location exists")
to_location = StorageLocation.all_objects.get_or_create(
file_storage=file_storage,
url=from_location.url,
connector_name=STORAGE_LOCAL_CONNECTOR,
)[0]
missing_data.append({
"connector_name":
from_location.connector_name,
"url":
from_location.url,
"data_id":
manager.data.id,
"to_storage_location_id":
to_location.id,
"from_storage_location_id":
from_location.id,
})
# Set last modified time so it does not get deleted.
from_location.last_update = now()
from_location.save()
return message.respond_ok(missing_data)
def test_handle_download_started_no_location(self):
obj = Message.command(
ExecutorProtocol.DOWNLOAD_STARTED,
{
"storage_location_id": -2,
"download_started_lock": True,
},
)
with self.assertRaises(StorageLocation.DoesNotExist):
self.processor.handle_download_started(obj, self.manager)
async def _get_response(self, peer_identity: PeerIdentity,
message: Message) -> Response:
"""Handle command received over 0MQ."""
# Logging and bootstraping are handled separately.
if message.command_name in ["log", "bootstrap", "liveness_probe"]:
try:
handler = getattr(self, f"handle_{message.command_name}")
return await database_sync_to_async(
handler, thread_sensitive=False)(message)
except Exception:
logger.exception(
__(
"Error handling command {} for peer {}.",
message.command_name,
peer_identity,
))
return message.respond("Error in handler.",
ResponseStatus.ERROR)
if peer_identity not in self.peers:
try:
peer = Processor(
peer_identity,
message.sequence_number,
self,
)
self.peers[peer_identity] = peer
except:
logger.exception("Error creating Processor instance.")
return message.respond("Error creating Processor instance",
ResponseStatus.ERROR)
peer = self.peers[peer_identity]
try:
response = await database_sync_to_async(
peer.process_command, thread_sensitive=False)(message)
except:
logger.exception("Error in process_command method.")
return message.respond("Error processing command",
ResponseStatus.ERROR)
else:
return response
def test_handle_download_aborted(self):
storage_location = StorageLocation.objects.create(
file_storage=self.file_storage,
connector_name="local",
status=StorageLocation.STATUS_UPLOADING,
)
obj = Message.command(ExecutorProtocol.DOWNLOAD_ABORTED, storage_location.id)
self.processor.handle_download_aborted(obj)
storage_location.refresh_from_db()
self.assertEqual(storage_location.status, StorageLocation.STATUS_PREPARING)
def test_handle_missing_data_locations_missing_storage_location(self):
obj = Message.command(ExecutorProtocol.MISSING_DATA_LOCATIONS, "")
parent = Data.objects.get(id=2)
child = Data.objects.get(id=1)
DataDependency.objects.create(
parent=parent, child=child, kind=DataDependency.KIND_IO
)
response = self.processor.handle_missing_data_locations(obj)
expected = Response(ResponseStatus.ERROR.value, "No storage location exists")
self.assertEqual(response, expected)
self.assertEqual(StorageLocation.all_objects.count(), 1)
def setUp(self):
self.communicator_mock = MagicMock()
self.missing_data = {
"connector_name": "S3",
"url": "transfer_url",
"data_id": 1,
"from_storage_location_id": 1,
"to_storage_location_id": 2,
}
self.COMMAND_DOWNLOAD_FINISHED = Message.command(
ExecutorProtocol.DOWNLOAD_FINISHED, 2
)
self.COMMAND_DOWNLOAD_ABORTED = Message.command(
ExecutorProtocol.DOWNLOAD_ABORTED, 2
)
self.COMMAND_GET_FILES = Message.command(
ExecutorProtocol.GET_FILES_TO_DOWNLOAD, 1
)
self.FILES_LIST = Response(ResponseStatus.OK.value, ["1", "dir/1"])
return super().setUp()
def test_transfer_downloaded(self):
commands = [
1,
(self.MISSING_DATA, self.MISSING_DATA_RESPONSE),
(Message.command("update_status", "PP"), self.RESULT_OK),
(self.DOWNLOAD_STARTED_LOCK, self.DOWNLOAD_FINISHED),
]
download_command = MagicMock(return_value=coroutine(True))
send_command = MagicMock(side_effect=partial(send, commands))
result = self._test_workflow(send_command, download_command, commands, True)
self.assertTrue(result)
download_command.assert_not_called()
def test_no_transfer(self):
self.MISSING_DATA_RESPONSE = Response(ResponseStatus.OK.value, [])
send_command = MagicMock(
side_effect=[
coroutine(self.MISSING_DATA_RESPONSE),
coroutine(self.RESULT_OK),
]
)
self.communicator_mock.send_command = send_command
run_async(transfer._transfer_data(self.communicator_mock))
send_command.assert_called_once_with(
Message.command("missing_data_locations", "")
)
def test_transfer_failed(self):
commands = [
1,
(self.MISSING_DATA, self.MISSING_DATA_RESPONSE),
(Message.command("update_status", "PP"), self.RESULT_OK),
(self.DOWNLOAD_STARTED_LOCK, self.DOWNLOAD_STARTED),
]
download_command = MagicMock(return_value=coroutine(False))
send_command = MagicMock(side_effect=partial(send, commands))
result = self._test_workflow(send_command, download_command, commands, True)
self.assertFalse(result)
download_command.assert_called_once_with(
self.missing_data[0], self.communicator_mock
)
async def send_single_message():
"""Open connection to listener and send single message."""
connection_string = f"{protocol}://{host}:{port}"
zmq_context = zmq.asyncio.Context.instance()
zmq_socket = zmq_context.socket(zmq.DEALER)
zmq_socket.setsockopt(zmq.IDENTITY, b"1")
zmq_socket.connect(connection_string)
communicator = ZMQCommunicator(zmq_socket,
"init_container <-> listener",
logger)
async with communicator:
await asyncio.ensure_future(
communicator.send_command(
Message.command("update_status", "PP")))
def test_handle_resolve_data_path(self):
"""Test data path resolwing."""
data = Data.objects.get(id=1)
Worker.objects.get_or_create(data=data, status=Worker.STATUS_PREPARING)
message = Message.command("resolve_data_path", data.pk)
response = self.manager.process_command(message)
assert response.message_data == str(constants.INPUTS_VOLUME)
connector_name = "local"
self.storage_location = StorageLocation.objects.create(
file_storage=self.file_storage,
connector_name=connector_name,
status="OK")
response = self.manager.process_command(message)
self.assertEqual(response.message_data, f"/data_{connector_name}")
async def terminate(self):
"""Send the terminate command to the worker.
Peer should terminate by itself and send finish message back to us.
"""
await database_sync_to_async(self._save_error
)("Processing was cancelled.")
# Ignore the possible timeout.
with suppress(RuntimeError):
await self._listener.communicator.send_command(
Message.command("terminate", ""),
peer_identity=str(self.data_id).encode(),
response_timeout=1,
)
def test_transfer_protocol_fail(self):
async def raise_exception(*args, **kwargs):
raise RuntimeError("Protocol error")
commands = [
1,
(self.MISSING_DATA, self.MISSING_DATA_RESPONSE),
(Message.command("update_status", "PP"), self.RESULT_OK),
(self.DOWNLOAD_STARTED_LOCK, self.DOWNLOAD_STARTED),
]
download_command = MagicMock(return_value=coroutine(True))
send_command = MagicMock(side_effect=raise_exception)
with self.assertRaises(RuntimeError):
self._test_workflow(send_command, download_command, commands, True)
download_command.assert_not_called()
def test_handle_download_started_ok_no_lock_preparing(self):
storage_location = StorageLocation.objects.create(
file_storage=self.file_storage, connector_name="local")
obj = Message.command(
ExecutorProtocol.DOWNLOAD_STARTED,
{
"storage_location_id": storage_location.id,
"download_started_lock": False,
},
)
response = self.processor.handle_download_started(obj, self.manager)
self.assertEqual(response,
Response(ResponseStatus.OK.value, "download_started"))
storage_location.refresh_from_db()
self.assertEqual(storage_location.status,
StorageLocation.STATUS_PREPARING)
def test_handle_missing_data_locations_none(self):
obj = Message.command(ExecutorProtocol.MISSING_DATA_LOCATIONS, "")
parent = Data.objects.get(id=2)
child = Data.objects.get(id=1)
DataDependency.objects.create(
parent=parent, child=child, kind=DataDependency.KIND_IO
)
StorageLocation.objects.create(
file_storage=parent.location,
connector_name="local",
status=StorageLocation.STATUS_DONE,
url="url",
)
response = self.processor.handle_missing_data_locations(obj)
expected = Response(ResponseStatus.OK.value, [])
self.assertEqual(response, expected)
self.assertEqual(StorageLocation.all_objects.count(), 2)
def test_handle_get_files_to_download(self):
obj = Message.command(ExecutorProtocol.GET_FILES_TO_DOWNLOAD,
self.storage_location.id)
response = self.processor.handle_get_files_to_download(
obj, self.manager)
expected = Response(
ResponseStatus.OK.value,
[{
"id": self.path.id,
"path": "test.me",
"size": -1,
"md5": "md5",
"crc32c": "crc",
"awss3etag": "aws",
"chunk_size": BaseStorageConnector.CHUNK_SIZE,
}],
)
self.assertEqual(response, expected)
def test_handle_download_finished(self):
storage_location = StorageLocation.objects.create(
file_storage=self.file_storage, connector_name="local"
)
obj = Message.command(ExecutorProtocol.DOWNLOAD_FINISHED, storage_location.id)
with patch(
"resolwe.storage.models.FileStorage.default_storage_location",
self.storage_location,
):
response = self.processor.handle_download_finished(obj)
self.assertEqual(response.response_status, ResponseStatus.OK)
storage_location.refresh_from_db()
self.assertEqual(storage_location.status, StorageLocation.STATUS_DONE)
self.assertEqual(storage_location.files.count(), 1)
file = storage_location.files.get()
self.assertEqual(file.path, "test.me")
self.assertEqual(file.md5, "md5")
self.assertEqual(file.crc32c, "crc")
self.assertEqual(file.awss3etag, "aws")
def process_command(self, message: Message) -> Response:
"""Process a single command from the peer.
This command is run in the database_sync_to_async so it is safe to
perform Django ORM operations inside.
Exceptions will be handler one level up and error response will be
sent in this case.
"""
# This worker must be in status processing or preparing.
# All messages from workers not in this status will be discarted and
# error will be returned.
if self.worker.status not in [
Worker.STATUS_PROCESSING,
Worker.STATUS_PREPARING,
]:
self._log_error(
f"Wrong worker status: {self.worker.status} for peer with id {self.data_id}."
)
return message.respond_error(
f"Wrong worker status: {self.worker.status}")
command_name = message.command_name
handler_name = f"handle_{command_name}"
handler = plugin_manager.get_handler(command_name)
if not handler:
error = f"No command handler for '{command_name}'."
self._log_error(error, save_to_data_object=False)
return message.respond_error(error)
# Read sequence number and refresh data object if it differs.
if self.expected_sequence_number != message.sequence_number:
try:
self.data.refresh_from_db()
self.worker.refresh_from_db()
except:
self._log_exception("Unable to refresh data object")
return message.respond_error(
"Unable to refresh the data object")
if self.worker.status != Worker.STATUS_PROCESSING:
self.worker.status = Worker.STATUS_PROCESSING
self.worker.save(update_fields=["status"])
if self.data.started is None:
self.data.started = now()
self.data.save(update_fields=["started"])
self.expected_sequence_number = message.sequence_number + 1
try:
with PrioritizedBatcher.global_instance():
result = handler(message, self)
# Set status of the response to ERROR when data object status
# is Data.STATUS_ERROR. Such response will trigger terminate
# procedure in the processing container and stop processing.
if self.data.status == Data.STATUS_ERROR:
result.type_data = ResponseStatus.ERROR.value
return result
except ValidationError as err:
error = (f"Validation error when saving Data object of process "
f"'{self.data.process.slug}' ({handler_name}): "
f"{err}")
self._log_exception(error)
return message.respond_error("Validation error")
except Exception as err:
error = f"Error in command handler '{handler_name}': {err}"
self._log_exception(error)
return message.respond_error(
f"Error in command handler '{handler_name}'")
def test_handle_download_aborted_missing_storage_location(self):
obj = Message.command(ExecutorProtocol.DOWNLOAD_ABORTED, -2)
response = self.processor.handle_download_aborted(obj, self.manager)
self.assertEqual(response.response_status, ResponseStatus.OK)
def test_handle_download_finished_missing_storage_location(self):
obj = Message.command(ExecutorProtocol.DOWNLOAD_FINISHED, -2)
with self.assertRaises(StorageLocation.DoesNotExist):
self.processor.handle_download_finished(obj, self.manager)
def test_handle_missing_data_locations_missing_data(self):
obj = Message.command(ExecutorProtocol.MISSING_DATA_LOCATIONS, "")
response = self.processor.handle_missing_data_locations(
obj, self.manager)
self.assertEqual(response, Response(ResponseStatus.OK.value, {}))
def handle_missing_data_locations(
self, message: Message, manager: "Processor"
) -> Response[Union[str, Dict[str, Dict[str, Any]]]]:
"""Handle an incoming request to get missing data locations."""
storage_name = "data"
filesystem_connector = None
filesystem_connectors = [
connector for connector in connectors.for_storage(storage_name)
if connector.mountable
]
if filesystem_connectors:
filesystem_connector = filesystem_connectors[0]
missing_data = dict()
dependencies = (Data.objects.filter(
children_dependency__child=manager.data,
children_dependency__kind=DataDependency.KIND_IO,
).exclude(location__isnull=True).exclude(
pk=manager.data.id).distinct())
for parent in dependencies:
file_storage = parent.location
# Is location available on some local connector?
if any(
file_storage.has_storage_location(
filesystem_connector.name)
for filesystem_connector in filesystem_connectors):
continue
from_location = file_storage.default_storage_location
if from_location is None:
manager._log_exception(
"No storage location exists (handle_get_missing_data_locations).",
extra={"file_storage_id": file_storage.id},
)
return message.respond_error("No storage location exists")
# When there exists at least one filesystem connector for the data
# storage download inputs to the shared storage.
missing_data_item = {
"data_id": parent.pk,
"from_connector": from_location.connector_name,
}
if filesystem_connector:
to_location = StorageLocation.all_objects.get_or_create(
file_storage=file_storage,
url=from_location.url,
connector_name=filesystem_connector.name,
)[0]
missing_data_item[
"from_storage_location_id"] = from_location.id
missing_data_item["to_storage_location_id"] = to_location.id
missing_data_item["to_connector"] = filesystem_connector.name
else:
missing_data_item["files"] = list(
ReferencedPath.objects.filter(
storage_locations=from_location).values())
missing_data[from_location.url] = missing_data_item
# Set last modified time so it does not get deleted.
from_location.last_update = now()
from_location.save()
return message.respond_ok(missing_data)
def test_handle_get_files_to_download_missing_storage_location(self):
obj = Message.command(ExecutorProtocol.GET_FILES_TO_DOWNLOAD, -2)
response = self.processor.handle_get_files_to_download(
obj, self.manager)
self.assertEqual(response, Response(ResponseStatus.OK.value, []))