Merge pull request #4642 from nilsvu/this_machine_py

Add `this_machine` Python function

CMakeLists.txt

@@ -189,6 +189,7 @@ spectre_include_directories(${CMAKE_BINARY_DIR}/src/Parallel)
 
 add_subdirectory(external)
 add_subdirectory(src)
+add_subdirectory(support)
 if(BUILD_TESTING)
   add_subdirectory(tests)
 endif()

docs/DevGuide/BuildSystem.md

@@ -247,6 +247,11 @@ cmake -D FLAG1=OPT1 ... -D FLAGN=OPTN <SPECTRE_ROOT>
   - Whether to keep the frame pointer. Needed for profiling or other cases
     where you need to be able to figure out what the call stack is.
     (default is `OFF`)
+- MACHINE
+  - Select a machine that we know how to run on, such as a particular
+    supercomputer. A file named MACHINE.yaml must exist in support/Machines and
+    a submit script template named MACHINE.sh must exist in
+    support/SubmitScripts.
 - MEMORY_ALLOCATOR
   - Set which memory allocator to use. If there are unexplained segfaults or
     other memory issues, it would be worth setting `MEMORY_ALLOCATOR=SYSTEM` to

support/CMakeLists.txt

@@ -0,0 +1,4 @@
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+
+add_subdirectory(Python)

support/Environments/wheeler_clang.sh

@@ -78,6 +78,7 @@ spectre_run_cmake() {
           -D BUILD_PYTHON_BINDINGS=ON \
           -D CMAKE_PREFIX_PATH="$PYTHON_HOME" \
           -D BOOTSTRAP_PY_DEPS=ON \
+          -D MACHINE=Wheeler \
           "$@" \
           $SPECTRE_HOME
 }

support/Environments/wheeler_gcc.sh

@@ -76,6 +76,7 @@ spectre_run_cmake() {
           -D BUILD_PYTHON_BINDINGS=ON \
           -D CMAKE_PREFIX_PATH="$PYTHON_HOME" \
           -D BOOTSTRAP_PY_DEPS=ON \
+          -D MACHINE=Wheeler \
           "$@" \
           $SPECTRE_HOME
 }

support/Machines/Wheeler.yaml

@@ -0,0 +1,13 @@
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+
+Machine:
+  Name: Wheeler
+  Description: |
+    Supercomputer at Caltech.
+    More information:
+    https://github.com/sxs-collaboration/WelcomeToSXS/wiki/Wheeler
+  # Use one of the 24 available cores for communication by default
+  DefaultProcsPerNode: 23
+  DefaultQueue: "productionQ"
+  DefaultTimeLimit: "1-00:00:00"

support/Python/CMakeLists.txt

@@ -0,0 +1,23 @@
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+
+option(MACHINE "Select a machine that we know how to run on, such as a \
+particular supercomputer" OFF)
+
+spectre_python_add_module(
+  support
+  PYTHON_FILES
+  Machines.py
+)
+
+if(MACHINE)
+  message(STATUS "Selected machine: ${MACHINE}")
+  configure_file(
+    ${CMAKE_SOURCE_DIR}/support/Machines/${MACHINE}.yaml
+    ${SPECTRE_PYTHON_PREFIX}/support/Machine.yaml
+    )
+  configure_file(
+    ${CMAKE_SOURCE_DIR}/support/SubmitScripts/${MACHINE}.sh
+    ${SPECTRE_PYTHON_PREFIX}/support/Submit.sh
+    )
+endif()

support/Python/Machines.py

