Add modules to jflow components

src/jflow/component.py

@@ -52,6 +52,7 @@ class Component(object):
         if isinstance(self.version, bytes):
             self.version = self.version.decode()
         self.batch_options = self.config_reader.get_component_batch_options(self.__class__.__name__)
+        self.modules = self.config_reader.get_component_modules(self.__class__.__name__)
         # in case of SGE, parse the cpu and memory parameter
         self.__cpu=None
         self.__memory=None
@@ -421,7 +422,7 @@ class Component(object):
                     for e in file :
                         commandline += ' %s $%s ' % (file.cmd_format, cpt)
                         cpt+=1
-            function = ShellFunction( commandline,  cmd_format='{EXE} {IN} {OUT}')
+            function = ShellFunction( commandline,  cmd_format='{EXE} {IN} {OUT}', modules=self.modules)
             function(inputs=inputs, outputs=outputs) 
         # weaver map abstraction
         elif abstraction == 'map' :
@@ -438,7 +439,7 @@ class Component(object):
                 if isinstance(file, ParameterList) :
                     outputs = file 
             
-            function = ShellFunction( commandline,  cmd_format='{EXE} {IN} {OUT}')
+            function = ShellFunction( commandline,  cmd_format='{EXE} {IN} {OUT}', modules=self.modules)
             exe = Map(function, inputs=inputs, outputs=outputs)
         
         # jflow multimap
@@ -450,7 +451,7 @@ class Component(object):
                 commandline += ' %s $%s ' % (file.cmd_format, cpt)
                 cpt+=1
             
-            function  = ShellFunction( commandline,  cmd_format='{EXE} {IN} {OUT}')
+            function  = ShellFunction( commandline,  cmd_format='{EXE} {IN} {OUT}', modules=self.modules)
             exe = MultiMap(function, inputs=inputs, outputs=outputs)
         # anything other than that will be considered errored
         else :
@@ -583,7 +584,7 @@ class Component(object):
                 cmd_format += " {IN}"
             if (isinstance(outputs, list) and len(outputs)>0) or (outputs is not None and outputs != []):
                 cmd_format += " {OUT}"
-        py_function = PythonFunction(function, add_path=add_path, cmd_format=cmd_format)
+        py_function = PythonFunction(function, add_path=add_path, cmd_format=cmd_format, modules=self.modules)
 
         
         new_inputs,includes_in = self.__generate_iolist(inputs, map)
@@ -600,7 +601,7 @@ class Component(object):
         
     def add_shell_execution(self, source, inputs=[], outputs=[], arguments=[], includes=[], 
                             cmd_format=None, map=False, shell=None, collect=False, local=False):
-        shell_function  = ShellFunction( source, shell=shell, cmd_format=cmd_format )
+        shell_function  = ShellFunction(source, shell=shell, cmd_format=cmd_format, modules=self.modules)
        
         # if abstraction is map or multimap
         if map :

src/jflow/config_reader.py

@@ -118,6 +118,12 @@ class JFlowConfigReader(object):
             return self.reader.get("components", component_class+".batch_options")
         except:
             return ""
+
+    def get_component_modules(self, component_class):
+        try:
+            return self.reader.get("components", component_class+".modules").split(",")
+        except:
+            return []
     
     def get_workflow_group(self, workflow_class):
         try:

src/weaver/function.py

@@ -185,7 +185,7 @@ class ShellFunction(ScriptFunction):
     }
     SHELL_DEFAULT = 'sh'
 
-    def __init__(self, source, shell=None, executable=None, cmd_format=None):
+    def __init__(self, source, shell=None, executable=None, cmd_format=None, modules=[]):
         if shell is None or not os.path.isabs(shell):
             if shell not in ShellFunction.SHELL_TABLE:
                 shell = ShellFunction.SHELL_DEFAULT
@@ -193,7 +193,9 @@ class ShellFunction(ScriptFunction):
         else:
             shell_path = shell
             shell = os.path.basename(shell)
-        source = '#!%s\n' % shell_path + source
+        source = '#!%s\n' % shell_path + \
+                 ((("module load " +  module) for module in modules) if len(modules) > 0 else "") + "\n" + source + \
+                 ((("module unload " + module) for module in modules) if len(modules) > 0 else "")
         ScriptFunction.__init__(self, source, executable, cmd_format)
 
 
@@ -246,7 +248,9 @@ def prepare_arguments(*args):
 {{2}}
 if __name__ == '__main__':
     (new_args, output_path) = prepare_arguments(*sys.argv[1:])
+    {{4}}
     outputs = {{3}}(*new_args)
+    {{5}}
     
     if not isinstance(outputs, list): outputs = [outputs]
     # in case the returned object is a table
@@ -262,7 +266,7 @@ if __name__ == '__main__':
             objh.close()
 '''.format(PYTHON_VERSION)
 
-    def __init__(self, function, add_path=None, executable=None, cmd_format=None):
+    def __init__(self, function, add_path=None, executable=None, cmd_format=None, modules=[]):
         # TODO: this doesn't work with Python3
         body = inspect.getsource(function)
         name = getfuncname(function)
@@ -280,8 +284,17 @@ if __name__ == '__main__':
         path = ""
         for apath in add_path:
             path += "sys.path.insert(0, '" + apath + "')\n"
+
+        modules_cmd = ""
+        modules_unload_cmd = ""
+        if len(modules) > 0:
+            for module in modules:
+                modules_cmd += "module load " + module + "; "
+                modules_unload_cmd += "module unload" + module + "; "
+            modules_cmd = "os.system(" + modules_cmd + ")"
+            modules_unload_cmd = "os.system(" + modules_unload_cmd + ")"
         
-        source = self.PYTHON_TEMPLATE.format(', '.join(imports), path, body, name)
+        source = self.PYTHON_TEMPLATE.format(', '.join(imports), path, body, name, modules_cmd, modules_unload_cmd)
         ScriptFunction.__init__(self, source, executable, cmd_format)