fix: distributing tests with pytest -n auto

Author: adi611

pydra/engine/run_pickled_function.py

@@ -3,7 +3,8 @@
 import sys
 
 def run_pickled():
-    with open('/pydra/pydra/engine/my_function.pkl', 'rb') as file:
+    file_path = sys.argv[1]
+    with open(file_path, 'rb') as file:
         loaded_function = pickle.load(file)
 
     result = loaded_function(rerun=False)

pydra/engine/run_pickled_function_2.py

@@ -3,11 +3,14 @@
 import sys
 
 def run_pickled():
-    with open('/pydra/pydra/engine/my_function.pkl', 'rb') as file:
+    file_path_1 = sys.argv[1]
+    file_path_2 = sys.argv[2]
+    file_path_3 = sys.argv[3]
+    with open(file_path_1, 'rb') as file:
         loaded_function = pickle.load(file)
-    with open('/pydra/pydra/engine/taskmain.pkl', 'rb') as file:
+    with open(file_path_2, 'rb') as file:
         taskmain = pickle.load(file)
-    with open('/pydra/pydra/engine/ind.pkl', 'rb') as file:
+    with open(file_path_3, 'rb') as file:
         ind = pickle.load(file)
 
     result = loaded_function(taskmain, ind, rerun=False)

pydra/engine/workers.py

@@ -923,22 +923,29 @@ def make_job(self, spec, attributes):
     async def exec_psij(self, runnable, rerun=False):
         import psij
         import pickle
+        import os
         self.psij = psij
         jex = psij.JobExecutor.get_instance('slurm')
-
+        
         if isinstance(runnable, TaskBase):
-            with open('/pydra/pydra/engine/my_function.pkl', 'wb') as file:
+            cache_dir = runnable.cache_dir
+            file_path = os.path.join(cache_dir, 'my_function.pkl')
+            with open(file_path, 'wb') as file:
                 pickle.dump(runnable._run, file)
-            spec = self.make_spec("python3.9", ["/pydra/pydra/engine/run_pickled_function.py"])
+            spec = self.make_spec("python3.9", ["/pydra/pydra/engine/run_pickled_function.py", file_path])
         else:  # it could be tuple that includes pickle files with tasks and inputs
+            cache_dir = runnable[-1].cache_dir
+            file_path_1 = os.path.join(cache_dir, 'my_function.pkl')
+            file_path_2 = os.path.join(cache_dir, 'taskmain.pkl')
+            file_path_3 = os.path.join(cache_dir, 'ind.pkl')
             ind, task_main_pkl, task_orig = runnable
-            with open('/pydra/pydra/engine/my_function.pkl', 'wb') as file:
+            with open(file_path_1, 'wb') as file:
                 pickle.dump(load_and_run, file)
-            with open('/pydra/pydra/engine/taskmain.pkl', 'wb') as file:
+            with open(file_path_2, 'wb') as file:
                 pickle.dump(task_main_pkl, file)
-            with open('/pydra/pydra/engine/ind.pkl', 'wb') as file:
+            with open(file_path_3, 'wb') as file:
                 pickle.dump(ind, file)
-            spec = self.make_spec("python3.9", ["/pydra/pydra/engine/run_pickled_function_2.py"])
+            spec = self.make_spec("python3.9", ["/pydra/pydra/engine/run_pickled_function_2.py", file_path_1, file_path_2, file_path_3])
 
         job = self.make_job(spec, None)
         jex.submit(job)