modules: use new `module` function instead of `get_module_cmd` (#8570)

Use new `module` function instead of `get_module_cmd`

Previously, Spack relied on either examining the bash `module()` function or using the `which` command to find the underlying executable for modules. More complicated module systems do not allow for the sort of simple analysis we were doing (see #6451).

Spack now uses the `module` function directly and copies environment changes from the resulting subprocess back into Spack. This should provide a future-proof implementation for changes to the logic underlying the module system on various HPC systems.

8 files changed, 118 insertions, 248 deletions

diff --git a/lib/spack/spack/operating_systems/cnl.py b/lib/spack/spack/operating_systems/cnl.py
index fc6c78b18a..7d21d6ef86 100644
--- a/lib/spack/spack/operating_systems/cnl.py
+++ b/lib/spack/spack/operating_systems/cnl.py
@@ -9,7 +9,7 @@ import llnl.util.tty as tty
 import llnl.util.multiproc as mp
 
 from spack.architecture import OperatingSystem
-from spack.util.module_cmd import get_module_cmd
+from spack.util.module_cmd import module
 
 
 class Cnl(OperatingSystem):
@@ -29,8 +29,7 @@ class Cnl(OperatingSystem):
         return self.name + str(self.version)
 
     def _detect_crayos_version(self):
-        modulecmd = get_module_cmd()
-        output = modulecmd("avail", "PrgEnv-", output=str, error=str)
+        output = module("avail", "PrgEnv-")
         matches = re.findall(r'PrgEnv-\w+/(\d+).\d+.\d+', output)
         major_versions = set(matches)
         latest_version = max(major_versions)
@@ -58,10 +57,7 @@ class Cnl(OperatingSystem):
             if not cmp_cls.PrgEnv_compiler:
                 tty.die('Must supply PrgEnv_compiler with PrgEnv')
 
-            modulecmd = get_module_cmd()
-
-            output = modulecmd(
-                'avail', cmp_cls.PrgEnv_compiler, output=str, error=str)
+            output = module('avail', cmp_cls.PrgEnv_compiler)
             version_regex = r'(%s)/([\d\.]+[\d])' % cmp_cls.PrgEnv_compiler
             matches = re.findall(version_regex, output)
             for name, version in matches:
diff --git a/lib/spack/spack/operating_systems/cray_frontend.py b/lib/spack/spack/operating_systems/cray_frontend.py
index 13e382030a..7b6359c5b8 100644
--- a/lib/spack/spack/operating_systems/cray_frontend.py
+++ b/lib/spack/spack/operating_systems/cray_frontend.py
@@ -6,7 +6,7 @@
 import os
 
 from spack.operating_systems.linux_distro import LinuxDistro
-from spack.util.module_cmd import get_module_cmd
+from spack.util.module_cmd import module
 
 
 class CrayFrontend(LinuxDistro):
@@ -41,10 +41,7 @@ class CrayFrontend(LinuxDistro):
             # into the PATH environment variable (i.e. the following modules:
             # 'intel', 'cce', 'gcc', etc.) will also be unloaded since they are
             # specified as prerequisites in the PrgEnv-* modulefiles.
-            modulecmd = get_module_cmd()
-            exec(compile(
-                modulecmd('unload', prg_env, output=str, error=os.devnull),
-                '<string>', 'exec'))
+            module('unload', prg_env)
 
         # Call the overridden method.
         clist = super(CrayFrontend, self).find_compilers(*paths)
diff --git a/lib/spack/spack/platforms/cray.py b/lib/spack/spack/platforms/cray.py
index 6814f5be8e..3499f0b164 100644
--- a/lib/spack/spack/platforms/cray.py
+++ b/lib/spack/spack/platforms/cray.py
@@ -11,7 +11,7 @@ from spack.util.executable import which
 from spack.architecture import Platform, Target, NoPlatformError
 from spack.operating_systems.cray_frontend import CrayFrontend
 from spack.operating_systems.cnl import Cnl
-from spack.util.module_cmd import get_module_cmd, unload_module
+from spack.util.module_cmd import module
 
 
 def _get_modules_in_modulecmd_output(output):
@@ -90,8 +90,8 @@ class Cray(Platform):
         # Unload these modules to prevent any silent linking or unnecessary
         # I/O profiling in the case of darshan.
         modules_to_unload = ["cray-mpich", "darshan", "cray-libsci", "altd"]
-        for module in modules_to_unload:
-            unload_module(module)
+        for mod in modules_to_unload:
+            module('unload', mod)
 
         env.set('CRAYPE_LINK_TYPE', 'dynamic')
         cray_wrapper_names = os.path.join(build_env_path, 'cray')
@@ -127,8 +127,7 @@ class Cray(Platform):
     def _avail_targets(self):
         '''Return a list of available CrayPE CPU targets.'''
         if getattr(self, '_craype_targets', None) is None:
-            module = get_module_cmd()
-            output = module('avail', '-t', 'craype-', output=str, error=str)
+            output = module('avail', '-t', 'craype-')
             craype_modules = _get_modules_in_modulecmd_output(output)
             self._craype_targets = targets = []
             _fill_craype_targets_from_modules(targets, craype_modules)
diff --git a/lib/spack/spack/test/build_environment.py b/lib/spack/spack/test/build_environment.py
index cc0cdad583..1927a76b38 100644
--- a/lib/spack/spack/test/build_environment.py
+++ b/lib/spack/spack/test/build_environment.py
@@ -205,7 +205,7 @@ def test_spack_paths_before_module_paths(
     assert paths.index(spack_path) < paths.index(module_path)
 
 
-def test_package_inheritance_module_setup(config, mock_packages):
+def test_package_inheritance_module_setup(config, mock_packages, working_env):
     s = spack.spec.Spec('multimodule-inheritance')
     s.concretize()
     pkg = s.package
@@ -217,8 +217,6 @@ def test_package_inheritance_module_setup(config, mock_packages):
     assert pkg.use_module_variable() == 'test_module_variable'
     assert os.environ['TEST_MODULE_VAR'] == 'test_module_variable'
 
-    os.environ.pop('TEST_MODULE_VAR')
-
 
 def test_set_build_environment_variables(
         config, mock_packages, working_env, monkeypatch,
diff --git a/lib/spack/spack/test/conftest.py b/lib/spack/spack/test/conftest.py
index b2d255c340..8107ffe122 100644
--- a/lib/spack/spack/test/conftest.py
+++ b/lib/spack/spack/test/conftest.py
@@ -142,7 +142,11 @@ def remove_whatever_it_is(path):
 def working_env():
     saved_env = os.environ.copy()
     yield
-    os.environ = saved_env
+    # os.environ = saved_env doesn't work
+    # it causes module_parsing::test_module_function to fail
+    # when it's run after any test using this fixutre
+    os.environ.clear()
+    os.environ.update(saved_env)
 
 
 @pytest.fixture(scope='function', autouse=True)
diff --git a/lib/spack/spack/test/flag_handlers.py b/lib/spack/spack/test/flag_handlers.py
index 1bf5042e30..094afe595a 100644
--- a/lib/spack/spack/test/flag_handlers.py
+++ b/lib/spack/spack/test/flag_handlers.py
@@ -17,7 +17,8 @@ from spack.pkgkit import inject_flags, env_flags, build_system_flags
 def temp_env():
     old_env = os.environ.copy()
     yield
-    os.environ = old_env
+    os.environ.clear()
+    os.environ.update(old_env)
 
 
 def add_o3_to_build_system_cflags(pkg, name, flags):
diff --git a/lib/spack/spack/test/module_parsing.py b/lib/spack/spack/test/module_parsing.py
index fc8111408f..0594724eaf 100644
--- a/lib/spack/spack/test/module_parsing.py
+++ b/lib/spack/spack/test/module_parsing.py
@@ -4,60 +4,75 @@
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
 import pytest
-import subprocess
 import os
+import spack
+
 from spack.util.module_cmd import (
+    module,
     get_path_from_module,
-    get_path_from_module_contents,
     get_path_arg_from_module_line,
-    get_module_cmd_from_bash,
-    get_module_cmd,
-    ModuleError)
+    get_path_from_module_contents
+)
+
+test_module_lines = ['prepend-path LD_LIBRARY_PATH /path/to/lib',
+                     'setenv MOD_DIR /path/to',
+                     'setenv LDFLAGS -Wl,-rpath/path/to/lib',
+                     'setenv LDFLAGS -L/path/to/lib',
+                     'prepend-path PATH /path/to/bin']
 
 
-env = os.environ.copy()
-env['LC_ALL'] = 'C'
-typeset_func = subprocess.Popen('module avail',
-                                env=env,
-                                stdout=subprocess.PIPE,
-                                stderr=subprocess.PIPE,
-                                shell=True)
-typeset_func.wait()
-typeset = typeset_func.stderr.read()
-MODULE_NOT_DEFINED = b'not found' in typeset
+@pytest.fixture
+def module_function_test_mode():
+    old_mode = spack.util.module_cmd._test_mode
+    spack.util.module_cmd._test_mode = True
+
+    yield
+
+    spack.util.module_cmd._test_mode = old_mode
 
 
 @pytest.fixture
-def save_env():
-    old_path = os.environ.get('PATH', None)
-    old_bash_func = os.environ.get('BASH_FUNC_module()', None)
+def save_module_func():
+    old_func = spack.util.module_cmd.module
 
     yield
 
-    if old_path:
-        os.environ['PATH'] = old_path
-    if old_bash_func:
-        os.environ['BASH_FUNC_module()'] = old_bash_func
+    spack.util.module_cmd.module = old_func
 
 
-def test_get_path_from_module(save_env):
-    lines = ['prepend-path LD_LIBRARY_PATH /path/to/lib',
-             'prepend-path CRAY_LD_LIBRARY_PATH /path/to/lib',
-             'setenv MOD_DIR /path/to',
-             'setenv LDFLAGS -Wl,-rpath/path/to/lib',
-             'setenv LDFLAGS -L/path/to/lib',
-             'prepend-path PATH /path/to/bin']
+def test_module_function_change_env(tmpdir, working_env,
+                                    module_function_test_mode):
+    src_file = str(tmpdir.join('src_me'))
+    with open(src_file, 'w') as f:
+        f.write('export TEST_MODULE_ENV_VAR=TEST_SUCCESS\n')
 
-    for line in lines:
-        module_func = '() { eval `echo ' + line + ' bash filler`\n}'
-        os.environ['BASH_FUNC_module()'] = module_func
-        path = get_path_from_module('mod')
-        assert path == '/path/to'
+    os.environ['NOT_AFFECTED'] = "NOT_AFFECTED"
+    module('load', src_file)
+
+    assert os.environ['TEST_MODULE_ENV_VAR'] == 'TEST_SUCCESS'
+    assert os.environ['NOT_AFFECTED'] == "NOT_AFFECTED"
 
-    os.environ['BASH_FUNC_module()'] = '() { eval $(echo fill bash $*)\n}'
-    path = get_path_from_module('mod')
 
-    assert path is None
+def test_module_function_no_change(tmpdir, module_function_test_mode):
+    src_file = str(tmpdir.join('src_me'))
+    with open(src_file, 'w') as f:
+        f.write('echo TEST_MODULE_FUNCTION_PRINT')
+
+    old_env = os.environ.copy()
+    text = module('show', src_file)
+
+    assert text == 'TEST_MODULE_FUNCTION_PRINT\n'
+    assert os.environ == old_env
+
+
+def test_get_path_from_module_faked(save_module_func):
+    for line in test_module_lines:
+        def fake_module(*args):
+            return line
+        spack.util.module_cmd.module = fake_module
+
+        path = get_path_from_module('mod')
+        assert path == '/path/to'
 
 
 def test_get_path_from_module_contents():
@@ -106,62 +121,3 @@ def test_get_argument_from_module_line():
     for bl in bad_lines:
         with pytest.raises(ValueError):
             get_path_arg_from_module_line(bl)
-
-
-@pytest.mark.skipif(MODULE_NOT_DEFINED, reason='Depends on defined module fn')
-def test_get_module_cmd_from_bash_using_modules():
-    module_list_proc = subprocess.Popen(['module list'],
-                                        stdout=subprocess.PIPE,
-                                        stderr=subprocess.STDOUT,
-                                        executable='/bin/bash',
-                                        shell=True)
-    module_list_proc.wait()
-    module_list = module_list_proc.stdout.read()
-
-    module_cmd = get_module_cmd_from_bash()
-    module_cmd_list = module_cmd('list', output=str, error=str)
-
-    # Lmod command reprints some env variables on every invocation.
-    # Test containment to avoid false failures on lmod systems.
-    assert module_list in module_cmd_list
-
-
-def test_get_module_cmd_from_bash_ticks(save_env):
-    os.environ['BASH_FUNC_module()'] = '() { eval `echo bash $*`\n}'
-
-    module_cmd = get_module_cmd()
-    module_cmd_list = module_cmd('list', output=str, error=str)
-
-    assert module_cmd_list == 'python list\n'
-
-
-def test_get_module_cmd_from_bash_parens(save_env):
-    os.environ['BASH_FUNC_module()'] = '() { eval $(echo fill sh $*)\n}'
-
-    module_cmd = get_module_cmd()
-    module_cmd_list = module_cmd('list', output=str, error=str)
-
-    assert module_cmd_list == 'fill python list\n'
-
-
-def test_get_module_cmd_fails(save_env):
-    os.environ.pop('BASH_FUNC_module()')
-    os.environ.pop('PATH')
-    with pytest.raises(ModuleError):
-        module_cmd = get_module_cmd(b'--norc')
-        module_cmd()  # Here to avoid Flake F841 on previous line
-
-
-def test_get_module_cmd_from_which(tmpdir, save_env):
-    f = tmpdir.mkdir('bin').join('modulecmd')
-    f.write('#!/bin/bash\n'
-            'echo $*')
-    f.chmod(0o770)
-
-    os.environ['PATH'] = str(tmpdir.join('bin')) + ':' + os.environ['PATH']
-    os.environ.pop('BASH_FUNC_module()')
-
-    module_cmd = get_module_cmd(b'--norc')
-    module_cmd_list = module_cmd('list', output=str, error=str)
-
-    assert module_cmd_list == 'python list\n'
diff --git a/lib/spack/spack/util/module_cmd.py b/lib/spack/spack/util/module_cmd.py
index be8607154c..9f29d76711 100644
--- a/lib/spack/spack/util/module_cmd.py
+++ b/lib/spack/spack/util/module_cmd.py
@@ -8,108 +8,53 @@ This module contains routines related to the module command for accessing and
 parsing environment modules.
 """
 import subprocess
-import re
 import os
-import llnl.util.tty as tty
-from spack.util.executable import which
-
-
-def get_module_cmd(bashopts=''):
-    try:
-        return get_module_cmd_from_bash(bashopts)
-    except ModuleError:
-        # Don't catch the exception this time; we have no other way to do it.
-        tty.warn("Could not detect module function from bash."
-                 " Trying to detect modulecmd from `which`")
-        try:
-            return get_module_cmd_from_which()
-        except ModuleError:
-            raise ModuleError('Spack requires modulecmd or a defined module'
-                              ' function. Make sure modulecmd is in your path'
-                              ' or the function "module" is defined in your'
-                              ' bash environment.')
-
-
-def get_module_cmd_from_which():
-    module_cmd = which('modulecmd')
-    if not module_cmd:
-        raise ModuleError('`which` did not find any modulecmd executable')
-    module_cmd.add_default_arg('python')
-
-    # Check that the executable works
-    module_cmd('list', output=str, error=str, fail_on_error=False)
-    if module_cmd.returncode != 0:
-        raise ModuleError('get_module_cmd cannot determine the module command')
-
-    return module_cmd
+import json
+import re
 
+import llnl.util.tty as tty
 
-def get_module_cmd_from_bash(bashopts=''):
-    # Find how the module function is defined in the environment
-    module_func = os.environ.get('BASH_FUNC_module()', None)
-    if module_func:
-        module_func = os.path.expandvars(module_func)
-    else:
-        module_func_proc = subprocess.Popen(['{0} typeset -f module | '
-                                             'envsubst'.format(bashopts)],
-                                            stdout=subprocess.PIPE,
-                                            stderr=subprocess.STDOUT,
-                                            executable='/bin/bash',
-                                            shell=True)
-        module_func_proc.wait()
-        module_func = module_func_proc.stdout.read()
-
-    # Find the portion of the module function that is evaluated
-    try:
-        find_exec = re.search(r'.*`(.*(:? bash | sh ).*)`.*', module_func)
-        exec_line = find_exec.group(1)
-    except BaseException:
-        try:
-            # This will fail with nested parentheses. TODO: expand regex.
-            find_exec = re.search(r'.*\(([^()]*(:? bash | sh )[^()]*)\).*',
-                                  module_func)
-            exec_line = find_exec.group(1)
-        except BaseException:
-            raise ModuleError('get_module_cmd cannot '
-                              'determine the module command from bash')
-
-    # Create an executable
-    args = exec_line.split()
-    module_cmd = which(args[0])
-    if module_cmd:
-        for arg in args[1:]:
-            if arg in ('bash', 'sh'):
-                module_cmd.add_default_arg('python')
-                break
-            else:
-                module_cmd.add_default_arg(arg)
+# This list is not exhaustive. Currently we only use load and unload
+# If we need another option that changes the environment, add it here.
+module_change_commands = ['load', 'swap', 'unload', 'purge', 'use', 'unuse']
+py_cmd = "'import os\nimport json\nprint(json.dumps(dict(os.environ)))'"
+
+# This is just to enable testing. I hate it but we can't find a better way
+_test_mode = False
+
+
+def module(*args):
+    module_cmd = 'module ' + ' '.join(args) + ' 2>&1'
+    if _test_mode:
+        tty.warn('module function operating in test mode')
+        module_cmd = ". %s 2>&1" % args[1]
+    if args[0] in module_change_commands:
+        # Do the module manipulation, then output the environment in JSON
+        # and read the JSON back in the parent process to update os.environ
+        module_cmd += ' >/dev/null; python -c %s' % py_cmd
+        module_p  = subprocess.Popen(module_cmd,
+                                     stdout=subprocess.PIPE,
+                                     stderr=subprocess.STDOUT,
+                                     shell=True)
+
+        # Cray modules spit out warnings that we cannot supress.
+        # This hack skips to the last output (the environment)
+        env_output = str(module_p.communicate()[0].decode())
+        print(env_output)
+        env = env_output.strip().split('\n')[-1]
+
+        # Update os.environ with new dict
+        env_dict = json.loads(env)
+        os.environ.clear()
+        os.environ.update(env_dict)
     else:
-        raise ModuleError('Could not create executable based on module'
-                          ' function.')
-
-    # Check that the executable works
-    module_cmd('list', output=str, error=str, fail_on_error=False)
-    if module_cmd.returncode != 0:
-        raise ModuleError('get_module_cmd cannot determine the module command'
-                          'from bash.')
-
-    return module_cmd
-
-
-def unload_module(mod):
-    """Takes a module name and unloads the module from the environment. It does
-    not check whether conflicts arise from the unloaded module"""
-    tty.debug("Unloading module: {0}".format(mod))
-
-    modulecmd = get_module_cmd()
-    unload_output = modulecmd('unload', mod, output=str, error=str)
-
-    try:
-        exec(compile(unload_output, '<string>', 'exec'))
-    except Exception:
-        tty.debug("Module unload output of {0}:\n{1}\n".format(
-            mod, unload_output))
-        raise
+        # Simply execute commands that don't change state and return output
+        module_p = subprocess.Popen(module_cmd,
+                                    stdout=subprocess.PIPE,
+                                    stderr=subprocess.STDOUT,
+                                    shell=True)
+        # Decode and str to return a string object in both python 2 and 3
+        return str(module_p.communicate()[0].decode())
 
 
 def load_module(mod):
@@ -117,37 +62,18 @@ def load_module(mod):
     load that module. It then loads the provided module. Depends on the
     modulecmd implementation of modules used in cray and lmod.
     """
-    tty.debug("Loading module: {0}".format(mod))
-
-    # Create an executable of the module command that will output python code
-    modulecmd = get_module_cmd()
-
     # Read the module and remove any conflicting modules
     # We do this without checking that they are already installed
     # for ease of programming because unloading a module that is not
     # loaded does nothing.
-    module_content = modulecmd('show', mod, output=str, error=str)
-    text = module_content.split()
-    try:
-        for i, word in enumerate(text):
-            if word == 'conflict':
-                unload_module(text[i + 1])
-    except Exception:
-        tty.debug("Module show output of {0}:\n{1}\n".format(
-            mod, module_content))
-        raise
+    text = module('show', mod).split()
+    for i, word in enumerate(text):
+        if word == 'conflict':
+            module('unload', text[i + 1])
 
     # Load the module now that there are no conflicts
     # Some module systems use stdout and some use stderr
-    load = modulecmd('load', mod, output=str, error='/dev/null')
-    if not load:
-        load = modulecmd('load', mod, error=str)
-
-    try:
-        exec(compile(load, '<string>', 'exec'))
-    except Exception:
-        tty.debug("Module load output of {0}:\n{1}\n".format(mod, load))
-        raise
+    module('load', mod)
 
 
 def get_path_arg_from_module_line(line):
@@ -172,11 +98,8 @@ def get_path_from_module(mod):
     """Inspects a TCL module for entries that indicate the absolute path
     at which the library supported by said module can be found.
     """
-    # Create a modulecmd executable
-    modulecmd = get_module_cmd()
-
     # Read the module
-    text = modulecmd('show', mod, output=str, error=str).split('\n')
+    text = module('show', mod).split('\n')
 
     p = get_path_from_module_contents(text, mod)
     if p and not os.path.exists(p):
@@ -229,7 +152,3 @@ def get_path_from_module_contents(text, module_name):
 
     # Unable to find module path
     return None
-
-
-class ModuleError(Exception):
-    """Raised the the module_cmd utility to indicate errors."""