improve psijworker performance - remove unnecessary pickle dumps

Author: adi611

pydra/engine/run_pickled.py

@@ -1,6 +1,6 @@
 import pickle
 import sys
-
+from pydra.engine.helpers import load_and_run
 
 def run_pickled(*file_paths, rerun=False):
     loaded_objects = []
@@ -11,8 +11,8 @@ def run_pickled(*file_paths, rerun=False):
 
     if len(loaded_objects) == 1:
         result = loaded_objects[0](rerun=rerun)
-    elif len(loaded_objects) == 3:
-        result = loaded_objects[0](loaded_objects[1], loaded_objects[2], rerun=rerun)
+    elif len(loaded_objects) == 2:
+        result = load_and_run(loaded_objects[0], loaded_objects[1], rerun=rerun)
     else:
         raise ValueError("Unsupported number of loaded objects")
 

pydra/engine/workers.py

@@ -995,15 +995,12 @@ async def exec_psij(self, runnable, rerun=False):
             spec = self.make_spec("python", [func_path, file_path])
         else:  # it could be tuple that includes pickle files with tasks and inputs
             cache_dir = runnable[-1].cache_dir
-            file_path_1 = cache_dir / "runnable_function.pkl"
-            file_path_2 = cache_dir / "taskmain.pkl"
-            file_path_3 = cache_dir / "ind.pkl"
+            file_path_1 = cache_dir / "taskmain.pkl"
+            file_path_2 = cache_dir / "ind.pkl"
             ind, task_main_pkl, task_orig = runnable
             with open(file_path_1, "wb") as file:
-                pickle.dump(load_and_run, file)
-            with open(file_path_2, "wb") as file:
                 pickle.dump(task_main_pkl, file)
-            with open(file_path_3, "wb") as file:
+            with open(file_path_2, "wb") as file:
                 pickle.dump(ind, file)
             func_path = absolute_path / "run_pickled.py"
             spec = self.make_spec(
@@ -1012,7 +1009,6 @@ async def exec_psij(self, runnable, rerun=False):
                     func_path,
                     file_path_1,
                     file_path_2,
-                    file_path_3,
                 ],
             )