Add base RunSettings support for slurm, pbs, and cobalt

conftest.py

@@ -108,6 +108,31 @@ def get_test_interface():
         global test_nic
         return test_nic
 
+    @staticmethod
+    def get_base_run_settings(exe, args, nodes=1, ntasks=1, **kwargs):
+        if test_launcher == "slurm":
+            run_args = {"--nodes": nodes,
+                        "--ntasks": ntasks,
+                        "--time": "00:10:00"}
+            run_args.update(kwargs)
+            settings = RunSettings(exe, args, run_command="srun", run_args=run_args)
+            return settings
+        if test_launcher == "pbs":
+            run_args = {"--pes": ntasks}
+            run_args.update(kwargs)
+            settings = RunSettings(exe, args, run_command="aprun", run_args=run_args)
+            return settings
+        if test_launcher == "cobalt":
+            run_args = {"--pes": ntasks}
+            run_args.update(kwargs)
+            settings = RunSettings(exe, args, run_command="aprun", run_args=run_args)
+            return settings
+        if test_launcher == "lsf":
+            raise SSUnsupportedError("SmartSim LSF launcher does not support custom run settings at this time.")
+        # TODO allow user to pick aprun vs MPIrun
+        return RunSettings(exe, args)
+
+
     @staticmethod
     def get_run_settings(exe, args, nodes=1, ntasks=1, **kwargs):
         if test_launcher == "slurm":

smartsim/launcher/cobalt/cobaltLauncher.py

@@ -30,11 +30,11 @@
 
 from ...constants import STATUS_CANCELLED, STATUS_COMPLETED
 from ...error import LauncherError, SSConfigError
-from ...settings import AprunSettings, CobaltBatchSettings, MpirunSettings
+from ...settings import AprunSettings, CobaltBatchSettings, MpirunSettings, RunSettings
 from ...utils import get_logger
 from ..launcher import WLMLauncher
 from ..pbs.pbsCommands import qdel, qstat
-from ..step import AprunStep, CobaltBatchStep, MpirunStep
+from ..step import AprunStep, CobaltBatchStep, LocalStep, MpirunStep
 from ..stepInfo import CobaltStepInfo
 from .cobaltParser import parse_cobalt_step_id, parse_cobalt_step_status, parse_qsub_out
 
@@ -80,6 +80,9 @@ def create_step(self, name, cwd, step_settings):
             if isinstance(step_settings, MpirunSettings):
                 step = MpirunStep(name, cwd, step_settings)
                 return step
+            if isinstance(step_settings, RunSettings):
+                step = LocalStep(name, cwd, step_settings)
+                return step
             raise TypeError(
                 f"RunSettings type {type(step_settings)} not supported by Cobalt"
             )

smartsim/launcher/pbs/pbsLauncher.py

@@ -28,10 +28,10 @@
 
 from ...constants import STATUS_CANCELLED, STATUS_COMPLETED
 from ...error import LauncherError, SSConfigError
-from ...settings import AprunSettings, MpirunSettings, QsubBatchSettings
+from ...settings import AprunSettings, MpirunSettings, QsubBatchSettings, RunSettings
 from ...utils import get_logger
 from ..launcher import WLMLauncher
-from ..step import AprunStep, MpirunStep, QsubBatchStep
+from ..step import AprunStep, LocalStep, MpirunStep, QsubBatchStep
 from ..stepInfo import PBSStepInfo
 from .pbsCommands import qdel, qstat
 from .pbsParser import parse_qstat_jobid, parse_step_id_from_qstat
@@ -76,6 +76,9 @@ def create_step(self, name, cwd, step_settings):
             if isinstance(step_settings, MpirunSettings):
                 step = MpirunStep(name, cwd, step_settings)
                 return step
+            if isinstance(step_settings, RunSettings):
+                step = LocalStep(name, cwd, step_settings)
+                return step
             raise TypeError(
                 f"RunSettings type {type(step_settings)} not supported by PBSPro"
             )

smartsim/launcher/slurm/slurmLauncher.py

@@ -29,10 +29,10 @@
 
 from ...constants import STATUS_CANCELLED
 from ...error import LauncherError, SSConfigError, SSUnsupportedError
-from ...settings import MpirunSettings, SbatchSettings, SrunSettings
+from ...settings import MpirunSettings, RunSettings, SbatchSettings, SrunSettings
 from ...utils import get_logger
 from ..launcher import WLMLauncher
-from ..step import MpirunStep, SbatchStep, SrunStep
+from ..step import LocalStep, MpirunStep, SbatchStep, SrunStep
 from ..stepInfo import SlurmStepInfo
 from .slurmCommands import sacct, scancel, sstat
 from .slurmParser import parse_sacct, parse_sstat_nodes, parse_step_id_from_sacct
@@ -77,6 +77,9 @@ def create_step(self, name, cwd, step_settings):
             if isinstance(step_settings, MpirunSettings):
                 step = MpirunStep(name, cwd, step_settings)
                 return step
