aiida-pythonjob 0.3.4__tar.gz → 0.4.0__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 (43) hide show
  1. {aiida_pythonjob-0.3.4/src/aiida_pythonjob.egg-info → aiida_pythonjob-0.4.0}/PKG-INFO +1 -2
  2. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/pyproject.toml +0 -2
  3. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/__init__.py +1 -2
  4. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/calculations/common.py +0 -8
  5. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/calculations/pyfunction.py +0 -3
  6. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/data/deserializer.py +1 -1
  7. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/data/serializer.py +31 -35
  8. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/decorator.py +0 -2
  9. aiida_pythonjob-0.4.0/src/aiida_pythonjob/launch.py +262 -0
  10. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/parsers/pythonjob.py +0 -2
  11. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/parsers/utils.py +5 -9
  12. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/utils.py +5 -6
  13. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0/src/aiida_pythonjob.egg-info}/PKG-INFO +1 -2
  14. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob.egg-info/entry_points.txt +0 -1
  15. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob.egg-info/requires.txt +0 -1
  16. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/tests/test_data.py +2 -10
  17. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/tests/test_pythonjob.py +2 -2
  18. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/tests/test_serializer.py +0 -17
  19. aiida_pythonjob-0.3.4/src/aiida_pythonjob/launch.py +0 -178
  20. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/LICENSE +0 -0
  21. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/README.md +0 -0
  22. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/setup.cfg +0 -0
  23. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/calculations/__init__.py +0 -0
  24. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/calculations/pythonjob.py +0 -0
  25. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/calculations/utils.py +0 -0
  26. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/config.py +0 -0
  27. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/data/__init__.py +0 -0
  28. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/data/atoms.py +0 -0
  29. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/data/common_data.py +0 -0
  30. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/data/jsonable_data.py +0 -0
  31. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/data/pickled_data.py +0 -0
  32. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/data/utils.py +0 -0
  33. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob/parsers/__init__.py +0 -0
  34. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob.egg-info/SOURCES.txt +0 -0
  35. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob.egg-info/dependency_links.txt +0 -0
  36. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/src/aiida_pythonjob.egg-info/top_level.txt +0 -0
  37. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/tests/test_create_env.py +0 -0
  38. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/tests/test_entry_points.py +0 -0
  39. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/tests/test_jsonable_data.py +0 -0
  40. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/tests/test_parser.py +0 -0
  41. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/tests/test_pickled_data.py +0 -0
  42. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/tests/test_pyfunction.py +0 -0
  43. {aiida_pythonjob-0.3.4 → aiida_pythonjob-0.4.0}/tests/test_utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: aiida-pythonjob
3
- Version: 0.3.4
3
+ Version: 0.4.0
4
4
  Summary: Run Python functions on a remote computer.
5
5
  Author-email: Xing Wang <xingwang1991@gmail.com>
6
6
  License: MIT License
@@ -38,7 +38,6 @@ Description-Content-Type: text/markdown
38
38
  License-File: LICENSE
39
39
  Requires-Dist: aiida-core<3,>=2.3
40
40
  Requires-Dist: ase
41
- Requires-Dist: cloudpickle
42
41
  Requires-Dist: node-graph>=0.3.0
43
42
  Provides-Extra: test
44
43
  Requires-Dist: pgtest>=1.3.1,~=1.3; extra == "test"
@@ -22,7 +22,6 @@ requires-python = ">=3.9"
22
22
  dependencies = [
23
23
  "aiida-core>=2.3,<3",
24
24
  "ase",
25
- "cloudpickle",
26
25
  "node-graph>=0.3.0",
27
26
  ]
28
27
 
@@ -53,7 +52,6 @@ Source = "https://github.com/aiidateam/aiida-pythonjob"
53
52
 
54
53
  [project.entry-points."aiida.data"]
55
54
  "pythonjob.jsonable_data" = "aiida_pythonjob.data.jsonable_data:JsonableData"
56
- "pythonjob.pickled_data" = "aiida_pythonjob.data.pickled_data:PickledData"
57
55
  "pythonjob.ase.atoms.Atoms" = "aiida_pythonjob.data.atoms:AtomsData"
58
56
  "pythonjob.builtins.NoneType" = "aiida_pythonjob.data.common_data:NoneData"
59
57
  "pythonjob.builtins.int" = "aiida.orm.nodes.data.int:Int"
@@ -1,6 +1,6 @@
1
1
  """AiiDA plugin that run Python function on remote computers."""
2
2
 
3
- __version__ = "0.3.4"
3
+ __version__ = "0.4.0"
4
4
 
5
5
  from node_graph import socket_spec as spec
6
6
 
@@ -10,7 +10,6 @@ from .launch import prepare_pyfunction_inputs, prepare_pythonjob_inputs
10
10
  from .parsers import PythonJobParser
11
11
 
12
12
  __all__ = (
13
- "PickledData",
14
13
  "PyFunction",
15
14
  "PythonJob",
16
15
  "PythonJobParser",
@@ -10,7 +10,6 @@ from aiida_pythonjob.data.deserializer import deserialize_to_raw_python_data
10
10
 
11
11
  # Attribute keys stored on ProcessNode.base.attributes
12
12
  ATTR_OUTPUTS_SPEC = "outputs_spec"
13
- ATTR_USE_PICKLE = "use_pickle"
14
13
  ATTR_SERIALIZERS = "serializers"
15
14
  ATTR_DESERIALIZERS = "deserializers"
16
15
 
@@ -28,12 +27,6 @@ def add_common_function_io(spec) -> None:
28
27
  required=False,
29
28
  help="Specification for the outputs.",
30
29
  )
31
- spec.input(
32
- "metadata.use_pickle",
33
- valid_type=bool,
34
- required=False,
35
- help="Allow pickling of function inputs and outputs.",
36
- )
37
30
  spec.input("process_label", valid_type=Str, serializer=to_aiida_type, required=False)
38
31
 
39
32
  spec.input_namespace("function_inputs", valid_type=Data, required=False)
@@ -115,7 +108,6 @@ class FunctionProcessMixin:
115
108
  def _setup_metadata(self, metadata: dict) -> None: # type: ignore[override]
116
109
  """Store common metadata on the ProcessNode and forward the rest."""
117
110
  self.node.base.attributes.set(ATTR_OUTPUTS_SPEC, metadata.pop("outputs_spec", {}))
118
- self.node.base.attributes.set(ATTR_USE_PICKLE, metadata.pop("use_pickle", False))
119
111
  self.node.base.attributes.set(ATTR_SERIALIZERS, metadata.pop("serializers", {}))
120
112
  self.node.base.attributes.set(ATTR_DESERIALIZERS, metadata.pop("deserializers", {}))
121
113
  super()._setup_metadata(metadata)
@@ -17,7 +17,6 @@ from aiida_pythonjob.calculations.common import (
17
17
  ATTR_DESERIALIZERS,
18
18
  ATTR_OUTPUTS_SPEC,
19
19
  ATTR_SERIALIZERS,
20
- ATTR_USE_PICKLE,
21
20
  FunctionProcessMixin,
22
21
  add_common_function_io,
23
22
  )
@@ -115,7 +114,6 @@ class PyFunction(FunctionProcessMixin, Process):
115
114
 
116
115
  # Parse & attach outputs
117
116
  outputs_spec = SocketSpec.from_dict(self.node.base.attributes.get(ATTR_OUTPUTS_SPEC) or {})
118
- use_pickle = self.node.base.attributes.get(ATTR_USE_PICKLE, False)
119
117
  serializers = self.node.base.attributes.get(ATTR_SERIALIZERS, {})
120
118
  outputs, exit_code = parse_outputs(
121
119
  results,
@@ -123,7 +121,6 @@ class PyFunction(FunctionProcessMixin, Process):
123
121
  exit_codes=self.exit_codes,
124
122
  logger=self.logger,
125
123
  serializers=serializers,
126
- use_pickle=use_pickle,
127
124
  )
128
125
  if exit_code:
129
126
  return exit_code
@@ -61,7 +61,7 @@ def deserialize_to_raw_python_data(
61
61
  ) -> Any:
62
62
  """Deserialize the AiiDA data node to an raw Python data."""
63
63
 
64
- deserializers = deserializers or {}
64
+ deserializers = deserializers or all_deserializers
65
65
 
66
66
  if isinstance(data, orm.Data):
67
67
  if hasattr(data, "value"):
@@ -96,18 +96,12 @@ def general_serializer(
96
96
  data: Any,
97
97
  serializers: dict | None = None,
98
98
  store: bool = True,
99
- use_pickle: bool | None = None,
100
99
  ) -> orm.Node:
101
100
  """
102
101
  Attempt to serialize the data to an AiiDA data node based on the preference from `config`:
103
102
  1) AiiDA data only, 2) JSON-serializable, 3) fallback to PickledData (if allowed).
104
103
  """
105
- from aiida_pythonjob.config import config
106
-
107
- # Merge user-provided config with defaults
108
- allow_json = config.get("allow_json", True)
109
- if use_pickle is None:
110
- use_pickle = config.get("use_pickle", False)
104
+ serializers = serializers or all_serializers
111
105
 
112
106
  # 1) If it is already an AiiDA node, just return it
113
107
  if isinstance(data, orm.Data):
