aiida-pythonjob 0.1.2__tar.gz → 0.1.4__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/PKG-INFO +10 -2
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/README.md +8 -1
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/gallery/autogen/how_to.py +57 -10
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/pyproject.toml +9 -1
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/__init__.py +1 -3
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/calculations/pythonjob.py +4 -40
- aiida_pythonjob-0.1.4/src/aiida_pythonjob/data/__init__.py +4 -0
- aiida_pythonjob-0.1.4/src/aiida_pythonjob/data/atoms.py +53 -0
- aiida_pythonjob-0.1.4/src/aiida_pythonjob/data/data_with_value.py +13 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/data/pickled_data.py +0 -12
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/data/serializer.py +1 -9
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/launch.py +33 -26
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/parsers/pythonjob.py +38 -32
- aiida_pythonjob-0.1.4/src/aiida_pythonjob/utils.py +255 -0
- aiida_pythonjob-0.1.4/tests/test_create_env.py +104 -0
- aiida_pythonjob-0.1.4/tests/test_data.py +63 -0
- aiida_pythonjob-0.1.4/tests/test_entry_points.py +61 -0
- aiida_pythonjob-0.1.4/tests/test_parser.py +112 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/tests/test_pythonjob.py +102 -30
- aiida_pythonjob-0.1.4/tests/test_utils.py +13 -0
- aiida_pythonjob-0.1.2/src/aiida_pythonjob/data/__init__.py +0 -4
- aiida_pythonjob-0.1.2/src/aiida_pythonjob/data/pickled_function.py +0 -145
- aiida_pythonjob-0.1.2/src/aiida_pythonjob/utils.py +0 -31
- aiida_pythonjob-0.1.2/tests/input.txt +0 -1
- aiida_pythonjob-0.1.2/tests/test_data.py +0 -23
- aiida_pythonjob-0.1.2/tests/test_parsers.py +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/.github/workflows/ci.yml +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/.github/workflows/python-publish.yml +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/.gitignore +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/.pre-commit-config.yaml +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/.readthedocs.yml +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/LICENSE +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/Makefile +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/environment.yml +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/gallery/autogen/GALLERY_HEADER.rst +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/make.bat +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/requirements.txt +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/conf.py +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/index.rst +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/installation.rst +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/tutorial/dft.ipynb +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/tutorial/html/atomization_energy.html +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/tutorial/html/pythonjob_eos_emt.html +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/tutorial/index.rst +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/examples/test_add.py +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/calculations/__init__.py +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/config.py +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/parsers/__init__.py +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/tests/conftest.py +0 -0
- {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/tests/inputs_folder/another_input.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.3
|
|
2
2
|
Name: aiida-pythonjob
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.4
|
|
4
4
|
Summary: Run Python functions on a remote computer.
|
|
5
5
|
Project-URL: Source, https://github.com/aiidateam/aiida-pythonjob
|
|
6
6
|
Author-email: Xing Wang <xingwang1991@gmail.com>
|
|
@@ -34,6 +34,7 @@ Classifier: Natural Language :: English
|
|
|
34
34
|
Classifier: Programming Language :: Python
|
|
35
35
|
Requires-Python: >=3.9
|
|
36
36
|
Requires-Dist: aiida-core<3,>=2.3
|
|
37
|
+
Requires-Dist: ase
|
|
37
38
|
Requires-Dist: cloudpickle
|
|
38
39
|
Requires-Dist: voluptuous
|
|
39
40
|
Provides-Extra: docs
|
|
@@ -55,7 +56,14 @@ Description-Content-Type: text/markdown
|
|
|
55
56
|
[](https://codecov.io/gh/aiidateam/aiida-pythonjob)
|
|
56
57
|
[](http://aiida-pythonjob.readthedocs.io/)
|
|
57
58
|
|
|
58
|
-
|
|
59
|
+
`PythonJob` allows users to run Python functions on a remote computer. It is designed to enable users from non-AiiDA communities to run their Python functions remotely and construct workflows with checkpoints, maintaining all data provenance. For instance, users can use ASE's calculator to run a DFT calculation on a remote computer directly.
|
|
60
|
+
|
|
61
|
+
## Key Features
|
|
62
|
+
|
|
63
|
+
1. **Remote Execution**: Seamlessly run Python functions on a remote computer.
|
|
64
|
+
2. **User-Friendly**: Designed for users who are not familiar with AiiDA, simplifying the process of remote execution.
|
|
65
|
+
3. **Workflow Management**: Construct workflows using WorkGraph with checkpoints, ensuring that intermediate states and results are preserved.
|
|
66
|
+
4. **Data Provenance**: Maintain comprehensive data provenance, tracking the full history and transformations of data.
|
|
59
67
|
|
|
60
68
|
|
|
61
69
|
|
|
@@ -4,7 +4,14 @@
|
|
|
4
4
|
[](https://codecov.io/gh/aiidateam/aiida-pythonjob)
|
|
5
5
|
[](http://aiida-pythonjob.readthedocs.io/)
|
|
6
6
|
|
|
7
|
-
|
|
7
|
+
`PythonJob` allows users to run Python functions on a remote computer. It is designed to enable users from non-AiiDA communities to run their Python functions remotely and construct workflows with checkpoints, maintaining all data provenance. For instance, users can use ASE's calculator to run a DFT calculation on a remote computer directly.
|
|
8
|
+
|
|
9
|
+
## Key Features
|
|
10
|
+
|
|
11
|
+
1. **Remote Execution**: Seamlessly run Python functions on a remote computer.
|
|
12
|
+
2. **User-Friendly**: Designed for users who are not familiar with AiiDA, simplifying the process of remote execution.
|
|
13
|
+
3. **Workflow Management**: Construct workflows using WorkGraph with checkpoints, ensuring that intermediate states and results are preserved.
|
|
14
|
+
4. **Data Provenance**: Maintain comprehensive data provenance, tracking the full history and transformations of data.
|
|
8
15
|
|
|
9
16
|
|
|
10
17
|
|
|
@@ -6,16 +6,60 @@ How to guides
|
|
|
6
6
|
|
|
7
7
|
|
|
8
8
|
######################################################################
|
|
9
|
-
#
|
|
10
|
-
#
|
|
9
|
+
# Preparing inputs for `PythonJob`
|
|
10
|
+
# --------------------------------
|
|
11
|
+
# The `prepare_pythonjob_inputs` function is available for setting up the
|
|
12
|
+
# inputs for a `PythonJob` calculation. This function simplifies the process
|
|
13
|
+
# of preparing and serializing data, and configuring the execution environment.
|
|
14
|
+
#
|
|
15
|
+
# - **Code**: You can specify the `computer` where the job will run, which will
|
|
16
|
+
# create a `python3@computer` code if it doesn't already exist. Alternatively,
|
|
17
|
+
# if the code has already been created, you can set the `code` directly.
|
|
18
|
+
#
|
|
19
|
+
# - **Data**: Use standard Python data types for input. The `prepare_pythonjob_inputs`
|
|
20
|
+
# function handles the conversion to AiiDA data. For serialization:
|
|
21
|
+
# - The function first searches for an AiiDA data entry point corresponding to the module
|
|
22
|
+
# and class names (e.g., `ase.atoms.Atoms`).
|
|
23
|
+
# - If a matching entry point exists, it is used for serialization.
|
|
24
|
+
# - If no match is found, the data is serialized into binary format using `PickledData`.
|
|
25
|
+
#
|
|
26
|
+
# - **Python Version**: Ensure the Python version on the remote computer matches the local environment.
|
|
27
|
+
# This is important since pickle is used for data storage and retrieval. Use **conda** to
|
|
28
|
+
# create and activate a virtual environment with the same Python version. Pass metadata
|
|
29
|
+
# to the scheduler to activate the environment during the job execution:
|
|
30
|
+
#
|
|
31
|
+
# .. code-block:: python
|
|
32
|
+
#
|
|
33
|
+
# metadata = {
|
|
34
|
+
# "options": {
|
|
35
|
+
# 'custom_scheduler_commands': 'module load anaconda\nconda activate py3.11\n',
|
|
36
|
+
# }
|
|
37
|
+
# }
|
|
38
|
+
#
|
|
39
|
+
# --------------------------------------------------
|
|
40
|
+
# Create a conda environment on the remote computer
|
|
41
|
+
# --------------------------------------------------
|
|
42
|
+
# One can use the `create_conda_env` function to create a conda environment
|
|
43
|
+
# on the remote computer. The function will create a conda environment with
|
|
44
|
+
# the specified packages and modules. The function will update the packages
|
|
45
|
+
# if the environment already exists.
|
|
46
|
+
#
|
|
47
|
+
# .. code-block:: python
|
|
48
|
+
#
|
|
49
|
+
# from aiida_pythonjob.utils import create_conda_env
|
|
50
|
+
# # create a conda environment on remote computer
|
|
51
|
+
# create_conda_env(
|
|
52
|
+
# "merlin6", # Remote computer
|
|
53
|
+
# "test_pythonjob", # Name of the conda environment
|
|
54
|
+
# modules=["anaconda"], # Modules to load (e.g., Anaconda)
|
|
55
|
+
# pip=["numpy", "matplotlib"], # Python packages to install via pip
|
|
56
|
+
# conda={ # Conda-specific settings
|
|
57
|
+
# "channels": ["conda-forge"], # Channels to use
|
|
58
|
+
# "dependencies": ["qe"] # Conda packages to install
|
|
59
|
+
# }
|
|
60
|
+
# )
|
|
11
61
|
#
|
|
12
|
-
# To run this tutorial, you need to load the AiiDA profile.
|
|
13
62
|
#
|
|
14
|
-
|
|
15
|
-
from aiida import load_profile
|
|
16
|
-
|
|
17
|
-
load_profile()
|
|
18
|
-
|
|
19
63
|
|
|
20
64
|
######################################################################
|
|
21
65
|
# Default outputs
|
|
@@ -24,8 +68,11 @@ load_profile()
|
|
|
24
68
|
# The default output of the function is `result`. The `PythonJob` task
|
|
25
69
|
# will store the result as one node in the database with the key `result`.
|
|
26
70
|
#
|
|
27
|
-
from aiida
|
|
28
|
-
from
|
|
71
|
+
from aiida import load_profile
|
|
72
|
+
from aiida.engine import run_get_node
|
|
73
|
+
from aiida_pythonjob import PythonJob, prepare_pythonjob_inputs
|
|
74
|
+
|
|
75
|
+
load_profile()
|
|
29
76
|
|
|
30
77
|
|
|
31
78
|
def add(x, y):
|
|
@@ -22,6 +22,7 @@ keywords = ["aiida", "plugin"]
|
|
|
22
22
|
requires-python = ">=3.9"
|
|
23
23
|
dependencies = [
|
|
24
24
|
"aiida-core>=2.3,<3",
|
|
25
|
+
"ase",
|
|
25
26
|
"cloudpickle",
|
|
26
27
|
"voluptuous"
|
|
27
28
|
]
|
|
@@ -46,7 +47,14 @@ Source = "https://github.com/aiidateam/aiida-pythonjob"
|
|
|
46
47
|
|
|
47
48
|
[project.entry-points."aiida.data"]
|
|
48
49
|
"pythonjob.pickled_data" = "aiida_pythonjob.data.pickled_data:PickledData"
|
|
49
|
-
"pythonjob.
|
|
50
|
+
"pythonjob.ase.atoms.Atoms" = "aiida_pythonjob.data.atoms:AtomsData"
|
|
51
|
+
"pythonjob.builtins.int" = "aiida.orm.nodes.data.int:Int"
|
|
52
|
+
"pythonjob.builtins.float" = "aiida.orm.nodes.data.float:Float"
|
|
53
|
+
"pythonjob.builtins.str" = "aiida.orm.nodes.data.str:Str"
|
|
54
|
+
"pythonjob.builtins.bool" = "aiida.orm.nodes.data.bool:Bool"
|
|
55
|
+
"pythonjob.builtins.list"="aiida_pythonjob.data.data_with_value:List"
|
|
56
|
+
"pythonjob.builtins.dict"="aiida_pythonjob.data.data_with_value:Dict"
|
|
57
|
+
|
|
50
58
|
|
|
51
59
|
[project.entry-points."aiida.calculations"]
|
|
52
60
|
"pythonjob.pythonjob" = "aiida_pythonjob.calculations.pythonjob:PythonJob"
|
|
@@ -1,16 +1,14 @@
|
|
|
1
1
|
"""AiiDA plugin that run Python function on remote computers."""
|
|
2
2
|
|
|
3
|
-
__version__ = "0.1.
|
|
3
|
+
__version__ = "0.1.4"
|
|
4
4
|
|
|
5
5
|
from .calculations import PythonJob
|
|
6
|
-
from .data import PickledData, PickledFunction
|
|
7
6
|
from .launch import prepare_pythonjob_inputs
|
|
8
7
|
from .parsers import PythonJobParser
|
|
9
8
|
|
|
10
9
|
__all__ = (
|
|
11
10
|
"PythonJob",
|
|
12
11
|
"PickledData",
|
|
13
|
-
"PickledFunction",
|
|
14
12
|
"prepare_pythonjob_inputs",
|
|
15
13
|
"PythonJobParser",
|
|
16
14
|
)
|
{aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/calculations/pythonjob.py
RENAMED
|
@@ -11,6 +11,7 @@ from aiida.common.folders import Folder
|
|
|
11
11
|
from aiida.engine import CalcJob, CalcJobProcessSpec
|
|
12
12
|
from aiida.orm import (
|
|
13
13
|
Data,
|
|
14
|
+
Dict,
|
|
14
15
|
FolderData,
|
|
15
16
|
List,
|
|
16
17
|
RemoteData,
|
|
@@ -19,8 +20,6 @@ from aiida.orm import (
|
|
|
19
20
|
to_aiida_type,
|
|
20
21
|
)
|
|
21
22
|
|
|
22
|
-
from aiida_pythonjob.data.pickled_function import PickledFunction, to_pickled_function
|
|
23
|
-
|
|
24
23
|
__all__ = ("PythonJob",)
|
|
25
24
|
|
|
26
25
|
|
|
@@ -42,31 +41,11 @@ class PythonJob(CalcJob):
|
|
|
42
41
|
:param spec: the calculation job process spec to define.
|
|
43
42
|
"""
|
|
44
43
|
super().define(spec)
|
|
45
|
-
spec.input(
|
|
46
|
-
"function",
|
|
47
|
-
valid_type=PickledFunction,
|
|
48
|
-
serializer=to_pickled_function,
|
|
49
|
-
required=False,
|
|
50
|
-
)
|
|
51
|
-
spec.input(
|
|
52
|
-
"function_source_code",
|
|
53
|
-
valid_type=Str,
|
|
54
|
-
serializer=to_aiida_type,
|
|
55
|
-
required=False,
|
|
56
|
-
)
|
|
57
|
-
spec.input("function_name", valid_type=Str, serializer=to_aiida_type, required=False)
|
|
44
|
+
spec.input("function_data", valid_type=Dict, serializer=to_aiida_type, required=False)
|
|
58
45
|
spec.input("process_label", valid_type=Str, serializer=to_aiida_type, required=False)
|
|
59
46
|
spec.input_namespace(
|
|
60
47
|
"function_inputs", valid_type=Data, required=False
|
|
61
48
|
) # , serializer=serialize_to_aiida_nodes)
|
|
62
|
-
spec.input(
|
|
63
|
-
"function_outputs",
|
|
64
|
-
valid_type=List,
|
|
65
|
-
default=lambda: List(),
|
|
66
|
-
required=False,
|
|
67
|
-
serializer=to_aiida_type,
|
|
68
|
-
help="The information of the output ports",
|
|
69
|
-
)
|
|
70
49
|
spec.input(
|
|
71
50
|
"parent_folder",
|
|
72
51
|
valid_type=(RemoteData, FolderData, SinglefileData),
|
|
@@ -146,7 +125,7 @@ class PythonJob(CalcJob):
|
|
|
146
125
|
if "process_label" in self.inputs:
|
|
147
126
|
return self.inputs.process_label.value
|
|
148
127
|
else:
|
|
149
|
-
data = self.
|
|
128
|
+
data = self.inputs.function_data.get_dict()
|
|
150
129
|
return f"PythonJob<{data['name']}>"
|
|
151
130
|
|
|
152
131
|
def on_create(self) -> None:
|
|
@@ -155,21 +134,6 @@ class PythonJob(CalcJob):
|
|
|
155
134
|
super().on_create()
|
|
156
135
|
self.node.label = self._build_process_label()
|
|
157
136
|
|
|
158
|
-
def get_function_data(self) -> dict[str, t.Any]:
|
|
159
|
-
"""Get the function data.
|
|
160
|
-
|
|
161
|
-
:returns: The function data.
|
|
162
|
-
"""
|
|
163
|
-
if "function" in self.inputs:
|
|
164
|
-
metadata = self.inputs.function.metadata
|
|
165
|
-
metadata["source_code"] = metadata["import_statements"] + "\n" + metadata["source_code_without_decorator"]
|
|
166
|
-
return metadata
|
|
167
|
-
else:
|
|
168
|
-
return {
|
|
169
|
-
"source_code": self.inputs.function_source_code.value,
|
|
170
|
-
"name": self.inputs.function_name.value,
|
|
171
|
-
}
|
|
172
|
-
|
|
173
137
|
def prepare_for_submission(self, folder: Folder) -> CalcInfo:
|
|
174
138
|
"""Prepare the calculation for submission.
|
|
175
139
|
|
|
@@ -192,7 +156,7 @@ class PythonJob(CalcJob):
|
|
|
192
156
|
parent_folder_name = self.inputs.parent_folder_name.value
|
|
193
157
|
else:
|
|
194
158
|
parent_folder_name = self._DEFAULT_PARENT_FOLDER_NAME
|
|
195
|
-
function_data = self.
|
|
159
|
+
function_data = self.inputs.function_data.get_dict()
|
|
196
160
|
# create python script to run the function
|
|
197
161
|
script = f"""
|
|
198
162
|
import pickle
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
from aiida.orm import Data
|
|
3
|
+
from ase import Atoms
|
|
4
|
+
from ase.db.row import atoms2dict
|
|
5
|
+
|
|
6
|
+
__all__ = ("AtomsData",)
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class AtomsData(Data):
|
|
10
|
+
"""Data to represent a ASE Atoms."""
|
|
11
|
+
|
|
12
|
+
_cached_atoms = None
|
|
13
|
+
|
|
14
|
+
def __init__(self, value=None, **kwargs):
|
|
15
|
+
"""Initialise a `AtomsData` node instance.
|
|
16
|
+
|
|
17
|
+
:param value: ASE Atoms instance to initialise the `AtomsData` node from
|
|
18
|
+
"""
|
|
19
|
+
atoms = value or Atoms()
|
|
20
|
+
super().__init__(**kwargs)
|
|
21
|
+
data, keys = self.atoms2dict(atoms)
|
|
22
|
+
self.base.attributes.set_many(data)
|
|
23
|
+
self.base.attributes.set("keys", keys)
|
|
24
|
+
|
|
25
|
+
@classmethod
|
|
26
|
+
def atoms2dict(cls, atoms):
|
|
27
|
+
data = atoms2dict(atoms)
|
|
28
|
+
data.pop("unique_id")
|
|
29
|
+
keys = list(data.keys())
|
|
30
|
+
formula = atoms.get_chemical_formula()
|
|
31
|
+
data = cls._convert_numpy_to_native(data)
|
|
32
|
+
data["formula"] = formula
|
|
33
|
+
data["symbols"] = atoms.get_chemical_symbols()
|
|
34
|
+
return data, keys
|
|
35
|
+
|
|
36
|
+
@classmethod
|
|
37
|
+
def _convert_numpy_to_native(cls, data):
|
|
38
|
+
"""Convert numpy types to Python native types for JSON compatibility."""
|
|
39
|
+
for key, value in data.items():
|
|
40
|
+
if isinstance(value, np.bool_):
|
|
41
|
+
data[key] = bool(value)
|
|
42
|
+
elif isinstance(value, np.ndarray):
|
|
43
|
+
data[key] = value.tolist()
|
|
44
|
+
elif isinstance(value, np.generic):
|
|
45
|
+
data[key] = value.item()
|
|
46
|
+
return data
|
|
47
|
+
|
|
48
|
+
@property
|
|
49
|
+
def value(self):
|
|
50
|
+
keys = self.base.attributes.get("keys")
|
|
51
|
+
data = self.base.attributes.get_many(keys)
|
|
52
|
+
data = dict(zip(keys, data))
|
|
53
|
+
return Atoms(**data)
|
|
@@ -7,18 +7,6 @@ import cloudpickle
|
|
|
7
7
|
from aiida import orm
|
|
8
8
|
|
|
9
9
|
|
|
10
|
-
class Dict(orm.Dict):
|
|
11
|
-
@property
|
|
12
|
-
def value(self):
|
|
13
|
-
return self.get_dict()
|
|
14
|
-
|
|
15
|
-
|
|
16
|
-
class List(orm.List):
|
|
17
|
-
@property
|
|
18
|
-
def value(self):
|
|
19
|
-
return self.get_list()
|
|
20
|
-
|
|
21
|
-
|
|
22
10
|
class PickledData(orm.Data):
|
|
23
11
|
"""Data to represent a pickled value using cloudpickle."""
|
|
24
12
|
|
|
@@ -33,12 +33,10 @@ def get_serializer_from_entry_points() -> dict:
|
|
|
33
33
|
eps.setdefault(key, [])
|
|
34
34
|
eps[key].append(ep)
|
|
35
35
|
|
|
36
|
-
# print("Time to load entry points: ", time.time() - ts)
|
|
37
36
|
# check if there are duplicates
|
|
38
37
|
for key, value in eps.items():
|
|
39
38
|
if len(value) > 1:
|
|
40
39
|
if key in serializers:
|
|
41
|
-
[ep for ep in value if ep.name == serializers[key]]
|
|
42
40
|
eps[key] = [ep for ep in value if ep.name == serializers[key]]
|
|
43
41
|
if not eps[key]:
|
|
44
42
|
raise ValueError(f"Entry point {serializers[key]} not found for {key}")
|
|
@@ -105,13 +103,7 @@ def general_serializer(data: Any, check_value=True) -> orm.Node:
|
|
|
105
103
|
new_node.store()
|
|
106
104
|
return new_node
|
|
107
105
|
except Exception:
|
|
108
|
-
|
|
109
|
-
try:
|
|
110
|
-
new_node = PickledData(data)
|
|
111
|
-
new_node.store()
|
|
112
|
-
return new_node
|
|
113
|
-
except Exception as e:
|
|
114
|
-
raise ValueError(f"Error in serializing {ep_key}: {e}")
|
|
106
|
+
raise ValueError(f"Error in storing data {ep_key}")
|
|
115
107
|
else:
|
|
116
108
|
# try to serialize the data as a PickledData
|
|
117
109
|
try:
|
|
@@ -1,39 +1,38 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
1
3
|
import inspect
|
|
2
4
|
import os
|
|
3
|
-
from typing import Any, Callable, Dict, Optional, Union
|
|
5
|
+
from typing import Any, Callable, Dict, List, Optional, Union
|
|
4
6
|
|
|
5
|
-
from aiida
|
|
7
|
+
from aiida import orm
|
|
6
8
|
|
|
7
|
-
from .data.pickled_function import PickledFunction
|
|
8
9
|
from .data.serializer import serialize_to_aiida_nodes
|
|
9
|
-
from .utils import get_or_create_code
|
|
10
|
+
from .utils import build_function_data, get_or_create_code
|
|
10
11
|
|
|
11
12
|
|
|
12
13
|
def prepare_pythonjob_inputs(
|
|
13
14
|
function: Optional[Callable[..., Any]] = None,
|
|
14
15
|
function_inputs: Optional[Dict[str, Any]] = None,
|
|
15
|
-
function_outputs: Optional[
|
|
16
|
-
code: Optional[AbstractCode] = None,
|
|
16
|
+
function_outputs: Optional[List[str | dict]] = None,
|
|
17
|
+
code: Optional[orm.AbstractCode] = None,
|
|
17
18
|
command_info: Optional[Dict[str, str]] = None,
|
|
18
|
-
computer: Union[str, Computer] = "localhost",
|
|
19
|
+
computer: Union[str, orm.Computer] = "localhost",
|
|
19
20
|
metadata: Optional[Dict[str, Any]] = None,
|
|
20
21
|
upload_files: Dict[str, str] = {},
|
|
21
22
|
process_label: Optional[str] = None,
|
|
22
|
-
|
|
23
|
+
function_data: dict | None = None,
|
|
23
24
|
**kwargs: Any,
|
|
24
25
|
) -> Dict[str, Any]:
|
|
25
26
|
pass
|
|
26
27
|
"""Prepare the inputs for PythonJob"""
|
|
27
28
|
|
|
28
|
-
if function is None and
|
|
29
|
-
raise ValueError("Either function or
|
|
30
|
-
if function is not None and
|
|
31
|
-
raise ValueError("Only one of function or
|
|
32
|
-
# if function is a function,
|
|
29
|
+
if function is None and function_data is None:
|
|
30
|
+
raise ValueError("Either function or function_data must be provided")
|
|
31
|
+
if function is not None and function_data is not None:
|
|
32
|
+
raise ValueError("Only one of function or function_data should be provided")
|
|
33
|
+
# if function is a function, inspect it and get the source code
|
|
33
34
|
if function is not None and inspect.isfunction(function):
|
|
34
|
-
|
|
35
|
-
if pickled_function is not None:
|
|
36
|
-
executor = pickled_function
|
|
35
|
+
function_data = build_function_data(function)
|
|
37
36
|
new_upload_files = {}
|
|
38
37
|
# change the string in the upload files to SingleFileData, or FolderData
|
|
39
38
|
for key, source in upload_files.items():
|
|
@@ -42,10 +41,12 @@ def prepare_pythonjob_inputs(
|
|
|
42
41
|
new_key = key.replace(".", "_dot_")
|
|
43
42
|
if isinstance(source, str):
|
|
44
43
|
if os.path.isfile(source):
|
|
45
|
-
new_upload_files[new_key] = SinglefileData(file=source)
|
|
44
|
+
new_upload_files[new_key] = orm.SinglefileData(file=source)
|
|
46
45
|
elif os.path.isdir(source):
|
|
47
|
-
new_upload_files[new_key] = FolderData(tree=source)
|
|
48
|
-
|
|
46
|
+
new_upload_files[new_key] = orm.FolderData(tree=source)
|
|
47
|
+
else:
|
|
48
|
+
raise ValueError(f"Invalid upload file path: {source}")
|
|
49
|
+
elif isinstance(source, (orm.SinglefileData, orm.FolderData)):
|
|
49
50
|
new_upload_files[new_key] = source
|
|
50
51
|
else:
|
|
51
52
|
raise ValueError(f"Invalid upload file type: {type(source)}, {source}")
|
|
@@ -54,11 +55,13 @@ def prepare_pythonjob_inputs(
|
|
|
54
55
|
command_info = command_info or {}
|
|
55
56
|
code = get_or_create_code(computer=computer, **command_info)
|
|
56
57
|
# get the source code of the function
|
|
57
|
-
function_name =
|
|
58
|
-
if
|
|
59
|
-
function_source_code =
|
|
58
|
+
function_name = function_data["name"]
|
|
59
|
+
if function_data.get("is_pickle", False):
|
|
60
|
+
function_source_code = (
|
|
61
|
+
function_data["import_statements"] + "\n" + function_data["source_code_without_decorator"]
|
|
62
|
+
)
|
|
60
63
|
else:
|
|
61
|
-
function_source_code = f"from {
|
|
64
|
+
function_source_code = f"from {function_data['module']} import {function_name}"
|
|
62
65
|
|
|
63
66
|
# serialize the kwargs into AiiDA Data
|
|
64
67
|
function_inputs = function_inputs or {}
|
|
@@ -66,12 +69,16 @@ def prepare_pythonjob_inputs(
|
|
|
66
69
|
# transfer the args to kwargs
|
|
67
70
|
inputs = {
|
|
68
71
|
"process_label": process_label or "PythonJob<{}>".format(function_name),
|
|
69
|
-
"
|
|
70
|
-
|
|
72
|
+
"function_data": orm.Dict(
|
|
73
|
+
{
|
|
74
|
+
"source_code": function_source_code,
|
|
75
|
+
"name": function_name,
|
|
76
|
+
"outputs": function_outputs or [],
|
|
77
|
+
}
|
|
78
|
+
),
|
|
71
79
|
"code": code,
|
|
72
80
|
"function_inputs": function_inputs,
|
|
73
81
|
"upload_files": new_upload_files,
|
|
74
|
-
"function_outputs": List(function_outputs),
|
|
75
82
|
"metadata": metadata or {},
|
|
76
83
|
**kwargs,
|
|
77
84
|
}
|
|
@@ -22,58 +22,64 @@ class PythonJobParser(Parser):
|
|
|
22
22
|
"""
|
|
23
23
|
import pickle
|
|
24
24
|
|
|
25
|
-
function_outputs = self.node.inputs.
|
|
25
|
+
function_outputs = self.node.inputs.function_data.get_dict()["outputs"]
|
|
26
26
|
if len(function_outputs) == 0:
|
|
27
27
|
function_outputs = [{"name": "result"}]
|
|
28
28
|
self.output_list = function_outputs
|
|
29
29
|
# first we remove nested outputs, e.g., "add_multiply.add"
|
|
30
30
|
top_level_output_list = [output for output in self.output_list if "." not in output["name"]]
|
|
31
|
-
exit_code = 0
|
|
32
31
|
try:
|
|
33
32
|
with self.retrieved.base.repository.open("results.pickle", "rb") as handle:
|
|
34
33
|
results = pickle.load(handle)
|
|
35
34
|
if isinstance(results, tuple):
|
|
36
35
|
if len(top_level_output_list) != len(results):
|
|
37
|
-
self.exit_codes.ERROR_RESULT_OUTPUT_MISMATCH
|
|
36
|
+
return self.exit_codes.ERROR_RESULT_OUTPUT_MISMATCH
|
|
38
37
|
for i in range(len(top_level_output_list)):
|
|
39
38
|
top_level_output_list[i]["value"] = self.serialize_output(results[i], top_level_output_list[i])
|
|
40
|
-
elif isinstance(results, dict)
|
|
39
|
+
elif isinstance(results, dict):
|
|
41
40
|
# pop the exit code if it exists
|
|
42
41
|
exit_code = results.pop("exit_code", 0)
|
|
43
|
-
|
|
44
|
-
if
|
|
42
|
+
if exit_code:
|
|
43
|
+
if isinstance(exit_code, dict):
|
|
44
|
+
exit_code = ExitCode(exit_code["status"], exit_code["message"])
|
|
45
|
+
elif isinstance(exit_code, int):
|
|
46
|
+
exit_code = ExitCode(exit_code)
|
|
47
|
+
return exit_code
|
|
48
|
+
if len(top_level_output_list) == 1:
|
|
49
|
+
# if output name in results, use it
|
|
50
|
+
if top_level_output_list[0]["name"] in results:
|
|
51
|
+
top_level_output_list[0]["value"] = self.serialize_output(
|
|
52
|
+
results.pop(top_level_output_list[0]["name"]),
|
|
53
|
+
top_level_output_list[0],
|
|
54
|
+
)
|
|
55
|
+
# if there are any remaining results, raise an warning
|
|
56
|
+
if len(results) > 0:
|
|
57
|
+
self.logger.warning(
|
|
58
|
+
f"Found extra results that are not included in the output: {results.keys()}"
|
|
59
|
+
)
|
|
60
|
+
# otherwise, we assume the results is the output
|
|
61
|
+
else:
|
|
62
|
+
top_level_output_list[0]["value"] = self.serialize_output(results, top_level_output_list[0])
|
|
63
|
+
elif len(top_level_output_list) > 1:
|
|
64
|
+
for output in top_level_output_list:
|
|
45
65
|
if output["name"] not in results:
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
if top_level_output_list[0]["name"] in results:
|
|
57
|
-
top_level_output_list[0]["value"] = self.serialize_output(
|
|
58
|
-
results[top_level_output_list[0]["name"]],
|
|
59
|
-
top_level_output_list[0],
|
|
60
|
-
)
|
|
61
|
-
# otherwise, we assume the results is the output
|
|
62
|
-
else:
|
|
63
|
-
top_level_output_list[0]["value"] = self.serialize_output(results, top_level_output_list[0])
|
|
66
|
+
if output.get("required", True):
|
|
67
|
+
return self.exit_codes.ERROR_MISSING_OUTPUT
|
|
68
|
+
else:
|
|
69
|
+
output["value"] = self.serialize_output(results.pop(output["name"]), output)
|
|
70
|
+
# if there are any remaining results, raise an warning
|
|
71
|
+
if len(results) > 0:
|
|
72
|
+
self.logger.warning(
|
|
73
|
+
f"Found extra results that are not included in the output: {results.keys()}"
|
|
74
|
+
)
|
|
75
|
+
|
|
64
76
|
elif len(top_level_output_list) == 1:
|
|
65
|
-
# otherwise, we assume the results is the output
|
|
77
|
+
# otherwise it returns a single value, we assume the results is the output
|
|
66
78
|
top_level_output_list[0]["value"] = self.serialize_output(results, top_level_output_list[0])
|
|
67
79
|
else:
|
|
68
|
-
|
|
80
|
+
return self.exit_codes.ERROR_RESULT_OUTPUT_MISMATCH
|
|
69
81
|
for output in top_level_output_list:
|
|
70
82
|
self.out(output["name"], output["value"])
|
|
71
|
-
if exit_code:
|
|
72
|
-
if isinstance(exit_code, dict):
|
|
73
|
-
exit_code = ExitCode(exit_code["status"], exit_code["message"])
|
|
74
|
-
elif isinstance(exit_code, int):
|
|
75
|
-
exit_code = ExitCode(exit_code)
|
|
76
|
-
return exit_code
|
|
77
83
|
except OSError:
|
|
78
84
|
return self.exit_codes.ERROR_READING_OUTPUT_FILE
|
|
79
85
|
except ValueError as exception:
|