+            if isinstance(step_settings, RunSettings):
+                step = LocalStep(name, cwd, step_settings)
+                return step
             raise SSUnsupportedError("RunSettings type not supported by Slurm")
         except SSConfigError as e:
             raise LauncherError("Step creation failed: " + str(e)) from None

tests/on_wlm/test_base_settings_on_wlm.py

@@ -0,0 +1,106 @@
+import time
+
+import pytest
+
+from smartsim import Experiment, constants
+from smartsim.settings.settings import RunSettings
+
+"""
+Test the launch and stop of models and ensembles using base
+RunSettings while on WLM.
+"""
+
+# retrieved from pytest fixtures
+if pytest.test_launcher not in pytest.wlm_options:
+    pytestmark = pytest.mark.skip(reason="Not testing WLM integrations")
+
+
+def test_model_on_wlm(fileutils, wlmutils):
+    launcher = wlmutils.get_test_launcher()
+    if launcher not in  ["pbs", "slurm", "cobalt"]:
+        pytest.skip("Test only runs on systems with PBSPro, Slurm, or Cobalt as WLM")
+
+    exp_name = "test-base-settings-model-launch"
+    exp = Experiment(exp_name, launcher=wlmutils.get_test_launcher())
+    test_dir = fileutils.make_test_dir(exp_name)
+
+    script = fileutils.get_test_conf_path("sleep.py")
+    settings1 = wlmutils.get_base_run_settings("python", f"{script} --time=5")
+    settings2 = wlmutils.get_base_run_settings("python", f"{script} --time=5")
+    M1 = exp.create_model("m1", path=test_dir, run_settings=settings1)
+    M2 = exp.create_model("m2", path=test_dir, run_settings=settings2)
+
+    # launch models twice to show that they can also be restarted
+    for _ in range(2):
+        exp.start(M1, M2, block=True)
+        statuses = exp.get_status(M1, M2)
+        assert all([stat == constants.STATUS_COMPLETED for stat in statuses])
+
+
+def test_model_stop_on_wlm(fileutils, wlmutils):
+    launcher = wlmutils.get_test_launcher()
+    if launcher not in  ["pbs", "slurm", "cobalt"]:
+        pytest.skip("Test only runs on systems with PBSPro, Slurm, or Cobalt as WLM")
+
+    exp_name = "test-base-settings-model-stop"
+    exp = Experiment(exp_name, launcher=wlmutils.get_test_launcher())
+    test_dir = fileutils.make_test_dir(exp_name)
+
+    script = fileutils.get_test_conf_path("sleep.py")
+    settings1 = wlmutils.get_base_run_settings("python", f"{script} --time=5")
+    settings2 = wlmutils.get_base_run_settings("python", f"{script} --time=5")
+    M1 = exp.create_model("m1", path=test_dir, run_settings=settings1)
+    M2 = exp.create_model("m2", path=test_dir, run_settings=settings2)
+
+    # stop launched models
+    exp.start(M1, M2, block=False)
+    time.sleep(2)
+    exp.stop(M1, M2)
+    assert M1.name in exp._control._jobs.completed
+    assert M2.name in exp._control._jobs.completed
+    statuses = exp.get_status(M1, M2)
+    assert all([stat == constants.STATUS_CANCELLED for stat in statuses])
+
+
+def test_ensemble_on_wlm(fileutils, wlmutils):
+    launcher = wlmutils.get_test_launcher()
+    if launcher not in  ["pbs", "slurm", "cobalt"]:
+        pytest.skip("Test only runs on systems with PBSPro, Slurm, or Cobalt as WLM")
+
+    exp_name = "test-base-settings-ensemble-launch"
+    exp = Experiment(exp_name, launcher=wlmutils.get_test_launcher())
+    test_dir = fileutils.make_test_dir(exp_name)
+
+    script = fileutils.get_test_conf_path("sleep.py")
+    settings = wlmutils.get_base_run_settings("python", f"{script} --time=5")
+    ensemble = exp.create_ensemble("ensemble", run_settings=settings, replicas=2)
+    ensemble.set_path(test_dir)
+
+    # launch ensemble twice to show that it can also be restarted
+    for _ in range(2):
+        exp.start(ensemble, block=True)
+        statuses = exp.get_status(ensemble)
+        assert all([stat == constants.STATUS_COMPLETED for stat in statuses])
+
+
+def test_ensemble_stop_on_wlm(fileutils, wlmutils):
+    launcher = wlmutils.get_test_launcher()
+    if launcher not in  ["pbs", "slurm", "cobalt"]:
+        pytest.skip("Test only runs on systems with PBSPro, Slurm, or Cobalt as WLM")
+
+    exp_name = "test-base-settings-ensemble-launch"
+    exp = Experiment(exp_name, launcher=wlmutils.get_test_launcher())
+    test_dir = fileutils.make_test_dir(exp_name)
+
+    script = fileutils.get_test_conf_path("sleep.py")
+    settings = wlmutils.get_base_run_settings("python", f"{script} --time=5")
+    ensemble = exp.create_ensemble("ensemble", run_settings=settings, replicas=2)
+    ensemble.set_path(test_dir)
+
+    # stop launched ensemble
+    exp.start(ensemble, block=False)
+    time.sleep(2)
+    exp.stop(ensemble)
+    statuses = exp.get_status(ensemble)
+    assert all([stat == constants.STATUS_CANCELLED for stat in statuses])
+    assert all([m.name in exp._control._jobs.completed for m in ensemble])

