reduce number of module load commands used to set up build environment

easybuild/framework/easyblock.py

@@ -92,7 +92,7 @@
 from easybuild.tools.config import get_failed_install_logs_path, get_log_filename, get_repository, get_repositorypath
 from easybuild.tools.config import install_path, log_path, package_path, source_paths, source_paths_data
 from easybuild.tools.config import DATA, SOFTWARE
-from easybuild.tools.environment import restore_env, sanitize_env
+from easybuild.tools.environment import copy_current_env, restore_env, sanitize_env
 from easybuild.tools.filetools import CHECKSUM_TYPE_SHA256
 from easybuild.tools.filetools import adjust_permissions, apply_patch, back_up_file, change_dir, check_lock, clean_dir
 from easybuild.tools.filetools import compute_checksum, convert_name, copy_dir, copy_file, create_lock
@@ -2163,6 +2163,22 @@ def install_extensions_sequential(self, install=True):
 
         exts_cnt = len(self.ext_instances)
 
+        # determine build environment, and cache it
+        if self.dry_run:
+            self.dry_run_msg("defining build environment based on toolchain (options) and dependencies...")
+        else:
+            with self.fake_module_environment(with_build_deps=True):
+                self.log.debug("List of loaded modules: %s", self.modules_tool.list())
+                # don't reload modules for toolchain, there is no need
+                # since they will be loaded already by the fake module
+                self.toolchain.prepare(onlymod=self.cfg['onlytcmod'], deps=self.cfg.dependencies(),
+                                       silent=True, loadmod=False,
+                                       rpath_filter_dirs=self.rpath_filter_dirs,
+                                       rpath_include_dirs=self.rpath_include_dirs,
+                                       rpath_wrappers_dir=self.rpath_wrappers_dir)
+
+                build_env = copy_current_env()
+
         for idx, ext in enumerate(self.ext_instances):
             self.log.info("Starting extension %s", ext.name)
 
@@ -2183,33 +2199,25 @@ def install_extensions_sequential(self, install=True):
                 msg = "\n* installing extension %s %s using '%s' easyblock\n" % tup
                 self.dry_run_msg(msg)
 
-            if self.dry_run:
-                self.dry_run_msg("defining build environment based on toolchain (options) and dependencies...")
-
             # actual installation of the extension
-            if install and not self.dry_run:
-                with self.fake_module_environment(with_build_deps=True):
-                    self.log.debug("List of loaded modules: %s", self.modules_tool.list())
-                    # don't reload modules for toolchain, there is no need
-                    # since they will be loaded already by the fake module
-                    ext.toolchain.prepare(onlymod=self.cfg['onlytcmod'], deps=self.cfg.dependencies(),
-                                          silent=True, loadmod=False,
-                                          rpath_filter_dirs=self.rpath_filter_dirs,
-                                          rpath_include_dirs=self.rpath_include_dirs,
-                                          rpath_wrappers_dir=self.rpath_wrappers_dir)
-                    try:
-                        ext.install_extension_substep("pre_install_extension")
-                        with self.module_generator.start_module_creation():
-                            txt = ext.install_extension_substep("install_extension")
-                        if txt:
-                            self.module_extra_extensions += txt
-                        ext.install_extension_substep("post_install_extension")
-                    finally:
-                        ext_duration = datetime.now() - start_time
-                        if ext_duration.total_seconds() >= 1:
-                            print_msg("\t... (took %s)", time2str(ext_duration), log=self.log, silent=self.silent)
-                        elif self.logdebug or build_option('trace'):
-                            print_msg("\t... (took < 1 sec)", log=self.log, silent=self.silent)
+            elif install:
+
+                # restore build environment for this extension
+                restore_env(build_env, log_changes=False)
+
+                try:
+                    ext.install_extension_substep("pre_install_extension")
+                    with self.module_generator.start_module_creation():
+                        txt = ext.install_extension_substep("install_extension")
+                    if txt:
+                        self.module_extra_extensions += txt
+                    ext.install_extension_substep("post_install_extension")
+                finally:
+                    ext_duration = datetime.now() - start_time
+                    if ext_duration.total_seconds() >= 1:
+                        print_msg("\t... (took %s)", time2str(ext_duration), log=self.log, silent=self.silent)
+                    elif self.logdebug or build_option('trace'):
+                        print_msg("\t... (took < 1 sec)", log=self.log, silent=self.silent)
 
             self.update_exts_progress_bar(progress_info, progress_size=1)
 

easybuild/tools/environment.py

@@ -76,6 +76,13 @@ def get_changes():
     return _changes
 
 
+def copy_current_env():
+    """
+    Copy current environment, and return it.
+    """
+    return copy.deepcopy(os.environ)
+
+
 def setvar(key, value, verbose=True, log_changes=True):
     """
     put key in the environment with value

easybuild/tools/modules.py

@@ -1127,8 +1127,8 @@ def load(self, modules, mod_paths=None, purge=False, init_env=None, allow_reload
 
         if not allow_reload:
             modules = set(modules) - set(self.loaded_modules())
-        for mod in modules:
-            self.run_module('load', mod)
+
+        self.run_module('load', *modules)
 
     def unload(self, modules, log_changes=None, *, hide_output=None):
         """

