Merge pull request #1674 from pyiron/debug_decorator

Debug decorator - fix delayed output for python functions
pyiron · Oct 22, 2024 · 0163935 · 0163935
2 parents b311b18 + 2ffe25f
commit 0163935
Show file tree

Hide file tree

Showing 4 changed files with 214 additions and 120 deletions.
diff --git a/pyiron_base/project/decorator.py b/pyiron_base/project/decorator.py
@@ -18,6 +18,7 @@ def pyiron_job(
     new_hdf: bool = True,
     output_file_lst: list = [],
     output_key_lst: list = [],
+    list_length: Optional[int] = None,
 ):
     """
     Decorator to create a pyiron job object from any python function
@@ -76,6 +77,7 @@ def get_delayed_object(
             delayed=True,
             output_file_lst=output_file_lst,
             output_key_lst=output_key_lst,
+            list_length=list_length,
             **kwargs,
         )
         delayed_job_object._server = Server(**pyiron_resource_dict)
@@ -87,13 +89,13 @@ def function(
             *args, pyiron_project: Project, pyiron_resource_dict: dict = {}, **kwargs
         ):
             resource_default_dict = {
-                "host": None,
-                "queue": None,
-                "cores": 1,
-                "threads": 1,
-                "gpus": None,
-                "run_mode": "modal",
-                "new_hdf": True,
+                "host": host,
+                "queue": queue,
+                "cores": cores,
+                "threads": threads,
+                "gpus": gpus,
+                "run_mode": run_mode,
+                "new_hdf": new_hdf,
             }
             return get_delayed_object(
                 *args,

diff --git a/pyiron_base/project/delayed.py b/pyiron_base/project/delayed.py
@@ -258,7 +258,10 @@ def draw(self):
         draw(node_dict=node_dict, edge_lst=edge_lst)
 
     def get_python_result(self):
-        return getattr(self._result.output, self._output_key)
+        if isinstance(self._result, dict):
+            return self._result[self._output_key]
+        else:
+            return getattr(self._result.output, self._output_key)
 
     def get_file_result(self):
         return getattr(self._result.files, self._output_file)

diff --git a/tests/unit/flex/test_pythonfunctioncontainer.py b/tests/unit/flex/test_pythonfunctioncontainer.py
@@ -24,6 +24,10 @@ def my_function_exe(a_lst, b_lst, executor):
     return [future.result() for future in future_lst]
 
 
+def function_with_dict(a, b=1, c=3):
+    return {"a": a, "b": b, "c": c}
+
+
 class TestPythonFunctionContainer(TestWithProject):
     def test_as_job(self):
         job = self.project.wrap_python_function(my_function)
@@ -238,3 +242,22 @@ def test_delayed(self):
         nodes_dict, edges_lst = d.get_graph()
         self.assertEqual(len(nodes_dict), 6)
         self.assertEqual(len(edges_lst), 6)
+
+    def test_delayed_dict(self):
+        job_1 = self.project.wrap_python_function(
+            python_function=function_with_dict,
+            a=6,
+            b=4,
+            c=3,
+            delayed=True,
+            output_key_lst=["a", "b", "c"],
+        )
+        job_2 = self.project.wrap_python_function(
+            python_function=function_with_dict,
+            a=job_1.output.a,
+            b=5,
+            c=job_1.output.c,
+            delayed=True,
+            output_key_lst=["a", "b", "c"],
+        )
+        self.assertEqual(job_2.pull(), {"a": 6, "b": 5, "c": 3})