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.
Files changed (50) hide show
  1. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/PKG-INFO +10 -2
  2. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/README.md +8 -1
  3. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/gallery/autogen/how_to.py +57 -10
  4. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/pyproject.toml +9 -1
  5. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/__init__.py +1 -3
  6. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/calculations/pythonjob.py +4 -40
  7. aiida_pythonjob-0.1.4/src/aiida_pythonjob/data/__init__.py +4 -0
  8. aiida_pythonjob-0.1.4/src/aiida_pythonjob/data/atoms.py +53 -0
  9. aiida_pythonjob-0.1.4/src/aiida_pythonjob/data/data_with_value.py +13 -0
  10. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/data/pickled_data.py +0 -12
  11. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/data/serializer.py +1 -9
  12. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/launch.py +33 -26
  13. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/parsers/pythonjob.py +38 -32
  14. aiida_pythonjob-0.1.4/src/aiida_pythonjob/utils.py +255 -0
  15. aiida_pythonjob-0.1.4/tests/test_create_env.py +104 -0
  16. aiida_pythonjob-0.1.4/tests/test_data.py +63 -0
  17. aiida_pythonjob-0.1.4/tests/test_entry_points.py +61 -0
  18. aiida_pythonjob-0.1.4/tests/test_parser.py +112 -0
  19. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/tests/test_pythonjob.py +102 -30
  20. aiida_pythonjob-0.1.4/tests/test_utils.py +13 -0
  21. aiida_pythonjob-0.1.2/src/aiida_pythonjob/data/__init__.py +0 -4
  22. aiida_pythonjob-0.1.2/src/aiida_pythonjob/data/pickled_function.py +0 -145
  23. aiida_pythonjob-0.1.2/src/aiida_pythonjob/utils.py +0 -31
  24. aiida_pythonjob-0.1.2/tests/input.txt +0 -1
  25. aiida_pythonjob-0.1.2/tests/test_data.py +0 -23
  26. aiida_pythonjob-0.1.2/tests/test_parsers.py +0 -0
  27. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/.github/workflows/ci.yml +0 -0
  28. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/.github/workflows/python-publish.yml +0 -0
  29. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/.gitignore +0 -0
  30. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/.pre-commit-config.yaml +0 -0
  31. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/.readthedocs.yml +0 -0
  32. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/LICENSE +0 -0
  33. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/Makefile +0 -0
  34. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/environment.yml +0 -0
  35. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/gallery/autogen/GALLERY_HEADER.rst +0 -0
  36. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/make.bat +0 -0
  37. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/requirements.txt +0 -0
  38. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/conf.py +0 -0
  39. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/index.rst +0 -0
  40. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/installation.rst +0 -0
  41. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/tutorial/dft.ipynb +0 -0
  42. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/tutorial/html/atomization_energy.html +0 -0
  43. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/tutorial/html/pythonjob_eos_emt.html +0 -0
  44. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/docs/source/tutorial/index.rst +0 -0
  45. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/examples/test_add.py +0 -0
  46. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/calculations/__init__.py +0 -0
  47. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/config.py +0 -0
  48. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/src/aiida_pythonjob/parsers/__init__.py +0 -0
  49. {aiida_pythonjob-0.1.2 → aiida_pythonjob-0.1.4}/tests/conftest.py +0 -0
  50. {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.2
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
  [![codecov](https://codecov.io/gh/aiidateam/aiida-pythonjob/branch/main/graph/badge.svg)](https://codecov.io/gh/aiidateam/aiida-pythonjob)
56
57
  [![Docs status](https://readthedocs.org/projects/aiida-pythonjob/badge)](http://aiida-pythonjob.readthedocs.io/)
57
58
 
58
- Efficiently design and manage flexible workflows with AiiDA, featuring an interactive GUI, checkpoints, provenance tracking, error-resistant, and remote execution capabilities.
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
  [![codecov](https://codecov.io/gh/aiidateam/aiida-pythonjob/branch/main/graph/badge.svg)](https://codecov.io/gh/aiidateam/aiida-pythonjob)
5
5
  [![Docs status](https://readthedocs.org/projects/aiida-pythonjob/badge)](http://aiida-pythonjob.readthedocs.io/)
6
6
 
7
- Efficiently design and manage flexible workflows with AiiDA, featuring an interactive GUI, checkpoints, provenance tracking, error-resistant, and remote execution capabilities.
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
- # Introduction
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.engine import run_get_node # noqa: E402
28
- from aiida_pythonjob import PythonJob, prepare_pythonjob_inputs # noqa: E402
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.pickled_function" = "aiida_pythonjob.data.pickled_function:PickledFunction"
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.2"
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
  )
@@ -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.get_function_data()
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.get_function_data()
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,4 @@
1
+ from .pickled_data import PickledData
2
+ from .serializer import general_serializer, serialize_to_aiida_nodes
3
+
4
+ __all__ = ("PickledData", "serialize_to_aiida_nodes", "general_serializer")
@@ -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)
@@ -0,0 +1,13 @@
1
+ from aiida import orm
2
+
3
+
4
+ class Dict(orm.Dict):
5
+ @property
6
+ def value(self):
7
+ return self.get_dict()
8
+
9
+
10
+ class List(orm.List):
11
+ @property
12
+ def value(self):
13
+ return self.get_list()
@@ -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
- # try to serialize the value as a PickledData
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.orm import AbstractCode, Computer, FolderData, List, SinglefileData, Str
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[Dict[str, Any]] = None,
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
- pickled_function: Optional[PickledFunction] = None,
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 pickled_function is None:
29
- raise ValueError("Either function or pickled_function must be provided")
30
- if function is not None and pickled_function is not None:
31
- raise ValueError("Only one of function or pickled_function should be provided")
32
- # if function is a function, convert it to a PickledFunction
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
- executor = PickledFunction.build_callable(function)
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
- elif isinstance(source, (SinglefileData, FolderData)):
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 = executor["name"]
58
- if executor.get("is_pickle", False):
59
- function_source_code = executor["import_statements"] + "\n" + executor["source_code_without_decorator"]
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 {executor['module']} import {function_name}"
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
- "function_source_code": Str(function_source_code),
70
- "function_name": Str(function_name),
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.function_outputs.get_list()
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) and len(top_level_output_list) > 1:
39
+ elif isinstance(results, dict):
41
40
  # pop the exit code if it exists
42
41
  exit_code = results.pop("exit_code", 0)
43
- for output in top_level_output_list:
44
- if output.get("required", False):
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
- self.exit_codes.ERROR_MISSING_OUTPUT
47
- output["value"] = self.serialize_output(results.pop(output["name"]), output)
48
- # if there are any remaining results, raise an warning
49
- if results:
50
- self.logger.warning(
51
- f"Found extra results that are not included in the output: {results.keys()}"
52
- )
53
- elif isinstance(results, dict) and len(top_level_output_list) == 1:
54
- exit_code = results.pop("exit_code", 0)
55
- # if output name in results, use it
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
- raise ValueError("The number of results does not match the number of outputs.")
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: