aiida-pythonjob 0.4.0__tar.gz → 0.4.2__tar.gz

{aiida_pythonjob-0.4.0/src/aiida_pythonjob.egg-info → aiida_pythonjob-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiida-pythonjob
-Version: 0.4.0
+Version: 0.4.2
 Summary: Run Python functions on a remote computer.
 Author-email: Xing Wang <xingwang1991@gmail.com>
 License: MIT License
@@ -38,7 +38,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: aiida-core<3,>=2.3
 Requires-Dist: ase
-Requires-Dist: node-graph>=0.3.0
+Requires-Dist: node-graph>=0.3.9
 Provides-Extra: test
 Requires-Dist: pgtest>=1.3.1,~=1.3; extra == "test"
 Requires-Dist: coverage~=7.0; extra == "test"

{aiida_pythonjob-0.4.0 → aiida_pythonjob-0.4.2}/pyproject.toml

@@ -22,7 +22,7 @@ requires-python = ">=3.9"
 dependencies = [
     "aiida-core>=2.3,<3",
     "ase",
-    "node-graph>=0.3.0",
+    "node-graph>=0.3.9",
 ]
 
 [project.optional-dependencies]

{aiida_pythonjob-0.4.0 → aiida_pythonjob-0.4.2}/src/aiida_pythonjob/__init__.py

@@ -1,6 +1,6 @@
 """AiiDA plugin that run Python function on remote computers."""
 
-__version__ = "0.4.0"
+__version__ = "0.4.2"
 
 from node_graph import socket_spec as spec
 

{aiida_pythonjob-0.4.0 → aiida_pythonjob-0.4.2}/src/aiida_pythonjob/calculations/pythonjob.py

@@ -143,6 +143,12 @@ class PythonJob(FunctionProcessMixin, CalcJob):
             invalidates_cache=True,
             message="The script failed for an unknown reason.\n{exception}\n{traceback}",
         )
+        spec.exit_code(
+            329,
+            "ERROR_IMPORT_MPI4PY_FAILED",
+            invalidates_cache=True,
+            message="Trying to run with MPI support, but importing mpi4py failed.\n{exception}\n{traceback}",
+        )
 
     @override
     def _setup_db_record(self) -> None:
@@ -242,6 +248,7 @@ class PythonJob(FunctionProcessMixin, CalcJob):
             pickled_function=pickled_function,
             source_code=source_code,
             function_name=function_name,
+            withmpi=self.inputs.metadata.options.get("withmpi", False),
         )
 
         # Write the script
@@ -283,7 +290,10 @@ class PythonJob(FunctionProcessMixin, CalcJob):
         local_copy_list.append((file_data.uuid, file_data.filename, filename))
 
         codeinfo = CodeInfo()
-        codeinfo.stdin_name = self.options.input_filename
+        if self.options.get("withmpi", False):
+            codeinfo.cmdline_params = [self.options.input_filename]
+        else:
+            codeinfo.stdin_name = self.options.input_filename
         codeinfo.stdout_name = self.options.output_filename
         codeinfo.code_uuid = self.inputs.code.uuid
 