tests/on_wlm/test_simple_base_settings_on_wlm.py

@@ -0,0 +1,97 @@
+import time
+
+import pytest
+
+from smartsim import Experiment, constants
+from smartsim.settings.settings import RunSettings
+
+"""
+Test the launch and stop of simple models and ensembles that use base
+RunSettings while on WLM.
+"""
+
+# retrieved from pytest fixtures
+if pytest.test_launcher not in pytest.wlm_options:
+    pytestmark = pytest.mark.skip(reason="Not testing WLM integrations")
+
+
+def test_simple_model_on_wlm(fileutils, wlmutils):
+    launcher = wlmutils.get_test_launcher()
+    if launcher not in  ["pbs", "slurm", "cobalt"]:
+        pytest.skip("Test only runs on systems with PBSPro, Slurm, or Cobalt as WLM")
+
+    exp_name = "test-simplebase-settings-model-launch"
+    exp = Experiment(exp_name, launcher=wlmutils.get_test_launcher())
+    test_dir = fileutils.make_test_dir(exp_name)
+
+    script = fileutils.get_test_conf_path("sleep.py")
+    settings = RunSettings("python", exe_args=f"{script} --time=5")
+    M = exp.create_model("m", path=test_dir, run_settings=settings)
+
+    # launch model twice to show that it can also be restarted
+    for _ in range(2):
+        exp.start(M, block=True)
+        assert exp.get_status(M)[0] == constants.STATUS_COMPLETED
+
+
+def test_simple_model_stop_on_wlm(fileutils, wlmutils):
+    launcher = wlmutils.get_test_launcher()
+    if launcher not in  ["pbs", "slurm", "cobalt"]:
+        pytest.skip("Test only runs on systems with PBSPro, Slurm, or Cobalt as WLM")
+
+    exp_name = "test-simplebase-settings-model-stop"
+    exp = Experiment(exp_name, launcher=wlmutils.get_test_launcher())
+    test_dir = fileutils.make_test_dir(exp_name)
+
+    script = fileutils.get_test_conf_path("sleep.py")
+    settings = RunSettings("python", exe_args=f"{script} --time=5")
+    M = exp.create_model("m", path=test_dir, run_settings=settings)
+
+    # stop launched model
+    exp.start(M, block=False)
+    time.sleep(2)
+    exp.stop(M)
+    assert M.name in exp._control._jobs.completed
+    assert exp.get_status(M)[0] == constants.STATUS_CANCELLED
+
+
+def test_simple_ensemble_on_wlm(fileutils, wlmutils):
+    launcher = wlmutils.get_test_launcher()
+    if launcher not in  ["pbs", "slurm", "cobalt"]:
+        pytest.skip("Test only runs on systems with PBSPro, Slurm, or Cobalt as WLM")
+
+    exp_name = "test-simple-base-settings-ensemble-launch"
+    exp = Experiment(exp_name, launcher=wlmutils.get_test_launcher())
+    test_dir = fileutils.make_test_dir(exp_name)
+
+    script = fileutils.get_test_conf_path("sleep.py")
+    settings = RunSettings("python", exe_args=f"{script} --time=5")
+    ensemble = exp.create_ensemble("ensemble", run_settings=settings, replicas=1)
+    ensemble.set_path(test_dir)
+
+    # launch ensemble twice to show that it can also be restarted
+    for _ in range(2):
+        exp.start(ensemble, block=True)
+        assert exp.get_status(ensemble)[0] == constants.STATUS_COMPLETED
+
+
+def test_simple_ensemble_stop_on_wlm(fileutils, wlmutils):
+    launcher = wlmutils.get_test_launcher()
+    if launcher not in  ["pbs", "slurm", "cobalt"]:
+        pytest.skip("Test only runs on systems with PBSPro, Slurm, or Cobalt as WLM")
+
+    exp_name = "test-simple-base-settings-ensemble-stop"
+    exp = Experiment(exp_name, launcher=wlmutils.get_test_launcher())
+    test_dir = fileutils.make_test_dir(exp_name)
+
+    script = fileutils.get_test_conf_path("sleep.py")
+    settings = RunSettings("python", exe_args=f"{script} --time=5")
+    ensemble = exp.create_ensemble("ensemble", run_settings=settings, replicas=1)
+    ensemble.set_path(test_dir)
+
+    # stop launched ensemble
+    exp.start(ensemble, block=False)
+    time.sleep(2)
+    exp.stop(ensemble)
+    assert exp.get_status(ensemble)[0] == constants.STATUS_CANCELLED
+    assert ensemble.models[0].name in exp._control._jobs.completed