Merge pull request #4759 from ESMCI/adds-query_config-compiler-info

jasonb5 · web-flow · commit 36cf5e4170fd · 2025-03-17T17:51:08.000-07:00
Adds --compiler argument to query_config. This argument should be used when querying a specific machine e.g. --machines frontier. When specified, modules and environment_variables that match the compiler will be printed. Additional requirements e.g. DEBUG="TRUE" will be printed above sections that require them. Additionally for environment_variables there is an attempt to resolve any CIME specific templates e.g. $ENV{HOME}. Test suite: Test baseline: n/a Test namelist changes: n/a Test status: n/a Fixes #4743 User interface changes?: Update gh-pages html (Y/N)?:
diff --git a/CIME/XML/machines.py b/CIME/XML/machines.py
@@ -1,17 +1,45 @@
 """
 Interface to the config_machines.xml file.  This class inherits from GenericXML.py
 """
+
 from CIME.XML.standard_module_setup import *
 from CIME.XML.generic_xml import GenericXML
 from CIME.XML.files import Files
-from CIME.utils import convert_to_unknown_type, get_cime_config
+from CIME.utils import CIMEError, expect, convert_to_unknown_type, get_cime_config
 
+import re
+import logging
 import socket
+from functools import partial
 from pathlib import Path
 
 logger = logging.getLogger(__name__)
 
 
