Merge pull request #1265 from Libensemble/feature/merge_platform_specs_detection

platform_specs is merged with detected platform

Author: shuds13

libensemble/executors/mpi_runner.py

@@ -37,7 +37,7 @@ def __init__(self, run_command="mpiexec", platform_info=None):
         self.arg_nnodes = ("--LIBE_NNODES_ARG_EMPTY",)
         self.arg_ppn = ("--LIBE_PPN_ARG_EMPTY",)
         self.default_mpi_options = None
-        self.default_gpu_arg = None
+        self.default_gpu_args = None
         self.default_gpu_arg_type = None
         self.platform_info = platform_info
 
@@ -126,16 +126,32 @@ def _set_gpu_env_var(self, wresources, task, gpus_per_node, gpus_env):
 
     def _local_runner_set_gpus(self, task, wresources, extra_args, gpus_per_node, ppn):
         """Set default GPU setting for MPI runner"""
-        if self.default_gpu_arg is not None:
-            arg_type = self.default_gpu_arg_type
+
+        arg_type = self.default_gpu_arg_type
+        if arg_type is not None:
             gpu_value = gpus_per_node // ppn if arg_type == "option_gpus_per_task" else gpus_per_node
-            gpu_setting_name = self.default_gpu_arg
+            gpu_setting_name = self.default_gpu_args[arg_type]
+            jassert(gpu_setting_name is not None, f"No default gpu_setting_name for {arg_type}")
             extra_args = self._set_gpu_cli_option(wresources, extra_args, gpu_setting_name, gpu_value)
         else:
             gpus_env = "CUDA_VISIBLE_DEVICES"
             self._set_gpu_env_var(wresources, task, gpus_per_node, gpus_env)
         return extra_args
 
+    def _get_default_arg(self, gpu_setting_type):
+        """Return default setting for the given gpu_setting_type if it exists, else error"""
+        jassert(
+            gpu_setting_type in ["option_gpus_per_node", "option_gpus_per_task"],
+            f"Unrecognized gpu_setting_type {gpu_setting_type}",
+        )
+        jassert(
+            self.default_gpu_args is not None,
+            "The current MPI runner has no default command line option for setting GPUs",
+        )
+        gpu_setting_name = self.default_gpu_args[gpu_setting_type]
+        jassert(gpu_setting_name is not None, f"No default GPU setting for {gpu_setting_type}")
+        return gpu_setting_name
+
     def _assign_gpus(self, task, resources, nprocs, nnodes, ppn, ngpus, extra_args, match_procs_to_gpus):
         """Assign GPU resources to slots, limited by ngpus if present.
 
@@ -199,7 +215,7 @@ def _assign_gpus(self, task, resources, nprocs, nnodes, ppn, ngpus, extra_args,
 
         elif gpu_setting_type in ["option_gpus_per_node", "option_gpus_per_task"]:
             gpu_value = gpus_per_node // ppn if gpu_setting_type == "option_gpus_per_task" else gpus_per_node
-            gpu_setting_name = self.platform_info.get("gpu_setting_name", self.default_gpu_arg)
+            gpu_setting_name = self.platform_info.get("gpu_setting_name", self._get_default_arg(gpu_setting_type))
             extra_args = self._set_gpu_cli_option(wresources, extra_args, gpu_setting_name, gpu_value)
 
         elif gpu_setting_type == "env":
@@ -319,7 +335,7 @@ def __init__(self, run_command="mpirun", platform_info=None):
         self.arg_nnodes = ("--LIBE_NNODES_ARG_EMPTY",)
         self.arg_ppn = ("--ppn", "-ppn")
         self.default_mpi_options = None
-        self.default_gpu_arg = None
+        self.default_gpu_args = None
         self.default_gpu_arg_type = None
         self.platform_info = platform_info
 
@@ -343,7 +359,7 @@ def __init__(self, run_command="mpirun", platform_info=None):
         self.arg_nnodes = ("--LIBE_NNODES_ARG_EMPTY",)
         self.arg_ppn = ("-npernode",)
         self.default_mpi_options = None
-        self.default_gpu_arg = None
+        self.default_gpu_args = None
         self.default_gpu_arg_type = None
         self.platform_info = platform_info
         self.mpi_command = [
@@ -388,7 +404,7 @@ def __init__(self, run_command="aprun", platform_info=None):
         self.arg_nnodes = ("--LIBE_NNODES_ARG_EMPTY",)
         self.arg_ppn = ("-N",)
         self.default_mpi_options = None
-        self.default_gpu_arg = None
+        self.default_gpu_args = None
         self.default_gpu_arg_type = None
         self.platform_info = platform_info
         self.mpi_command = [
@@ -410,7 +426,7 @@ def __init__(self, run_command="mpiexec", platform_info=None):
         self.arg_nnodes = ("--LIBE_NNODES_ARG_EMPTY",)
         self.arg_ppn = ("-cores",)
         self.default_mpi_options = None
-        self.default_gpu_arg = None
+        self.default_gpu_args = None
         self.default_gpu_arg_type = None
         self.platform_info = platform_info
         self.mpi_command = [
@@ -431,8 +447,9 @@ def __init__(self, run_command="srun", platform_info=None):
         self.arg_nnodes = ("-N", "--nodes")
         self.arg_ppn = ("--ntasks-per-node",)
         self.default_mpi_options = "--exact"
-        self.default_gpu_arg = "--gpus-per-task"
         self.default_gpu_arg_type = "option_gpus_per_task"
+        self.default_gpu_args = {"option_gpus_per_task": "--gpus-per-task", "option_gpus_per_node": "--gpus-per-node"}
+
         self.platform_info = platform_info
         self.mpi_command = [
             self.run_command,
@@ -453,8 +470,8 @@ def __init__(self, run_command="jsrun", platform_info=None):
         self.arg_nnodes = ("--LIBE_NNODES_ARG_EMPTY",)
         self.arg_ppn = ("-r",)
         self.default_mpi_options = None
-        self.default_gpu_arg = "-g"
         self.default_gpu_arg_type = "option_gpus_per_task"
+        self.default_gpu_args = {"option_gpus_per_task": "-g", "option_gpus_per_node": None}
 
         self.platform_info = platform_info
         self.mpi_command = [self.run_command, "-n {num_procs}", "-r {procs_per_node}", "{extra_args}"]

libensemble/resources/platforms.py

@@ -279,53 +279,54 @@ class Known_platforms(BaseModel):
 
 # Dictionary of known systems (or system partitions) detectable by domain name
 detect_systems = {
-    "crusher.olcf.ornl.gov": Crusher,
-    "frontier.olcf.ornl.gov": Frontier,
-    "hostmgmt.cm.aurora.alcf.anl.gov": Aurora,
-    "hsn.cm.polaris.alcf.anl.gov": Polaris,
-    "spock.olcf.ornl.gov": Spock,
-    "summit.olcf.ornl.gov": Summit,  # Need to detect gpu count
+    "crusher.olcf.ornl.gov": "crusher",
+    "frontier.olcf.ornl.gov": "frontier",
+    "hostmgmt.cm.aurora.alcf.anl.gov": "aurora",
+    "hsn.cm.polaris.alcf.anl.gov": "polaris",
+    "spock.olcf.ornl.gov": "spock",
+    "summit.olcf.ornl.gov": "summit",  # Need to detect gpu count
 }
 
 
 def known_envs():
     """Detect system by environment variables"""
-    platform_info = {}
+    name = None
     if os.environ.get("NERSC_HOST") == "perlmutter":
         if "gpu_" in os.environ.get("SLURM_JOB_PARTITION"):
-            platform_info = specs_dump(PerlmutterGPU(), by_alias=True)
+            name = "perlmutter_g"
         else:
-            platform_info = specs_dump(PerlmutterCPU(), by_alias=True)
-    return platform_info
+            name = "perlmutter_c"
+    return name
 
 
 def known_system_detect(cmd="hostname -d"):
     """Detect known systems
 
-    This function attempts to detect if on a known system, but users
-    should specify systems to be sure.
+    This function attempts to detect if on a known system, and
+    returns the name of the system as a string.
     """
     run_cmd = cmd.split()
-    platform_info = {}
+    name = None
     try:
         domain_name = subprocess.check_output(run_cmd).decode().rstrip()
-        platform_info = specs_dump(detect_systems[domain_name](), by_alias=True)
+        name = detect_systems[domain_name]
     except Exception:
-        platform_info = known_envs()
-    return platform_info
+        name = known_envs()
+    return name
 
 
 def get_platform(libE_specs):
     """Return platform as a dictionary from relevant libE_specs option.
 
     For internal use, return a platform as a dictionary from either
-    platform name or platform_specs.
+    platform name or platform_specs or auto-detection.
 
-    If both platform and platform_spec fields are present, any fields in
-    platform_specs are added or overwrite fields in the known platform.
+    If a platform is given or detected and platform_spec fields are present,
+    any fields in platform_specs are added to or overwrite fields in the known
+    platform.
     """
-
-    name = libE_specs.get("platform") or os.environ.get("LIBE_PLATFORM")
+    platform_info = {}
+    name = libE_specs.get("platform") or os.environ.get("LIBE_PLATFORM") or known_system_detect()
     if name:
         try:
             known_platforms = specs_dump(Known_platforms(), exclude_none=True)
@@ -340,9 +341,5 @@ def get_platform(libE_specs):
                 platform_info[k] = v
     elif libE_specs.get("platform_specs"):
         platform_info = libE_specs["platform_specs"]
-    else:
-        # See if in detection list
-        platform_info = known_system_detect()
-
-    platform_info = {k: v for k, v in platform_info.items() if v is not None}
+        platform_info = {k: v for k, v in platform_info.items() if v is not None}
     return platform_info

libensemble/tests/unit_tests/test_platform.py

@@ -1,6 +1,8 @@
 import pytest
 
+from libensemble.utils.misc import specs_dump
 from libensemble.resources.platforms import PlatformException, get_platform, known_system_detect
+from libensemble.resources.platforms import Known_platforms
 
 my_spec = {
     "mpi_runner": "srun",
@@ -10,15 +12,12 @@
 
 summit_spec = {
     "mpi_runner": "jsrun",
-    "runner_name": None,
     "cores_per_node": 42,
     "logical_cores_per_node": 168,
     "gpus_per_node": 6,
-    "tiles_per_gpu": None,
     "gpu_setting_type": "option_gpus_per_task",
     "gpu_setting_name": "-g",
     "scheduler_match_slots": False,
-    "gpu_env_fallback": None,
 }
 
 
@@ -60,7 +59,6 @@ def test_platform_known():
 def test_platform_specs():
     """Test known platform and platform_specs supplied"""
     from libensemble.specs import LibeSpecs
-    from libensemble.utils.misc import specs_dump
 
     exp = my_spec
     libE_specs = {"platform_specs": my_spec}
@@ -85,16 +83,18 @@ def test_platform_specs():
 
 
 def test_known_sys_detect():
+    known_platforms = specs_dump(Known_platforms(), exclude_none=True)
     get_sys_cmd = "echo summit.olcf.ornl.gov"  # Overrides default "hostname -d"
-    platform_info = known_system_detect(cmd=get_sys_cmd)
+    name = known_system_detect(cmd=get_sys_cmd)
+    platform_info = known_platforms[name]
     assert platform_info == summit_spec, f"Summit spec does not match expected ({platform_info})"
 
     # Try unknown system
     get_sys_cmd = "echo madeup.system"  # Overrides default "hostname -d"
-    platform_info = known_system_detect(cmd=get_sys_cmd)
+    name = known_system_detect(cmd=get_sys_cmd)
     assert (
-        platform_info == {}
-    ), f"Expected known_system_detect to return empty dict for unknown system ({platform_info})"
+        name is None
+    ), f"Expected known_system_detect to return None ({name})"
 
 
 if __name__ == "__main__":