Update envmod.setup_user_modules to inspect PATH and LOADEDMODULES be…

…fore/after loading modules

payu/envmod.py

@@ -129,14 +129,17 @@ def lib_update(required_libs, lib_name):
 
 
 def setup_user_modules(user_modules, user_modulepaths):
-    """Run module use + load commands for user-defined modules"""
+    """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
+        return (None, None)
 
     # Add user-defined directories to MODULEPATH
     for modulepath in user_modulepaths:
@@ -147,18 +150,26 @@ def setup_user_modules(user_modules, user_modulepaths):
 
         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
-        module_subcommand = f"avail --terse {modulefile}"
-        output = run_cmd(module_cmd(module_subcommand)).stderr
+        output = run_module_cmd("avail --terse", modulefile).stderr
 
-        # Extract out the modulefiles available
+        # 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:
-            # Modules are used for finding model executable paths - so check
-            # for unique module -TODO: Could be a warning rather than an error?
             raise ValueError(
                 f"There are multiple modules available for {modulefile}:\n" +
                 f"{output}\n{MULTIPLE_MODULES_HELP}")
@@ -170,34 +181,18 @@ def setup_user_modules(user_modules, user_modulepaths):
         # 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(':')))
 
-def env_var_set_by_modules(user_modules, env_var):
-    """Return an environment variable post loading only user-defined modules
-    - this is used for getting $PATH for searching for the model executable"""
-    if 'MODULESHOME' not in os.environ:
-        print('payu: warning: No Environment Modules found; skipping '
-              f'inspecting user module changes to ${env_var}')
-        return
-
-    # Note: Using subprocess shell to isolate changes to environment
-    load_commands = [f'load {module}' for module in user_modules]
-    commands = ['purge'] + load_commands
-    module_cmds = [f"eval `{module_cmd(c)}`" for c in commands]
-    module_cmds += [f'echo ${env_var}']
-    command = ' && '.join(module_cmds)
-    output = run_cmd(command)
-
-    # Extract out $env_var from output
-    output.check_returncode()
-    lines = output.stdout.strip().split('\n')
-    return lines[-1]
-
-
-def module_cmd(command):
-    """Format module subcommand using modulecmd"""
-    return f"{os.environ['MODULESHOME']}/bin/modulecmd bash {command}"
+    return (loaded_modules, paths)
 
 
-def run_cmd(command):
-    """Wrapper around subprocess command that captures output"""
+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,7 +132,7 @@ def __init__(self, lab, reproduce=False, force=False):
 
         self.run_id = None
 
-        self.user_modules_path = None
+        self.user_modules_paths = None
 
     def init_models(self):
 
@@ -227,23 +227,15 @@ def setup_modules(self):
         envmod.setup()
 
         # Get user modules info from config
-        self.user_modulepaths = self.config.get('modules', {}).get('use', [])
-        self.user_modules = self.config.get('modules', {}).get('load', [])
-
-        # Run module use + load commands for user-defined modules
-        envmod.setup_user_modules(self.user_modules, self.user_modulepaths)
-
-        # Get paths and loaded modules post loading only the user modules
-        self.user_modules_path = envmod.env_var_set_by_modules(
-            self.user_modules, 'PATH'
-        )
-
-        # Store list of all modules loaded by user-modules
-        self.loaded_user_modules = envmod.env_var_set_by_modules(
-            self.user_modules, 'LOADEDMODULES'
-        )
-        if self.loaded_user_modules is not None:
-            self.loaded_user_modules = self.loaded_user_modules.split(':')
+        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

payu/models/model.py

@@ -184,15 +184,11 @@ def expand_executable_path(self, exec):
             return exec
 
         # Check if path set by loading user modules has been defined
-        module_added_path = self.expt.user_modules_path
-        if module_added_path is None:
+        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 = []
-        elif module_added_path == '':
-            module_added_paths = []
-        else:
-            module_added_paths = module_added_path.split(':')
 
         # Search for exe inside paths added to $PATH by user-defined modules
         exec_paths = []