spack module_cmd: set LD_LIBRARY_PATH for python in subshell. (#16827)

* Ensure python runs with proper ld_library_path within Spack regardless of environment

3 files changed, 46 insertions, 8 deletions

diff --git a/lib/spack/spack/main.py b/lib/spack/spack/main.py
index 3677764cb6..bdeef049ff 100644
--- a/lib/spack/spack/main.py
+++ b/lib/spack/spack/main.py
@@ -92,6 +92,7 @@ required_command_properties = ['level', 'section', 'description']
 
 #: Recorded directory where spack command was originally invoked
 spack_working_dir = None
+spack_ld_library_path = os.environ.get('LD_LIBRARY_PATH', '')
 
 
 def set_working_dir():
diff --git a/lib/spack/spack/test/util/executable.py b/lib/spack/spack/test/util/executable.py
index 26719e94d1..8baae3b453 100644
--- a/lib/spack/spack/test/util/executable.py
+++ b/lib/spack/spack/test/util/executable.py
@@ -4,18 +4,19 @@
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
 import sys
+import os
 
 import llnl.util.filesystem as fs
-
+import spack
 import spack.util.executable as ex
 from spack.hooks.sbang import filter_shebangs_in_directory
 
 
-def test_read_unicode(tmpdir):
+def test_read_unicode(tmpdir, working_env):
     script_name = 'print_unicode.py'
 
     with tmpdir.as_cwd():
-
+        os.environ['LD_LIBRARY_PATH'] = spack.main.spack_ld_library_path
         # make a script that prints some unicode
         with open(script_name, 'w') as f:
             f.write('''#!{0}
diff --git a/lib/spack/spack/util/module_cmd.py b/lib/spack/spack/util/module_cmd.py
index 143ad3d43e..a5ebe73c84 100644
--- a/lib/spack/spack/util/module_cmd.py
+++ b/lib/spack/spack/util/module_cmd.py
@@ -13,12 +13,13 @@ import sys
 import json
 import re
 
+import spack
 import llnl.util.tty as tty
 
 # 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;import json;print(json.dumps(dict(os.environ)))'"
+py_cmd = 'import os;import json;print(json.dumps(dict(os.environ)))'
 _cmd_template = "'module ' + ' '.join(args) + ' 2>&1'"
 
 
@@ -27,7 +28,31 @@ def module(*args):
     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;' + sys.executable + ' -c %s' % py_cmd
+        # For python, we use the same python running the Spack process, because
+        # we can guarantee its existence. We have to do some LD_LIBRARY_PATH
+        # shenanigans to ensure python will run.
+
+        # LD_LIBRARY_PATH under which Spack ran
+        os.environ['SPACK_LD_LIBRARY_PATH'] = spack.main.spack_ld_library_path
+
+        # suppress output from module function
+        module_cmd += ' >/dev/null;'
+
+        # Capture the new LD_LIBRARY_PATH after `module` was run
+        module_cmd += 'export SPACK_NEW_LD_LIBRARY_PATH="$LD_LIBRARY_PATH";'
+
+        # Set LD_LIBRARY_PATH to value at Spack startup time to ensure that
+        # python executable finds its libraries
+        module_cmd += 'LD_LIBRARY_PATH="$SPACK_LD_LIBRARY_PATH" '
+
+        # Execute the python command
+        module_cmd += '%s -c "%s";' % (sys.executable, py_cmd)
+
+        # If LD_LIBRARY_PATH was set after `module`, dump the old value because
+        # we have since corrupted it to ensure python would run.
+        # dump SPACKIGNORE as a placeholder for parsing if LD_LIBRARY_PATH null
+        module_cmd += 'echo "${SPACK_NEW_LD_LIBRARY_PATH:-SPACKIGNORE}"'
+
         module_p  = subprocess.Popen(module_cmd,
                                      stdout=subprocess.PIPE,
                                      stderr=subprocess.STDOUT,
@@ -36,13 +61,24 @@ def module(*args):
 
         # 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())
-        env = env_output.strip().split('\n')[-1]
+        env_out = str(module_p.communicate()[0].decode()).strip().split('\n')
+
+        # The environment dumped as json
+        env_json = env_out[-2]
+        # Either the uncorrupted $LD_LIBRARY_PATH or SPACKIGNORE
+        new_ld_library_path = env_out[-1]
 
         # Update os.environ with new dict
-        env_dict = json.loads(env)
+        env_dict = json.loads(env_json)
         os.environ.clear()
         os.environ.update(env_dict)
+
+        # Override restored LD_LIBRARY_PATH with pre-python value
+        if new_ld_library_path == 'SPACKIGNORE':
+            os.environ.pop('LD_LIBRARY_PATH', None)
+        else:
+            os.environ['LD_LIBRARY_PATH'] = new_ld_library_path
+
     else:
         # Simply execute commands that don't change state and return output
         module_p = subprocess.Popen(module_cmd,