From 3d3cea1c870b9be391437023163ed2608e8121a1 Mon Sep 17 00:00:00 2001 From: Greg Becker Date: Fri, 10 May 2019 07:04:24 +0900 Subject: 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. --- lib/spack/spack/operating_systems/cnl.py | 10 +- lib/spack/spack/operating_systems/cray_frontend.py | 7 +- lib/spack/spack/platforms/cray.py | 9 +- lib/spack/spack/test/build_environment.py | 4 +- lib/spack/spack/test/conftest.py | 6 +- lib/spack/spack/test/flag_handlers.py | 3 +- lib/spack/spack/test/module_parsing.py | 148 ++++++----------- lib/spack/spack/util/module_cmd.py | 179 ++++++--------------- 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), - '', '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, '', '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, '', '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.""" -- cgit v1.2.3-60-g2f50