Increase functionality of try_toolchain

easybuild/framework/easyconfig/easyconfig.py

@@ -82,6 +82,9 @@
 # name of easyconfigs archive subdirectory
 EASYCONFIGS_ARCHIVE_DIR = '__archive__'
 
+# Available capabilities of toolchains
+CAPABILITIES = ['comp_family', 'mpi_family', 'blas_family', 'lapack_family', 'cuda']
+
 
 try:
     import autopep8
@@ -116,16 +119,16 @@ def toolchain_hierarchy_cache(func):
     cache = {}
 
     @functools.wraps(func)
-    def cache_aware_func(toolchain):
+    def cache_aware_func(toolchain, require_capabilities=False):
         """Look up toolchain hierarchy in cache first, determine and cache it if not available yet."""
-        cache_key = (toolchain['name'], toolchain['version'])
+        cache_key = (toolchain['name'], toolchain['version'], require_capabilities)
 
         # fetch from cache if available, cache it if it's not
         if cache_key in cache:
             _log.debug("Using cache to return hierarchy for toolchain %s: %s", str(toolchain), cache[cache_key])
             return cache[cache_key]
         else:
-            toolchain_hierarchy = func(toolchain)
+            toolchain_hierarchy = func(toolchain, require_capabilities)
             cache[cache_key] = toolchain_hierarchy
             return cache[cache_key]
 
@@ -135,7 +138,7 @@ def cache_aware_func(toolchain):
     return cache_aware_func
 
 
-def det_subtoolchain_version(current_tc, subtoolchain_name, optional_toolchains, cands):
+def det_subtoolchain_version(current_tc, subtoolchain_name, optional_toolchains, cands, require_capabilities=False):
     """
     Returns unique version for subtoolchain, in tc dict.
     If there is no unique version:
@@ -152,7 +155,7 @@ def det_subtoolchain_version(current_tc, subtoolchain_name, optional_toolchains,
 
     # dummy toolchain: bottom of the hierarchy
     if subtoolchain_name == DUMMY_TOOLCHAIN_NAME:
-        if build_option('add_dummy_to_minimal_toolchains'):
+        if build_option('add_dummy_to_minimal_toolchains') and not require_capabilities:
             subtoolchain_version = ''
     elif len(uniq_subtc_versions) == 1:
         subtoolchain_version = list(uniq_subtc_versions)[0]
@@ -169,7 +172,7 @@ def det_subtoolchain_version(current_tc, subtoolchain_name, optional_toolchains,
 
 
 @toolchain_hierarchy_cache
-def get_toolchain_hierarchy(parent_toolchain):
+def get_toolchain_hierarchy(parent_toolchain, require_capabilities=False):
     """
     Determine list of subtoolchains for specified parent toolchain.
     Result starts with the most minimal subtoolchains first, ends with specified toolchain.
@@ -199,8 +202,9 @@ def get_toolchain_hierarchy(parent_toolchain):
     optional_toolchains = set(tc_class.NAME for tc_class in all_tc_classes if getattr(tc_class, 'OPTIONAL', False))
     composite_toolchains = set(tc_class.NAME for tc_class in all_tc_classes if len(tc_class.__bases__) > 1)
 
-    # the parent toolchain is at the top of the hierarchy
-    toolchain_hierarchy = [parent_toolchain]
+    # the parent toolchain is at the top of the hierarchy, we need a copy so that adding capabilities (below) doesn't
+    # affect the original object
+    toolchain_hierarchy = [copy.copy(parent_toolchain)]
     # use a queue to handle a breadth-first-search of the hierarchy,
     # which is required to take into account the potential for multiple subtoolchains
     bfs_queue = [parent_toolchain]
@@ -268,14 +272,30 @@ def get_toolchain_hierarchy(parent_toolchain):
         cands = [c for c in cands if c['name'] in subtoolchain_names]
 
         for subtoolchain_name in subtoolchain_names:
-            subtoolchain_version = det_subtoolchain_version(current_tc, subtoolchain_name, optional_toolchains, cands)
+            subtoolchain_version = det_subtoolchain_version(current_tc, subtoolchain_name, optional_toolchains, cands,
+                                                            require_capabilities=require_capabilities)
             # add to hierarchy and move to next
             if subtoolchain_version is not None and subtoolchain_name not in visited:
                 tc = {'name': subtoolchain_name, 'version': subtoolchain_version}
                 toolchain_hierarchy.insert(0, tc)
                 bfs_queue.insert(0, tc)
                 visited.add(subtoolchain_name)
 
+    # also add toolchain capabilities
+    if require_capabilities:
+        for toolchain in toolchain_hierarchy:
+            toolchain_class, _ = search_toolchain(toolchain['name'])
+            tc = toolchain_class(version=toolchain['version'])
+            for capability in CAPABILITIES:
+                # cuda is the special case which doesn't have a family attribute
+                if capability == 'cuda':
+                    # use None rather than False, useful to have it consistent with the rest
+                    toolchain['cuda'] = ('CUDA_CC' in tc.variables) or None
+                else:
+                    if hasattr(tc, capability):
+                        toolchain[capability] = getattr(tc, capability)()
+
+
     _log.info("Found toolchain hierarchy for toolchain %s: %s", parent_toolchain, toolchain_hierarchy)
     return toolchain_hierarchy