@@ -119,7 +113,6 @@ def general_serializer(
119
113
  # 3) check entry point
120
114
  data_type = type(data)
121
115
  ep_key = f"{data_type.__module__}.{data_type.__name__}"
122
- serializers = serializers or {}
123
116
  if ep_key in serializers:
124
117
  try:
125
118
  serializer = import_from_path(serializers[ep_key])
@@ -131,30 +124,33 @@ def general_serializer(
131
124
  error_traceback = traceback.format_exc()
132
125
  raise ValueError(f"Error in serializing {ep_key}: {error_traceback}")
133
126
 
134
- # check if we can JSON-serialize the data
135
- if allow_json:
136
- try:
137
- node = JsonableData(data)
138
- if store:
139
- node.store()
140
- return node
141
- except (TypeError, ValueError):
142
- # print(f"Error in JSON-serializing {type(data).__name__}")
143
- pass
144
-
145
- # fallback to pickling
146
- if use_pickle:
147
- from .pickled_data import PickledData
148
-
149
- try:
150
- new_node = PickledData(data)
151
- if store:
152
- new_node.store()
153
- return new_node
154
- except Exception as e:
155
- raise ValueError(f"Error in pickling {type(data).__name__}: {e}")
156
-
157
- raise ValueError(
158
- f"Cannot serialize type={type(data).__name__}. No suitable method found "
159
- f"(json_allowed={allow_json}, pickle_allowed={use_pickle})."
160
- )
127
+ try:
128
+ node = JsonableData(data)
129
+ if store:
130
+ node.store()
131
+ return node
132
+ except (TypeError, ValueError):
133
+ suggestions = [
134
+ "How to fix:",
135
+ "1) Register a type-specific AiiDA Data class as an `aiida.data` entry point "
136
+ "(recommended for domain objects).",
137
+ " Example in `pyproject.toml`:",
138
+ ' [project.entry-points."aiida.data"]',
139
+ f' myplugin.{ep_key} = "myplugin.data.mytype:MyTypeData"',
140
+ " where `MyTypeData` is a subclass of `aiida.orm.Data` that knows how to store your object.",
141
+ "",
142
+ "2) Or make the class JSON-serializable so `JsonableData` can handle it by implementing:",
143
+ " - `to_dict()` / `as_dict()` (any one) returning only JSON-friendly structures, and",
144
+ " - `from_dict(cls, dct)` / `fromdict(cls, dct)` to rebuild the object later.",
145
+ "",
146
+ "3) Or pass an ad-hoc serializer function via the `serializers` argument:",
147
+ f" general_serializer(obj, serializers={{'{ep_key}': 'my_pkg.mod:to_aiida_node'}})",
148
+ " where `to_aiida_node(obj)` returns an `aiida.orm.Data` instance.",
149
+ ]
150
+ raise ValueError(
151
+ (
152
+ "Cannot serialize the provided object.\n\n"
153
+ f"Type: {ep_key}\n"
154
+ f"Tried entry-point key: '{ep_key}' — not found in provided serializers.\n" + "\n".join(suggestions)
155
+ )
156
+ )
@@ -22,7 +22,6 @@ LOGGER = logging.getLogger(__name__)
22
22
  def pyfunction(
23
23
  inputs: t.Optional[SocketSpec | List[str]] = None,
24
24
  outputs: t.Optional[t.List[SocketSpec | List[str]]] = None,
25
- use_pickle: bool | None = None,
26
25
  ) -> t.Callable[[FunctionType], FunctionType]:
27
26
  """The base function decorator to create a FunctionProcess out of a normal python function.
28
27
 
@@ -83,7 +82,6 @@ def pyfunction(
83
82
  deserializers=deserializers,
84
83
  serializers=serializers,
85
84
  register_pickle_by_value=register_pickle_by_value,
86
- use_pickle=use_pickle,
87
85
  )
88
86
 
89
87
  process = PyFunction(inputs=process_inputs, runner=runner)
@@ -0,0 +1,262 @@
1
+ from __future__ import annotations
2
+
3
+ import inspect
4
+ import os
5
+ import types
6
+ from typing import Any, Callable, Dict, Optional, Tuple, Union
7
+
8
+ from aiida import orm
9
+ from node_graph.node_spec import BaseHandle
10
+ from node_graph.socket_spec import infer_specs_from_callable
11
+
12
+ from aiida_pythonjob.data.deserializer import all_deserializers
13
+ from aiida_pythonjob.data.serializer import all_serializers
14
+
15
+ from .utils import build_function_data, get_or_create_code, serialize_ports
16
+
17
+
18
+ def _unwrap_callable(func: Any) -> Callable[..., Any] | None:
19
+ """
20
+ Return a plain Python function from several supported wrappers.
21
+ Returns None if func is None. Raises for unsupported types.
22
+ """
23
+ if func is None:
24
+ return None
25
+ if isinstance(func, BaseHandle) and hasattr(func, "_func"):
26
+ return func._func
27
+ if getattr(func, "is_process_function", False):
28
+ # aiida process_function wrapper (e.g., calcfunction/workfunction)
29
+ return func.func
30
+ if inspect.isfunction(func):
31
+ return func
32
+ if isinstance(func, types.BuiltinFunctionType):
33
+ raise NotImplementedError("Built-in functions are not supported yet.")
34
+ raise ValueError(f"Invalid function type: {type(func)!r}")
35
+
36
+
37
+ def _validate_inputs_against_signature(func: Callable[..., Any], inputs: dict) -> None:
38
+ """Raise ValueError if inputs do not bind to func's signature."""
39
+ sig = inspect.signature(func)
40
+ try:
41
+ sig.bind(**inputs)
42
+ except TypeError as e:
43
+ raise ValueError(f"Invalid function inputs: {e}") from e
44
+
45
+
46
+ def _merge_registry(overrides: dict | None, base: dict) -> dict:
47
+ """Shallow-merge (user overrides win)."""
48
+ return {**base, **(overrides or {})}
49
+
50
+
51
+ def _normalize_upload_files(
52
+ upload_files: Dict[str, Union[str, orm.SinglefileData, orm.FolderData]] | None,
53
+ ) -> Dict[str, Union[orm.SinglefileData, orm.FolderData]]:
54
+ """
55
+ Convert string paths to AiiDA SinglefileData/FolderData and sanitize keys.
56
+ """
57
+ result: Dict[str, Union[orm.SinglefileData, orm.FolderData]] = {}
58
+ if not upload_files:
59
+ return result
60
+
61
+ for key, source in upload_files.items():
62
+ # Only alphanumeric + underscore in keys; also make dots explicit
63
+ new_key = key.replace(".", "_dot_")
64
+
65
+ if isinstance(source, str):
66
+ if os.path.isfile(source):
67
+ result[new_key] = orm.SinglefileData(file=source)
68
+ elif os.path.isdir(source):
69
+ result[new_key] = orm.FolderData(tree=source)
70
+ else:
71
+ raise ValueError(f"Invalid upload file path: {source!r}")
72
+ elif isinstance(source, (orm.SinglefileData, orm.FolderData)):
73
+ result[new_key] = source
74
+ else:
75
+ raise ValueError(f"Invalid upload file type: {type(source)}, value={source!r}")
76
+
77
+ return result
78
+
79
+
80
+ def _maybe_build_function_data(func: Callable[..., Any] | None, *, register_pickle_by_value: bool) -> dict | None:
81
+ """Build function_data if we have a Python function; else return None."""
82
+ if func is None:
83
+ return None
84
+ return build_function_data(func, register_pickle_by_value=register_pickle_by_value)
85
+
86
+
87
+ def _prepare_common(
88
+ *,
89
+ function: Optional[Callable[..., Any]],
90
+ function_data: Optional[dict],
91
+ function_inputs: Optional[Dict[str, Any]],
92
+ inputs_spec: Optional[type],
93
+ outputs_spec: Optional[type],
94
+ serializers: Optional[dict],
95
+ deserializers: Optional[dict],
96
+ register_pickle_by_value: bool,
97
+ validate_signature: bool,
98
+ ) -> Tuple[dict, dict, dict, dict]:
99
+ """
100
+ Shared logic used by both PyFunction and PythonJob preparations.
101
+
102
+ Returns:
103
+ (prepared_inputs, outputs_spec_dict, merged_serializers, merged_deserializers)
104
+ where prepared_inputs = {"function_data": ..., "function_inputs": ..., "metadata": {...}}
105
+ """
106
+ # Unwrap and normalize the function
107
+ fn = _unwrap_callable(function)
108
+
109
+ # Guard: either function or function_data must be present, but not both
110
+ if fn is None and function_data is None:
111
+ raise ValueError("Either `function` or `function_data` must be provided.")
112
+ if fn is not None and function_data is not None:
113
+ raise ValueError("Only one of `function` or `function_data` should be provided.")
114
+
115
+ # If we have a Python function, build function_data from source/pickle
116
+ if fn is not None:
117
+ function_data = _maybe_build_function_data(fn, register_pickle_by_value=register_pickle_by_value)
118
+
119
+ # Infer I/O specs
120
+ in_spec, out_spec = infer_specs_from_callable(fn, inputs=inputs_spec, outputs=outputs_spec)
121
+
122
+ # Merge serializer/deserializer registries (user wins)
123
+ merged_serializers = _merge_registry(serializers, all_serializers)
124
+ merged_deserializers = _merge_registry(deserializers, all_deserializers)
125
+
126
+ # Serialize inputs according to (possibly nested) input schema
127
+ py_inputs = function_inputs or {}
128
+ serialized_inputs = serialize_ports(
129
+ python_data=py_inputs,
130
+ port_schema=in_spec,
131
+ serializers=merged_serializers,
132
+ )
133
+
134
+ # Optional: validate against fn signature (bind) using the PROVIDED keys.
135
+ # Binding cares about names/arity, not the exact serialized types.
136
+ if validate_signature and fn is not None:
137
+ _validate_inputs_against_signature(fn, serialized_inputs)
138
+
139
+ metadata = {
140
+ "outputs_spec": out_spec.to_dict(),
141
+ "serializers": merged_serializers,
142
+ "deserializers": merged_deserializers,
143
+ }
144
+
145
+ prepared = {
146
+ "function_data": function_data,
147
+ "function_inputs": serialized_inputs,
148
+ "metadata": metadata,
149
+ }
150
+ return prepared, metadata["outputs_spec"], merged_serializers, merged_deserializers
151
+
152
+
153
+ def create_inputs(func: Callable[..., Any], *args: Any, **kwargs: Any) -> dict[str, Any]:
154
+ """
155
+ Create the input dictionary for calling a Python function by name-binding.
156
+ Positional args are mapped to positional parameters; **kwargs are merged on top.
157
+ Variable positional parameters (*args) are not supported.
158
+ """
159
+ inputs = dict(kwargs or {})
160
+ arguments = list(args)
161
+ for name, param in inspect.signature(func).parameters.items():
162
+ if param.kind in (param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD):
163
+ try:
164
+ inputs[name] = arguments.pop(0)
165
+ except IndexError:
166
+ pass
167
+ elif param.kind is param.VAR_POSITIONAL:
168
+ raise NotImplementedError("Variable positional arguments (*args) are not supported.")
169
+ return inputs
170
+
171
+
172
+ def prepare_pythonjob_inputs(
173
+ function: Optional[Callable[..., Any]] = None,
174
+ function_inputs: Optional[Dict[str, Any]] = None,
175
+ inputs_spec: Optional[type] = None,
176
+ outputs_spec: Optional[type] = None,
177
+ code: Optional[orm.AbstractCode] = None,
178
+ command_info: Optional[Dict[str, str]] = None,
179
+ computer: Union[str, orm.Computer] = "localhost",
180
+ metadata: Optional[Dict[str, Any]] = None,
181
+ upload_files: Optional[Dict[str, Union[str, orm.SinglefileData, orm.FolderData]]] = None,
182
+ process_label: Optional[str] = None,
183
+ function_data: dict | None = None,
184
+ deserializers: dict | None = None,
185
+ serializers: dict | None = None,
186
+ register_pickle_by_value: bool = False,
187
+ **kwargs: Any,
188
+ ) -> Dict[str, Any]:
189
+ """
190
+ Prepare the inputs for a PythonJob (runner that needs a Code and optional upload_files).
191
+ """
192
+ prepared, _, _, _ = _prepare_common(
193
+ function=function,
194
+ function_data=function_data,
195
+ function_inputs=function_inputs,
196
+ inputs_spec=inputs_spec,
197
+ outputs_spec=outputs_spec,
198
+ serializers=serializers,
199
+ deserializers=deserializers,
200
+ register_pickle_by_value=register_pickle_by_value,
201
+ validate_signature=(function is not None), # only when we actually got a function
202
+ )
203
+
204
+ # Files & Code specifics
205
+ new_upload_files = _normalize_upload_files(upload_files)
206
+ if code is None:
207
+ code = get_or_create_code(computer=computer, **(command_info or {}))
208
+
209
+ # Merge external metadata if provided
210
+ md = {**prepared["metadata"], **(metadata or {})}
211
+ prepared["metadata"] = md
212
+
213
+ inputs: Dict[str, Any] = {
214
+ **prepared,
215
+ "code": code,
216
+ "upload_files": new_upload_files,
217
+ **kwargs,
218
+ }
219
+ if process_label:
220
+ inputs["process_label"] = process_label
221
+ return inputs
222
+
223
+
224
+ def prepare_pyfunction_inputs(
225
+ function: Optional[Callable[..., Any]] = None,
226
+ function_inputs: Optional[Dict[str, Any]] = None,
227
+ inputs_spec: Optional[type] = None,
228
+ outputs_spec: Optional[type] = None,
229
+ metadata: Optional[Dict[str, Any]] = None,
230
+ process_label: Optional[str] = None,
231
+ function_data: dict | None = None,
232
+ deserializers: dict | None = None,
233
+ serializers: dict | None = None,
234
+ register_pickle_by_value: bool = False,
235
+ **kwargs: Any,
236
+ ) -> Dict[str, Any]:
237
+ """
238
+ Prepare the inputs for a local PyFunction (no Code/upload_files).
239
+ """
240
+ prepared, _, _, _ = _prepare_common(
241
+ function=function,
242
+ function_data=function_data,
243
+ function_inputs=function_inputs,
244
+ inputs_spec=inputs_spec,
245
+ outputs_spec=outputs_spec,
246
+ serializers=serializers,
247
+ deserializers=deserializers,
248
+ register_pickle_by_value=register_pickle_by_value,
249
+ validate_signature=False, # leave binding checks to the engine if desired
250
+ )
251
+
252
+ # Merge external metadata if provided
253
+ md = {**prepared["metadata"], **(metadata or {})}
254
+ prepared["metadata"] = md
255
+
256
+ inputs: Dict[str, Any] = {
257
+ **prepared,
258
+ **kwargs,
259
+ }
260
+ if process_label:
261
+ inputs["process_label"] = process_label
262
+ return inputs
@@ -25,7 +25,6 @@ class PythonJobParser(Parser):
25
25
 
26
26
  # Read outputs SocketSpec
27
27
  spec_dict = self.node.base.attributes.get("outputs_spec", {})
28
- use_pickle = self.node.base.attributes.get("use_pickle", False)
29
28
  self.outputs_spec = SocketSpec.from_dict(spec_dict)
30
29
 
31
30
  # load custom serializers
@@ -65,7 +64,6 @@ class PythonJobParser(Parser):
65
64
  exit_codes=self.exit_codes,
66
65
  logger=self.logger,
67
66
  serializers=self.serializers,
68
- use_pickle=use_pickle,
69
67
  )
70
68
  if exit_code:
71
69
  return exit_code
@@ -35,7 +35,6 @@ def parse_outputs(
35
35
  exit_codes,
36
36
  logger,
37
37
  serializers: Optional[Dict[str, str]] = None,
38
- use_pickle: bool = False,
39
38
  ) -> Tuple[Optional[Dict[str, Any]], Optional[ExitCode]]:
40
39
  """Validate & convert *results* according to *output_spec*.
41
40
 
@@ -60,7 +59,7 @@ def parse_outputs(
60
59
  for i, name in enumerate(names):
61
60
  child_spec = fields[name]
62
61
  val = results[i]
63
- outs[name] = serialize_ports(val, child_spec, serializers=serializers, use_pickle=use_pickle)
62
+ outs[name] = serialize_ports(val, child_spec, serializers=serializers)
64
63
  return outs, None
65
64
 
66
65
  # dict
@@ -85,21 +84,19 @@ def parse_outputs(
85
84
  ((only_name, only_spec),) = fields.items()
86
85
  # if user used the same key as port name, use that value;
87
86
  if only_name in results:
88
- outs[only_name] = serialize_ports(
89
- results.pop(only_name), only_spec, serializers=serializers, use_pickle=use_pickle
90
- )
87
+ outs[only_name] = serialize_ports(results.pop(only_name), only_spec, serializers=serializers)
91
88
  if results:
92
89
  logger.warning(f"Found extra results that are not included in the output: {list(results.keys())}")
93
90
  else:
94
91
  # else treat the entire dict as the value for that single port.
95
- outs[only_name] = serialize_ports(results, only_spec, serializers=serializers, use_pickle=use_pickle)
92
+ outs[only_name] = serialize_ports(results, only_spec, serializers=serializers)
96
93
  return outs, None
97
94
 
98
95
  # fixed fields
99
96
  for name, child_spec in fields.items():
100
97
  if name in remaining:
101
98
  value = remaining.pop(name)
102
- outs[name] = serialize_ports(value, child_spec, serializers=serializers, use_pickle=use_pickle)
99
+ outs[name] = serialize_ports(value, child_spec, serializers=serializers)
103
100
  else:
104
101
  # If the field is explicitly required -> invalid output
105
102
  required = getattr(child_spec.meta, "required", None)
@@ -115,7 +112,6 @@ def parse_outputs(
115
112
  value,
116
113
  item_spec or SocketSpec(identifier="node_graph.any"),
117
114
  serializers=serializers,
118
- use_pickle=use_pickle,
119
115
  )
120
116
  return outs, None
121
117
  # not dynamic -> leftovers are unexpected (warn but continue)
@@ -126,6 +122,6 @@ def parse_outputs(
126
122
  # single fixed output + non-dict/tuple scalar
127
123
  if len(fields) == 1 and not is_dyn:
128
124
  ((only_name, only_spec),) = fields.items()
129
- return {only_name: serialize_ports(results, only_spec, serializers=serializers, use_pickle=use_pickle)}, None
125
+ return {only_name: serialize_ports(results, only_spec, serializers=serializers)}, None
130
126
 
131
127
  return None, exit_codes.ERROR_RESULT_OUTPUT_MISMATCH
@@ -283,7 +283,6 @@ def serialize_ports(
283
283
  python_data: Any,
284
284
  port_schema: SocketSpec | Dict[str, Any],
285
285
  serializers: Optional[Dict[str, str]] = None,
286
- use_pickle: bool | None = None,
287
286
  ) -> Any:
288
287
  """Serialize raw Python data to AiiDA Data following a SocketSpec schema.
289
288
 
@@ -309,21 +308,21 @@ def serialize_ports(
309
308
  if key in fields:
310
309
  child_spec = fields[key]
311
310
  if child_spec.is_namespace():
312
- out[key] = serialize_ports(value, child_spec, serializers=serializers, use_pickle=use_pickle)
311
+ out[key] = serialize_ports(value, child_spec, serializers=serializers)
313
312
  else:
314
- out[key] = general_serializer(value, serializers=serializers, store=False, use_pickle=use_pickle)
313
+ out[key] = general_serializer(value, serializers=serializers, store=False)
315
314
  elif (is_dyn and item_spec is not None) or allow_extra:
316
315
  schema = item_spec if (is_dyn and item_spec is not None) else catch_schema
317
316
  if schema.is_namespace():
318
- out[key] = serialize_ports(value, schema, serializers=serializers, use_pickle=use_pickle)
317
+ out[key] = serialize_ports(value, schema, serializers=serializers)
319
318
  else:
320
- out[key] = general_serializer(value, serializers=serializers, store=False, use_pickle=use_pickle)
319
+ out[key] = general_serializer(value, serializers=serializers, store=False)
321
320
  else:
322
321
  raise ValueError(f"Unexpected key '{key}' for namespace '{name}' (not dynamic).")
323
322
  return out
324
323
 
325
324
  # Leaf
326
- return general_serializer(python_data, serializers=serializers, store=False, use_pickle=use_pickle)
325
+ return general_serializer(python_data, serializers=serializers, store=False)
327
326
 
328
327
 
329
328
  def deserialize_ports(
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: aiida-pythonjob
3
- Version: 0.3.4
3
+ Version: 0.4.0
4
4
  Summary: Run Python functions on a remote computer.
5
5
  Author-email: Xing Wang <xingwang1991@gmail.com>
6
6
  License: MIT License
@@ -38,7 +38,6 @@ Description-Content-Type: text/markdown
38
38
  License-File: LICENSE
39
39
  Requires-Dist: aiida-core<3,>=2.3
40
40
  Requires-Dist: ase
41
- Requires-Dist: cloudpickle
42
41
  Requires-Dist: node-graph>=0.3.0
43
42
  Provides-Extra: test
44
43
  Requires-Dist: pgtest>=1.3.1,~=1.3; extra == "test"
@@ -16,7 +16,6 @@ pythonjob.numpy.float32 = aiida.orm.nodes.data.float:Float
16
16
  pythonjob.numpy.float64 = aiida.orm.nodes.data.float:Float
17
17
  pythonjob.numpy.int64 = aiida.orm.nodes.data.int:Int
18
18
  pythonjob.numpy.ndarray = aiida.orm.nodes.data.array.array.ArrayData
19
- pythonjob.pickled_data = aiida_pythonjob.data.pickled_data:PickledData
20
19
 
21
20
  [aiida.parsers]
22
21
  pythonjob.pythonjob = aiida_pythonjob.parsers.pythonjob:PythonJobParser
@@ -1,6 +1,5 @@
1
1
  aiida-core<3,>=2.3
2
2
  ase
3
- cloudpickle
4
3
  node-graph>=0.3.0
5
4
 
6
5
  [docs]
@@ -30,22 +30,14 @@ def test_typing():
30
30
 
31
31
  def test_python_job():
32
32
  """Test a simple python node."""
33
- from aiida_pythonjob.config import config
34
- from aiida_pythonjob.data.pickled_data import PickledData
35
33
  from aiida_pythonjob.data.serializer import serialize_to_aiida_nodes
36
34
 
37
35
  inputs = {"a": 1, "b": 2.0, "c": set()}
38
36
  with pytest.raises(
39
37
  ValueError,
40
- match="Cannot serialize type=set. No suitable method found",
38
+ match="Cannot serialize the provided object.",
41
39
  ):
42
- new_inputs = serialize_to_aiida_nodes(inputs, serializers=all_serializers)
43
- # Allow pickling
44
- config["use_pickle"] = True
45
- new_inputs = serialize_to_aiida_nodes(inputs, serializers=all_serializers)
46
- assert isinstance(new_inputs["a"], aiida.orm.Int)
47
- assert isinstance(new_inputs["b"], aiida.orm.Float)
48
- assert isinstance(new_inputs["c"], PickledData)
40
+ serialize_to_aiida_nodes(inputs, serializers=all_serializers)
49
41
 
50
42
 
51
43
  def test_atoms_data():
@@ -14,11 +14,11 @@ def test_validate_inputs(fixture_localhost):
14
14
  def add(x, y):
15
15
  return x + y
16
16
 
17
- with pytest.raises(ValueError, match="Either function or function_data must be provided"):
17
+ with pytest.raises(ValueError, match="Either `function` or `function_data` must be provided."):
18
18
  prepare_pythonjob_inputs(
19
19
  function_inputs={"x": 1, "y": 2},
20
20
  )
21
- with pytest.raises(ValueError, match="Only one of function or function_data should be provided"):
21
+ with pytest.raises(ValueError, match="Only one of `function` or `function_data` should be provided."):
22
22
  prepare_pythonjob_inputs(
23
23
  function=add,
24
24
  function_data={"module_path": "math", "name": "sqrt", "is_pickle": False},
@@ -58,20 +58,3 @@ def test_serialize_json():
58
58
 
59
59
  serialized_data = general_serializer(data, serializers=all_serializers)
60
60
  assert isinstance(serialized_data, JsonableData)
61
-
62
-
63
- def test_serialize_pickle():
64
- from aiida_pythonjob.config import config
65
- from aiida_pythonjob.data.pickled_data import PickledData
66
- from aiida_pythonjob.data.serializer import general_serializer
67
-
68
- data = NonJsonableData("a", 1)
69
- config["use_pickle"] = False
70
- with pytest.raises(
71
- ValueError,
72
- match="Cannot serialize type=NonJsonableData. No suitable method found",
73
- ):
74
- general_serializer(data, serializers=all_serializers)
75
- config["use_pickle"] = True
76
- serialized_data = general_serializer(data, serializers=all_serializers)
77
- assert isinstance(serialized_data, PickledData)
@@ -1,178 +0,0 @@
1
- from __future__ import annotations
2
-
3
- import inspect
4
- import os
5
- from typing import Any, Callable, Dict, Optional, Union
6
-
7
- from aiida import orm
8
- from node_graph.node_spec import BaseHandle
9
- from node_graph.socket_spec import infer_specs_from_callable
10
-
11
- from aiida_pythonjob.data.deserializer import all_deserializers
12
- from aiida_pythonjob.data.serializer import all_serializers
13
-
14
- from .utils import build_function_data, get_or_create_code, serialize_ports
15
-
16
-
17
- def validate_inputs(func, inputs: dict):
18
- sig = inspect.signature(func)
19
-
20
- try:
21
- # Bind the provided inputs to the function's signature
22
- sig.bind(**inputs)
23
- except TypeError as e:
24
- return False, str(e)
25
-
26
- return True, "Inputs are valid."
27
-
28
-
29
- def prepare_pythonjob_inputs(
30
- function: Optional[Callable[..., Any]] = None,
31
- function_inputs: Optional[Dict[str, Any]] = None,
32
- inputs_spec: Optional[type] = None,
33
- outputs_spec: Optional[type] = None,
34
- code: Optional[orm.AbstractCode] = None,
35
- command_info: Optional[Dict[str, str]] = None,
36
- computer: Union[str, orm.Computer] = "localhost",
37
- metadata: Optional[Dict[str, Any]] = None,
38
- upload_files: Dict[str, str] = {},
39
- process_label: Optional[str] = None,
40
- function_data: dict | None = None,
41
- deserializers: dict | None = None,
42
- serializers: dict | None = None,
43
- register_pickle_by_value: bool = False,
44
- use_pickle: bool | None = None,
45
- **kwargs: Any,
46
- ) -> Dict[str, Any]:
47
- """Prepare the inputs for PythonJob"""
48
-
49
- if function is None and function_data is None:
50
- raise ValueError("Either function or function_data must be provided")
51
- if function is not None and function_data is not None:
52
- raise ValueError("Only one of function or function_data should be provided")
53
- if isinstance(function, BaseHandle):
54
- function = function._func
55
- # if function is a function, inspect it and get the source code
56
- if function is not None and inspect.isfunction(function):
57
- function_data = build_function_data(function, register_pickle_by_value=register_pickle_by_value)
58
- new_upload_files = {}
59
- # change the string in the upload files to SingleFileData, or FolderData
60
- for key, source in upload_files.items():
61
- # only alphanumeric and underscores are allowed in the key
62
- # replace all "." with "_dot_"
63
- new_key = key.replace(".", "_dot_")
64
- if isinstance(source, str):
65
- if os.path.isfile(source):
66
- new_upload_files[new_key] = orm.SinglefileData(file=source)
67
- elif os.path.isdir(source):
68
- new_upload_files[new_key] = orm.FolderData(tree=source)
69
- else:
70
- raise ValueError(f"Invalid upload file path: {source}")
71
- elif isinstance(source, (orm.SinglefileData, orm.FolderData)):
72
- new_upload_files[new_key] = source
73
- else:
74
- raise ValueError(f"Invalid upload file type: {type(source)}, {source}")
75
- if code is None:
76
- command_info = command_info or {}
77
- code = get_or_create_code(computer=computer, **command_info)
78
- in_spec, out_spec = infer_specs_from_callable(function, inputs=inputs_spec, outputs=outputs_spec)
79
- metadata = metadata or {}
80
- metadata["outputs_spec"] = out_spec.to_dict()
81
- # serialize kwargs against the (nested) input schema
82
- serializers = {**all_serializers, **(serializers or {})}
83
- deserializers = {**all_deserializers, **(deserializers or {})}
84
- function_inputs = function_inputs or {}
85
- function_inputs = serialize_ports(
86
- python_data=function_inputs, port_schema=in_spec, serializers=serializers, use_pickle=use_pickle
87
- )
88
- if function is not None:
89
- valid, msg = validate_inputs(function, function_inputs)
90
- if not valid:
91
- raise ValueError(f"Invalid function inputs: {msg}")
92
- metadata["serializers"] = serializers
93
- metadata["deserializers"] = deserializers
94
- inputs = {
95
- "function_data": function_data,
96
- "code": code,
97
- "function_inputs": function_inputs,
98
- "upload_files": new_upload_files,
99
- "metadata": metadata,
100
- **kwargs,
101
- }
102
- if process_label:
103
- inputs["process_label"] = process_label
104
- return inputs
105
-
106
-
107
- def create_inputs(func, *args: Any, **kwargs: Any) -> dict[str, Any]:
108
- """Create the input dictionary for the ``FunctionProcess``."""
109
- # The complete input dictionary consists of the keyword arguments...
110
- inputs = dict(kwargs or {})
111
- arguments = list(args)
112
- for name, parameter in inspect.signature(func).parameters.items():
113
- if parameter.kind in [parameter.POSITIONAL_ONLY, parameter.POSITIONAL_OR_KEYWORD]:
114
- try:
115
- inputs[name] = arguments.pop(0)
116
- except IndexError:
117
- pass
118
- elif parameter.kind is parameter.VAR_POSITIONAL:
119
- raise NotImplementedError("Variable positional arguments are not yet supported")
120
-
121
- return inputs
122
-
123
-
124
- def prepare_pyfunction_inputs(
125
- function: Optional[Callable[..., Any]] = None,
126
- function_inputs: Optional[Dict[str, Any]] = None,
127
- inputs_spec: Optional[type] = None,
128
- outputs_spec: Optional[type] = None,
129
- metadata: Optional[Dict[str, Any]] = None,
130
- process_label: Optional[str] = None,
131
- function_data: dict | None = None,
132
- deserializers: dict | None = None,
133
- serializers: dict | None = None,
134
- register_pickle_by_value: bool = False,
135
- use_pickle: bool | None = None,
136
- **kwargs: Any,
137
- ) -> Dict[str, Any]:
138
- """Prepare the inputs for PyFunction."""
139
- import types
140
-
141
- if function is None and function_data is None:
142
- raise ValueError("Either function or function_data must be provided")
143
- if function is not None and function_data is not None:
144
- raise ValueError("Only one of function or function_data should be provided")
145
- if isinstance(function, BaseHandle):
146
- function = function._func
147
- elif hasattr(function, "is_process_function") and function.is_process_function:
148
- function = function.func
149
- # if function is a function, inspect it and get the source code
150
- if function is not None:
151
- if inspect.isfunction(function):
152
- function_data = build_function_data(function, register_pickle_by_value=register_pickle_by_value)
153
- elif isinstance(function, types.BuiltinFunctionType):
154
- raise NotImplementedError("Built-in functions are not supported yet")
155
- else:
156
- raise ValueError("Invalid function type")
157
- # spec
158
- in_spec, out_spec = infer_specs_from_callable(function, inputs=inputs_spec, outputs=outputs_spec)
159
- metadata = metadata or {}
160
- metadata["outputs_spec"] = out_spec.to_dict()
161
- # serialize the kwargs into AiiDA Data
162
- serializers = {**all_serializers, **(serializers or {})}
163
- deserializers = {**all_deserializers, **(deserializers or {})}
164
- function_inputs = function_inputs or {}
165
- function_inputs = serialize_ports(
166
- python_data=function_inputs, port_schema=in_spec, serializers=serializers, use_pickle=use_pickle
167
- )
168
- metadata["serializers"] = serializers
169
- metadata["deserializers"] = deserializers
170
- inputs = {
171
- "function_data": function_data,
172
- "function_inputs": function_inputs,
173
- "metadata": metadata,
174
- **kwargs,
175
- }
176
- if process_label:
177
- inputs["process_label"] = process_label
178
- return inputs
File without changes