diff --git a/toolchain/bootstrap/modules.sh b/toolchain/bootstrap/modules.sh index d68811b2f..0cdc4d48a 100644 --- a/toolchain/bootstrap/modules.sh +++ b/toolchain/bootstrap/modules.sh @@ -85,6 +85,87 @@ __extract() { __combine "$(grep -E "^$1\s+" toolchain/modules | sed "s/^$1\s\+//")" } +# A toolchain/modules line is a list of module names optionally followed by +# KEY=value assignments. Two shapes pull in opposite directions: +# +# * several assignments on one line -- CC=nvc CXX=nvc++ FC=nvfortran +# * a single value that itself contains spaces -- linker/compiler flags +# +# Splitting on whitespace truncates the latter; splitting on the first '=' +# merges the former. Both readers below therefore share one rule: a word shaped +# like an identifier followed by '=' opens a new assignment, and every word +# after it belongs to that assignment's value until the next such word. Words +# ahead of the first assignment are module names. + +# Suppress globbing while a line is word-split, so a value such as -Wl,* cannot +# pick up filenames from the working directory. Records the caller's setting. +__noglob_push() { + __noglob_prev=0 + case "$-" in *f*) __noglob_prev=1 ;; esac + set -f +} + +__noglob_pop() { + [ "$__noglob_prev" -eq 0 ] && set +f + unset __noglob_prev + return 0 +} + +# Echo only the module names on a line, dropping assignments and their values. +__module_words() { + local _word _out="" _in_assignment=0 + + __noglob_push + for _word in $1; do + if [[ "$_word" =~ ^[A-Za-z_][A-Za-z0-9_]*= ]]; then + _in_assignment=1 + fi + [ "$_in_assignment" -eq 0 ] && _out="$_out $_word" + done + __noglob_pop + + echo "${_out# }" +} + +# Export the KEY=value assignments on one toolchain/modules line. +# +# Values are still passed through eval so that "$VAR" references to previously +# exported variables keep expanding, as they always have. +__export_assignments() { + local _entry="$1" _word _acc="" _key _val + + __noglob_push + + __flush() { + [ -z "$_acc" ] && return 0 + _key="${_acc%%=*}" + _val="${_acc#*=}" + # One round of expansion (so "$VAR" references still work), then export + # the result directly -- re-evaluating it would word-split the value + # again, which is the bug this function exists to avoid. The expansion + # goes through the positional parameters rather than echo, so a value + # such as -n or -e is not mistaken for an option to echo. + eval "set -- $_val" + _val="$*" + log " \$ export $_key=$_val" + export "$_key=$_val" + _acc="" + } + + for _word in $_entry; do + if [[ "$_word" =~ ^[A-Za-z_][A-Za-z0-9_]*= ]]; then + __flush + _acc="$_word" + else + _acc="$_acc $_word" + fi + done + __flush + unset -f __flush + + __noglob_pop +} + COMPUTER="$(__extract "$u_c")" if [[ -z "$COMPUTER" ]]; then @@ -107,8 +188,18 @@ else module purge > /dev/null 2>&1 fi -ELEMENTS="$(__extract "$u_c-all") $(__extract "$u_c-$cg")" -MODULES=`echo "$ELEMENTS" | tr ' ' '\n' | grep -v = | xargs` +# Collect the module names a line at a time. __extract() concatenates every +# matching line, which would let one line's assignments swallow the next line's +# module names, so the assignment-aware split has to happen per line. +MODULES="" +for _suffix in "all" "$cg"; do + while IFS= read -r _entry; do + _entry="$(__module_words "$_entry")" + [ -n "$_entry" ] && MODULES="$MODULES $_entry" + done < <(grep -E "^$u_c-$_suffix\s+" toolchain/modules | sed "s/^$u_c-$_suffix\s\+//") +done +MODULES="$(echo "$MODULES" | xargs)" +unset _suffix _entry log " $ module load $MODULES" if ! module load $MODULES; then @@ -122,8 +213,7 @@ fi for _suffix in "all" "$cg"; do while IFS= read -r _entry; do if echo "$_entry" | grep -q '='; then - log " $ export $(eval "echo \"$_entry\"")" - eval "export $_entry" + __export_assignments "$_entry" fi done < <(grep -E "^$u_c-$_suffix\s+" toolchain/modules | sed "s/^$u_c-$_suffix\s\+//") done diff --git a/toolchain/mfc/bootstrap_tests/__init__.py b/toolchain/mfc/bootstrap_tests/__init__.py new file mode 100644 index 000000000..ae61419a4 --- /dev/null +++ b/toolchain/mfc/bootstrap_tests/__init__.py @@ -0,0 +1 @@ +"""Tests for the shell bootstrap scripts under toolchain/bootstrap/.""" diff --git a/toolchain/mfc/bootstrap_tests/test_modules_env.py b/toolchain/mfc/bootstrap_tests/test_modules_env.py new file mode 100644 index 000000000..d1bbd5f35 --- /dev/null +++ b/toolchain/mfc/bootstrap_tests/test_modules_env.py @@ -0,0 +1,249 @@ +"""Tests for the toolchain/modules line reader in toolchain/bootstrap/modules.sh. + +`. ./mfc.sh load` reads each `toolchain/modules` line twice: once to collect the +module names it should `module load`, and once to export the KEY=value +assignments. Those lines carry two shapes that pull in opposite directions: + + * several assignments on one line -- `CC=nvc CXX=nvc++ FC=nvfortran` + * a single value that itself contains spaces -- linker/compiler flags + +Both readers have to agree on where an assignment starts and ends, or a value's +words leak into the module list. The tests below drive the real shell functions +out of modules.sh rather than a copy of them, so they fail if the implementation +regresses. +""" + +import re +import shlex +import shutil +import subprocess +import textwrap +from pathlib import Path + +import pytest + +MODULES_SH = Path(__file__).resolve().parents[2] / "bootstrap" / "modules.sh" + +pytestmark = pytest.mark.skipif(shutil.which("bash") is None, reason="bash is required to exercise modules.sh") + + +def _extract_function(name): + """Return the source text of a shell function defined in modules.sh. + + modules.sh is meant to be `source`d with arguments and runs an interactive + loader at import time, so the function under test is lifted out on its own. + """ + text = MODULES_SH.read_text() + match = re.search(rf"^{name}\(\) \{{.*?^\}}", text, re.MULTILINE | re.DOTALL) + assert match, f"{name}() not found in {MODULES_SH}" + return match.group(0) + + +def _helpers(*names): + """Return the source of the named functions plus the glob-guard helpers.""" + return "\n".join(_extract_function(n) for n in ("__noglob_push", "__noglob_pop", *names)) + + +def _load(entry, preset_env=None): + """Run one `toolchain/modules` entry through the exporter. + + Returns the resulting environment as a dict, so a test can assert on the + variables the entry was supposed to set. + """ + preamble = "\n".join(f"export {k}={shlex.quote(v)}" for k, v in (preset_env or {}).items()) + script = textwrap.dedent( + """\ + log() {{ :; }} # modules.sh logging stub + {preamble} + {function_src} + __export_assignments {entry} + env + """ + ).format(preamble=preamble, function_src=_helpers("__export_assignments"), entry=_quote(entry)) + + proc = subprocess.run(["bash", "-c", script], capture_output=True, text=True, check=True) + env = {} + for line in proc.stdout.splitlines(): + key, sep, value = line.partition("=") + if sep: + env[key] = value + return env + + +def _module_words(entry): + """Return the module names `. ./mfc.sh load` would pass to `module load`.""" + script = textwrap.dedent( + """\ + {function_src} + __module_words {entry} + """ + ).format(function_src=_helpers("__module_words"), entry=_quote(entry)) + + proc = subprocess.run(["bash", "-c", script], capture_output=True, text=True, check=True) + return proc.stdout.strip() + + +def _quote(value): + """Single-quote a string for safe interpolation into the bash script.""" + return "'" + value.replace("'", "'\\''") + "'" + + +@pytest.mark.parametrize( + "entry,expected", + [ + # Every multi-assignment shape currently in toolchain/modules. + ( + "MFC_CUDA_CC=70,75,80 NVHPC_CUDA_HOME=$CUDA_HOME CC=nvc CXX=nvc++ FC=nvfortran", + {"MFC_CUDA_CC": "70,75,80", "NVHPC_CUDA_HOME": "/opt/cuda", "CC": "nvc", "CXX": "nvc++", "FC": "nvfortran"}, + ), + ("CC=nvc CXX=nvc++ FC=nvfortran", {"CC": "nvc", "CXX": "nvc++", "FC": "nvfortran"}), + ("MPICC=mpiicc MPICXX=mpiicpc MPIFC=mpiifort", {"MPICC": "mpiicc", "MPICXX": "mpiicpc", "MPIFC": "mpiifort"}), + ('UCX_NET_DEVICES="mlx5_4:1,mlx5_7:1"', {"UCX_NET_DEVICES": "mlx5_4:1,mlx5_7:1"}), + ('PYTHONPATH=""', {"PYTHONPATH": ""}), + ("NVHPC_CUDA_HOME=$CUDA_HOME", {"NVHPC_CUDA_HOME": "/opt/cuda"}), + ], +) +def test_existing_module_lines_still_load(entry, expected): + """Assignment shapes already in toolchain/modules keep their old meaning.""" + env = _load(entry, preset_env={"CUDA_HOME": "/opt/cuda"}) + for key, value in expected.items(): + assert env.get(key) == value + + +def test_multi_word_value_survives_intact(): + """A value with spaces is exported whole, not truncated to its first word. + + Written unquoted, `CRAY_CCE_LLD_ARGS` used to keep only the first flag: the + rest were treated as further names to export and the failure went to stderr, + which `./mfc.sh load` routinely discards. The build then succeeds with a + compiler workaround silently missing. + """ + flags = "-plugin-opt=-mattr=-mai-insts -plugin-opt=-disable-promote-alloca-to-vector" + env = _load(f"CRAY_CCE_LLD_ARGS={flags}") + assert env["CRAY_CCE_LLD_ARGS"] == flags + + +def test_quoted_multi_word_value_survives_intact(): + """Quoting the value in toolchain/modules works too, and drops the quotes.""" + env = _load('CRAY_CCE_LLD_ARGS="-O2 -g"') + assert env["CRAY_CCE_LLD_ARGS"] == "-O2 -g" + + +def test_multi_word_value_followed_by_another_assignment(): + """A space-carrying value ends where the next NAME= begins.""" + env = _load("CFLAGS=-O2 -march=native CC=gcc") + assert env["CFLAGS"] == "-O2 -march=native" + assert env["CC"] == "gcc" + + +def test_variable_reference_inside_a_multi_word_value(): + """Values are still expanded once, so they can reference earlier exports.""" + env = _load("LDFLAGS=-L$CUDA_HOME/lib64 -lcudart", preset_env={"CUDA_HOME": "/opt/cuda"}) + assert env["LDFLAGS"] == "-L/opt/cuda/lib64 -lcudart" + + +def test_glob_characters_in_a_value_are_not_expanded(): + """A value such as `-Wl,*` must not pick up filenames from the cwd.""" + env = _load("MYFLAGS=-Wl,* -O2") + assert env["MYFLAGS"] == "-Wl,* -O2" + + +@pytest.mark.parametrize("value", ["-n", "-e", "-E", "-n -e trailing"]) +def test_value_that_looks_like_an_echo_option_survives(value): + """A value starting with -n/-e must be exported, not read as an option. + + Expanding through `echo` made bash's builtin treat these as flags, so + `MYVAR=-n` exported an empty string. Expansion now runs through the + positional parameters instead. + """ + env = _load(f"MYVAR={value}") + assert env["MYVAR"] == value + + +def test_return_status_is_zero_when_the_caller_already_set_noglob(): + """`set -f` in the caller must not make the loader report failure.""" + script = textwrap.dedent( + """\ + log() {{ :; }} + {function_src} + set -f + __export_assignments 'CC=gcc' + echo "status=$?" + """ + ).format(function_src=_helpers("__export_assignments")) + proc = subprocess.run(["bash", "-c", script], capture_output=True, text=True, check=True) + assert "status=0" in proc.stdout + + +def test_every_assignment_line_in_toolchain_modules_loads_cleanly(): + """No line shipped in toolchain/modules may error out or lose a value.""" + modules_file = MODULES_SH.resolve().parents[1] / "modules" + entries = [] + for line in modules_file.read_text().splitlines(): + match = re.match(r"^[a-z][a-z0-9]*-(?:all|cpu|gpu)\s+(.*)$", line) + if match and "=" in match.group(1): + entries.append(match.group(1)) + + assert entries, "expected at least one KEY=value line in toolchain/modules" + + for entry in entries: + env = _load(entry, preset_env={"CUDA_HOME": "/opt/cuda", "OLCF_AFAR_ROOT": "/sw/afar"}) + for name in re.findall(r"(?:^|\s)([A-Za-z_][A-Za-z0-9_]*)=", entry): + assert name in env, f"{name} was not exported by: {entry}" + + +# The module-name side of the same line. `module load` receives whatever these +# return, and an unknown name aborts `. ./mfc.sh load` outright, so a value word +# leaking into this list is a hard failure rather than a cosmetic one. + + +@pytest.mark.parametrize( + "entry,expected", + [ + ("python cmake/3.22.2", "python cmake/3.22.2"), + ("nvhpc/22.9 cuda/11.7 openmpi/4.0.5-nvhpc22.9", "nvhpc/22.9 cuda/11.7 openmpi/4.0.5-nvhpc22.9"), + ("CC=nvc CXX=nvc++ FC=nvfortran", ""), + ("python cmake CC=nvc CXX=nvc++", "python cmake"), + ], +) +def test_module_names_are_separated_from_assignments(entry, expected): + """Module names come first; everything from the first NAME= on is a value.""" + assert _module_words(entry) == expected + + +@pytest.mark.parametrize( + "entry", + [ + # Words of a multi-word value that happen to lack '=' of their own. + "LDFLAGS=-L$CUDA_HOME/lib64 -lcudart", + "CFLAGS=-O2 -march=native CC=gcc", + 'CRAY_CCE_LLD_ARGS="-O2 -g"', + "MYFLAGS=-Wl,* -O2", + ], +) +def test_value_words_are_never_passed_to_module_load(entry): + """A value's own words must not be mistaken for module names. + + The classification step used to drop every word containing '=' and keep the + rest, so `LDFLAGS=-L$CUDA_HOME/lib64 -lcudart` handed `-lcudart` to + `module load`. That fails and trips `error "Failed to load modules."`, + aborting the loader before any variable is exported. + """ + assert _module_words(entry) == "" + + +def test_module_names_in_toolchain_modules_are_unchanged(): + """Every line shipped today yields exactly the module list it always did.""" + modules_file = MODULES_SH.resolve().parents[1] / "modules" + checked = 0 + for line in modules_file.read_text().splitlines(): + match = re.match(r"^[a-z][a-z0-9]*-(?:all|cpu|gpu)\s+(.*)$", line) + if not match: + continue + entry = match.group(1) + # What the pre-existing implementation produced: drop '='-bearing words. + expected = " ".join(word for word in entry.split() if "=" not in word) + assert _module_words(entry) == expected, entry + checked += 1 + + assert checked, "expected at least one module line in toolchain/modules"