Get module executables from path loaded by environment modules (#439)

Get module executables from path loaded by environment modules

- Run module use/load commands for user modules (specified in config.yaml) in Experiment.setup()
- Inspect changes to PATH before/after loading user modules
- Move setting executable paths in Model class to Experiment.setup()
- In Collate job, load any user modules, and expand path for collate executable

payu/envmod.py

@@ -17,6 +17,22 @@
 DEFAULT_BASEPATH = '/opt/Modules'
 DEFAULT_VERSION = 'v4.3.0'
 
+MODULE_NOT_FOUND_HELP = """ To fix module not being found:
+- Check module name and version in config.yaml (listed under `modules: load:`)
+- If module is found in a module directory, ensure this path is listed in
+config.yaml under `modules: use:`, or run `module use` command prior to running
+payu commands.
+"""
+
+MULTIPLE_MODULES_HELP = """ To fix having multiple modules available:
+- Add version to the module in config.yaml (under `modules: load:`)
+- Modify module directories in config.yaml (under `modules: use:`)
+- Or modify module directories in user environment by using module use/unuse
+commands, e.g.:
+    $ module use dir # Add dir to $MODULEPATH
+    $ module unuse dir # Remove dir from $MODULEPATH
+"""
+
 
 def setup(basepath=DEFAULT_BASEPATH):
     """Set the environment modules used by the Environment Module system."""
@@ -109,4 +125,74 @@ def lib_update(required_libs, lib_name):
             return '{0}/{1}'.format(mod_name, mod_version)
 
     # If there are no libraries, return an empty string
-    return ''
+    return ''
+
+
+def setup_user_modules(user_modules, user_modulepaths):
+    """Run module use + load commands for user-defined modules. Return
+    tuple containing a set of loaded modules and paths added to
+    LOADEDMODULES and PATH environment variable, as result of
+    loading user-modules"""
+
+    if 'MODULESHOME' not in os.environ:
+        print(
+            'payu: warning: No Environment Modules found; ' +
+            'skipping running module use/load commands for any module ' +
+            'directories/modulefiles defined in config.yaml')
+        return (None, None)
+
+    # Add user-defined directories to MODULEPATH
+    for modulepath in user_modulepaths:
+        if not os.path.isdir(modulepath):
+            raise ValueError(
+                f"Module directory is not found: {modulepath}" +
+                "\n Check paths listed under `modules: use:` in config.yaml")
+
+        module('use', modulepath)
+
+    # First un-load all user modules, if they are loaded, so can store
+    # LOADEDMODULES and PATH to compare to later
+    for modulefile in user_modules:
+        if run_module_cmd("is-loaded", modulefile).returncode == 0:
+            module('unload', modulefile)
+    previous_loaded_modules = os.environ.get('LOADEDMODULES', '')
+    previous_path = os.environ.get('PATH', '')
+
+    for modulefile in user_modules:
+        # Check module exists and there is not multiple available
+        output = run_module_cmd("avail --terse", modulefile).stderr
+
+        # Extract out the modulefiles available - strip out lines like:
+        # /apps/Modules/modulefiles:
+        modules = [line for line in output.strip().splitlines()
+                   if not (line.startswith('/') and line.endswith(':'))]
+
+        # Modules are used for finding model executable paths - so check
+        # for unique module
+        if len(modules) > 1:
+            raise ValueError(
+                f"There are multiple modules available for {modulefile}:\n" +
+                f"{output}\n{MULTIPLE_MODULES_HELP}")
+        elif len(modules) == 0:
+            raise ValueError(
+                f"Module is not found: {modulefile}\n{MODULE_NOT_FOUND_HELP}"
+            )
+
+        # Load module
+        module('load', modulefile)
+
+    # Create a set of paths and modules loaded by user modules
+    loaded_modules = os.environ.get('LOADEDMODULES', '')
+    path = os.environ.get('PATH', '')
+    loaded_modules = set(loaded_modules.split(':')).difference(
+        previous_loaded_modules.split(':'))
+    paths = set(path.split(':')).difference(set(previous_path.split(':')))
+
+    return (loaded_modules, paths)
+
+
+def run_module_cmd(subcommand, *args):
+    """Wrapper around subprocess module command that captures output"""
+    modulecmd = f"{os.environ['MODULESHOME']}/bin/modulecmd bash"
+    command = f"{modulecmd} {subcommand} {' '.join(args)}"
+    return subprocess.run(command, shell=True, text=True, capture_output=True)

payu/experiment.py

@@ -132,6 +132,8 @@ def __init__(self, lab, reproduce=False, force=False):
 
         self.run_id = None
 
+        self.user_modules_paths = None
+
     def init_models(self):
 
         self.model_name = self.config.get('model')
@@ -220,9 +222,22 @@ def set_stacksize(self, stacksize):
         resource.setrlimit(resource.RLIMIT_STACK,
                            (stacksize, resource.RLIM_INFINITY))
 
-    def load_modules(self):
-        # NOTE: This function is increasingly irrelevant, and may be removable.
+    def setup_modules(self):
+        """Setup modules and get paths added to $PATH by user-modules"""
+        envmod.setup()
+
+        # Get user modules info from config
+        user_modulepaths = self.config.get('modules', {}).get('use', [])
+        user_modules = self.config.get('modules', {}).get('load', [])
+
+        # Run module use + load commands for user-defined modules, and
+        # get a set of paths and loaded modules added by loading the modules
+        loaded_mods, paths = envmod.setup_user_modules(user_modules,
+                                                       user_modulepaths)
+        self.user_modules_paths = paths
+        self.loaded_user_modules = [] if loaded_mods is None else loaded_mods
 
+    def load_modules(self):
         # Scheduler
         sched_modname = self.config.get('scheduler', 'pbs')
         self.modules.add(sched_modname)
@@ -245,18 +260,14 @@ def load_modules(self):
             if len(mod) > 0:
                 print('mod '+mod)
                 mod_base = mod.split('/')[0]
-                if mod_base not in core_modules:
+                if (mod_base not in core_modules and
+                        mod not in self.loaded_user_modules):
                     envmod.module('unload', mod)
 
         # Now load model-dependent modules
         for mod in self.modules:
             envmod.module('load', mod)
 
-        # User-defined modules
-        user_modules = self.config.get('modules', {}).get('load', [])
-        for mod in user_modules:
-            envmod.module('load', mod)
-
         envmod.module('list')
 
         for prof in self.profilers:
@@ -414,6 +425,11 @@ def setup(self, force_archive=False):
 
         make_symlink(self.work_path, self.work_sym_path)
 
+        # Set up executable paths - first search through paths added by modules
+        self.setup_modules()
+        for model in self.models:
+            model.setup_executable_paths()
+
         # Set up all file manifests
         self.manifest.setup()
 
@@ -453,13 +469,6 @@ def setup(self, force_archive=False):
             self.get_restarts_to_prune()
 
     def run(self, *user_flags):
-        # XXX: This was previously done in reversion
-        envmod.setup()
-
-        # Add any user-defined module dir(s) to MODULEPATH
-        for module_dir in self.config.get('modules', {}).get('use', []):
-            envmod.module('use', module_dir)
-
         self.load_modules()
 
         f_out = open(self.stdout_fname, 'w')
@@ -804,6 +813,9 @@ def archive(self, force_prune_restarts=False):
             self.postprocess()
 
     def collate(self):
+        # Setup modules - load user-defined modules
+        self.setup_modules()
+
         for model in self.models:
             model.collate()
 

payu/models/fms.py

@@ -76,8 +76,7 @@ def fms_collate(model):
     mpi = collate_config.get('mpi', False)
 
     if mpi:
-        # Must use envmod to be able to load mpi modules for collation
-        envmod.setup()
+        # Load mpi modules for collation
         model.expt.load_modules()
         default_exe = 'mppnccombine-fast'
     else:
@@ -92,8 +91,7 @@ def fms_collate(model):
                 mppnc_path = os.path.join(model.expt.lab.bin_path, f)
                 break
     else:
-        if not os.path.isabs(mppnc_path):
-            mppnc_path = os.path.join(model.expt.lab.bin_path, mppnc_path)
+        mppnc_path = model.expand_executable_path(mppnc_path)
 
     assert mppnc_path, 'No mppnccombine program found'
 

payu/models/model.py

@@ -82,21 +82,6 @@ def set_model_pathnames(self):
         self.work_output_path = self.work_path
         self.work_init_path = self.work_path
 
-        self.exec_prefix = self.config.get('exe_prefix', '')
-        self.exec_name = self.config.get('exe', self.default_exec)
-        if self.exec_name:
-            # By default os.path.join will not prepend the lab bin_path
-            # to an absolute path
-            self.exec_path = os.path.join(self.expt.lab.bin_path,
-                                          self.exec_name)
-        else:
-            self.exec_path = None
-        if self.exec_path:
-            # Make exec_name consistent for models with fully qualified path.
-            # In all cases it will just be the name of the executable without a
-            # path
-            self.exec_name = os.path.basename(self.exec_path)
-
     def set_local_pathnames(self):
 
         # This is the path relative to the control directory, required for
@@ -129,12 +114,6 @@ def set_local_pathnames(self):
                 os.path.relpath(self.work_init_path, self.expt.work_path)
             )
         )