@@ -0,0 +1,92 @@
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+"""Support for host machines, such as supercomputers.
+
+Machines are defined as YAML files in 'support/Machines/'. To add support for a
+new machine, add a YAML file that defines a `Machine:` key with the attributes
+listed in the `Machine` class below.
+
+To select a machine, specify the `MACHINE` option when configuring the CMake
+build.
+"""
+
+import os
+from dataclasses import dataclass
+
+# functools.cache was added in Py 3.9. Fall back to 'lru_cache' in earlier
+# versions, which is pretty much the same but slightly slower.
+try:
+    from functools import cache
+except ImportError:
+    from functools import lru_cache
+    cache = lru_cache(maxsize=None)
+
+import yaml
+
+
+@dataclass(frozen=True)
+class Machine(yaml.YAMLObject):
+    """A machine we know how to run on, such as a particular supercomputer.
+
+    Many configuration options for job submission are hardcoded in the submit
+    script for the machine (such as the total number of available cores per
+    node). Here we provide parameters that we want to adjust when scheduling
+    jobs, but that have sensible defaults based on our experience running on the
+    machine. Parameters are set to 'None' when they have no sensible default for
+    the machine.
+
+    Attributes:
+      Name: A short name for the machine. Must match the YAML file name.
+      Description: A description of the machine. Give some basic context and
+        any information that may help people get started using the machine.
+        Provide links to wiki pages, signup pages, etc., for additional
+        information.
+      DefaultProcsPerNode: Default number of worker threads spawned per node.
+        It is often advised to leave one core per node or socket free for
+        communication, so this might be the number of cores or hyperthreads
+        per node minus one.
+      DefaultQueue: Default queue that jobs are submitted to. On Slurm systems
+        you can see the available queues with `sinfo`.
+      DefaultTimeLimit: Default wall time limit for submitted jobs. For
+        acceptable formats, see: https://slurm.schedmd.com/sbatch.html#OPT_time
+    """
+    yaml_tag = '!Machine'
+    yaml_loader = yaml.SafeLoader
+    # The YAML machine files can have these attributes:
+    Name: str
+    Description: str
+    DefaultProcsPerNode: int
+    DefaultQueue: str
+    DefaultTimeLimit: str
+
+
+# Parse YAML machine files as Machine objects
+yaml.SafeLoader.add_path_resolver('!Machine', ['Machine'], dict)
+
+
+class UnknownMachineError(Exception):
+    """Indicates we were unsuccessful in identifying the current machine"""
+    pass
+
+
+@cache
+def this_machine(machinefile_path=os.path.join(os.path.dirname(__file__),
+                                               'Machine.yaml')) -> Machine:
+    """Determine the machine we are running on.
+
+    Raises `UnknownMachineError` if no machine was selected. Specify the
+    `MACHINE` option in the CMake build configuration to select a machine, or
+    pass the `machinefile_path` argument to this function.
+
+    Arguments:
+      machinefile_path: Path to a YAML file that describes the current machine.
+        Defaults to the machine selected in the CMake build configuration.
+    """
+    if not os.path.exists(machinefile_path):
+        raise UnknownMachineError(
+            "No machine was selected. Specify the 'MACHINE' option when "
+            "configuring the build with CMake. If you are running on a new "
+            "machine, please add it to 'support/Machines/'. The machine file "
+            f"was expected at the following path:\n  {machinefile_path}")
+    with open(machinefile_path, 'r') as open_machinefile:
+        return yaml.safe_load(open_machinefile)['Machine']

support/Python/__main__.py

@@ -5,6 +5,7 @@
 import logging
 import rich.logging
 import rich.traceback
+from spectre.support.Machines import this_machine, UnknownMachineError
 
 logger = logging.getLogger(__name__)
 
@@ -71,11 +72,29 @@ def get_command(self, ctx, name):
         raise NotImplementedError(f"The command '{name}' is not implemented.")
 
 
+def print_machine(ctx, param, value):
+    if not value or ctx.resilient_parsing:
+        return
+    try:
+        machine = this_machine()
+        click.echo(machine.Name)
+        ctx.exit(1)
+    except UnknownMachineError as exc:
+        click.echo(exc)
+        ctx.exit()
+
+
 # Set up CLI entry point
 @click.group(context_settings=dict(help_option_names=["-h", "--help"]),
              help=f"SpECTRE version: {SPECTRE_VERSION}",
              cls=Cli)
 @click.version_option(version=SPECTRE_VERSION, message="%(version)s")
+@click.option('--machine',
+              is_flag=True,
+              expose_value=False,
+              is_eager=True,
+              callback=print_machine,
+              help="Show the machine we're running on and exit.")
 @click.option('--debug',
               'log_level',
               flag_value=logging.DEBUG,

tests/support/Python/CMakeLists.txt

@@ -9,3 +9,9 @@ if (BUILD_PYTHON_BINDINGS)
     PASS_REGULAR_EXPRESSION "${SPECTRE_VERSION}"
     LABELS "python")
 endif()
+
+spectre_add_python_bindings_test(
+  "support.Machines"
+  Test_Machines.py
+  "Machines"
+  None)

tests/support/Python/Test_Machines.py

@@ -0,0 +1,37 @@
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+
+import os
+import unittest
+
+import yaml
+from spectre.Informer import unit_test_build_path
+from spectre.support.Machines import Machine, UnknownMachineError, this_machine
+
+
+class TestMachines(unittest.TestCase):
+    def setUp(self):
+        self.machinefile_path = os.path.join(unit_test_build_path(),
+                                             'TestMachine.yaml')
+        yaml.safe_dump(
+            dict(Machine=dict(Name="TestMachine",
+                              Description="Just for testing",
+                              DefaultProcsPerNode=15,
+                              DefaultQueue="production",
+                              DefaultTimeLimit="1-00:00:00")),
+            open(self.machinefile_path, 'w'))
+
+    def test_this_machine(self):
+        with self.assertRaises(UnknownMachineError):
+            this_machine('NonexistentMachinefile.yaml')
+        machine = this_machine(self.machinefile_path)
+        self.assertIsInstance(machine, Machine)
+        self.assertEqual(machine.Name, "TestMachine")
+        self.assertEqual(machine.Description, "Just for testing")
+        self.assertEqual(machine.DefaultProcsPerNode, 15)
+        self.assertEqual(machine.DefaultQueue, "production")
+        self.assertEqual(machine.DefaultTimeLimit, "1-00:00:00")
+
+
+if __name__ == '__main__':
+    unittest.main(verbosity=2)