easybuild/tools/toolchain/toolchain.py

@@ -570,7 +570,7 @@ def _check_dependencies(self, dependencies, check_modules=True):
         # current $MODULEPATH without loading the prior dependencies in a module hierarchy
         # (e.g. OpenMPI module may only be available after loading GCC module);
         # when actually loading the modules for the dependencies, the *short* module name is used,
-        # see _load_dependencies_modules()
+        # see _load_modules_toolchain_and_deps()
         dep_mod_names = [dep['full_mod_name'] for dep in dependencies]
 
         # check whether modules exist
@@ -641,75 +641,82 @@ def _simulated_load_dependency_module(self, name, version, metadata, verbose=Fal
         for var, value in env_vars.items():
             setvar(var, value, verbose=verbose)
 
-    def _load_toolchain_module(self, silent=False):
-        """Load toolchain module."""
-
-        tc_mod = self.det_short_module_name()
-
-        if self.dry_run:
-            dry_run_msg("Loading toolchain module...\n", silent=silent)
-
-            # load toolchain module, or simulate load of toolchain components if it is not available
-            if self.modules_tool.exist([tc_mod], skip_avail=True)[0]:
-                self.modules_tool.load([tc_mod])
-                dry_run_msg("module load %s" % tc_mod, silent=silent)
-            else:
-                # first simulate loads for toolchain dependencies, if required information is available
-                for tcdep in self.tcdeps:
-                    modname = tcdep['short_mod_name']
-                    dry_run_msg("module load %s [SIMULATED]" % modname, silent=silent)
-                    # 'use '$EBROOTNAME' as value for dep install prefix (looks nice in dry run output)
-                    deproot = '$%s' % get_software_root_env_var_name(tcdep['name'])
-                    self._simulated_load_dependency_module(tcdep['name'], tcdep['version'], {'prefix': deproot})
+    def _load_modules_toolchain_and_deps(self, silent=False):
+        """
+        Load modules for toolchain + dependencies, and handle special cases like external modules.
 
-                dry_run_msg("module load %s [SIMULATED]" % tc_mod, silent=silent)
-                # use name of $EBROOT* env var as value for $EBROOT* env var (results in sensible dry run output)
-                tcroot = '$%s' % get_software_root_env_var_name(self.name)
-                self._simulated_load_dependency_module(self.name, self.version, {'prefix': tcroot})
+        :param silent: boolean indicating whether or not to stay silent
+        """
+        system_toolchain = self.is_system_toolchain()
+        if system_toolchain:
+            self.log.info("Loading dependencies using system toolchain...")
         else:
-            # make sure toolchain is available using short module name by running 'module use' on module path subdir
-            if self.init_modpaths:
-                mod_path_suffix = build_option('suffix_modules_path')
-                for modpath in self.init_modpaths:
-                    modpath = os.path.join(install_path('mod'), mod_path_suffix, modpath)
-                    if os.path.exists(modpath):
-                        self.modules_tool.prepend_module_path(modpath)
-
-            # load modules for all dependencies
-            self.log.debug("Loading module for toolchain: %s", tc_mod)
-            trace_msg("loading toolchain module: " + tc_mod)
-            self.modules_tool.load([tc_mod])
-
-        # append toolchain module to list of modules
-        self.modules.append(tc_mod)
-
-    def _load_dependencies_modules(self, silent=False):
-        """Load modules for dependencies, and handle special cases like external modules."""
+            self.log.debug("Loading toolchain module and dependencies...")
+
+        tc_mod = None if system_toolchain else self.det_short_module_name()
         dep_mods = [dep['short_mod_name'] for dep in self.dependencies]
 
         if self.dry_run:
+            if not system_toolchain:
+                dry_run_msg("Loading toolchain module...\n", silent=silent)
+
+                # load toolchain module, or simulate load of toolchain components if it is not available
+                if self.modules_tool.exist([tc_mod], skip_avail=True)[0]:
+                    self.modules_tool.load([tc_mod])
+                    dry_run_msg(f"module load {tc_mod}", silent=silent)
+                else:
+                    # first simulate loads for toolchain dependencies, if required information is available
+                    for tc_dep in self.tcdeps:
+                        mod_name = tcdep['short_mod_name']
+                        dry_run_msg(f"module load {mod_name} [SIMULATED]", silent=silent)
+                        # 'use '$EBROOTNAME' as value for dep install prefix (looks nice in dry run output)
+                        deproot = '$' + get_software_root_env_var_name(tc_dep['name'])
+                        self._simulated_load_dependency_module(tc_dep['name'], tc_dep['version'], {'prefix': deproot})
+
+                    dry_run_msg("module load %s [SIMULATED]" % tc_mod, silent=silent)
+                    # use name of $EBROOT* env var as value for $EBROOT* env var (results in sensible dry run output)
+                    tcroot = '$' + get_software_root_env_var_name(self.name)
+                    self._simulated_load_dependency_module(self.name, self.version, {'prefix': tcroot})
+
             dry_run_msg("\nLoading modules for dependencies...\n", silent=silent)
 
             mods_exist = self.modules_tool.exist(dep_mods)
 
             # load available modules for dependencies, simulate load for others
+            mods_to_load = []
             for dep, dep_mod_exists in zip(self.dependencies, mods_exist):
                 mod_name = dep['short_mod_name']
                 if dep_mod_exists:
-                    self.modules_tool.load([mod_name])
-                    dry_run_msg("module load %s" % mod_name, silent=silent)
+                    mods_to_load.append(mod_name)
+                    dry_run_msg(f"module load {mod_name}", silent=silent)
                 else:
-                    dry_run_msg("module load %s [SIMULATED]" % mod_name, silent=silent)
+                    dry_run_msg("module load {mod_name} [SIMULATED]", silent=silent)
                     # 'use '$EBROOTNAME' as value for dep install prefix (looks nice in dry run output)
                     if not dep['external_module']:
-                        deproot = '$%s' % get_software_root_env_var_name(dep['name'])
+                        deproot = '$' + get_software_root_env_var_name(dep['name'])
                         self._simulated_load_dependency_module(dep['name'], dep['version'], {'prefix': deproot})
-        else:
-            # load modules for all dependencies
-            self.log.debug("Loading modules for dependencies: %s", dep_mods)
-            self.modules_tool.load(dep_mods)
 
-            if self.dependencies:
+            self.modules_tool.load(mods_to_load)
+
+        else:
+            mods_to_load = []
+            if not system_toolchain:
+                # make sure toolchain is available using short module name by running 'module use' on module path subdir
+                if self.init_modpaths:
+                    mod_path_suffix = build_option('suffix_modules_path')
+                    for modpath in self.init_modpaths:
+                        modpath = os.path.join(install_path('mod'), mod_path_suffix, modpath)
+                        if os.path.exists(modpath):
+                            # FIXME
+                            self.modules_tool.prepend_module_path(modpath)
+
+                self.log.debug(f"Loading module for toolchain: {tc_mod}")
+                trace_msg(f"loading toolchain module: {tc_mod}")
+                mods_to_load.append(tc_mod)
+
+            self.log.debug(f"Loading modules for dependencies: {' '.join(dep_mods)}")
+            mods_to_load.extend(dep_mods)
+            if dep_mods:
                 build_dep_mods = [dep['short_mod_name'] for dep in self.dependencies if dep['build_only']]
                 if build_dep_mods:
                     trace_msg("loading modules for build dependencies:")
@@ -726,8 +733,11 @@ def _load_dependencies_modules(self, silent=False):
                 else:
                     trace_msg("(no (runtime) dependencies specified)")
 
-        # append dependency modules to list of modules
-        self.modules.extend(dep_mods)
+            # actually load modules, all in one go
+            self.modules_tool.load(mods_to_load)
+
+        # append toolchain + dependency modules to list of modules
+        self.modules.extend([tc_mod] + dep_mods)
 
         # define $EBROOT* and $EBVERSION* for external modules, if metadata is available
         for dep in [d for d in self.dependencies if d['external_module']]:
@@ -749,24 +759,18 @@ def _load_modules(self, silent=False):
             raise EasyBuildError("No modules tool defined in Toolchain instance.")
 
         if not self._toolchain_exists() and not self.dry_run:
-            raise EasyBuildError("No module found for toolchain: %s", self.mod_short_name)
+            raise EasyBuildError(f"No module found for toolchain: {self.mod_short_name}")
 
-        if self.is_system_toolchain():
-            self.log.info("Loading dependencies using system toolchain...")
-            self._load_dependencies_modules(silent=silent)
-        else:
-            # load the toolchain and dependencies modules
-            self.log.debug("Loading toolchain module and dependencies...")
-            self._load_toolchain_module(silent=silent)
-            self._load_dependencies_modules(silent=silent)
+        # load modules for toolchain + specified (build) dependencies
+        self._load_modules_toolchain_and_deps(silent=silent)
 
         # include list of loaded modules in dry run output
         if self.dry_run:
             loaded_mods = self.modules_tool.list()
             dry_run_msg("\nFull list of loaded modules:", silent=silent)
             if loaded_mods:
                 for i, mod_name in enumerate([m['mod_name'] for m in loaded_mods]):
-                    dry_run_msg("  %d) %s" % (i + 1, mod_name), silent=silent)
+                    dry_run_msg(f"  {i + 1}) {mod_name}", silent=silent)
             else:
                 dry_run_msg("  (none)", silent=silent)
             dry_run_msg('', silent=silent)

test/framework/toy_build.py

@@ -3649,7 +3649,6 @@ def post_build_and_install_loop_hook(ecs):
             echo toy
             in module-write hook hook for {mod_name}
             installing of extension bar is done!
-            in module-write hook hook for {mod_name}
             pre_run_shell_cmd_hook triggered for ' gcc toy.c -o toy '
             ' gcc toy.c -o toy  && copy_toy_file toy copy_of_toy' command failed (exit code 127), but I fixed it!
             installing of extension toy is done!