-        if self.exec_path:
-            # Local path in work directory
-            self.exec_path_local = os.path.join(
-                self.work_path_local,
-                os.path.basename(self.exec_path)
-            )
 
     def set_input_paths(self):
         if len(self.expt.models) == 1:
@@ -198,6 +177,55 @@ def get_prior_restart_files(self):
             print("No prior restart files found: {error}".format(error=str(e)))
             return []
 
+    def expand_executable_path(self, exec):
+        """Given an executable, return the expanded executable path"""
+        # Check if exe is already an absolute path
+        if os.path.isabs(exec):
+            return exec
+
+        # Check if path set by loading user modules has been defined
+        module_added_paths = self.expt.user_modules_paths
+        if module_added_paths is None:
+            print("payu: warning: Skipping searching for model executable " +
+                  "in $PATH set by user modules")
+            module_added_paths = []
+
+        # Search for exe inside paths added to $PATH by user-defined modules
+        exec_paths = []
+        for path in module_added_paths:
+            exec_path = os.path.join(path, exec)
+            if os.path.exists(exec_path) and os.access(exec_path, os.X_OK):
+                exec_paths.append(exec_path)
+
+        if len(exec_paths) > 1:
+            raise ValueError(
+                f"Executable {exec} found in multiple $PATH paths added by " +
+                f"user-defined modules in `config.yaml`. Paths: {exec_paths}")
+        elif len(exec_paths) == 1:
+            return exec_paths[0]
+
+        # Else prepend the lab bin path to exec
+        return os.path.join(self.expt.lab.bin_path, exec)
+
+    def setup_executable_paths(self):
+        """Set model executable paths"""
+        self.exec_prefix = self.config.get('exe_prefix', '')
+        self.exec_name = self.config.get('exe', self.default_exec)
+        self.exec_path = None
+        if self.exec_name:
+            self.exec_path = self.expand_executable_path(self.exec_name)
+
+            # Make exec_name consistent for models with fully qualified path.
+            # In all cases it will just be the name of the executable without a
+            # path
+            self.exec_name = os.path.basename(self.exec_path)
+
+            # Local path in work directory
+            self.exec_path_local = os.path.join(
+                self.work_path_local,
+                os.path.basename(self.exec_path)
+            )
+
     def setup_configuration_files(self):
         """Copy configuration and optional configuration files from control
          path to work path"""
@@ -339,6 +367,7 @@ def collate(self):
         raise NotImplementedError
 
     def build_model(self):
+        self.setup_executable_paths()
 
         if not self.repo_url:
             return