test

Author: adi611

pydra/engine/run_pickled_function.py

@@ -0,0 +1,14 @@
+import pickle
+import pydra
+import sys
+
+def run_pickled():
+    with open('/pydra/pydra/engine/my_function.pkl', 'rb') as file:
+        loaded_function = pickle.load(file)
+
+    result = loaded_function(rerun=False)
+    
+    print(f'Result: {result}')
+
+if __name__ == '__main__':
+    run_pickled()

pydra/engine/run_pickled_function_2.py

@@ -0,0 +1,18 @@
+import pickle
+import pydra
+import sys
+
+def run_pickled():
+    with open('/pydra/pydra/engine/my_function.pkl', 'rb') as file:
+        loaded_function = pickle.load(file)
+    with open('/pydra/pydra/engine/taskmain.pkl', 'rb') as file:
+        taskmain = pickle.load(file)
+    with open('/pydra/pydra/engine/ind.pkl', 'rb') as file:
+        ind = pickle.load(file)
+
+    result = loaded_function(taskmain, ind, rerun=False)
+
+    print(f'Result: {result}')    
+
+if __name__ == '__main__':
+    run_pickled()

pydra/engine/workers.py

@@ -892,10 +892,69 @@ def close(self):
         pass
 
 
+class PsijWorker(Worker):
+    def __init__(self, **kwargs):
+        """Initialize worker."""
+        try:
+            import psij
+        except ImportError:
+            logger.critical("Please install psij.")
+            raise
+        logger.debug("Initialize PsijWorker")
+
+    def run_el(self, interface, rerun=False, **kwargs):
+        """Run a task."""
+        return self.exec_psij(interface, rerun=rerun)
+
+    def make_spec(self, cmd=None, arg=None):
+        spec = self.psij.JobSpec()
+        spec.executable = cmd
+        spec.arguments = arg
+        spec.stdout_path = '/pydra/pydra/engine/demo.stdout'
+        spec.stderr_path = '/pydra/pydra/engine/demo.stderr'
+
+        return spec
+
+    def make_job(self, spec, attributes):
+        job = self.psij.Job()
+        job.spec = spec
+        return job
+
+    async def exec_psij(self, runnable, rerun=False):
+        import psij
+        import pickle
+        self.psij = psij
+        jex = psij.JobExecutor.get_instance('slurm')
+
+        if isinstance(runnable, TaskBase):
+            with open('/pydra/pydra/engine/my_function.pkl', 'wb') as file:
+                pickle.dump(runnable._run, file)
+            spec = self.make_spec("python3.9", ["/pydra/pydra/engine/run_pickled_function.py"])
+        else:  # it could be tuple that includes pickle files with tasks and inputs
+            ind, task_main_pkl, task_orig = runnable
+            with open('/pydra/pydra/engine/my_function.pkl', 'wb') as file:
+                pickle.dump(load_and_run, file)
+            with open('/pydra/pydra/engine/taskmain.pkl', 'wb') as file:
+                pickle.dump(task_main_pkl, file)
+            with open('/pydra/pydra/engine/ind.pkl', 'wb') as file:
+                pickle.dump(ind, file)
+            spec = self.make_spec("python3.9", ["/pydra/pydra/engine/run_pickled_function_2.py"])
+
+        job = self.make_job(spec, None)
+        jex.submit(job)
+        job.wait()
+        
+        return
+
+    def close(self):
+        """Finalize the internal pool of tasks."""
+        pass
+
 WORKERS = {
     "serial": SerialWorker,
     "cf": ConcurrentFuturesWorker,
     "slurm": SlurmWorker,
     "dask": DaskWorker,
     "sge": SGEWorker,
+    "psij": PsijWorker,
 }