Fix recursive module find for upstream dependencies (spack#11304)

"spack module tcl find -r <spec>" (and equivalents for other module
systems) was failing when a dependency was installed in an upstream
Spack instance. This updates the module index to handle locating
module files for upstream Spack installations (encapsulating the
logic in a new class called UpstreamModuleIndex); the updated index
handles the case where a Spack installation has multiple upstream
instances.

Note that if a module is not available locally but we are using the
local package, then we shouldn't use a module (i.e. if the package is
also installed upstream, and there is a module file for it, Spack
should not use that module). Likewise, if we are instance X using
upstreams Y and Z like X->Y->Z, and if we are using a package from
instance Y, then we should only use a module from instance Y. This
commit includes tests to check that this is handled properly.

lib/spack/spack/cmd/modules/__init__.py

@@ -15,6 +15,7 @@
 import spack.cmd
 import spack.modules
 import spack.repo
+import spack.modules.common
 
 import spack.cmd.common.arguments as arguments
 
@@ -162,49 +163,24 @@ def loads(module_type, specs, args, out=sys.stdout):
 
 
 def find(module_type, specs, args):
-    """Returns the module file "use" name if there's a single match. Raises
-    error messages otherwise.
-    """
-
-    spec = one_spec_or_raise(specs)
-
-    if spec.package.installed_upstream:
-        module = spack.modules.common.upstream_module(spec, module_type)
-        if module:
-            print(module.path)
-        return
+    """Retrieve paths or use names of module files"""
 
-    # Check if the module file is present
-    def module_exists(spec):
-        writer = spack.modules.module_types[module_type](spec)
-        return os.path.isfile(writer.layout.filename)
+    single_spec = one_spec_or_raise(specs)
 
-    if not module_exists(spec):
-        msg = 'Even though {1} is installed, '
-        msg += 'no {0} module has been generated for it.'
-        tty.die(msg.format(module_type, spec))
-
-    # Check if we want to recurse and load all dependencies. In that case
-    # modify the list of specs adding all the dependencies in post order
     if args.recurse_dependencies:
-        specs = [
-            item for item in spec.traverse(order='post', cover='nodes')
-            if module_exists(item)
-        ]
-
-    # ... and if it is print its use name or full-path if requested
-    def module_str(specs):
-        modules = []
-        for x in specs:
-            writer = spack.modules.module_types[module_type](x)
-            if args.full_path:
-                modules.append(writer.layout.filename)
-            else:
-                modules.append(writer.layout.use_name)
-
-        return ' '.join(modules)
-
-    print(module_str(specs))
+        specs_to_retrieve = list(
+            single_spec.traverse(order='post', cover='nodes',
+                                 deptype=('link', 'run')))
+    else:
+        specs_to_retrieve = [single_spec]
+
+    try:
+        modules = [spack.modules.common.get_module(module_type, spec,
+                                                   args.full_path)
+                   for spec in specs_to_retrieve]
+    except spack.modules.common.ModuleNotFoundError as e:
+        tty.die(e.message)
+    print(' '.join(modules))
 
 
 def rm(module_type, specs, args):

lib/spack/spack/modules/common.py

@@ -238,6 +238,16 @@ def generate_module_index(root, modules):
         syaml.dump(index, index_file, default_flow_style=False)
 
 
+def _generate_upstream_module_index():
+    module_indices = read_module_indices()
+
+    return UpstreamModuleIndex(spack.store.db, module_indices)
+
+
+upstream_module_index = llnl.util.lang.Singleton(
+    _generate_upstream_module_index)
+
+
 ModuleIndexEntry = collections.namedtuple(
     'ModuleIndexEntry', ['path', 'use_name'])
 
@@ -247,38 +257,90 @@ def read_module_index(root):
     if not os.path.exists(index_path):
         return {}
     with open(index_path, 'r') as index_file:
