|
| 1 | +########################################################################### |
| 2 | +# Copyright (c), The AiiDA team. All rights reserved. # |
| 3 | +# This file is part of the AiiDA code. # |
| 4 | +# # |
| 5 | +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # |
| 6 | +# For further information on the license, see the LICENSE.txt file # |
| 7 | +# For further information please visit http://www.aiida.net # |
| 8 | +########################################################################### |
| 9 | +"""Implementation of UnstashCalculation.""" |
| 10 | + |
| 11 | +import json |
| 12 | +from pathlib import Path |
| 13 | + |
| 14 | +from aiida import orm |
| 15 | +from aiida.common import AIIDA_LOGGER |
| 16 | +from aiida.common.datastructures import CalcInfo, CodeInfo, UnstashTargetMode |
| 17 | +from aiida.engine import CalcJob |
| 18 | + |
| 19 | +from .stash import StashCalculation |
| 20 | + |
| 21 | +EXEC_LOGGER = AIIDA_LOGGER.getChild('UnstashCalculation') |
| 22 | + |
| 23 | + |
| 24 | +class UnstashCalculation(CalcJob): |
| 25 | + """ |
| 26 | + Utility to unstash files from a remote folder. |
| 27 | +
|
| 28 | + An example of how the input should look like: |
| 29 | +
|
| 30 | + .. code-block:: python |
| 31 | +
|
| 32 | + inputs = { |
| 33 | + 'metadata': { |
| 34 | + 'computer': Computer.collection.get(label="localhost"), |
| 35 | + 'options': { |
| 36 | + 'resources': {'num_machines': 1}, |
| 37 | + 'unstash': { |
| 38 | + 'unstash_target_mode': UnstashTargetMode.NewNode.value, |
| 39 | + 'source_list': ['aiida.in', '_aiidasubmit.sh'], # also accepts ['*'] |
| 40 | + }, |
| 41 | + }, |
| 42 | + }, |
| 43 | + 'source_node': <RemoteStashData>, |
| 44 | + 'code': <MY_CODE> # only in case of `type(source_node)==RemoteStashCustomData` |
| 45 | + } |
| 46 | + """ |
| 47 | + |
| 48 | + def __init__(self, *args, **kwargs): |
| 49 | + super().__init__(*args, **kwargs) |
| 50 | + |
| 51 | + @classmethod |
| 52 | + def define(cls, spec): |
| 53 | + super().define(spec) |
| 54 | + |
| 55 | + spec.input( |
| 56 | + 'source_node', |
| 57 | + valid_type=orm.RemoteStashData, |
| 58 | + required=True, |
| 59 | + help='', |
| 60 | + ) |
| 61 | + |
| 62 | + spec.inputs['metadata']['computer'].required = True |
| 63 | + spec.inputs['metadata']['options']['unstash'].required = True |
| 64 | + spec.inputs['metadata']['options']['unstash']['unstash_target_mode'].required = True |
| 65 | + spec.inputs['metadata']['options']['resources'].default = { |
| 66 | + 'num_machines': 1, |
| 67 | + 'num_mpiprocs_per_machine': 1, |
| 68 | + } |
| 69 | + |
| 70 | + spec.inputs['metadata']['options']['input_filename'].default = 'aiida.in' |
| 71 | + spec.inputs['metadata']['options']['output_filename'].default = 'aiida.out' |
| 72 | + |
| 73 | + def prepare_for_submission(self, folder): |
| 74 | + if self.inputs.source_node.computer.uuid != self.inputs.metadata.computer.uuid: |
| 75 | + EXEC_LOGGER.warning( |
| 76 | + 'YOUR SETTING MIGHT RESULT IN A SILENT FAILURE!' |
| 77 | + ' The computer of the source node and the computer of the calculation are strongly advised be the same.' |
| 78 | + ' However, it is not mandatory,' |
| 79 | + ' in order to support the case that original computer somehow is not usable, anymore.' |
| 80 | + ' E.g. the original computer was configured for ``core.torque``, but the HPC has move to SLURM,' |
| 81 | + ' so you had to create a new computer configured with ``core.slurm``,' |
| 82 | + " and you'll need a job submission to do this." |
| 83 | + ) |
| 84 | + source_node = self.inputs.get('source_node') |
| 85 | + unstash_target_mode = self.inputs.metadata.options.unstash.get('unstash_target_mode') |
| 86 | + |
| 87 | + calc_info = CalcInfo() |
| 88 | + |
| 89 | + if isinstance(source_node, orm.RemoteStashCustomData): |
| 90 | + if unstash_target_mode == UnstashTargetMode.OriginalPlace.value: |
| 91 | + |
| 92 | + def traverse(node_): |
| 93 | + for link in node_.base.links.get_incoming(): |
| 94 | + if (isinstance(link.node, CalcJob) and not isinstance(link.node, StashCalculation)) or ( |
| 95 | + isinstance(link.node, orm.RemoteData) |
| 96 | + ): |
| 97 | + return link.node |
| 98 | + return traverse(link.node) |
| 99 | + return None |
| 100 | + |
| 101 | + stashed_calculation_node = traverse(source_node) |
| 102 | + |
| 103 | + if not stashed_calculation_node: |
| 104 | + raise ValueError( |
| 105 | + 'Your stash node is not connected to any calcjob node, cannot find the source path.' |
| 106 | + ) |
| 107 | + |
| 108 | + target_path = stashed_calculation_node.get_remote_path() |
| 109 | + else: # UnstashTargetMode.NewRemoteData.value |
| 110 | + computer = self.inputs.metadata.get('computer') |
| 111 | + with computer.get_transport() as transport: |
| 112 | + remote_user = transport.whoami_async() |
| 113 | + remote_working_directory = computer.get_workdir().format(username=remote_user) |
| 114 | + |
| 115 | + # The following line is set at calcjob::presubmit, but we need it here |
| 116 | + calc_info_uuid = str(self.node.uuid) |
| 117 | + |
| 118 | + # This is normally done in execmanager::upload_calculation, |
| 119 | + # however unfortunatly is not modular and I had to copy-paste the logic here |
| 120 | + target_path = Path(remote_working_directory).joinpath( |
| 121 | + calc_info_uuid[:2], calc_info_uuid[2:4], calc_info_uuid[4:] |
| 122 | + ) |
| 123 | + |
| 124 | + with folder.open(self.options.input_filename, 'w', encoding='utf8') as handle: |
| 125 | + stash_dict = { |
| 126 | + 'source_path': self.inputs.source_node.target_basepath, |
| 127 | + 'source_list': self.inputs.metadata.options.unstash.get('source_list'), |
| 128 | + 'target_base': str(target_path), |
| 129 | + } |
| 130 | + stash_json = json.dumps(stash_dict) |
| 131 | + handle.write(f'{stash_json}\n') |
| 132 | + |
| 133 | + code_info = CodeInfo() |
| 134 | + code_info.stdin_name = self.options.input_filename |
| 135 | + code_info.stdout_name = self.options.output_filename |
| 136 | + |
| 137 | + if 'code' in self.inputs: |
| 138 | + code_info.code_uuid = self.inputs.code.uuid |
| 139 | + else: |
| 140 | + raise ValueError( |
| 141 | + f"Input 'code' is required for `UnstashTargetMode.{UnstashTargetMode(unstash_target_mode)}` mode." |
| 142 | + ) |
| 143 | + |
| 144 | + calc_info.codes_info = [code_info] |
| 145 | + calc_info.retrieve_list = [self.options.output_filename] |
| 146 | + calc_info.local_copy_list = [] |
| 147 | + calc_info.remote_copy_list = [] |
| 148 | + calc_info.remote_symlink_list = [] |
| 149 | + |
| 150 | + else: |
| 151 | + if 'code' in self.inputs: |
| 152 | + raise ValueError( |
| 153 | + f"Input 'code' cannot be used for `UnstashTargetMode.{UnstashTargetMode(unstash_target_mode)}`" |
| 154 | + ' mode. This UnStash mode is performed on the login node, ' |
| 155 | + 'no submission is planned therefore no code is needed.' |
| 156 | + ) |
| 157 | + |
| 158 | + calc_info.skip_submit = True |
| 159 | + |
| 160 | + calc_info.codes_info = [] |
| 161 | + calc_info.retrieve_list = [] |
| 162 | + calc_info.local_copy_list = [] |
| 163 | + calc_info.remote_copy_list = [] |
| 164 | + calc_info.remote_symlink_list = [] |
| 165 | + |
| 166 | + return calc_info |
0 commit comments