Accommodate separate lmod, tcl modules.yaml's

lib/jcsda-emc/spack-stack/stack/cmd/stack_cmds/create.py

@@ -142,6 +142,13 @@ def setup_env_parser(subparser):
         help="Include upstream environment (/path/to/spack-stack-x.y.z/envs/unified-env/install)",
     )
 
+    subparser.add_argument(
+        "--modulesys",
+        type=str,
+        choices=["lmod", "tcl"],
+        help="Override choice of tcl vs. lmod config (default is based on site config)",
+    )
+
     subparser.add_argument(
         "--site", type=str, required=False, default=default_site(), help=site_help()
     )
@@ -193,6 +200,7 @@ def dict_from_args(args):
     dict["base_packages"] = args.packages
     dict["dir"] = args.dir
     dict["upstreams"] = args.upstream
+    dict["modulesys"] = args.modulesys
 
     return dict
 

lib/jcsda-emc/spack-stack/stack/stack_env.py

@@ -96,6 +96,7 @@ def __init__(self, **kwargs):
         self.install_prefix = kwargs.get("install_prefix", None)
         self.mirror = kwargs.get("mirror", None)
         self.upstreams = kwargs.get("upstreams", None)
+        self.modulesys = kwargs.get("modulesys", None)
 
         if not self.name:
             # site = self.site if self.site else 'default'
@@ -108,11 +109,31 @@ def env_dir(self):
     def add_includes(self, includes):
         self.includes.extend(includes)
 
+    def get_lmod_or_tcl(self, site_configs_dir):
+        site_modules_yaml_path = os.path.join(site_configs_dir, "modules.yaml")
+        with open(site_modules_yaml_path, "r") as f:
+            site_modules_yaml = syaml.load_config(f)
+        lmod_or_tcl_list = site_modules_yaml["modules"]["default"]["enable"]
+        assert lmod_or_tcl_list and (
+            set(lmod_or_tcl_list) in ({"tcl"}, {"lmod"})
+        ), """Set one and only one value ('lmod' or 'tcl') under 'modules:default:enable'
+           in site modules.yaml, or use '--modulesys {tcl,lmod}'"""
+        return lmod_or_tcl_list[0]
+
     def _copy_common_includes(self):
         """Copy common directory into environment"""
         self.includes.append("common")
         env_common_dir = os.path.join(self.env_dir(), "common")
-        shutil.copytree(common_path, env_common_dir)
+        shutil.copytree(
+            common_path, env_common_dir, ignore=shutil.ignore_patterns("modules_*.yaml")
+        )
+        if self.modulesys:
+            lmod_or_tcl = self.modulesys
+        else:
+            lmod_or_tcl = self.get_lmod_or_tcl(self.site_configs_dir())
+        common_modules_yaml_path = os.path.join(common_path, "modules_%s.yaml" % lmod_or_tcl)
+        destination = os.path.join(env_common_dir, "modules.yaml")
+        shutil.copy(common_modules_yaml_path, destination)
 
     def site_configs_dir(self):
         site_configs_dir = os.path.join(site_path, self.site)
@@ -127,6 +148,23 @@ def _copy_site_includes(self):
         self.includes.append(site_name)
         env_site_dir = os.path.join(self.env_dir(), site_name)
         shutil.copytree(self.site_configs_dir(), env_site_dir)
+        # Update site modules.yaml if user overrides default module system
+        if not self.modulesys:
+            return
+        lmod_or_tcl = self.modulesys
+        site_modules_yaml_path = os.path.join(env_site_dir, "modules.yaml")
+        with open(site_modules_yaml_path, "r") as f:
+            site_modules_yaml = syaml.load_config(f)
+        current_sys = site_modules_yaml["modules"]["default"]["enable"][0]
+        if lmod_or_tcl == current_sys:
+            return
+        logging.info("Updating site modules.yaml to reflect env module system override setting\n")
+        site_modules_yaml["modules"]["default"]["enable"][0] = lmod_or_tcl
+        current_config = site_modules_yaml["modules"]["default"][current_sys]
+        site_modules_yaml["modules"]["default"][lmod_or_tcl] = current_config
+        del site_modules_yaml["modules"]["default"][current_sys]
+        with open(site_modules_yaml_path, "w") as f:
+            syaml.dump_config(site_modules_yaml, f)
 
     def _copy_package_includes(self):
         """Overwrite base packages in environment common dir"""

lib/jcsda-emc/spack-stack/tests/test_stack_create.py

@@ -100,3 +100,50 @@ def test_containers(container, spec):
         test_dir,
         "--overwrite",
     )
+
+
+@pytest.mark.extension("stack")
+@pytest.mark.filterwarnings("ignore::UserWarning")
+def test_modulesys():
+    modsystems = {"lmod", "tcl"}
+    for modulesys in modsystems:
+        stack_create(
+            "create",
+            "env",
+            "--site",
+            "hera",
+            "--name",
+            "modulesys_test",
+            "--dir",
+            test_dir,
+            "--overwrite",
+            "--modulesys",
+            modulesys,
+        )
+    modules_yaml_path = os.path.join(test_dir, "modulesys_test", "common", "modules.yaml")
+    with open(modules_yaml_path, "r") as f:
+        modules_yaml_txt = f.read()
+    assert "%s:" % modulesys in modules_yaml_txt
+    assert "%s:" % list(modsystems.difference(modulesys))[0] not in modules_yaml_txt
+
+
+@pytest.mark.extension("stack")
+@pytest.mark.filterwarnings("ignore::UserWarning")
+def test_upstream():
+    stack_create(
+        "create",
+        "env",
+        "--site",
+        "hera",
+        "--name",
+        "upstream_test",
+        "--dir",
+        test_dir,
+        "--overwrite",
+        "--upstream",
+        "/test/path/to/upstream/env",
+    )
+    spack_yaml_path = os.path.join(test_dir, "upstream_test", "spack.yaml")
+    with open(spack_yaml_path, "r") as f:
+        spack_yaml_txt = f.read()
+    assert "install_tree: /test/path/to/upstream/env" in spack_yaml_txt