+def match_value_by_attribute_regex(element, attribute_name, value):
+    """Checks element contains attribute whose pattern matches a value.
+
+    If the element does not have the attribute it's considered a match.
+
+    Args:
+        element (CIME.XML.generic_xml._Element): XML element to check attributes.
+        attribute_name (str): Name of attribute with regex value.
+        value (str): Value that is matched against attributes regex value.
+
+    Returns:
+        bool: True if attribute regex matches the target value otherwise False.
+    """
+    attribute_value = element.attrib.get(attribute_name, None)
+
+    return (
+        True
+        if value is None
+        or attribute_value is None
+        or re.match(attribute_value, value) is not None
+        else False
+    )
+
+
 class Machines(GenericXML):
     def __init__(
         self,
@@ -484,30 +512,179 @@ def set_value(self, vid, value, subgroup=None, ignore_type=True):
         # A temporary cache only
         self.custom_settings[vid] = value
 
-    def print_values(self):
-        # write out machines
-        machines = self.get_children("machine")
-        logger.info("Machines")
-        for machine in machines:
-            name = self.get(machine, "MACH")
-            desc = self.get_child("DESC", root=machine)
-            os_ = self.get_child("OS", root=machine)
-            compilers = self.get_child("COMPILERS", root=machine)
-            max_tasks_per_node = self.get_child("MAX_TASKS_PER_NODE", root=machine)
-            max_mpitasks_per_node = self.get_child(
-                "MAX_MPITASKS_PER_NODE", root=machine
+    def print_values(self, compiler=None):
+        """Prints machine values.
+
+        Args:
+            compiler (str, optional): Name of the compiler to print extra details for. Defaults to None.
+        """
+        current = self.probe_machine_name(False)
+
+        if self.machine_node is None:
+            for machine in self.get_children("machine"):
+                self._print_machine_values(machine, current)
+        else:
+            self._print_machine_values(self.machine_node, current, compiler)
+
+    def _print_machine_values(self, machine, current=None, compiler=None):
+        """Prints a machines details.
+
+        Args:
+            machine (CIME.XML.machines.Machine): Machine object.
+            current (str, optional): Name of the current machine. Defaults to None.
+            compiler (str, optional): If not None, then modules and environment variables matching compiler are printed. Defaults to None.
+
+        Raises:
+            CIMEError: If `compiler` is not valid.
+        """
+        name = self.get(machine, "MACH")
+        if current is not None and current == name:
+            name = f"{name} (current)"
+        desc = self.text(self.get_child("DESC", root=machine))
+        os_ = self.text(self.get_child("OS", root=machine))
+
+        compilers = self.text(self.get_child("COMPILERS", root=machine))
+        if compiler is not None and compiler not in compilers.split(","):
+            raise CIMEError(
+                f"Compiler {compiler!r} is not a valid choice from ({compilers})"
             )
-            max_gpus_per_node = self.get_child("MAX_GPUS_PER_NODE", root=machine)
 
-            print("  {} : {} ".format(name, self.text(desc)))
-            print("      os             ", self.text(os_))
-            print("      compilers      ", self.text(compilers))
-            if max_mpitasks_per_node is not None:
-                print("      pes/node       ", self.text(max_mpitasks_per_node))
-            if max_tasks_per_node is not None:
-                print("      max_tasks/node ", self.text(max_tasks_per_node))
-            if max_gpus_per_node is not None:
-                print("      max_gpus/node ", self.text(max_gpus_per_node))
+        mpilibs_nodes = self._get_children_filter_attribute_regex(
+            "MPILIBS", "compiler", compiler, root=machine
+        )
+        mpilibs = set([y for x in mpilibs_nodes for y in self.text(x).split(",")])
+
+        max_tasks_per_node = self.text(
+            self.get_child("MAX_TASKS_PER_NODE", root=machine)
+        )
+        max_mpitasks_per_node = self.text(
+            self.get_child("MAX_MPITASKS_PER_NODE", root=machine)
+        )
+        max_gpus_per_node = self.get_optional_child("MAX_GPUS_PER_NODE", root=machine)
+        max_gpus_per_node_text = (
+            self.text(max_gpus_per_node) if max_gpus_per_node else 0
+        )
+
+        if compiler is not None:
+            name = f"{name} ({compiler})"
+
+        print("  {} : {} ".format(name, desc))
+        print("      os             ", os_)
+        print("      compilers      ", compilers)
+        print("      mpilibs        ", ",".join(mpilibs))
+        print("      pes/node       ", max_mpitasks_per_node)
+        print("      max_tasks/node ", max_tasks_per_node)
+        print("      max_gpus/node  ", max_gpus_per_node_text)
+        print("")
+
+        if compiler is not None:
+            module_system_node = self.get_child("module_system", root=machine)
+
+            def command_formatter(node):
+                if node.text is None:
+                    return f"{node.attrib['name']}"
+                else:
+                    return f"{node.attrib['name']} {node.text}"
+
+            print("    Module commands:")
+            for requirements, commands in self._filter_children_by_compiler(
+                "modules", "command", compiler, command_formatter, module_system_node
+            ):
+                indent = "" if requirements == "" else "  "
+                if requirements != "":
+                    print(f"      (with {requirements})")
+                for x in commands:
+                    print(f"      {indent}{x}")
+            print("")
+
+            def env_formatter(node, machines=None):
+                return f"{node.attrib['name']}: {machines._get_resolved_environment_variable(node.text)}"
+
+            print("    Environment variables:")
+            for requirements, variables in self._filter_children_by_compiler(
+                "environment_variables",
+                "env",
+                compiler,
+                partial(env_formatter, machines=self),
+                machine,
+            ):
+                indent = "" if requirements == "" else "  "
+                if requirements != "":
+                    print(f"      (with {requirements})")
+                for x in variables:
+                    print(f"      {indent}{x}")
+
+    def _filter_children_by_compiler(self, parent, child, compiler, formatter, root):
+        """Filters parent nodes and returns requirements and children of filtered nodes.
+
+        Example of a yielded values:
+
+        "mpilib=openmpi DEBUG=true", ["HOME: /home/dev", "NETCDF_C_PATH: ../netcdf"]
+
+        Args:
+            parent (str): Name of the nodes to filter.
+            child (str): Name of the children nodes from filtered parent nodes.
+            compiler (str): Name of the compiler that will be matched against the regex.
+            formatter (function): Function to format the child nodes from the parents that match.
+            root (CIME.XML.generic_xml._Element): Root node to filter parent nodes from.
+
+        Yields:
+            str, list: Requirements for parent node and list of formated child nodes.
+        """
+        nodes = self._get_children_filter_attribute_regex(
+            parent, "compiler", compiler, root=root
+        )
+
+        for x in nodes:
+            attrib = {**x.attrib}
+            attrib.pop("compiler", None)
+
+            requirements = " ".join([f"{y}={z!r}" for y, z in attrib.items()])
+            values = [formatter(y) for y in self.get_children(child, root=x)]
+
+            yield requirements, values
+
+    def _get_children_filter_attribute_regex(self, name, attribute_name, value, root):
+        """Filter children nodes using regex.
+
+        Uses regex from attribute of children nodes to match a value.
+
+        Args:
+            name (str): Name of the children nodes.
+            attribute_name (str): Name of the attribute on the child nodes to build regex from.
+            value (str): Value that is matched using regex from attribute.
+            root (CIME.XML.generic_xml._Element): Root node to query children nodes from.
+
+        Returns:
+            list: List of children whose regex attribute matches the value.
+        """
+        return [
+            x
+            for x in self.get_children(name, root=root)
+            if match_value_by_attribute_regex(x, attribute_name, value)
+        ]
+
+    def _get_resolved_environment_variable(self, text):
+        """Attempts to resolve machines environment variable.
+
+        Args:
+            text (str): Environment variable value.
+
+        Returns:
+            str: Resolved value or error message.
+        """
+        if text is None:
+            return ""
+
+        try:
+            value = self.get_resolved_value(text, allow_unresolved_envvars=True)
+        except Exception as e:
+            return f"Failed to resolve {text!r} with: {e!s}"
+
+        if value == text and "$" in text:
+            value = f"Failed to resolve {text!r}"
+
+        return value
 
     def return_values(self):
         """return a dictionary of machine info
diff --git a/CIME/scripts/query_config.py b/CIME/scripts/query_config.py
@@ -96,6 +96,11 @@ def parse_command_line(description):
         help="Query machines for model. If not value is passed, all machines will be printed.",
     )
 
+    config_group.add_argument(
+        "--compiler",
+        help="Prints compiler details when combined with --machines",
+    )
+
     output_group = parser.add_argument_group("Output options")
 
     output_group.add_argument(
@@ -133,6 +138,18 @@ def parse_command_line(description):
     if not any([kwargs[x] for x in ["grids", "compsets", "components", "machines"]]):
         parser.print_help(sys.stderr)
 
+    if kwargs["compiler"] is not None and (
+        kwargs["machines"] is None or kwargs["machines"] == "all"
+    ):
+        parser.print_help(sys.stderr)
+
+        print("")
+        print(
+            "The --compiler argument must be used when specifying a machine with --machines <name>"
+        )
+
+        sys.exit(1)
+
     kwargs["files"] = files[kwargs["driver"]]
 
     return kwargs
@@ -313,7 +330,7 @@ def _query_component_settings(component, files, xml=False, all_components=False,
         component.print_values()
 
 
-def query_machines(files, machines, xml, **_):
+def query_machines(files, machines, xml, compiler, **_):
     config_file = files.get_value("MACHINES_SPEC_FILE")
     utils.expect(
         os.path.isfile(config_file),
@@ -334,68 +351,33 @@ def query_machines(files, machines, xml, **_):
                 )
             )
     else:
-        print_machine_values(xml_machines, machines)
+        print_machine_values(xml_machines, compiler, machines)
 
 
 def print_machine_values(
-    machine, machine_name="all"
+    machine,
+    compiler,
+    machine_name="all",
 ):  # pylint: disable=arguments-differ
-    # set flag to look for single machine
-    if "all" not in machine_name:
-        single_machine = True
-        if machine_name == "current":
-            machine_name = machine.probe_machine_name(warn=False)
+    """Prints machine values
+
+    Args:
+        machine (CIME.XML.machines.Machine): Machine object.
+        machine_name (str, optional): Which machine to print values for, can be "all", "current", or specific name. Defaults to "all".
+    """
+    if machine_name == "current":
+        machine_name = machine.probe_machine_name(False)
+
+    if machine_name == "all":
+        machine_names = machine.list_available_machines()
     else:
-        single_machine = False
-
-    # if we can't find the specified machine
-    if single_machine and machine_name is None:
-        files = Files()
-        config_file = files.get_value("MACHINES_SPEC_FILE")
-        print("Machine is not listed in config file: {}".format(config_file))
-    else:  # write out machines
-        if single_machine:
-            machine_names = [machine_name]
-        else:
-            machine_names = machine.list_available_machines()
-        print("Machine(s)\n")
-        for name in machine_names:
-            machine.set_machine(name)
-            desc = machine.text(machine.get_child("DESC"))
-            os_ = machine.text(machine.get_child("OS"))
-            compilers = machine.text(machine.get_child("COMPILERS"))
-            mpilibnodes = machine.get_children("MPILIBS", root=machine.machine_node)
-            mpilibs = []
-            for node in mpilibnodes:
-                mpilibs.extend(machine.text(node).split(","))
-            # This does not include the possible depedancy of mpilib on compiler
-            # it simply provides a list of mpilibs available on the machine
-            mpilibs = list(set(mpilibs))
-            max_tasks_per_node = machine.text(machine.get_child("MAX_TASKS_PER_NODE"))
-            mpitasks_node = machine.get_optional_child(
-                "MAX_MPITASKS_PER_NODE", root=machine.machine_node
-            )
-            max_mpitasks_per_node = (
-                machine.text(mpitasks_node) if mpitasks_node else max_tasks_per_node
-            )
-            max_gpus_node = machine.get_optional_child(
-                "MAX_GPUS_PER_NODE", root=machine.machine_node
-            )
-            max_gpus_per_node = machine.text(max_gpus_node) if max_gpus_node else 0
-
-            current_machine = machine.probe_machine_name(warn=False)
-            name += " (current)" if current_machine and current_machine in name else ""
-            print("  {} : {} ".format(name, desc))
-            print("      os             ", os_)
-            print("      compilers      ", compilers)
-            print("      mpilibs        ", mpilibs)
-            if max_mpitasks_per_node is not None:
-                print("      pes/node       ", max_mpitasks_per_node)
-            if max_tasks_per_node is not None:
-                print("      max_tasks/node ", max_tasks_per_node)
-            if max_gpus_per_node is not None:
-                print("      max_gpus/node ", max_gpus_per_node)
-            print("")
+        machine_names = [machine_name]
+
+    print("Machine(s)\n")
+
+    for name in machine_names:
+        machine.set_machine(name)
+        machine.print_values(compiler)
 
 
 if __name__ == "__main__":
diff --git a/CIME/tests/test_unit_xml_machines.py b/CIME/tests/test_unit_xml_machines.py