{aiida_pythonjob-0.4.0 → aiida_pythonjob-0.4.2}/src/aiida_pythonjob/calculations/utils.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 
 def generate_script_py(
-    pickled_function: bytes | None, source_code: str | None, function_name: str = "user_function"
+    pickled_function: bytes | None, source_code: str | None, function_name: str = "user_function", withmpi: bool = False
 ) -> str:
     """
     Generate the script.py content as a single string with robust exception handling.
@@ -36,6 +36,23 @@ def generate_script_py(
         "        write_error_file('IMPORT_CLOUDPICKLE_FAILED', e, traceback.format_exc())",
         "        sys.exit(1)",
         "",
+    ]
+
+    if withmpi:
+        script_lines += [
+            "    # Attempt to import mpi4py",
+            "    try:",
+            "        from mpi4py import MPI",
+            "    except ImportError as e:",
+            "        write_error_file('IMPORT_MPI4PY_FAILED', e, traceback.format_exc())",
+            "        sys.exit(1)",
+            "",
+            "    # MPI initialization",
+            "    RANK = MPI.COMM_WORLD.Get_rank()",
+            "",
+        ]
+
+    script_lines += [
         "    # 2) Attempt to unpickle the inputs",
         "    try:",
         "        with open('inputs.pickle', 'rb') as handle:",
@@ -87,8 +104,19 @@ def generate_script_py(
         "",
         "    # 5) Attempt to pickle the result",
         "    try:",
-        "        with open('results.pickle', 'wb') as handle:",
-        "            pickle.dump(result, handle)",
+    ]
+    if withmpi:
+        script_lines += [
+            "        if RANK == 0:",  # Only the root process saves the result
+            "            with open('results.pickle', 'wb') as handle:",
+            "                pickle.dump(result, handle)",
+        ]
+    else:
+        script_lines += [
+            "        with open('results.pickle', 'wb') as handle:",
+            "            pickle.dump(result, handle)",
+        ]
+    script_lines += [
         "    except Exception as e:",
         "        write_error_file('PICKLE_RESULTS_FAILED', e, traceback.format_exc())",
         "        sys.exit(1)",

{aiida_pythonjob-0.4.0 → aiida_pythonjob-0.4.2}/src/aiida_pythonjob/parsers/pythonjob.py

@@ -10,6 +10,7 @@ from .utils import parse_outputs
 # Map error_type from script.py to exit code label
 ERROR_TYPE_TO_EXIT_CODE = {
     "IMPORT_CLOUDPICKLE_FAILED": "ERROR_IMPORT_CLOUDPICKLE_FAILED",
+    "IMPORT_MPI4PY_FAILED": "ERROR_IMPORT_MPI4PY_FAILED",
     "UNPICKLE_INPUTS_FAILED": "ERROR_UNPICKLE_INPUTS_FAILED",
     "UNPICKLE_FUNCTION_FAILED": "ERROR_UNPICKLE_FUNCTION_FAILED",
     "FUNCTION_EXECUTION_FAILED": "ERROR_FUNCTION_EXECUTION_FAILED",

{aiida_pythonjob-0.4.0 → aiida_pythonjob-0.4.2}/src/aiida_pythonjob/utils.py

@@ -307,7 +307,10 @@ def serialize_ports(
         for key, value in python_data.items():
             if key in fields:
                 child_spec = fields[key]
-                if child_spec.is_namespace():
+                # Metadata is not serialized
+                if child_spec.meta.is_metadata:
+                    out[key] = value
+                elif child_spec.is_namespace():
                     out[key] = serialize_ports(value, child_spec, serializers=serializers)
                 else:
                     out[key] = general_serializer(value, serializers=serializers, store=False)

{aiida_pythonjob-0.4.0 → aiida_pythonjob-0.4.2/src/aiida_pythonjob.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiida-pythonjob
-Version: 0.4.0
+Version: 0.4.2
 Summary: Run Python functions on a remote computer.
 Author-email: Xing Wang <xingwang1991@gmail.com>
 License: MIT License
@@ -38,7 +38,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: aiida-core<3,>=2.3
 Requires-Dist: ase
-Requires-Dist: node-graph>=0.3.0
+Requires-Dist: node-graph>=0.3.9
 Provides-Extra: test
 Requires-Dist: pgtest>=1.3.1,~=1.3; extra == "test"
 Requires-Dist: coverage~=7.0; extra == "test"

{aiida_pythonjob-0.4.0 → aiida_pythonjob-0.4.2}/src/aiida_pythonjob.egg-info/requires.txt

@@ -1,6 +1,6 @@
 aiida-core<3,>=2.3
 ase
-node-graph>=0.3.0
+node-graph>=0.3.9
 
 [docs]
 sphinx_rtd_theme