-        yaml_content = syaml.load(index_file)
-        index = {}
-        yaml_index = yaml_content['module_index']
-        for dag_hash, module_properties in yaml_index.items():
-            index[dag_hash] = ModuleIndexEntry(
-                module_properties['path'],
-                module_properties['use_name'])
-        return index
+        return _read_module_index(index_file)
+
+
+def _read_module_index(str_or_file):
+    """Read in the mapping of spec hash to module location/name. For a given
+       Spack installation there is assumed to be (at most) one such mapping
+       per module type."""
+    yaml_content = syaml.load(str_or_file)
+    index = {}
+    yaml_index = yaml_content['module_index']
+    for dag_hash, module_properties in yaml_index.items():
+        index[dag_hash] = ModuleIndexEntry(
+            module_properties['path'],
+            module_properties['use_name'])
+    return index
 
 
 def read_module_indices():
-    module_type_to_indices = {}
     other_spack_instances = spack.config.get(
         'upstreams') or {}
 
+    module_indices = []
+
     for install_properties in other_spack_instances.values():
+        module_type_to_index = {}
         module_type_to_root = install_properties.get('modules', {})
         for module_type, root in module_type_to_root.items():
-            indices = module_type_to_indices.setdefault(module_type, [])
-            indices.append(read_module_index(root))
-
-    return module_type_to_indices
-
-
-module_type_to_indices = read_module_indices()
+            module_type_to_index[module_type] = read_module_index(root)
+        module_indices.append(module_type_to_index)
+
+    return module_indices
+
+
+class UpstreamModuleIndex(object):
+    """This is responsible for taking the individual module indices of all
+       upstream Spack installations and locating the module for a given spec
+       based on which upstream install it is located in."""
+    def __init__(self, local_db, module_indices):
+        self.local_db = local_db
+        self.upstream_dbs = local_db.upstream_dbs
+        self.module_indices = module_indices
+
+    def upstream_module(self, spec, module_type):
+        db_for_spec = self.local_db.db_for_spec_hash(spec.dag_hash())
+        if db_for_spec in self.upstream_dbs:
+            db_index = self.upstream_dbs.index(db_for_spec)
+        elif db_for_spec:
+            raise spack.error.SpackError(
+                "Unexpected: {0} is installed locally".format(spec))
+        else:
+            raise spack.error.SpackError(
+                "Unexpected: no install DB found for {0}".format(spec))
+        module_index = self.module_indices[db_index]
+        module_type_index = module_index.get(module_type, {})
+        if not module_type_index:
+            raise ModuleNotFoundError(
+                "No {0} modules associated with the Spack instance where"
+                " {1} is installed".format(module_type, spec))
+        if spec.dag_hash() in module_type_index:
+            return module_type_index[spec.dag_hash()]
+        else:
+            raise ModuleNotFoundError(
+                "No module is available for upstream package {0}".format(spec))
 
 
-def upstream_module(spec, module_type):
-    indices = module_type_to_indices[module_type]
-    for index in indices:
-        if spec.dag_hash() in index:
-            return index[spec.dag_hash()]
+def get_module(module_type, spec, get_full_path):
+    if spec.package.installed_upstream:
+        module = spack.modules.common.upstream_module_index.upstream_module(
+            spec, module_type)
+        if get_full_path:
+            return module.path
+        else:
+            return module.use_name
+    else:
+        writer = spack.modules.module_types[module_type](spec)
+        if not os.path.isfile(writer.layout.filename):
+            err_msg = "No module available for package {0} at {1}".format(
+                spec, writer.layout.filename
+            )
+            raise ModuleNotFoundError(err_msg)
+        if get_full_path:
+            return writer.layout.filename
+        else:
+            return writer.layout.use_name
 
 
 class BaseConfiguration(object):
@@ -773,6 +835,10 @@ class ModulesError(spack.error.SpackError):
     """Base error for modules."""
 
 
