async def transfer_inputs(communicator: BaseCommunicator):
"""Transfer missing input data.
:raises DataTransferError: on failure.
"""
inputs_connector = connectors["local"].duplicate()
inputs_connector.path = INPUTS_VOLUME
inputs_connector.name = "inputs"
inputs_connector.config["path"] = INPUTS_VOLUME
logger.debug(f"Inputs connector: {inputs_connector}")
logger.debug("Transfering missing data.")
response = await communicator.send_command(
Message.command("get_inputs_no_shared_storage", ""))
try:
for base_url, (connector_name, files) in response.message_data.items():
await download_to_location(base_url, files,
connectors[connector_name],
inputs_connector)
except Exception as e:
error_message = f"Preparing input for url {base_url} from connector {connector_name} failed: {e}."
await communicator.send_command(
Message.command("process_log", {"error": error_message}))
await communicator.send_command(Message.command("finish", {}))
raise DataTransferError(error_message)
async def error(error_message: str, communicator: BaseCommunicator):
"""Error occured inside container.
Send the error and terminate the process.
"""
with suppress(Exception):
await communicator.send_command(
Message.command("process_log", {"error": error_message}))
with suppress(Exception):
await communicator.send_command(Message.command("finish", {}))
async def handle_terminate(self, message: Message,
identity: PeerIdentity) -> Response[str]:
"""Handle terminate command."""
response = await self.processing_communicator.send_command(
Message.command("terminate", ""))
response.uuid = message.uuid
return response
async def main():
"""Start the main program.
:raises RuntimeError: when runtime error occurs.
:raises asyncio.exceptions.CancelledError: when task is terminated.
"""
communicator = _get_communicator()
protocol = InitProtocol(communicator, logger)
communicate_task = asyncio.ensure_future(protocol.communicate())
# Initialize settings constants by bootstraping.
response = await communicator.send_command(
Message.command("bootstrap", (DATA_ID, "init")))
global_settings.initialize_constants(DATA_ID, response.message_data)
initialize_secrets(response.message_data[ExecutorFiles.SECRETS_DIR])
modify_connector_settings()
connectors.recreate_connectors()
try:
prepare_volumes()
await protocol.transfer_missing_data()
except PreviousDataExistsError:
logger.warning("Previous run data exists, aborting processing.")
raise
except Exception as exception:
message = f"Unexpected exception in init container: {exception}."
logger.exception(message)
await error(message, communicator)
raise
finally:
protocol.stop_communicate()
with suppress(asyncio.TimeoutError):
await asyncio.wait_for(communicate_task, timeout=10)
async def get_script(self) -> str:
"""Get the script from the listener."""
response = await self.communicator.send_command(
Message.command("get_script", ""))
if response.response_status == ResponseStatus.ERROR:
raise RuntimeError("Response status error while fetching script.")
return response.message_data
async def transfer_missing_data(self):
"""Transfer missing data.
:raises RuntimeError: when data transfer error occurs.
"""
await self.communicator.send_command(
Message.command("update_status", "PP"))
if DATA_ALL_VOLUME_SHARED:
transfering_coroutine = transfer_data(self.communicator)
else:
transfering_coroutine = transfer_inputs(self.communicator)
await transfering_coroutine
async def process_script(self, script: str) -> int:
"""Send the script to the processing container.
This method can be very long running as it waits for the return code
the processing container.
:returns: return code of the process running the script.
"""
try:
response = await self.communicator.send_command(
Message.command("process_script", script),
response_timeout=None)
return response.message_data
except asyncio.CancelledError:
return 1
async def transfer_missing_data(self):
"""Transfer missing data.
Log error re-raise exception on failure.
:raises: RuntimeError on failure.
"""
try:
await transfer_data(self.listener_communicator)
except RuntimeError:
with suppress(Exception):
await self.listener_communicator.send_command(
Message.command(
"process_log",
{"error": ["Error transfering missing data."]}))
raise
async def handle_upload_dirs(self, message: Message[list[str]],
identity: PeerIdentity) -> Response[str]:
"""Create directories and sent them to the listener.
This is needed in case empty dirs are referenced.
"""
directories = message.message_data
referenced_dirs = []
for directory in directories:
# Make sure all referenced directories exists if
# UPLOAD_CONNECTOR_NAME equals "local".
if UPLOAD_CONNECTOR_NAME == "local":
destination_dir = DATA_VOLUME / directory
destination_dir.mkdir(parents=True, exist_ok=True)
referenced_dirs.append({
"path": os.path.join(directory, ""),
"size": 0
})
return await self.listener_communicator.send_command(
Message.command("referenced_files", referenced_dirs))
async def collect_produced_files(self):
"""Collect files produced by the worker.
Keep only files that are referenced in the data model.
Log error re-raise exception on failure.
:raises: RuntimeError on failure.
"""
try:
logger.debug("Collecting files")
await collect_files(self.listener_communicator)
logger.debug("Collected files")
return True
except RuntimeError:
with suppress(Exception):
await self.listener_communicator.send_command(
Message.command(
"process_log",
{"error": ["Error collecting produced files."]},
))
return False
async def send_referenced_files(self, referenced_files):
"""Send referenced files to the listener."""
return await self.listener_communicator.send_command(
Message.command("referenced_files", referenced_files))
async def terminate(self):
"""Terminate the processing container."""
await self.communicator.send_command(Message.command("terminate", ""))
async def finish(self, return_code: int):
"""Send finish command."""
await self.communicator.send_command(
Message.command("finish", {"rc": return_code}))
async def start(self) -> int:
"""Start the main program."""
try:
return_code = 1
await self.start_processing_socket()
self.listener_communicator = await self.open_listener_connection()
try:
logger.debug("Waiting for the processing container to connect")
await asyncio.wait_for(
self.processing_container_connected.wait(),
PROCESSING_CONTAINER_TIMEOUT,
)
except asyncio.TimeoutError:
message = "Unable to connect to the processing container."
logger.critical(message)
with suppress(Exception):
await self.listener_communicator.send_command(
Message.command("process_log", {"error": [message]}))
sys.exit(1)
logger.debug("Connected to the processing container.")
listener = ListenerProtocol(self.listener_communicator,
self.processing_communicator)
processing = ProcessingProtocol(self.processing_communicator,
self.listener_communicator)
try:
# Start listening for messages from the communication and the
# processing container.
listener_task = asyncio.ensure_future(listener.communicate())
processing_task = asyncio.ensure_future(
processing.communicate())
listener_task.add_done_callback(self._communicator_stopped)
processing_task.add_done_callback(self._communicator_stopped)
await self.listener_communicator.send_command(
Message.command("update_status", "PR"))
await self.transfer_missing_data()
script = await listener.get_script()
self._process_script_task = asyncio.create_task(
processing.process_script(script))
return_code = await self._process_script_task
self._process_script_task = None
except RuntimeError:
logger.exception("Error processing script.")
with suppress(Exception):
await self.listener_communicator.send_command(
Message.command(
"process_log",
{
"error":
["Runtime error in communication container."]
},
))
except Exception:
logger.exception("While running communication container")
finally:
if not KEEP_DATA:
purge_secrets()
if not await self.collect_produced_files():
if return_code == 0:
return_code = 1
# Notify listener that the processing is finished.
with suppress(Exception):
await listener.finish(return_code)
listener.stop_communicate()
processing.stop_communicate()
# Wait for up to 10 seconds to close the tasks.
with suppress(asyncio.TimeoutError):
await asyncio.wait_for(asyncio.gather(listener_task,
processing_task),
timeout=10)
return return_code
async def get_script(self) -> str:
"""Get the script from the listener."""
response = await self.communicator.send_command(
Message.command("get_script", ""))
return response.message_data
async def start(self) -> int:
"""Start the main program."""
try:
return_code = 1
logger.debug("Starting upload thread")
upload_thread = Uploader(self, asyncio.get_running_loop())
upload_thread.start()
# Wait up to 60 seconds for uploader to get ready.
if not upload_thread.ready.wait(60):
logger.error("Upload thread failed to start, terminating.")
raise RuntimeError("Upload thread failed to start.")
await self.start_processing_socket()
self.listener_communicator = await self.open_listener_connection()
try:
logger.debug("Waiting for the processing container to connect")
await asyncio.wait_for(
self.processing_container_connected.wait(),
constants.CONTAINER_TIMEOUT,
)
except asyncio.TimeoutError:
message = "Unable to connect to the processing container."
logger.critical(message)
with suppress(Exception):
await self.listener_communicator.send_command(
Message.command("process_log", {"error": [message]}))
sys.exit(1)
logger.debug("Connected to the processing container.")
listener = ListenerProtocol(self.listener_communicator,
self.processing_communicator)
processing = ProcessingProtocol(self.processing_communicator,
self.listener_communicator)
try:
# Start listening for messages from the communication and the
# processing container.
listener_task = asyncio.ensure_future(listener.communicate())
# Initialize settings constants by bootstraping.
response = await self.listener_communicator.send_command(
Message.command("bootstrap", (DATA_ID, "communication")))
global_settings.initialize_constants(DATA_ID,
response.message_data)
modify_connector_settings()
# Recreate connectors with received settings.
connectors.recreate_connectors()
set_default_storage_connectors()
processing_task = asyncio.ensure_future(
processing.communicate())
listener_task.add_done_callback(self._communicator_stopped)
processing_task.add_done_callback(self._communicator_stopped)
await self.listener_communicator.send_command(
Message.command("update_status", "PR"))
script = await listener.get_script()
self._process_script_task = asyncio.create_task(
processing.process_script(script))
return_code = await self._process_script_task
self._process_script_task = None
except RuntimeError as runtime_exception:
logger.exception("Error processing script.")
with suppress(Exception):
await self.listener_communicator.send_command(
Message.command(
"process_log",
{
"error": [
"Runtime error in communication container: "
f"{runtime_exception}."
]
},
))
except Exception:
logger.exception("While running communication container")
finally:
logger.debug("Terminating upload thread.")
upload_thread.terminate()
upload_thread.join()
if not KEEP_DATA:
purge_secrets()
# Notify listener that the processing is finished.
try:
await listener.finish(return_code)
except RuntimeError:
logger.exception("Error sending finish command.")
except:
logger.exception("Unknown error sending finish command.")
listener.stop_communicate()
processing.stop_communicate()
# Wait for up to 10 seconds to close the tasks.
with suppress(asyncio.TimeoutError):
await asyncio.wait_for(asyncio.gather(listener_task,
processing_task),
timeout=10)
return return_code
"""
subpath = global_settings.LOCATION_SUBPATH
directories = message.message_data
referenced_dirs = []
for directory in directories:
if storage_connectors := STORAGE_CONNECTOR.get("data"):
if mounted_connector := storage_connectors[1]:
destination_dir = mounted_connector.path / subpath / directory
destination_dir.mkdir(parents=True, exist_ok=True)
referenced_dirs.append({
"path": os.path.join(directory, ""),
"size": 0
})
return await self.listener_communicator.send_command(
Message.command("referenced_files", referenced_dirs))
async def process_script(self, script: str) -> int:
"""Send the script to the processing container.
This method can be very long running as it waits for the return code
the processing container.
:returns: return code of the process running the script.
"""
try:
response = await self.communicator.send_command(
Message.command("process_script", script),
response_timeout=None)
return response.message_data
except asyncio.CancelledError: