Feature/detect perlmutter login nodes

libensemble/resources/platforms.py

@@ -8,6 +8,7 @@
 option or the environment variable ``LIBE_PLATFORM``.
 """
 
+import logging
 import os
 import subprocess
 from typing import Optional
@@ -16,6 +17,10 @@
 
 from libensemble.utils.misc import specs_dump
 
+logger = logging.getLogger(__name__)
+# To change logging level for just this module
+# logger.setLevel(logging.DEBUG)
+
 
 class PlatformException(Exception):
     """Platform module exception"""
@@ -269,6 +274,7 @@ class Known_platforms(BaseModel):
     generic_rocm: GenericROCm = GenericROCm()
     crusher: Crusher = Crusher()
     frontier: Frontier = Frontier()
+    perlmutter: Perlmutter = Perlmutter()
     perlmutter_c: PerlmutterCPU = PerlmutterCPU()
     perlmutter_g: PerlmutterGPU = PerlmutterGPU()
     polaris: Polaris = Polaris()
@@ -292,10 +298,15 @@ def known_envs():
     """Detect system by environment variables"""
     name = None
     if os.environ.get("NERSC_HOST") == "perlmutter":
-        if "gpu_" in os.environ.get("SLURM_JOB_PARTITION"):
-            name = "perlmutter_g"
+        partition = os.environ.get("SLURM_JOB_PARTITION")
+        if partition:
+            if "gpu_" in partition:
+                name = "perlmutter_g"
+            else:
+                name = "perlmutter_c"
         else:
-            name = "perlmutter_c"
+            name = "perlmutter"
+            logger.manager_warning("Perlmutter detected, but no compute partition detected. Are you on login nodes?")
     return name
 
 

libensemble/tests/unit_tests/test_platform.py

@@ -1,6 +1,12 @@
 import pytest
 
-from libensemble.resources.platforms import Known_platforms, PlatformException, get_platform, known_system_detect
+from libensemble.resources.platforms import (
+    Known_platforms,
+    PlatformException,
+    get_platform,
+    known_envs,
+    known_system_detect,
+)
 from libensemble.utils.misc import specs_dump
 
 my_spec = {
@@ -20,8 +26,12 @@
 }
 
 
-def test_platform_empty():
+def test_platform_empty(monkeypatch):
     """Test no platform options supplied"""
+
+    # Ensure NERSC_HOST not set
+    monkeypatch.delenv("NERSC_HOST", raising=False)
+
     exp = {}
     libE_specs = {}
     platform_info = get_platform(libE_specs)
@@ -55,10 +65,13 @@ def test_platform_known():
     assert platform_info == exp, f"platform_info does not match expected: {platform_info}"
 
 
-def test_platform_specs():
+def test_platform_specs(monkeypatch):
     """Test known platform and platform_specs supplied"""
     from libensemble.specs import LibeSpecs
 
+    # Ensure NERSC_HOST not set
+    monkeypatch.delenv("NERSC_HOST", raising=False)
+
     exp = my_spec
     libE_specs = {"platform_specs": my_spec}
     platform_info = get_platform(libE_specs)
@@ -81,7 +94,12 @@ def test_platform_specs():
     assert specs_dump(LS.platform_specs, exclude_none=True) == exp, "Conversion isn't as expected"
 
 
-def test_known_sys_detect():
+def test_known_sys_detect(monkeypatch):
+    """Test detection of known system"""
+
+    # Ensure NERSC_HOST not set
+    monkeypatch.delenv("NERSC_HOST", raising=False)
+
     known_platforms = specs_dump(Known_platforms(), exclude_none=True)
     get_sys_cmd = "echo summit.olcf.ornl.gov"  # Overrides default "hostname -d"
     name = known_system_detect(cmd=get_sys_cmd)
@@ -94,9 +112,26 @@ def test_known_sys_detect():
     assert name is None, f"Expected known_system_detect to return None ({name})"
 
 
+def test_env_sys_detect(monkeypatch):
+    """Test detection of system partitions"""
+    monkeypatch.setenv("NERSC_HOST", "other_host")
+    monkeypatch.setenv("SLURM_JOB_PARTITION", "cpu_test_partition")
+    name = known_envs()
+    assert name is None
+    monkeypatch.setenv("NERSC_HOST", "perlmutter")
+
+    monkeypatch.setenv("SLURM_JOB_PARTITION", "gpu_test_partition")
+    name = known_envs()
+    assert name == "perlmutter_g"
+
+    monkeypatch.setenv("SLURM_JOB_PARTITION", "cpu_test_partition")
+    name = known_envs()
+    assert name == "perlmutter_c"
+
+    monkeypatch.delenv("SLURM_JOB_PARTITION", raising=False)
+    name = known_envs()
+    assert name == "perlmutter"
+
+
 if __name__ == "__main__":
-    test_platform_empty()
-    test_unknown_platform()
-    test_platform_known()
-    test_platform_specs()
-    test_known_sys_detect()
+    pytest.main([__file__])