+class ModuleNotFoundError(ModulesError):
+    """Raised when a module cannot be found for a spec"""
+
+
 class DefaultTemplateNotDefined(AttributeError, ModulesError):
     """Raised if the attribute 'default_template' has not been specified
     in the derived classes.

lib/spack/spack/test/modules/common.py

@@ -3,12 +3,17 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
-import pytest
 import os
 import stat
+import pytest
+import collections
+
 import spack.spec
-import spack.modules.common
 import spack.modules.tcl
+from spack.modules.common import (
+    UpstreamModuleIndex, ModuleNotFoundError)
+
+import spack.error
 
 
 def test_update_dictionary_extending_list():
@@ -70,3 +75,115 @@ def test_modules_written_with_proper_permissions(mock_module_filename,
 
     assert mock_package_perms & os.stat(
         mock_module_filename).st_mode == mock_package_perms
+
+
+class MockDb(object):
+    def __init__(self, db_ids, spec_hash_to_db):
+        self.upstream_dbs = db_ids
+        self.spec_hash_to_db = spec_hash_to_db
+
+    def db_for_spec_hash(self, spec_hash):
+        return self.spec_hash_to_db.get(spec_hash)
+
+
+class MockSpec(object):
+    def __init__(self, unique_id):
+        self.unique_id = unique_id
+
+    def dag_hash(self):
+        return self.unique_id
+
+
+def test_upstream_module_index():
+    s1 = MockSpec('spec-1')
+    s2 = MockSpec('spec-2')
+    s3 = MockSpec('spec-3')
+    s4 = MockSpec('spec-4')
+
+    tcl_module_index = """\
+module_index:
+  {0}:
+    path: /path/to/a
+    use_name: a
+""".format(s1.dag_hash())
+
+    module_indices = [
+        {
+            'tcl': spack.modules.common._read_module_index(tcl_module_index)
+        },
+        {}
+    ]
+
+    dbs = [
+        'd0',
+        'd1'
+    ]
+
+    mock_db = MockDb(
+        dbs,
+        {
+            s1.dag_hash(): 'd0',
+            s2.dag_hash(): 'd1',
+            s3.dag_hash(): 'd0'
+        }
+    )
+    upstream_index = UpstreamModuleIndex(mock_db, module_indices)
+
+    m1 = upstream_index.upstream_module(s1, 'tcl')
+    assert m1.path == '/path/to/a'
+
+    # No modules are defined for the DB associated with s2
+    with pytest.raises(ModuleNotFoundError):
+        upstream_index.upstream_module(s2, 'tcl')
+
+    # Modules are defined for the index associated with s1, but none are
+    # defined for the requested type
+    with pytest.raises(ModuleNotFoundError):
+        upstream_index.upstream_module(s1, 'lmod')
+
+    # A module is registered with a DB and the associated module index has
+    # modules of the specified type defined, but not for the requested spec
+    with pytest.raises(ModuleNotFoundError):
+        upstream_index.upstream_module(s3, 'tcl')
+
+    # The spec isn't recorded as installed in any of the DBs
+    with pytest.raises(spack.error.SpackError):
+        upstream_index.upstream_module(s4, 'tcl')
+
+
+def test_get_module_upstream():
+    s1 = MockSpec('spec-1')
+
+    tcl_module_index = """\
+module_index:
+  {0}:
+    path: /path/to/a
+    use_name: a
+""".format(s1.dag_hash())
+
+    module_indices = [
+        {},
+        {
+            'tcl': spack.modules.common._read_module_index(tcl_module_index)
+        }
+    ]
+
+    dbs = ['d0', 'd1']
+
+    mock_db = MockDb(
+        dbs,
+        {s1.dag_hash(): 'd1'}
+    )
+    upstream_index = UpstreamModuleIndex(mock_db, module_indices)
+
+    MockPackage = collections.namedtuple('MockPackage', ['installed_upstream'])
+    setattr(s1, "package", MockPackage(True))
+
+    try:
+        old_index = spack.modules.common.upstream_module_index
+        spack.modules.common.upstream_module_index = upstream_index
+
+        m1_path = spack.modules.common.get_module('tcl', s1, True)
+        assert m1_path == '/path/to/a'
+    finally:
+        spack.modules.common.upstream_module_index = old_index

lib/spack/spack/test/modules/tcl.py

@@ -3,8 +3,8 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
-
 import pytest
+
 import spack.modules.common
 import spack.modules.tcl
 import spack.spec