the_three_commits_in_april

Author: khsrali

pyproject.toml

@@ -68,6 +68,7 @@ requires-python = '>=3.9'
 'core.stash' = 'aiida.calculations.stash:StashCalculation'
 'core.templatereplacer' = 'aiida.calculations.templatereplacer:TemplatereplacerCalculation'
 'core.transfer' = 'aiida.calculations.transfer:TransferCalculation'
+'core.unstash' = 'aiida.calculations.unstash:UnstashCalculation'
 
 [project.entry-points.'aiida.calculations.importers']
 'core.arithmetic.add' = 'aiida.calculations.importers.arithmetic.add:ArithmeticAddCalculationImporter'

src/aiida/calculations/stash.py

@@ -54,7 +54,7 @@ class StashCalculation(CalcJob):
 
         inputs = {
             'metadata': {
-                'computer': Computer.collection.get(label="localhost"),
+                'computer': load_computer(label="localhost"),
                 'options': {
                     'resources': {'num_machines': 1},
                     'stash': {
@@ -64,7 +64,8 @@ class StashCalculation(CalcJob):
                     },
                 },
             },
-            'source_node': node_1,
+            'source_node': <RemoteData_NODE>,
+            'code': <MY_CODE>
         }
     """
 
@@ -91,6 +92,7 @@ def define(cls, spec):
             'num_machines': 1,
             'num_mpiprocs_per_machine': 1,
         }
+
         spec.inputs['metadata']['options']['input_filename'].default = 'aiida.in'
         spec.inputs['metadata']['options']['output_filename'].default = 'aiida.out'
 

src/aiida/calculations/unstash.py

@@ -0,0 +1,123 @@
+###########################################################################
+# Copyright (c), The AiiDA team. All rights reserved.                     #
+# This file is part of the AiiDA code.                                    #
+#                                                                         #
+# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
+# For further information on the license, see the LICENSE.txt file        #
+# For further information please visit http://www.aiida.net               #
+###########################################################################
+""""""
+
+from aiida import orm
+from aiida.common.datastructures import CalcInfo, CodeInfo, UnStashMode
+from aiida.engine import CalcJob
+
+
+class UnStashCalculation(CalcJob):
+    """
+    Utility to unstash files from a remote folder.
+
+    An example of how the input should look like:
+
+    .. code-block:: python
+
+        inputs = {
+            'metadata': {
+                'computer': Computer.collection.get(label="localhost"),
+                'options': {
+                    'resources': {'num_machines': 1},
+                    'unstash': {
+                        'unstash_mode': UnStashMode.OriginalPlace.value,
+                        'unstash_mode': UnStashMode.NewFolderData.value,
+                    },
+                },
+            },
+            'source_node': node_1,
+            'custom_command': 'rsync -av aiida.in _aiidasubmit.sh /scratch/my_stashing/',
+        }
+    """
+
+
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @classmethod
+    def define(cls, spec):
+        super().define(spec)
+
+        spec.input(
+            'source_node',
+            valid_type=(orm.RemoteStashCompressedData, orm.RemoteStashCustomData, orm.RemoteStashFolderData),
+            required=True,
+            help='',
+        )
+
+        spec.inputs['metadata']['options']['input_filename'].default = 'aiida.in'
+        spec.inputs['metadata']['options']['output_filename'].default = 'aiida.out'
+        # Code is irrelevant for this calculation.
+        spec.inputs.pop('code', None)
+
+        # Ideally one could use the same computer as the one of the `source_node`.
+        # However, if another computer has access to the directory, we don't want to restrict.`
+        spec.inputs['metadata']['computer'].required = True
+        spec.inputs['metadata']['options']['resources'].default = {
+            'num_machines': 1,
+            'num_mpiprocs_per_machine': 1,
+        }
+
+    def prepare_for_submission(self, folder):
+        source_node = self.inputs.get('source_node')
+        unstash_mode = self.inputs.metadata.options.stash.get('unstash_mode')
+
+
+        calc_info = CalcInfo()
+
+        if isinstance(source_node, (orm.RemoteStashCompressedData, orm.RemoteStashFolderData)):
+            calc_info.skip_submit = True
+
+            calc_info.codes_info = []
+            calc_info.retrieve_list = []
+            calc_info.local_copy_list = []
+            calc_info.remote_copy_list = []
+            calc_info.remote_symlink_list = []
+
+        elif isinstance(source_node, orm.RemoteStashCustomData):
+            custom_command = self.inputs.metadata.options.stash.get('custom_command')
+
+            if custom_command is None:
+                raise ValueError("Input 'custom_command' is required for `StashMode.CUSTOM_SCRIPT` mode.")
+
+            if unstash_mode == UnStashMode.OriginalPlace.value:
+                src_of_the_src= orm.load_node(source_node.source_uuid)
+                working_directory = src_of_the_src.get_remote_path()  # type: ignore[union-attr]
+            elif unstash_mode == UnStashMode.NewFolderData.value:
+                # TODO:
+                # Execmanager should make REmoteDAta
+                # Generate a path, so the submitted job can put the data directly there.
+                # Give that path to execmanager to generate and save a RemoteData node.
+
+
+
+            working_directory = source_node.get_remote_path()  # type: ignore[union-attr]
+            change_dir = f'cd {working_directory}\n'
+
+            with folder.open(self.options.input_filename, 'w', encoding='utf8') as handle:
+                the_scripts = change_dir + '\n' + custom_command
+                handle.write(the_scripts)
+
+            code_info = CodeInfo()
+            code_info.stdin_name = self.options.input_filename
+            code_info.stdout_name = self.options.output_filename
+
+            calc_info.codes_info = [code_info]
+            calc_info.retrieve_list = [self.options.output_filename]
+            calc_info.local_copy_list = []
+            calc_info.remote_copy_list = []
+            calc_info.remote_symlink_list = []
+
+            # The stashed node is going to be created by ``execmanager``, once the job is finished.
+
+
+
+        return calc_info

src/aiida/common/__init__.py

@@ -78,6 +78,7 @@
     'StoringNotAllowed',
     'TestsNotAllowedError',
     'TransportTaskException',
+    'UnStashMode',
     'UniquenessError',
     'UnsupportedSpeciesError',
     'ValidationError',

src/aiida/common/datastructures.py

@@ -15,7 +15,7 @@
 
 from .extendeddicts import DefaultFieldsAttributeDict
 
-__all__ = ('CalcInfo', 'CalcJobState', 'CodeInfo', 'CodeRunMode', 'StashMode')
+__all__ = ('CalcInfo', 'CalcJobState', 'CodeInfo', 'CodeRunMode', 'StashMode', 'UnStashMode')
 
 
 class StashMode(Enum):
@@ -29,6 +29,13 @@ class StashMode(Enum):
     SUBMIT_CUSTOM_CODE = 'submit_custom_code'
 
 
+class UnStashMode(Enum):
+    """Mode to use when unstashing files."""
+
+    NewFolderData = 'new_folder_data'
+    OriginalPlace = 'original_place'
+
+
 class CalcJobState(Enum):
     """The sub state of a CalcJobNode while its Process is in an active state (i.e. Running or Waiting)."""
 

src/aiida/engine/daemon/execmanager.py

@@ -491,6 +491,30 @@ async def stash_calculation(calculation: CalcJobNode, transport: Transport) -> N
 
     ###
 
+    if stash_mode == StashMode.SUBMIT_CUSTOM_CODE.value:
+        if calculation.process_type != 'aiida.calculations:core.stash':
+            EXEC_LOGGER.warning(
+                f'Stashing as {StashMode.SUBMIT_CUSTOM_CODE.value}'
+                ' is only possible through job submission. Stashing skipped.'
+            )
+            # Note we could easily support it via `transport.exec_command_wait_async`
+            # However, that may confuse users, as it's done in a different manner than job-submission
+            # So just to stay safe, we decided not to provide this feature.
+            return
+
+        remote_stash = RemoteStashCustomData(
+            computer=calculation.computer,
+            stash_mode=StashMode(stash_mode),
+            target_basepath=str(target_base),
+            source_list=source_list,
+        )
+
+        remote_stash.store()
+        remote_stash.base.links.add_incoming(calculation, link_type=LinkType.CREATE, link_label='remote_stash')
+        return
+
+    ###
+
     if stash_mode == StashMode.COPY.value:
         target_basepath = target_base / uuid[:2] / uuid[2:4] / uuid[4:]
 

src/aiida/engine/processes/calcjobs/calcjob.py

@@ -121,14 +121,28 @@ def validate_stash_options(stash_options: Any, _: Any) -> Optional[str]:
     source_list = stash_options.get('source_list', None)
     stash_mode = stash_options.get('stash_mode', None)
 
-    if not isinstance(target_base, str) or not os.path.isabs(target_base):
-        return f'`metadata.options.stash.target_base` should be an absolute filepath, got: {target_base}'
+    if stash_mode is not StashMode.CUSTOM_SCRIPT.value:
+        if not isinstance(target_base, str) or not os.path.isabs(target_base):
+            return f'`metadata.options.stash.target_base` should be an absolute filepath, got: {target_base}'
 
-    if not isinstance(source_list, (list, tuple)) or any(
-        not isinstance(src, str) or os.path.isabs(src) for src in source_list
-    ):
-        port = 'metadata.options.stash.source_list'
-        return f'`{port}` should be a list or tuple of relative filepaths, got: {source_list}'
+        if not isinstance(source_list, (list, tuple)) or any(
+            not isinstance(src, str) or os.path.isabs(src) for src in source_list
+        ):
+            port = 'metadata.options.stash.source_list'
+            return f'`{port}` should be a list or tuple of relative filepaths, got: {source_list}'
+
+        if stash_options.get('custom_command', None) is not None:
+            return f'`metadata.options.stash.target_base` is not an input for {stash_mode} stashing mode'
+    else:
+        custom_command = stash_options.get('custom_command', None)
+        if not isinstance(custom_command, str):
+            return f'`metadata.options.stash.custom_command` should be a string, got: {custom_command}'
+
+        if stash_options.get('target_base', None) is not None:
+            return '`metadata.options.stash.target_base` is not an input for `StashMode.CUSTOM_SCRIPT` stashing mode'
+
+        if stash_options.get('source_list', None) is not None:
+            return '`metadata.options.stash.source_list` is not an input for `StashMode.CUSTOM_SCRIPT` stashing mode'
 
     try:
         StashMode(stash_mode)
@@ -438,6 +452,12 @@ def define(cls, spec: CalcJobProcessSpec) -> None:  # type: ignore[override]
             required=False,
             help='Whether to follow symlinks while stashing or not, specific to StashMode.COMPRESS_* enums',
         )
+        spec.input(
+            'metadata.options.stash.custom_command',
+            valid_type=str,
+            required=False,
+            help='',
+        )
         spec.output(
             'remote_folder',
             valid_type=orm.RemoteData,

src/aiida/engine/processes/calcjobs/tasks.py

@@ -44,6 +44,7 @@
 UPDATE_COMMAND = 'update'
 RETRIEVE_COMMAND = 'retrieve'
 STASH_COMMAND = 'stash'
+UN_STASH_COMMAND = 'unstash'
 KILL_COMMAND = 'kill'
 
 RETRY_INTERVAL_OPTION = 'transport.task_retry_initial_interval'
@@ -521,6 +522,12 @@ async def execute(self) -> plumpy.process_states.State:  # type: ignore[override
                 else:
                     result = self.submit()
 
+            elif self._command == UN_STASH_COMMAND:
+                if node.get_option('unstash') is not None:
+                    await self._launch_task(task_unstash_job, node, transport_queue)
+
+                result = self.retrieve(monitor_result=self._monitor_result)
+
             elif self._command == SUBMIT_COMMAND:
                 result = await self._launch_task(task_submit_job, node, transport_queue)
 
@@ -555,6 +562,7 @@ async def execute(self) -> plumpy.process_states.State:  # type: ignore[override
             elif self._command == STASH_COMMAND:
                 if node.get_option('stash'):
                     await self._launch_task(task_stash_job, node, transport_queue)
+
                 result = self.retrieve(monitor_result=self._monitor_result)
 
             elif self._command == RETRIEVE_COMMAND:
@@ -667,6 +675,13 @@ def stash(self, monitor_result: CalcJobMonitorResult | None = None) -> 'Waiting'
             ProcessState.WAITING, None, msg=msg, data={'command': STASH_COMMAND, 'monitor_result': monitor_result}
         )
 
+    def unstash(self, monitor_result: CalcJobMonitorResult | None = None) -> 'Waiting':
+        """Return the `Waiting` state that will `unstash` the `CalcJob`."""
+        msg = 'Waiting to unstash'
+        return self.create_state(  # type: ignore[return-value]
+            ProcessState.WAITING, None, msg=msg, data={'command': UN_STASH_COMMAND, 'monitor_result': monitor_result}
+        )
+
     def retrieve(self, monitor_result: CalcJobMonitorResult | None = None) -> 'Waiting':
         """Return the `Waiting` state that will `retrieve` the `CalcJob`."""
         msg = 'Waiting to retrieve'

tests/orm/test_fields/fields_aiida.data.core.remote.stash.custom.RemoteStashCustomData.yml

@@ -20,3 +20,4 @@ stash_mode: QbField('stash_mode', dtype=<enum 'StashMode'>, is_attribute=True)
 target_basepath: QbStrField('target_basepath', dtype=<class 'str'>, is_attribute=True)
 user: QbNumericField('user', dtype=typing.Optional[int], is_attribute=False)
 uuid: QbStrField('uuid', dtype=typing.Optional[str], is_attribute=False)
+