From 42f7ff1a04727ab276590656e89516bea7e5320f Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Mon, 17 Jun 2024 18:25:51 -0500 Subject: [PATCH 01/34] Builders are generic, have a format method --- smartsim/settings/builders/launch/dragon.py | 2 +- smartsim/settings/builders/launch/local.py | 2 +- smartsim/settings/builders/launch/lsf.py | 2 +- smartsim/settings/builders/launch/mpi.py | 2 +- smartsim/settings/builders/launch/pals.py | 2 +- smartsim/settings/builders/launch/slurm.py | 2 +- smartsim/settings/builders/launchArgBuilder.py | 16 +++++++++++++--- 7 files changed, 19 insertions(+), 9 deletions(-) diff --git a/smartsim/settings/builders/launch/dragon.py b/smartsim/settings/builders/launch/dragon.py index 1ca0a244d..6043d6934 100644 --- a/smartsim/settings/builders/launch/dragon.py +++ b/smartsim/settings/builders/launch/dragon.py @@ -37,7 +37,7 @@ logger = get_logger(__name__) -class DragonArgBuilder(LaunchArgBuilder): +class DragonArgBuilder(LaunchArgBuilder[t.Any]): # TODO: come back and fix this def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Dragon.value diff --git a/smartsim/settings/builders/launch/local.py b/smartsim/settings/builders/launch/local.py index 595514f15..cdcd22a1c 100644 --- a/smartsim/settings/builders/launch/local.py +++ b/smartsim/settings/builders/launch/local.py @@ -37,7 +37,7 @@ logger = get_logger(__name__) -class LocalArgBuilder(LaunchArgBuilder): +class LocalArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Local.value diff --git a/smartsim/settings/builders/launch/lsf.py b/smartsim/settings/builders/launch/lsf.py index 2c72002e5..293f33281 100644 --- a/smartsim/settings/builders/launch/lsf.py +++ b/smartsim/settings/builders/launch/lsf.py @@ -37,7 +37,7 @@ logger = get_logger(__name__) -class JsrunArgBuilder(LaunchArgBuilder): +class JsrunArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Lsf.value diff --git a/smartsim/settings/builders/launch/mpi.py b/smartsim/settings/builders/launch/mpi.py index 1331be317..2fa68450c 100644 --- a/smartsim/settings/builders/launch/mpi.py +++ b/smartsim/settings/builders/launch/mpi.py @@ -37,7 +37,7 @@ logger = get_logger(__name__) -class _BaseMPIArgBuilder(LaunchArgBuilder): +class _BaseMPIArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def _reserved_launch_args(self) -> set[str]: """Return reserved launch arguments.""" return {"wd", "wdir"} diff --git a/smartsim/settings/builders/launch/pals.py b/smartsim/settings/builders/launch/pals.py index 051409c29..77a907e7a 100644 --- a/smartsim/settings/builders/launch/pals.py +++ b/smartsim/settings/builders/launch/pals.py @@ -37,7 +37,7 @@ logger = get_logger(__name__) -class PalsMpiexecArgBuilder(LaunchArgBuilder): +class PalsMpiexecArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Pals.value diff --git a/smartsim/settings/builders/launch/slurm.py b/smartsim/settings/builders/launch/slurm.py index 80d3d6be2..e0edd74dd 100644 --- a/smartsim/settings/builders/launch/slurm.py +++ b/smartsim/settings/builders/launch/slurm.py @@ -39,7 +39,7 @@ logger = get_logger(__name__) -class SlurmArgBuilder(LaunchArgBuilder): +class SlurmArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Slurm.value diff --git a/smartsim/settings/builders/launchArgBuilder.py b/smartsim/settings/builders/launchArgBuilder.py index bb1f389f3..8a839a5c8 100644 --- a/smartsim/settings/builders/launchArgBuilder.py +++ b/smartsim/settings/builders/launchArgBuilder.py @@ -37,7 +37,10 @@ logger = get_logger(__name__) -class LaunchArgBuilder(ABC): +_T = t.TypeVar("_T") + + +class LaunchArgBuilder(ABC, t.Generic[_T]): """Abstract base class that defines all generic launcher argument methods that are not supported. It is the responsibility of child classes for each launcher to translate @@ -50,12 +53,14 @@ def __init__(self, launch_args: t.Dict[str, str | None] | None) -> None: @abstractmethod def launcher_str(self) -> str: """Get the string representation of the launcher""" - pass @abstractmethod def set(self, arg: str, val: str | None) -> None: """Set the launch arguments""" - pass + + @abstractmethod + def finalize(self, exe: ExecutableLike, env: dict[str, str | None]) -> _T: + """Prepare an entity for launch using the built options""" def format_launch_args(self) -> t.Union[t.List[str], None]: """Build formatted launch arguments""" @@ -90,3 +95,8 @@ def format_env_vars( def __str__(self) -> str: # pragma: no-cover string = f"\nLaunch Arguments:\n{fmt_dict(self._launch_args)}" return string + + +class ExecutableLike(t.Protocol): + @abstractmethod + def as_program_arguments(self) -> t.Sequence[str]: ... From f76ee4fe161a85ab4b5cf38d57e83111b3b08f79 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Tue, 18 Jun 2024 16:32:26 -0500 Subject: [PATCH 02/34] Impl slurm --- smartsim/settings/builders/launch/slurm.py | 13 +++++++ tests/temp_tests/test_settings/conftest.py | 38 +++++++++++++++++++ .../test_settings/test_slurmLauncher.py | 38 +++++++++++++++++++ 3 files changed, 89 insertions(+) create mode 100644 tests/temp_tests/test_settings/conftest.py diff --git a/smartsim/settings/builders/launch/slurm.py b/smartsim/settings/builders/launch/slurm.py index e0edd74dd..3210c6665 100644 --- a/smartsim/settings/builders/launch/slurm.py +++ b/smartsim/settings/builders/launch/slurm.py @@ -36,6 +36,9 @@ from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder +if t.TYPE_CHECKING: + from smartsim.settings.builders.launchArgBuilder import ExecutableLike + logger = get_logger(__name__) @@ -315,3 +318,13 @@ def set(self, key: str, value: str | None) -> None: if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") self._launch_args[key] = value + + def finalize( + self, exe: ExecutableLike, env: dict[str, str | None] + ) -> t.Sequence[str]: + return ( + "srun", + *(self.format_launch_args() or ()), + "--", + *exe.as_program_arguments(), + ) diff --git a/tests/temp_tests/test_settings/conftest.py b/tests/temp_tests/test_settings/conftest.py new file mode 100644 index 000000000..ad1fa0f4a --- /dev/null +++ b/tests/temp_tests/test_settings/conftest.py @@ -0,0 +1,38 @@ +# BSD 2-Clause License +# +# Copyright (c) 2021-2024, Hewlett Packard Enterprise +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import pytest + +from smartsim.settings.builders import launchArgBuilder as launch + + +@pytest.fixture +def echo_executable_like(): + class _ExeLike(launch.ExecutableLike): + def as_program_arguments(self): + return ("echo", "hello", "world") + + return _ExeLike() diff --git a/tests/temp_tests/test_settings/test_slurmLauncher.py b/tests/temp_tests/test_settings/test_slurmLauncher.py index c5e9b5b62..bfa7dd9e1 100644 --- a/tests/temp_tests/test_settings/test_slurmLauncher.py +++ b/tests/temp_tests/test_settings/test_slurmLauncher.py @@ -4,6 +4,8 @@ from smartsim.settings.builders.launch.slurm import SlurmArgBuilder from smartsim.settings.launchCommand import LauncherType +pytestmark = pytest.mark.group_a + def test_launcher_str(): """Ensure launcher_str returns appropriate value""" @@ -253,3 +255,39 @@ def test_set_het_groups(monkeypatch): assert slurmLauncher._arg_builder._launch_args["het-group"] == "3,2" with pytest.raises(ValueError): slurmLauncher.launch_args.set_het_group([4]) + + +@pytest.mark.parametrize( + "args, expected", + ( + pytest.param({}, ("srun", "--", "echo", "hello", "world"), id="Empty Args"), + pytest.param( + {"N": "1"}, + ("srun", "-N", "1", "--", "echo", "hello", "world"), + id="Short Arg", + ), + pytest.param( + {"nodes": "1"}, + ("srun", "--nodes=1", "--", "echo", "hello", "world"), + id="Long Arg", + ), + pytest.param( + {"v": None}, + ("srun", "-v", "--", "echo", "hello", "world"), + id="Short Arg (No Value)", + ), + pytest.param( + {"verbose": None}, + ("srun", "--verbose", "--", "echo", "hello", "world"), + id="Long Arg (No Value)", + ), + pytest.param( + {"nodes": "1", "n": "123"}, + ("srun", "--nodes=1", "-n", "123", "--", "echo", "hello", "world"), + id="Short and Long Args", + ), + ), +) +def test_formatting_launch_args(echo_executable_like, args, expected): + cmd = SlurmArgBuilder(args).finalize(echo_executable_like, {}) + assert tuple(cmd) == expected From 3027721ae9ad0aa8e74de677f4e3c9acb3ac4d38 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Tue, 18 Jun 2024 16:34:22 -0500 Subject: [PATCH 03/34] Impl local --- smartsim/settings/builders/launch/local.py | 8 ++++++++ tests/temp_tests/test_settings/test_localLauncher.py | 7 +++++++ 2 files changed, 15 insertions(+) diff --git a/smartsim/settings/builders/launch/local.py b/smartsim/settings/builders/launch/local.py index cdcd22a1c..78e2dd28f 100644 --- a/smartsim/settings/builders/launch/local.py +++ b/smartsim/settings/builders/launch/local.py @@ -34,6 +34,9 @@ from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder +if t.TYPE_CHECKING: + from smartsim.settings.builders.launchArgBuilder import ExecutableLike + logger = get_logger(__name__) @@ -72,3 +75,8 @@ def set(self, key: str, value: str | None) -> None: if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") self._launch_args[key] = value + + def finalize( + self, exe: ExecutableLike, env: dict[str, str | None] + ) -> t.Sequence[str]: + return exe.as_program_arguments() diff --git a/tests/temp_tests/test_settings/test_localLauncher.py b/tests/temp_tests/test_settings/test_localLauncher.py index 1ee7b9d87..4eb314a8b 100644 --- a/tests/temp_tests/test_settings/test_localLauncher.py +++ b/tests/temp_tests/test_settings/test_localLauncher.py @@ -4,6 +4,8 @@ from smartsim.settings.builders.launch.local import LocalArgBuilder from smartsim.settings.launchCommand import LauncherType +pytestmark = pytest.mark.group_a + def test_launcher_str(): """Ensure launcher_str returns appropriate value""" @@ -110,3 +112,8 @@ def test_format_env_vars(): localLauncher = LaunchSettings(launcher=LauncherType.Local, env_vars=env_vars) assert isinstance(localLauncher._arg_builder, LocalArgBuilder) assert localLauncher.format_env_vars() == ["A=a", "B=", "C=", "D=12"] + + +def test_formatting_returns_original_exe(echo_executable_like): + cmd = LocalArgBuilder({}).finalize(echo_executable_like, {}) + assert tuple(cmd) == ("echo", "hello", "world") From 6266b759f243440f71577093771a23cfb36e5cbc Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Tue, 18 Jun 2024 17:31:01 -0500 Subject: [PATCH 04/34] Impl jsrun --- smartsim/settings/builders/launch/lsf.py | 13 +++++++ .../test_settings/test_lsfLauncher.py | 38 +++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/smartsim/settings/builders/launch/lsf.py b/smartsim/settings/builders/launch/lsf.py index 293f33281..7fd3217ab 100644 --- a/smartsim/settings/builders/launch/lsf.py +++ b/smartsim/settings/builders/launch/lsf.py @@ -34,6 +34,9 @@ from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder +if t.TYPE_CHECKING: + from smartsim.settings.builders.launchArgBuilder import ExecutableLike + logger = get_logger(__name__) @@ -115,3 +118,13 @@ def set(self, key: str, value: str | None) -> None: if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") self._launch_args[key] = value + + def finalize( + self, exe: ExecutableLike, env: dict[str, str | None] + ) -> t.Sequence[str]: + return ( + "jsrun", + *(self.format_launch_args() or ()), + "--", + *exe.as_program_arguments(), + ) diff --git a/tests/temp_tests/test_settings/test_lsfLauncher.py b/tests/temp_tests/test_settings/test_lsfLauncher.py index 4c4260ac5..592c80ce7 100644 --- a/tests/temp_tests/test_settings/test_lsfLauncher.py +++ b/tests/temp_tests/test_settings/test_lsfLauncher.py @@ -4,6 +4,8 @@ from smartsim.settings.builders.launch.lsf import JsrunArgBuilder from smartsim.settings.launchCommand import LauncherType +pytestmark = pytest.mark.group_a + def test_launcher_str(): """Ensure launcher_str returns appropriate value""" @@ -56,3 +58,39 @@ def test_launch_args(): "--np=100", ] assert formatted == result + + +@pytest.mark.parametrize( + "args, expected", + ( + pytest.param({}, ("jsrun", "--", "echo", "hello", "world"), id="Empty Args"), + pytest.param( + {"n": "1"}, + ("jsrun", "-n", "1", "--", "echo", "hello", "world"), + id="Short Arg", + ), + pytest.param( + {"nrs": "1"}, + ("jsrun", "--nrs=1", "--", "echo", "hello", "world"), + id="Long Arg", + ), + pytest.param( + {"v": None}, + ("jsrun", "-v", "--", "echo", "hello", "world"), + id="Short Arg (No Value)", + ), + pytest.param( + {"verbose": None}, + ("jsrun", "--verbose", "--", "echo", "hello", "world"), + id="Long Arg (No Value)", + ), + pytest.param( + {"tasks_per_rs": "1", "n": "123"}, + ("jsrun", "--tasks_per_rs=1", "-n", "123", "--", "echo", "hello", "world"), + id="Short and Long Args", + ), + ), +) +def test_formatting_launch_args(echo_executable_like, args, expected): + cmd = JsrunArgBuilder(args).finalize(echo_executable_like, {}) + assert tuple(cmd) == expected From bbb515297498985f929535b3475a52a8484527ab Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Tue, 18 Jun 2024 18:15:11 -0500 Subject: [PATCH 05/34] Impl mpi{run,exec}, orterun REBASEME: moar mpi REBASEME: fmt mpi --- smartsim/settings/builders/launch/mpi.py | 46 +++++++++++-------- smartsim/settings/builders/launch/pals.py | 13 ++++++ .../test_settings/test_mpiLauncher.py | 46 +++++++++++++++++++ .../test_settings/test_palsLauncher.py | 38 +++++++++++++++ 4 files changed, 125 insertions(+), 18 deletions(-) diff --git a/smartsim/settings/builders/launch/mpi.py b/smartsim/settings/builders/launch/mpi.py index 2fa68450c..b012a4271 100644 --- a/smartsim/settings/builders/launch/mpi.py +++ b/smartsim/settings/builders/launch/mpi.py @@ -34,6 +34,9 @@ from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder +if t.TYPE_CHECKING: + from smartsim.settings.builders.launchArgBuilder import ExecutableLike + logger = get_logger(__name__) @@ -215,36 +218,43 @@ def set(self, key: str, value: str | None) -> None: class MpiArgBuilder(_BaseMPIArgBuilder): - def __init__( - self, - launch_args: t.Dict[str, str | None] | None, - ) -> None: - super().__init__(launch_args) - def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Mpirun.value + def finalize( + self, exe: ExecutableLike, env: dict[str, str | None] + ) -> t.Sequence[str]: + return ("mpirun", *self.format_launch_args(), "--", *exe.as_program_arguments()) -class MpiexecArgBuilder(_BaseMPIArgBuilder): - def __init__( - self, - launch_args: t.Dict[str, str | None] | None, - ) -> None: - super().__init__(launch_args) +class MpiexecArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Mpiexec.value + def finalize( + self, exe: ExecutableLike, env: dict[str, str | None] + ) -> t.Sequence[str]: + return ( + "mpiexec", + *self.format_launch_args(), + "--", + *exe.as_program_arguments(), + ) -class OrteArgBuilder(_BaseMPIArgBuilder): - def __init__( - self, - launch_args: t.Dict[str, str | None] | None, - ) -> None: - super().__init__(launch_args) +class OrteArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Orterun.value + + def finalize( + self, exe: ExecutableLike, env: dict[str, str | None] + ) -> t.Sequence[str]: + return ( + "orterun", + *self.format_launch_args(), + "--", + *exe.as_program_arguments(), + ) diff --git a/smartsim/settings/builders/launch/pals.py b/smartsim/settings/builders/launch/pals.py index 77a907e7a..af8cd7706 100644 --- a/smartsim/settings/builders/launch/pals.py +++ b/smartsim/settings/builders/launch/pals.py @@ -34,6 +34,9 @@ from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder +if t.TYPE_CHECKING: + from smartsim.settings.builders.launchArgBuilder import ExecutableLike + logger = get_logger(__name__) @@ -149,3 +152,13 @@ def set(self, key: str, value: str | None) -> None: if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") self._launch_args[key] = value + + def finalize( + self, exe: ExecutableLike, env: dict[str, str | None] + ) -> t.Sequence[str]: + return ( + "mpiexec", + *(self.format_launch_args() or ()), + "--", + *exe.as_program_arguments(), + ) diff --git a/tests/temp_tests/test_settings/test_mpiLauncher.py b/tests/temp_tests/test_settings/test_mpiLauncher.py index 815f0c5c1..9b651c220 100644 --- a/tests/temp_tests/test_settings/test_mpiLauncher.py +++ b/tests/temp_tests/test_settings/test_mpiLauncher.py @@ -10,6 +10,8 @@ ) from smartsim.settings.launchCommand import LauncherType +pytestmark = pytest.mark.group_a + @pytest.mark.parametrize( "launcher", @@ -205,3 +207,47 @@ def test_invalid_hostlist_format(launcher): mpiSettings.launch_args.set_hostlist([5]) with pytest.raises(TypeError): mpiSettings.launch_args.set_hostlist(5) + + +@pytest.mark.parametrize( + "cls, cmd", + ( + pytest.param(MpiArgBuilder, "mpirun", id="w/ mpirun"), + pytest.param(MpiexecArgBuilder, "mpiexec", id="w/ mpiexec"), + pytest.param(OrteArgBuilder, "orterun", id="w/ orterun"), + ), +) +@pytest.mark.parametrize( + "args, expected", + ( + pytest.param({}, ("--", "echo", "hello", "world"), id="Empty Args"), + pytest.param( + {"n": "1"}, + ("--n", "1", "--", "echo", "hello", "world"), + id="Short Arg", + ), + pytest.param( + {"host": "myhost"}, + ("--host", "myhost", "--", "echo", "hello", "world"), + id="Long Arg", + ), + pytest.param( + {"v": None}, + ("--v", "--", "echo", "hello", "world"), + id="Short Arg (No Value)", + ), + pytest.param( + {"verbose": None}, + ("--verbose", "--", "echo", "hello", "world"), + id="Long Arg (No Value)", + ), + pytest.param( + {"n": "1", "host": "myhost"}, + ("--n", "1", "--host", "myhost", "--", "echo", "hello", "world"), + id="Short and Long Args", + ), + ), +) +def test_formatting_launch_args(echo_executable_like, cls, cmd, args, expected): + fmt = cls(args).finalize(echo_executable_like, {}) + assert tuple(fmt) == (cmd,) + expected diff --git a/tests/temp_tests/test_settings/test_palsLauncher.py b/tests/temp_tests/test_settings/test_palsLauncher.py index 01cbea2ed..a0bc7821c 100644 --- a/tests/temp_tests/test_settings/test_palsLauncher.py +++ b/tests/temp_tests/test_settings/test_palsLauncher.py @@ -4,6 +4,8 @@ from smartsim.settings.builders.launch.pals import PalsMpiexecArgBuilder from smartsim.settings.launchCommand import LauncherType +pytestmark = pytest.mark.group_a + def test_launcher_str(): """Ensure launcher_str returns appropriate value""" @@ -67,3 +69,39 @@ def test_invalid_hostlist_format(): palsLauncher.launch_args.set_hostlist([5]) with pytest.raises(TypeError): palsLauncher.launch_args.set_hostlist(5) + + +@pytest.mark.parametrize( + "args, expected", + ( + pytest.param({}, ("mpiexec", "--", "echo", "hello", "world"), id="Empty Args"), + pytest.param( + {"n": "1"}, + ("mpiexec", "--n", "1", "--", "echo", "hello", "world"), + id="Short Arg", + ), + pytest.param( + {"host": "myhost"}, + ("mpiexec", "--host", "myhost", "--", "echo", "hello", "world"), + id="Long Arg", + ), + pytest.param( + {"v": None}, + ("mpiexec", "--v", "--", "echo", "hello", "world"), + id="Short Arg (No Value)", + ), + pytest.param( + {"verbose": None}, + ("mpiexec", "--verbose", "--", "echo", "hello", "world"), + id="Long Arg (No Value)", + ), + pytest.param( + {"n": "1", "host": "myhost"}, + ("mpiexec", "--n", "1", "--host", "myhost", "--", "echo", "hello", "world"), + id="Short and Long Args", + ), + ), +) +def test_formatting_launch_args(echo_executable_like, args, expected): + cmd = PalsMpiexecArgBuilder(args).finalize(echo_executable_like, {}) + assert tuple(cmd) == expected From b96a1605aef7b591f918d5f2d2ed0372b4ff3e9a Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Tue, 18 Jun 2024 18:32:24 -0500 Subject: [PATCH 06/34] Impl aprun --- smartsim/settings/builders/launch/alps.py | 15 +++++++- .../test_settings/test_alpsLauncher.py | 38 +++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/smartsim/settings/builders/launch/alps.py b/smartsim/settings/builders/launch/alps.py index a527cafac..538ddb00e 100644 --- a/smartsim/settings/builders/launch/alps.py +++ b/smartsim/settings/builders/launch/alps.py @@ -34,10 +34,13 @@ from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder +if t.TYPE_CHECKING: + from smartsim.settings.builders.launchArgBuilder import ExecutableLike + logger = get_logger(__name__) -class AprunArgBuilder(LaunchArgBuilder): +class AprunArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def _reserved_launch_args(self) -> set[str]: """Return reserved launch arguments.""" return {"wdir"} @@ -213,3 +216,13 @@ def set(self, key: str, value: str | None) -> None: if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") self._launch_args[key] = value + + def finalize( + self, exe: ExecutableLike, env: dict[str, str | None] + ) -> t.Sequence[str]: + return ( + "aprun", + *(self.format_launch_args() or ()), + "--", + *exe.as_program_arguments(), + ) diff --git a/tests/temp_tests/test_settings/test_alpsLauncher.py b/tests/temp_tests/test_settings/test_alpsLauncher.py index 7f9a4c3b9..7fa95cb6d 100644 --- a/tests/temp_tests/test_settings/test_alpsLauncher.py +++ b/tests/temp_tests/test_settings/test_alpsLauncher.py @@ -4,6 +4,8 @@ from smartsim.settings.builders.launch.alps import AprunArgBuilder from smartsim.settings.launchCommand import LauncherType +pytestmark = pytest.mark.group_a + def test_launcher_str(): """Ensure launcher_str returns appropriate value""" @@ -147,3 +149,39 @@ def test_invalid_exclude_hostlist_format(): alpsLauncher.launch_args.set_excluded_hosts([5]) with pytest.raises(TypeError): alpsLauncher.launch_args.set_excluded_hosts(5) + + +@pytest.mark.parametrize( + "args, expected", + ( + pytest.param({}, ("aprun", "--", "echo", "hello", "world"), id="Empty Args"), + pytest.param( + {"N": "1"}, + ("aprun", "-N", "1", "--", "echo", "hello", "world"), + id="Short Arg", + ), + pytest.param( + {"cpus-per-pe": "1"}, + ("aprun", "--cpus-per-pe=1", "--", "echo", "hello", "world"), + id="Long Arg", + ), + pytest.param( + {"q": None}, + ("aprun", "-q", "--", "echo", "hello", "world"), + id="Short Arg (No Value)", + ), + pytest.param( + {"quiet": None}, + ("aprun", "--quiet", "--", "echo", "hello", "world"), + id="Long Arg (No Value)", + ), + pytest.param( + {"N": "1", "cpus-per-pe": "123"}, + ("aprun", "-N", "1", "--cpus-per-pe=123", "--", "echo", "hello", "world"), + id="Short and Long Args", + ), + ), +) +def test_formatting_launch_args(echo_executable_like, args, expected): + cmd = AprunArgBuilder(args).finalize(echo_executable_like, {}) + assert tuple(cmd) == expected From dc1cf0ba60d9054c2ccb7b79386d4fdf2b7572dc Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Thu, 20 Jun 2024 14:18:10 -0500 Subject: [PATCH 07/34] Impl dragon --- smartsim/settings/builders/launch/dragon.py | 30 ++++++++++++++- .../test_settings/test_dragonLauncher.py | 38 ++++++++++++++++++- 2 files changed, 65 insertions(+), 3 deletions(-) diff --git a/smartsim/settings/builders/launch/dragon.py b/smartsim/settings/builders/launch/dragon.py index 6043d6934..bc5f2b528 100644 --- a/smartsim/settings/builders/launch/dragon.py +++ b/smartsim/settings/builders/launch/dragon.py @@ -26,18 +26,23 @@ from __future__ import annotations +import os import typing as t +from smartsim._core.schemas.dragonRequests import DragonRunRequest from smartsim.log import get_logger from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder +if t.TYPE_CHECKING: + from smartsim.settings.builders.launchArgBuilder import ExecutableLike + logger = get_logger(__name__) -class DragonArgBuilder(LaunchArgBuilder[t.Any]): # TODO: come back and fix this +class DragonArgBuilder(LaunchArgBuilder[DragonRunRequest]): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Dragon.value @@ -54,7 +59,7 @@ def set_tasks_per_node(self, tasks_per_node: int) -> None: :param tasks_per_node: number of tasks per node """ - self.set("tasks-per-node", str(tasks_per_node)) + self.set("tasks_per_node", str(tasks_per_node)) def set(self, key: str, value: str | None) -> None: """Set the launch arguments""" @@ -62,3 +67,24 @@ def set(self, key: str, value: str | None) -> None: if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") self._launch_args[key] = value + + def finalize( + self, exe: ExecutableLike, env: dict[str, str | None] + ) -> DragonRunRequest: + exe_, *args = exe.as_program_arguments() + return DragonRunRequest( + exe=exe_, + exe_args=args, + path=os.getcwd(), # FIXME: Currently this is hard coded because + # the schema requires it, but in future, + # it is almost certainly necessary that + # this will need to be injected by the + # user or by us to have the command + # execute next to any generated files. A + # similar problem exists for the other + # settings. + # TODO: Find a way to inject this path + env=env, + current_env=dict(os.environ), + **self._launch_args, + ) diff --git a/tests/temp_tests/test_settings/test_dragonLauncher.py b/tests/temp_tests/test_settings/test_dragonLauncher.py index d21a21c59..004090eef 100644 --- a/tests/temp_tests/test_settings/test_dragonLauncher.py +++ b/tests/temp_tests/test_settings/test_dragonLauncher.py @@ -1,9 +1,12 @@ import pytest +from smartsim._core.schemas.dragonRequests import DragonRunRequest from smartsim.settings import LaunchSettings from smartsim.settings.builders.launch.dragon import DragonArgBuilder from smartsim.settings.launchCommand import LauncherType +pytestmark = pytest.mark.group_a + def test_launcher_str(): """Ensure launcher_str returns appropriate value""" @@ -16,7 +19,7 @@ def test_launcher_str(): [ pytest.param("set_nodes", (2,), "2", "nodes", id="set_nodes"), pytest.param( - "set_tasks_per_node", (2,), "2", "tasks-per-node", id="set_tasks_per_node" + "set_tasks_per_node", (2,), "2", "tasks_per_node", id="set_tasks_per_node" ), ], ) @@ -25,3 +28,36 @@ def test_dragon_class_methods(function, value, flag, result): assert isinstance(dragonLauncher._arg_builder, DragonArgBuilder) getattr(dragonLauncher.launch_args, function)(*value) assert dragonLauncher.launch_args._launch_args[flag] == result + + +NOT_SET = object() + + +@pytest.mark.parametrize("nodes", (NOT_SET, 20, 40)) +@pytest.mark.parametrize("tasks_per_node", (NOT_SET, 1, 20)) +def test_formatting_launch_args_into_request( + echo_executable_like, nodes, tasks_per_node +): + builder = DragonArgBuilder({}) + if nodes is not NOT_SET: + builder.set_nodes(nodes) + if tasks_per_node is not NOT_SET: + builder.set_tasks_per_node(tasks_per_node) + req = builder.finalize(echo_executable_like, {}) + + args = dict( + (k, v) + for k, v in ( + ("nodes", nodes), + ("tasks_per_node", tasks_per_node), + ) + if v is not NOT_SET + ) + expected = DragonRunRequest( + exe="echo", exe_args=["hello", "world"], path="/tmp", env={}, **args + ) + + assert req.nodes == expected.nodes + assert req.tasks_per_node == expected.tasks_per_node + assert req.hostlist == expected.hostlist + assert req.pmi_enabled == expected.pmi_enabled From e5553426e9084d1584e3d8a0d75ccb301736476e Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Thu, 20 Jun 2024 16:09:39 -0500 Subject: [PATCH 08/34] Type errors supressed for now --- smartsim/entity/entity.py | 3 +++ smartsim/settings/launchSettings.py | 11 +++++++++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/smartsim/entity/entity.py b/smartsim/entity/entity.py index 2f4b651f9..5304d2c29 100644 --- a/smartsim/entity/entity.py +++ b/smartsim/entity/entity.py @@ -26,6 +26,9 @@ import typing as t +if t.TYPE_CHECKING: + import smartsim.settings.base.RunSettings + class TelemetryConfiguration: """A base class for configuraing telemetry production behavior on diff --git a/smartsim/settings/launchSettings.py b/smartsim/settings/launchSettings.py index a9e5e8103..f512fcc34 100644 --- a/smartsim/settings/launchSettings.py +++ b/smartsim/settings/launchSettings.py @@ -67,8 +67,15 @@ def launcher(self) -> str: return self._launcher.value @property - def launch_args(self) -> LaunchArgBuilder: + def launch_args(self) -> LaunchArgBuilder[t.Any]: """Return the launch argument translator.""" + # FIXME: We _REALLY_ need to make the `LaunchSettings` class generic at + # on `_arg_builder` if we are expecting users to call specific + # `set_*` methods defined specificially on each of the + # subclasses. Otherwise we have no way of showing what methods + # are available at intellisense/static analysis/compile time. + # This whole object basically resolves to being one step removed + # from `Any` typed!! return self._arg_builder @launch_args.setter @@ -88,7 +95,7 @@ def env_vars(self, value: t.Dict[str, str]) -> None: """Set the environment variables.""" self._env_vars = copy.deepcopy(value) - def _get_arg_builder(self, launch_args: StringArgument | None) -> LaunchArgBuilder: + def _get_arg_builder(self, launch_args: StringArgument | None) -> LaunchArgBuilder[t.Any]: """Map the Launcher to the LaunchArgBuilder""" if self._launcher == LauncherType.Slurm: return SlurmArgBuilder(launch_args) From 91611258ad9790269db0dadb2d69c4f93d03deda Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Sat, 22 Jun 2024 17:20:53 -0500 Subject: [PATCH 09/34] Add a dispatcher class to send built settings to a launcher --- smartsim/settings/builders/launch/alps.py | 2 + smartsim/settings/builders/launch/local.py | 2 + smartsim/settings/builders/launch/lsf.py | 2 + smartsim/settings/builders/launch/mpi.py | 4 + smartsim/settings/builders/launch/pals.py | 2 + smartsim/settings/builders/launch/slurm.py | 3 +- smartsim/settings/dispatch.py | 164 ++++++++++++++++++ tests/temp_tests/test_settings/conftest.py | 30 +++- .../temp_tests/test_settings/test_dispatch.py | 142 +++++++++++++++ 9 files changed, 349 insertions(+), 2 deletions(-) create mode 100644 smartsim/settings/dispatch.py create mode 100644 tests/temp_tests/test_settings/test_dispatch.py diff --git a/smartsim/settings/builders/launch/alps.py b/smartsim/settings/builders/launch/alps.py index 538ddb00e..97e83a6c1 100644 --- a/smartsim/settings/builders/launch/alps.py +++ b/smartsim/settings/builders/launch/alps.py @@ -29,6 +29,7 @@ import typing as t from smartsim.log import get_logger +from smartsim.settings.dispatch import default_dispatcher, ShellLauncher from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType @@ -40,6 +41,7 @@ logger = get_logger(__name__) +@default_dispatcher.dispatch(to_launcher=ShellLauncher) class AprunArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def _reserved_launch_args(self) -> set[str]: """Return reserved launch arguments.""" diff --git a/smartsim/settings/builders/launch/local.py b/smartsim/settings/builders/launch/local.py index 78e2dd28f..2cab0df75 100644 --- a/smartsim/settings/builders/launch/local.py +++ b/smartsim/settings/builders/launch/local.py @@ -29,6 +29,7 @@ import typing as t from smartsim.log import get_logger +from smartsim.settings.dispatch import default_dispatcher, ShellLauncher from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType @@ -40,6 +41,7 @@ logger = get_logger(__name__) +@default_dispatcher.dispatch(to_launcher=ShellLauncher) class LocalArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/lsf.py b/smartsim/settings/builders/launch/lsf.py index 7fd3217ab..118946ca2 100644 --- a/smartsim/settings/builders/launch/lsf.py +++ b/smartsim/settings/builders/launch/lsf.py @@ -29,6 +29,7 @@ import typing as t from smartsim.log import get_logger +from smartsim.settings.dispatch import default_dispatcher, ShellLauncher from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType @@ -40,6 +41,7 @@ logger = get_logger(__name__) +@default_dispatcher.dispatch(to_launcher=ShellLauncher) class JsrunArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/mpi.py b/smartsim/settings/builders/launch/mpi.py index b012a4271..d2a1f3170 100644 --- a/smartsim/settings/builders/launch/mpi.py +++ b/smartsim/settings/builders/launch/mpi.py @@ -29,6 +29,7 @@ import typing as t from smartsim.log import get_logger +from smartsim.settings.dispatch import default_dispatcher, ShellLauncher from ...common import set_check_input from ...launchCommand import LauncherType @@ -217,6 +218,7 @@ def set(self, key: str, value: str | None) -> None: self._launch_args[key] = value +@default_dispatcher.dispatch(to_launcher=ShellLauncher) class MpiArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" @@ -228,6 +230,7 @@ def finalize( return ("mpirun", *self.format_launch_args(), "--", *exe.as_program_arguments()) +@default_dispatcher.dispatch(to_launcher=ShellLauncher) class MpiexecArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" @@ -244,6 +247,7 @@ def finalize( ) +@default_dispatcher.dispatch(to_launcher=ShellLauncher) class OrteArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/pals.py b/smartsim/settings/builders/launch/pals.py index af8cd7706..ff1755fb6 100644 --- a/smartsim/settings/builders/launch/pals.py +++ b/smartsim/settings/builders/launch/pals.py @@ -29,6 +29,7 @@ import typing as t from smartsim.log import get_logger +from smartsim.settings.dispatch import default_dispatcher, ShellLauncher from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType @@ -40,6 +41,7 @@ logger = get_logger(__name__) +@default_dispatcher.dispatch(to_launcher=ShellLauncher) class PalsMpiexecArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/slurm.py b/smartsim/settings/builders/launch/slurm.py index 3210c6665..1cdbea3f7 100644 --- a/smartsim/settings/builders/launch/slurm.py +++ b/smartsim/settings/builders/launch/slurm.py @@ -31,6 +31,7 @@ import typing as t from smartsim.log import get_logger +from smartsim.settings.dispatch import default_dispatcher, ShellLauncher from ...common import set_check_input from ...launchCommand import LauncherType @@ -41,7 +42,7 @@ logger = get_logger(__name__) - +@default_dispatcher.dispatch(to_launcher=ShellLauncher) class SlurmArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py new file mode 100644 index 000000000..27b247593 --- /dev/null +++ b/smartsim/settings/dispatch.py @@ -0,0 +1,164 @@ +# BSD 2-Clause License +# +# Copyright (c) 2021-2024, Hewlett Packard Enterprise +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import subprocess as sp +import uuid +import typing as t + +from smartsim.error import errors + +if t.TYPE_CHECKING: + from typing_extensions import Self + from smartsim.experiment import Experiment + from smartsim.settings.builders import LaunchArgBuilder + +_T = t.TypeVar("_T") +_T_contra = t.TypeVar("_T_contra", contravariant=True) + +JobID = t.NewType("JobID", uuid.UUID) + + +def create_job_id() -> JobID: + return JobID(uuid.uuid4()) + + +class LauncherLike(t.Protocol[_T_contra]): + def start(self, launchable: _T_contra) -> JobID: ... + @classmethod + def create(cls, exp: Experiment) -> Self: ... + + +@t.final +class Dispatcher: + """A class capable of deciding which launcher type should be used to launch + a given settings builder type. + """ + + def __init__( + self, + *, + dispatch_registry: ( + t.Mapping[type[LaunchArgBuilder[t.Any]], type[LauncherLike[t.Any]]] | None + ) = None, + ) -> None: + self._dispatch_registry = ( + dict(dispatch_registry) if dispatch_registry is not None else {} + ) + + def copy(self) -> Self: + """Create a shallow copy of the Dispatcher""" + return type(self)(dispatch_registry=self._dispatch_registry) + + @t.overload + def dispatch( + self, + args: None = ..., + *, + to_launcher: type[LauncherLike[_T]], + allow_overwrite: bool = ..., + ) -> t.Callable[[type[LaunchArgBuilder[_T]]], type[LaunchArgBuilder[_T]]]: ... + @t.overload + def dispatch( + self, + args: type[LaunchArgBuilder[_T]], + *, + to_launcher: type[LauncherLike[_T]], + allow_overwrite: bool = ..., + ) -> None: ... + def dispatch( + self, + args: type[LaunchArgBuilder[_T]] | None = None, + *, + to_launcher: type[LauncherLike[_T]], + allow_overwrite: bool = False, + ) -> t.Callable[[type[LaunchArgBuilder[_T]]], type[LaunchArgBuilder[_T]]] | None: + """A type safe way to add a mapping of settings builder to launcher to + handle the settings at launch time. + """ + + def register( + args_: type[LaunchArgBuilder[_T]], / + ) -> type[LaunchArgBuilder[_T]]: + if args_ in self._dispatch_registry and not allow_overwrite: + launcher_type = self._dispatch_registry[args_] + raise TypeError( + f"{args_.__name__} has already been registered to be " + f"launched with {launcher_type}" + ) + self._dispatch_registry[args_] = to_launcher + return args_ + + if args is not None: + register(args) + return None + return register + + def get_launcher_for( + self, args: LaunchArgBuilder[_T] | type[LaunchArgBuilder[_T]], / + ) -> type[LauncherLike[_T]]: + """Find a type of launcher that is registered as being able to launch + the output of the provided builder + """ + if not isinstance(args, type): + args = type(args) + launcher_type = self._dispatch_registry.get(args, None) + if launcher_type is None: + raise TypeError( + f"{type(self).__name__} {self} has no launcher type to " + f"dispatch to for argument builder of type {args}" + ) + # Note the sleight-of-hand here: we are secretly casting a type of + # `LauncherLike[Any]` to `LauncherLike[_T]`. This is safe to do if all + # entries in the mapping were added using a type safe method (e.g. + # `Dispatcher.dispatch`), but if a user were to supply a custom + # dispatch registry or otherwise modify the registry THIS IS NOT + # NECESSARILY 100% TYPE SAFE!! + return launcher_type + + +default_dispatcher: t.Final = Dispatcher() + + +class ShellLauncher: + """Mock launcher for launching/tracking simple shell commands + + TODO: this is probably all we need for a "local" launcher, but probably + best to move this to a `smartsim._core.launcher` module/submodule + """ + + def __init__(self) -> None: + self._launched: dict[JobID, sp.Popen[bytes]] = {} + + def start(self, launchable: t.Sequence[str]) -> JobID: + id_ = create_job_id() + self._launched[id_] = sp.Popen(launchable) + return id_ + + @classmethod + def create(cls, exp: Experiment) -> Self: + return cls() diff --git a/tests/temp_tests/test_settings/conftest.py b/tests/temp_tests/test_settings/conftest.py index ad1fa0f4a..334431fce 100644 --- a/tests/temp_tests/test_settings/conftest.py +++ b/tests/temp_tests/test_settings/conftest.py @@ -25,8 +25,10 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import pytest +from unittest.mock import Mock from smartsim.settings.builders import launchArgBuilder as launch +from smartsim.settings import dispatch @pytest.fixture @@ -35,4 +37,30 @@ class _ExeLike(launch.ExecutableLike): def as_program_arguments(self): return ("echo", "hello", "world") - return _ExeLike() + yield _ExeLike() + + +@pytest.fixture +def settings_builder(): + class _SettingsBuilder(launch.LaunchArgBuilder): + def launcher_str(self): + return "Mock Settings Builder" + + def set(self, arg, val): ... + def finalize(self, exe, env): + return Mock() + + yield _SettingsBuilder({}) + + +@pytest.fixture +def launcher_like(): + class _LuancherLike(dispatch.LauncherLike): + def start(self, launchable): + return dispatch.create_job_id() + + @classmethod + def create(cls, exp): + return cls() + + yield _LuancherLike() diff --git a/tests/temp_tests/test_settings/test_dispatch.py b/tests/temp_tests/test_settings/test_dispatch.py new file mode 100644 index 000000000..4dbd25c3f --- /dev/null +++ b/tests/temp_tests/test_settings/test_dispatch.py @@ -0,0 +1,142 @@ +# BSD 2-Clause License +# +# Copyright (c) 2021-2024, Hewlett Packard Enterprise +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import pytest +from smartsim.settings import dispatch + +import contextlib + + +def test_declaritive_form_dispatch_declaration(launcher_like, settings_builder): + d = dispatch.Dispatcher() + assert type(settings_builder) == d.dispatch(to_launcher=type(launcher_like))( + type(settings_builder) + ) + assert d._dispatch_registry == {type(settings_builder): type(launcher_like)} + + +def test_imperative_form_dispatch_declaration(launcher_like, settings_builder): + d = dispatch.Dispatcher() + assert None == d.dispatch(type(settings_builder), to_launcher=type(launcher_like)) + assert d._dispatch_registry == {type(settings_builder): type(launcher_like)} + + +def test_dispatchers_from_same_registry_do_not_cross_polute( + launcher_like, settings_builder +): + some_starting_registry = {} + d1 = dispatch.Dispatcher(dispatch_registry=some_starting_registry) + d2 = dispatch.Dispatcher(dispatch_registry=some_starting_registry) + assert ( + d1._dispatch_registry == d2._dispatch_registry == some_starting_registry == {} + ) + assert ( + d1._dispatch_registry is not d2._dispatch_registry is not some_starting_registry + ) + + d2.dispatch(type(settings_builder), to_launcher=type(launcher_like)) + assert d1._dispatch_registry == {} + assert d2._dispatch_registry == {type(settings_builder): type(launcher_like)} + + +def test_copied_dispatchers_do_not_cross_pollute(launcher_like, settings_builder): + some_starting_registry = {} + d1 = dispatch.Dispatcher(dispatch_registry=some_starting_registry) + d2 = d1.copy() + assert ( + d1._dispatch_registry == d2._dispatch_registry == some_starting_registry == {} + ) + assert ( + d1._dispatch_registry is not d2._dispatch_registry is not some_starting_registry + ) + + d2.dispatch(type(settings_builder), to_launcher=type(launcher_like)) + assert d1._dispatch_registry == {} + assert d2._dispatch_registry == {type(settings_builder): type(launcher_like)} + + +@pytest.mark.parametrize( + "add_dispatch, expected_ctx", + ( + pytest.param( + lambda d, s, l: d.dispatch(s, to_launcher=l), + pytest.raises(TypeError, match="has already been registered"), + id="Imperative -- Disallowed implicitly", + ), + pytest.param( + lambda d, s, l: d.dispatch(s, to_launcher=l, allow_overwrite=True), + contextlib.nullcontext(), + id="Imperative -- Allowed with flag", + ), + pytest.param( + lambda d, s, l: d.dispatch(to_launcher=l)(s), + pytest.raises(TypeError, match="has already been registered"), + id="Declarative -- Disallowed implicitly", + ), + pytest.param( + lambda d, s, l: d.dispatch(to_launcher=l, allow_overwrite=True)(s), + contextlib.nullcontext(), + id="Declarative -- Allowed with flag", + ), + ), +) +def test_dispatch_overwriting( + add_dispatch, expected_ctx, launcher_like, settings_builder +): + registry = {type(settings_builder): type(launcher_like)} + d = dispatch.Dispatcher(dispatch_registry=registry) + with expected_ctx: + add_dispatch(d, type(settings_builder), type(launcher_like)) + + +@pytest.mark.parametrize( + "map_settings", + ( + pytest.param(type, id="From settings type"), + pytest.param(lambda s: s, id="From settings instance"), + ), +) +def test_dispatch_can_retrieve_launcher_to_dispatch_to( + map_settings, launcher_like, settings_builder +): + registry = {type(settings_builder): type(launcher_like)} + d = dispatch.Dispatcher(dispatch_registry=registry) + assert type(launcher_like) == d.get_launcher_for(map_settings(settings_builder)) + + +@pytest.mark.parametrize( + "map_settings", + ( + pytest.param(type, id="From settings type"), + pytest.param(lambda s: s, id="From settings instance"), + ), +) +def test_dispatch_raises_if_settings_type_not_registered( + map_settings, launcher_like, settings_builder +): + d = dispatch.Dispatcher(dispatch_registry={}) + with pytest.raises(TypeError, match="no launcher type to dispatch to"): + d.get_launcher_for(map_settings(settings_builder)) From d765718ea5c491581dc49bf7d49cda860d5efac3 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Mon, 24 Jun 2024 12:11:23 -0500 Subject: [PATCH 10/34] Isort/Black --- smartsim/settings/builders/launch/alps.py | 2 +- smartsim/settings/builders/launch/dragon.py | 19 ++++++++++--------- smartsim/settings/builders/launch/local.py | 2 +- smartsim/settings/builders/launch/lsf.py | 2 +- smartsim/settings/builders/launch/mpi.py | 2 +- smartsim/settings/builders/launch/pals.py | 2 +- smartsim/settings/builders/launch/slurm.py | 3 ++- smartsim/settings/dispatch.py | 5 ++--- smartsim/settings/launchSettings.py | 4 +++- tests/temp_tests/test_settings/conftest.py | 5 +++-- .../temp_tests/test_settings/test_dispatch.py | 5 ++++- 11 files changed, 29 insertions(+), 22 deletions(-) diff --git a/smartsim/settings/builders/launch/alps.py b/smartsim/settings/builders/launch/alps.py index 97e83a6c1..8f425dacc 100644 --- a/smartsim/settings/builders/launch/alps.py +++ b/smartsim/settings/builders/launch/alps.py @@ -29,7 +29,7 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import default_dispatcher, ShellLauncher +from smartsim.settings.dispatch import ShellLauncher, default_dispatcher from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType diff --git a/smartsim/settings/builders/launch/dragon.py b/smartsim/settings/builders/launch/dragon.py index bc5f2b528..1c8f1ac78 100644 --- a/smartsim/settings/builders/launch/dragon.py +++ b/smartsim/settings/builders/launch/dragon.py @@ -75,15 +75,16 @@ def finalize( return DragonRunRequest( exe=exe_, exe_args=args, - path=os.getcwd(), # FIXME: Currently this is hard coded because - # the schema requires it, but in future, - # it is almost certainly necessary that - # this will need to be injected by the - # user or by us to have the command - # execute next to any generated files. A - # similar problem exists for the other - # settings. - # TODO: Find a way to inject this path + # FIXME: Currently this is hard coded because + # the schema requires it, but in future, + # it is almost certainly necessary that + # this will need to be injected by the + # user or by us to have the command + # execute next to any generated files. A + # similar problem exists for the other + # settings. + # TODO: Find a way to inject this path + path=os.getcwd(), env=env, current_env=dict(os.environ), **self._launch_args, diff --git a/smartsim/settings/builders/launch/local.py b/smartsim/settings/builders/launch/local.py index 2cab0df75..23c5d75f0 100644 --- a/smartsim/settings/builders/launch/local.py +++ b/smartsim/settings/builders/launch/local.py @@ -29,7 +29,7 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import default_dispatcher, ShellLauncher +from smartsim.settings.dispatch import ShellLauncher, default_dispatcher from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType diff --git a/smartsim/settings/builders/launch/lsf.py b/smartsim/settings/builders/launch/lsf.py index 118946ca2..e7b22cb47 100644 --- a/smartsim/settings/builders/launch/lsf.py +++ b/smartsim/settings/builders/launch/lsf.py @@ -29,7 +29,7 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import default_dispatcher, ShellLauncher +from smartsim.settings.dispatch import ShellLauncher, default_dispatcher from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType diff --git a/smartsim/settings/builders/launch/mpi.py b/smartsim/settings/builders/launch/mpi.py index d2a1f3170..6bcde18da 100644 --- a/smartsim/settings/builders/launch/mpi.py +++ b/smartsim/settings/builders/launch/mpi.py @@ -29,7 +29,7 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import default_dispatcher, ShellLauncher +from smartsim.settings.dispatch import ShellLauncher, default_dispatcher from ...common import set_check_input from ...launchCommand import LauncherType diff --git a/smartsim/settings/builders/launch/pals.py b/smartsim/settings/builders/launch/pals.py index ff1755fb6..c8bdf2432 100644 --- a/smartsim/settings/builders/launch/pals.py +++ b/smartsim/settings/builders/launch/pals.py @@ -29,7 +29,7 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import default_dispatcher, ShellLauncher +from smartsim.settings.dispatch import ShellLauncher, default_dispatcher from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType diff --git a/smartsim/settings/builders/launch/slurm.py b/smartsim/settings/builders/launch/slurm.py index 1cdbea3f7..11e9a7b15 100644 --- a/smartsim/settings/builders/launch/slurm.py +++ b/smartsim/settings/builders/launch/slurm.py @@ -31,7 +31,7 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import default_dispatcher, ShellLauncher +from smartsim.settings.dispatch import ShellLauncher, default_dispatcher from ...common import set_check_input from ...launchCommand import LauncherType @@ -42,6 +42,7 @@ logger = get_logger(__name__) + @default_dispatcher.dispatch(to_launcher=ShellLauncher) class SlurmArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def launcher_str(self) -> str: diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index 27b247593..db2094030 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -27,13 +27,12 @@ from __future__ import annotations import subprocess as sp -import uuid import typing as t - -from smartsim.error import errors +import uuid if t.TYPE_CHECKING: from typing_extensions import Self + from smartsim.experiment import Experiment from smartsim.settings.builders import LaunchArgBuilder diff --git a/smartsim/settings/launchSettings.py b/smartsim/settings/launchSettings.py index f512fcc34..17dbfa7a2 100644 --- a/smartsim/settings/launchSettings.py +++ b/smartsim/settings/launchSettings.py @@ -95,7 +95,9 @@ def env_vars(self, value: t.Dict[str, str]) -> None: """Set the environment variables.""" self._env_vars = copy.deepcopy(value) - def _get_arg_builder(self, launch_args: StringArgument | None) -> LaunchArgBuilder[t.Any]: + def _get_arg_builder( + self, launch_args: StringArgument | None + ) -> LaunchArgBuilder[t.Any]: """Map the Launcher to the LaunchArgBuilder""" if self._launcher == LauncherType.Slurm: return SlurmArgBuilder(launch_args) diff --git a/tests/temp_tests/test_settings/conftest.py b/tests/temp_tests/test_settings/conftest.py index 334431fce..ebf361e97 100644 --- a/tests/temp_tests/test_settings/conftest.py +++ b/tests/temp_tests/test_settings/conftest.py @@ -24,11 +24,12 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -import pytest from unittest.mock import Mock -from smartsim.settings.builders import launchArgBuilder as launch +import pytest + from smartsim.settings import dispatch +from smartsim.settings.builders import launchArgBuilder as launch @pytest.fixture diff --git a/tests/temp_tests/test_settings/test_dispatch.py b/tests/temp_tests/test_settings/test_dispatch.py index 4dbd25c3f..ccd1e81cd 100644 --- a/tests/temp_tests/test_settings/test_dispatch.py +++ b/tests/temp_tests/test_settings/test_dispatch.py @@ -24,10 +24,13 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +import contextlib + import pytest + from smartsim.settings import dispatch -import contextlib +pytestmark = pytest.mark.group_a def test_declaritive_form_dispatch_declaration(launcher_like, settings_builder): From 01393c7602688e0a99993563127064a4959674fa Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Wed, 26 Jun 2024 16:38:27 -0500 Subject: [PATCH 11/34] Env: dict -> Mapping --- smartsim/settings/builders/launch/alps.py | 2 +- smartsim/settings/builders/launch/dragon.py | 2 +- smartsim/settings/builders/launch/local.py | 2 +- smartsim/settings/builders/launch/lsf.py | 2 +- smartsim/settings/builders/launch/mpi.py | 6 +++--- smartsim/settings/builders/launch/pals.py | 2 +- smartsim/settings/builders/launch/slurm.py | 2 +- smartsim/settings/builders/launchArgBuilder.py | 2 +- 8 files changed, 10 insertions(+), 10 deletions(-) diff --git a/smartsim/settings/builders/launch/alps.py b/smartsim/settings/builders/launch/alps.py index 8f425dacc..f1a196e7c 100644 --- a/smartsim/settings/builders/launch/alps.py +++ b/smartsim/settings/builders/launch/alps.py @@ -220,7 +220,7 @@ def set(self, key: str, value: str | None) -> None: self._launch_args[key] = value def finalize( - self, exe: ExecutableLike, env: dict[str, str | None] + self, exe: ExecutableLike, env: t.Mapping[str, str | None] ) -> t.Sequence[str]: return ( "aprun", diff --git a/smartsim/settings/builders/launch/dragon.py b/smartsim/settings/builders/launch/dragon.py index 1c8f1ac78..0d0062bf2 100644 --- a/smartsim/settings/builders/launch/dragon.py +++ b/smartsim/settings/builders/launch/dragon.py @@ -69,7 +69,7 @@ def set(self, key: str, value: str | None) -> None: self._launch_args[key] = value def finalize( - self, exe: ExecutableLike, env: dict[str, str | None] + self, exe: ExecutableLike, env: t.Mapping[str, str | None] ) -> DragonRunRequest: exe_, *args = exe.as_program_arguments() return DragonRunRequest( diff --git a/smartsim/settings/builders/launch/local.py b/smartsim/settings/builders/launch/local.py index 23c5d75f0..64770f696 100644 --- a/smartsim/settings/builders/launch/local.py +++ b/smartsim/settings/builders/launch/local.py @@ -79,6 +79,6 @@ def set(self, key: str, value: str | None) -> None: self._launch_args[key] = value def finalize( - self, exe: ExecutableLike, env: dict[str, str | None] + self, exe: ExecutableLike, env: t.Mapping[str, str | None] ) -> t.Sequence[str]: return exe.as_program_arguments() diff --git a/smartsim/settings/builders/launch/lsf.py b/smartsim/settings/builders/launch/lsf.py index e7b22cb47..e1a03ef3b 100644 --- a/smartsim/settings/builders/launch/lsf.py +++ b/smartsim/settings/builders/launch/lsf.py @@ -122,7 +122,7 @@ def set(self, key: str, value: str | None) -> None: self._launch_args[key] = value def finalize( - self, exe: ExecutableLike, env: dict[str, str | None] + self, exe: ExecutableLike, env: t.Mapping[str, str | None] ) -> t.Sequence[str]: return ( "jsrun", diff --git a/smartsim/settings/builders/launch/mpi.py b/smartsim/settings/builders/launch/mpi.py index 6bcde18da..6eac12a24 100644 --- a/smartsim/settings/builders/launch/mpi.py +++ b/smartsim/settings/builders/launch/mpi.py @@ -225,7 +225,7 @@ def launcher_str(self) -> str: return LauncherType.Mpirun.value def finalize( - self, exe: ExecutableLike, env: dict[str, str | None] + self, exe: ExecutableLike, env: t.Mapping[str, str | None] ) -> t.Sequence[str]: return ("mpirun", *self.format_launch_args(), "--", *exe.as_program_arguments()) @@ -237,7 +237,7 @@ def launcher_str(self) -> str: return LauncherType.Mpiexec.value def finalize( - self, exe: ExecutableLike, env: dict[str, str | None] + self, exe: ExecutableLike, env: t.Mapping[str, str | None] ) -> t.Sequence[str]: return ( "mpiexec", @@ -254,7 +254,7 @@ def launcher_str(self) -> str: return LauncherType.Orterun.value def finalize( - self, exe: ExecutableLike, env: dict[str, str | None] + self, exe: ExecutableLike, env: t.Mapping[str, str | None] ) -> t.Sequence[str]: return ( "orterun", diff --git a/smartsim/settings/builders/launch/pals.py b/smartsim/settings/builders/launch/pals.py index c8bdf2432..d21edc8bd 100644 --- a/smartsim/settings/builders/launch/pals.py +++ b/smartsim/settings/builders/launch/pals.py @@ -156,7 +156,7 @@ def set(self, key: str, value: str | None) -> None: self._launch_args[key] = value def finalize( - self, exe: ExecutableLike, env: dict[str, str | None] + self, exe: ExecutableLike, env: t.Mapping[str, str | None] ) -> t.Sequence[str]: return ( "mpiexec", diff --git a/smartsim/settings/builders/launch/slurm.py b/smartsim/settings/builders/launch/slurm.py index 11e9a7b15..1125c2611 100644 --- a/smartsim/settings/builders/launch/slurm.py +++ b/smartsim/settings/builders/launch/slurm.py @@ -322,7 +322,7 @@ def set(self, key: str, value: str | None) -> None: self._launch_args[key] = value def finalize( - self, exe: ExecutableLike, env: dict[str, str | None] + self, exe: ExecutableLike, env: t.Mapping[str, str | None] ) -> t.Sequence[str]: return ( "srun", diff --git a/smartsim/settings/builders/launchArgBuilder.py b/smartsim/settings/builders/launchArgBuilder.py index 8a839a5c8..b125046cd 100644 --- a/smartsim/settings/builders/launchArgBuilder.py +++ b/smartsim/settings/builders/launchArgBuilder.py @@ -59,7 +59,7 @@ def set(self, arg: str, val: str | None) -> None: """Set the launch arguments""" @abstractmethod - def finalize(self, exe: ExecutableLike, env: dict[str, str | None]) -> _T: + def finalize(self, exe: ExecutableLike, env: t.Mapping[str, str | None]) -> _T: """Prepare an entity for launch using the built options""" def format_launch_args(self) -> t.Union[t.List[str], None]: From e504ccfcee2e2fcb7e21c365d8d2c0c631631682 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Wed, 26 Jun 2024 16:44:19 -0500 Subject: [PATCH 12/34] Organize dispatch file, call out work to do --- smartsim/settings/dispatch.py | 36 +++++++++++++++++++++-------------- smartsim/types.py | 32 +++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 14 deletions(-) create mode 100644 smartsim/types.py diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index db2094030..013ae6319 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -30,6 +30,8 @@ import typing as t import uuid +from smartsim.types import LaunchedJobID + if t.TYPE_CHECKING: from typing_extensions import Self @@ -39,18 +41,6 @@ _T = t.TypeVar("_T") _T_contra = t.TypeVar("_T_contra", contravariant=True) -JobID = t.NewType("JobID", uuid.UUID) - - -def create_job_id() -> JobID: - return JobID(uuid.uuid4()) - - -class LauncherLike(t.Protocol[_T_contra]): - def start(self, launchable: _T_contra) -> JobID: ... - @classmethod - def create(cls, exp: Experiment) -> Self: ... - @t.final class Dispatcher: @@ -143,6 +133,21 @@ def get_launcher_for( default_dispatcher: t.Final = Dispatcher() +# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> +# TODO: move these to a common module under `smartsim._core.launcher` +# ----------------------------------------------------------------------------- + + +def create_job_id() -> LaunchedJobID: + return LaunchedJobID(uuid.uuid4()) + + +class LauncherLike(t.Protocol[_T_contra]): + def start(self, launchable: _T_contra) -> LaunchedJobID: ... + @classmethod + def create(cls, exp: Experiment) -> Self: ... + + class ShellLauncher: """Mock launcher for launching/tracking simple shell commands @@ -151,9 +156,9 @@ class ShellLauncher: """ def __init__(self) -> None: - self._launched: dict[JobID, sp.Popen[bytes]] = {} + self._launched: dict[LaunchedJobID, sp.Popen[bytes]] = {} - def start(self, launchable: t.Sequence[str]) -> JobID: + def start(self, launchable: t.Sequence[str]) -> LaunchedJobID: id_ = create_job_id() self._launched[id_] = sp.Popen(launchable) return id_ @@ -161,3 +166,6 @@ def start(self, launchable: t.Sequence[str]) -> JobID: @classmethod def create(cls, exp: Experiment) -> Self: return cls() + + +# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< diff --git a/smartsim/types.py b/smartsim/types.py new file mode 100644 index 000000000..84eb31a85 --- /dev/null +++ b/smartsim/types.py @@ -0,0 +1,32 @@ +# BSD 2-Clause License +# +# Copyright (c) 2021-2024, Hewlett Packard Enterprise +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import typing as t +import uuid + +LaunchedJobID = t.NewType("LaunchedJobID", uuid.UUID) From 910b2d903f3c7d0b6f5a9ab7591f5a10dee5ea5d Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Wed, 26 Jun 2024 17:30:52 -0500 Subject: [PATCH 13/34] Wire up dispatch dragon builder to dragon launcher --- smartsim/_core/config/config.py | 5 +- .../_core/launcher/dragon/dragonConnector.py | 25 +++-- .../_core/launcher/dragon/dragonLauncher.py | 66 +++++++++---- smartsim/_core/utils/helpers.py | 8 +- smartsim/experiment.py | 95 ++++++++++++------- smartsim/settings/__init__.py | 2 + smartsim/settings/builders/launch/dragon.py | 26 ++--- smartsim/settings/dispatch.py | 6 +- smartsim/types.py | 3 +- 9 files changed, 155 insertions(+), 81 deletions(-) diff --git a/smartsim/_core/config/config.py b/smartsim/_core/config/config.py index 374457f3a..1012129e9 100644 --- a/smartsim/_core/config/config.py +++ b/smartsim/_core/config/config.py @@ -161,10 +161,7 @@ def dragon_dotenv(self) -> Path: @property def dragon_server_path(self) -> t.Optional[str]: - return os.getenv( - "SMARTSIM_DRAGON_SERVER_PATH", - os.getenv("SMARTSIM_DRAGON_SERVER_PATH_EXP", None), - ) + return os.getenv("SMARTSIM_DRAGON_SERVER_PATH", None) @property def dragon_server_timeout(self) -> int: diff --git a/smartsim/_core/launcher/dragon/dragonConnector.py b/smartsim/_core/launcher/dragon/dragonConnector.py index 0cd68c24e..6ce0e257f 100644 --- a/smartsim/_core/launcher/dragon/dragonConnector.py +++ b/smartsim/_core/launcher/dragon/dragonConnector.py @@ -57,6 +57,11 @@ ) from ...utils.network import find_free_port, get_best_interface_and_address +if t.TYPE_CHECKING: + from typing_extensions import Self + + from smartsim.experiment import Experiment + logger = get_logger(__name__) _SchemaT = t.TypeVar("_SchemaT", bound=t.Union[DragonRequest, DragonResponse]) @@ -69,21 +74,29 @@ class DragonConnector: to start a Dragon server and communicate with it. """ - def __init__(self) -> None: + def __init__(self, path: str | os.PathLike[str]) -> None: self._context: zmq.Context[t.Any] = zmq.Context.instance() self._context.setsockopt(zmq.REQ_CORRELATE, 1) self._context.setsockopt(zmq.REQ_RELAXED, 1) self._authenticator: t.Optional[zmq.auth.thread.ThreadAuthenticator] = None config = get_config() self._reset_timeout(config.dragon_server_timeout) + + # TODO: We should be able to make these "non-optional" + # by simply moving the impl of + # `DragonConnectior.connect_to_dragon` to this method. This is + # fine as we expect the that method should only be called once + # without hitting a guard clause. self._dragon_head_socket: t.Optional[zmq.Socket[t.Any]] = None self._dragon_head_process: t.Optional[subprocess.Popen[bytes]] = None # Returned by dragon head, useful if shutdown is to be requested # but process was started by another connector self._dragon_head_pid: t.Optional[int] = None - self._dragon_server_path = config.dragon_server_path + self._dragon_server_path = _resolve_dragon_path(path) logger.debug(f"Dragon Server path was set to {self._dragon_server_path}") self._env_vars: t.Dict[str, str] = {} + + # TODO: Remove! in theory this is unreachable if self._dragon_server_path is None: raise SmartSimError( "DragonConnector could not find the dragon server path. " @@ -293,8 +306,7 @@ def connect_to_dragon(self) -> None: "Establishing connection with Dragon server or starting a new one..." ) - path = _resolve_dragon_path(self._dragon_server_path) - + path = self._dragon_server_path self._connect_to_existing_server(path) if self.is_connected: return @@ -520,8 +532,9 @@ def _dragon_cleanup( def _resolve_dragon_path(fallback: t.Union[str, "os.PathLike[str]"]) -> Path: - dragon_server_path = get_config().dragon_server_path or os.path.join( - fallback, ".smartsim", "dragon" + config = get_config() + dragon_server_path = config.dragon_server_path or os.path.join( + fallback, config.dragon_default_subdir ) dragon_server_paths = dragon_server_path.split(":") if len(dragon_server_paths) > 1: diff --git a/smartsim/_core/launcher/dragon/dragonLauncher.py b/smartsim/_core/launcher/dragon/dragonLauncher.py index 17b47e309..829828224 100644 --- a/smartsim/_core/launcher/dragon/dragonLauncher.py +++ b/smartsim/_core/launcher/dragon/dragonLauncher.py @@ -29,6 +29,8 @@ import os import typing as t +from smartsim.types import LaunchedJobID + from ...._core.launcher.stepMapping import StepMap from ....error import LauncherError, SmartSimError from ....log import get_logger @@ -42,6 +44,7 @@ from ....status import SmartSimStatus from ...schemas import ( DragonRunRequest, + DragonRunRequestView, DragonRunResponse, DragonStopRequest, DragonStopResponse, @@ -55,6 +58,11 @@ from ..stepInfo import StepInfo from .dragonConnector import DragonConnector, _SchemaT +if t.TYPE_CHECKING: + from typing_extensions import Self + + from smartsim.experiment import Experiment + logger = get_logger(__name__) @@ -72,9 +80,9 @@ class DragonLauncher(WLMLauncher): the Job Manager to interact with it. """ - def __init__(self) -> None: + def __init__(self, server_path: str | os.PathLike[str]) -> None: super().__init__() - self._connector = DragonConnector() + self._connector = DragonConnector(server_path) """Connector used to start and interact with the Dragon server""" self._slurm_launcher = SlurmLauncher() """Slurm sub-launcher, used only for batch jobs""" @@ -119,6 +127,19 @@ def add_step_to_mapping_table(self, name: str, step_map: StepMap) -> None: ) sublauncher.add_step_to_mapping_table(name, sublauncher_step_map) + @classmethod + def create(cls, exp: Experiment) -> Self: + self = cls(exp.exp_path) + self._connector.connect_to_dragon() # TODO: protected access + return self + + def start(self, req_args: DragonRunRequestView) -> LaunchedJobID: + self._connector.load_persisted_env() + merged_env = self._connector.merge_persisted_env(os.environ.copy()) + req = DragonRunRequest(**dict(req_args), current_env=merged_env) + res = _assert_schema_type(self._connector.send_request(req), DragonRunResponse) + return LaunchedJobID(res.step_id) + def run(self, step: Step) -> t.Optional[str]: """Run a job step through Slurm @@ -165,27 +186,21 @@ def run(self, step: Step) -> t.Optional[str]: run_args = step.run_settings.run_args req_env = step.run_settings.env_vars self._connector.load_persisted_env() - merged_env = self._connector.merge_persisted_env(os.environ.copy()) nodes = int(run_args.get("nodes", None) or 1) tasks_per_node = int(run_args.get("tasks-per-node", None) or 1) - response = _assert_schema_type( - self._connector.send_request( - DragonRunRequest( - exe=cmd[0], - exe_args=cmd[1:], - path=step.cwd, - name=step.name, - nodes=nodes, - tasks_per_node=tasks_per_node, - env=req_env, - current_env=merged_env, - output_file=out, - error_file=err, - ) - ), - DragonRunResponse, + step_id = self.start( + DragonRunRequestView( + exe=cmd[0], + exe_args=cmd[1:], + path=step.cwd, + name=step.name, + nodes=nodes, + tasks_per_node=tasks_per_node, + env=req_env, + output_file=out, + error_file=err, + ) ) - step_id = str(response.step_id) else: # pylint: disable-next=consider-using-with out_strm = open(out, "w+", encoding="utf-8") @@ -319,3 +334,14 @@ def _assert_schema_type(obj: object, typ: t.Type[_SchemaT], /) -> _SchemaT: if not isinstance(obj, typ): raise TypeError(f"Expected schema of type `{typ}`, but got {type(obj)}") return obj + + +# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> +# TODO: Remove this registry and move back to builder file after fixing +# circular import +# ----------------------------------------------------------------------------- +from smartsim.settings.dispatch import default_dispatcher +from smartsim.settings.builders.launch.dragon import DragonArgBuilder + +default_dispatcher.dispatch(DragonArgBuilder, to_launcher=DragonLauncher) +# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< diff --git a/smartsim/_core/utils/helpers.py b/smartsim/_core/utils/helpers.py index 70f52bc4e..646b086dd 100644 --- a/smartsim/_core/utils/helpers.py +++ b/smartsim/_core/utils/helpers.py @@ -27,6 +27,8 @@ """ A file of helper functions for SmartSim """ +from __future__ import annotations + import base64 import collections.abc import os @@ -45,6 +47,7 @@ from types import FrameType +_T = t.TypeVar("_T") _TSignalHandlerFn = t.Callable[[int, t.Optional["FrameType"]], object] @@ -122,7 +125,6 @@ def expand_exe_path(exe: str) -> str: # which returns none if not found in_path = which(exe) - print(f"hmm what is this: {in_path}") if not in_path: if os.path.isfile(exe) and os.access(exe, os.X_OK): return os.path.abspath(exe) @@ -412,6 +414,10 @@ def is_crayex_platform() -> bool: return result.is_cray +def first(predicate: t.Callable[[_T], bool], iterable: t.Iterable[_T]) -> _T | None: + return next((item for item in iterable if predicate(item)), None) + + @t.final class SignalInterceptionStack(collections.abc.Collection[_TSignalHandlerFn]): """Registers a stack of callables to be called when a signal is diff --git a/smartsim/experiment.py b/smartsim/experiment.py index 5ffc6102e..07095461a 100644 --- a/smartsim/experiment.py +++ b/smartsim/experiment.py @@ -26,15 +26,21 @@ # pylint: disable=too-many-lines +from __future__ import annotations + +import itertools import os import os.path as osp +import textwrap import typing as t from os import environ, getcwd from tabulate import tabulate from smartsim._core.config import CONFIG +from smartsim._core.utils.helpers import first from smartsim.error.errors import SSUnsupportedError +from smartsim.settings.dispatch import default_dispatcher from smartsim.status import SmartSimStatus from ._core import Controller, Generator, Manifest, previewrenderer @@ -49,7 +55,12 @@ from .error import SmartSimError from .log import ctx_exp_path, get_logger, method_contextualizer from .settings import BatchSettings, Container, RunSettings -from .wlm import detect_launcher + +if t.TYPE_CHECKING: + from smartsim.launchable.job import Job + from smartsim.settings.builders import LaunchArgBuilder + from smartsim.settings.dispatch import Dispatcher, LauncherLike + from smartsim.types import LaunchedJobID logger = get_logger(__name__) @@ -101,8 +112,9 @@ class Experiment: def __init__( self, name: str, - exp_path: t.Optional[str] = None, - launcher: str = "local", + exp_path: str | None = None, + *, + settings_dispatcher: Dispatcher = default_dispatcher, ): """Initialize an Experiment instance. @@ -110,7 +122,7 @@ def __init__( local launcher, which will start all Experiment created instances on the localhost. - Example of initializing an Experiment with the local launcher + Example of initializing an Experiment .. highlight:: python .. code-block:: python @@ -143,10 +155,6 @@ def __init__( :param name: name for the ``Experiment`` :param exp_path: path to location of ``Experiment`` directory - :param launcher: type of launcher being used, options are "slurm", "pbs", - "lsf", "sge", or "local". If set to "auto", - an attempt will be made to find an available launcher - on the system. """ self.name = name if exp_path: @@ -160,28 +168,45 @@ def __init__( self.exp_path = exp_path - self._launcher = launcher.lower() - - if self._launcher == "auto": - self._launcher = detect_launcher() - if self._launcher == "cobalt": - raise SSUnsupportedError("Cobalt launcher is no longer supported.") - - if launcher == "dragon": - self._set_dragon_server_path() - - self._control = Controller(launcher=self._launcher) + # TODO: Remove this! The contoller is becoming obsolete + self._control = Controller(launcher="local") + self._dispatcher = settings_dispatcher + self._active_launchers: set[LauncherLike[t.Any]] = set() self.fs_identifiers: t.Set[str] = set() self._telemetry_cfg = ExperimentTelemetryConfiguration() - def _set_dragon_server_path(self) -> None: - """Set path for dragon server through environment varialbes""" - if not "SMARTSIM_DRAGON_SERVER_PATH" in environ: - environ["SMARTSIM_DRAGON_SERVER_PATH_EXP"] = osp.join( - self.exp_path, CONFIG.dragon_default_subdir + def start_jobs(self, *jobs: Job) -> tuple[LaunchedJobID, ...]: + """WIP: replacemnt method to launch jobs using the new API""" + + if not jobs: + raise TypeError( + f"{type(self).__name__}.start_jobs() missing at least 1 required " + "positional argument" ) + def _start(job: Job) -> LaunchedJobID: + builder = job.launch_settings.launch_args + launcher_type = self._dispatcher.get_launcher_for(builder) + launcher = first( + lambda launcher: type(launcher) is launcher_type, + self._active_launchers, + ) + if launcher is None: + launcher = launcher_type.create(self) + self._active_launchers.add(launcher) + # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + # FIXME: Opting out of type check here. Fix this later!! + # TODO: Very much dislike that we have to pass in attrs off of `job` + # into `builder`, which is itself an attr of an attr of `job`. + # Why is `Job` not generic based on launch arg builder? + # --------------------------------------------------------------------- + finalized = builder.finalize(job.entity, job.launch_settings.env_vars) + # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< + return launcher.start(finalized) + + return tuple(map(_start, jobs)) + @_contextualize def start( self, @@ -477,7 +502,7 @@ def preview( """ # Retrieve any active feature store jobs - active_fsjobs = self._control.active_active_feature_store_jobs + active_fsjobs = self._control.active_feature_store_jobs preview_manifest = Manifest(*args) @@ -490,10 +515,6 @@ def preview( active_fsjobs, ) - @property - def launcher(self) -> str: - return self._launcher - @_contextualize def summary(self, style: str = "github") -> str: """Return a summary of the ``Experiment`` @@ -551,11 +572,19 @@ def _launch_summary(self, manifest: Manifest) -> None: :param manifest: Manifest of deployables. """ + launcher_list = "\n".join(str(launcher) for launcher in self._active_launchers) + # ^^^^^^^^^^^^^ + # TODO: make this a nicer string + summary = textwrap.dedent(f"""\ + + + === Launch Summary === + Experiment: {self.name} + Experiment Path: {self.exp_path} + Launchers: + {textwrap.indent(" - ", launcher_list)} + """) - summary = "\n\n=== Launch Summary ===\n" - summary += f"Experiment: {self.name}\n" - summary += f"Experiment Path: {self.exp_path}\n" - summary += f"Launcher: {self._launcher}\n" if manifest.applications: summary += f"Applications: {len(manifest.applications)}\n" diff --git a/smartsim/settings/__init__.py b/smartsim/settings/__init__.py index e0313f341..e904240af 100644 --- a/smartsim/settings/__init__.py +++ b/smartsim/settings/__init__.py @@ -24,6 +24,8 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +import typing as t + from .baseSettings import BaseSettings from .batchSettings import BatchSettings from .launchSettings import LaunchSettings diff --git a/smartsim/settings/builders/launch/dragon.py b/smartsim/settings/builders/launch/dragon.py index 0d0062bf2..8fb2ebc48 100644 --- a/smartsim/settings/builders/launch/dragon.py +++ b/smartsim/settings/builders/launch/dragon.py @@ -29,7 +29,7 @@ import os import typing as t -from smartsim._core.schemas.dragonRequests import DragonRunRequest +from smartsim._core.schemas.dragonRequests import DragonRunRequestView from smartsim.log import get_logger from ...common import StringArgument, set_check_input @@ -42,7 +42,7 @@ logger = get_logger(__name__) -class DragonArgBuilder(LaunchArgBuilder[DragonRunRequest]): +class DragonArgBuilder(LaunchArgBuilder[DragonRunRequestView]): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Dragon.value @@ -70,22 +70,22 @@ def set(self, key: str, value: str | None) -> None: def finalize( self, exe: ExecutableLike, env: t.Mapping[str, str | None] - ) -> DragonRunRequest: + ) -> DragonRunRequestView: exe_, *args = exe.as_program_arguments() - return DragonRunRequest( + return DragonRunRequestView( exe=exe_, exe_args=args, - # FIXME: Currently this is hard coded because - # the schema requires it, but in future, - # it is almost certainly necessary that - # this will need to be injected by the - # user or by us to have the command - # execute next to any generated files. A - # similar problem exists for the other - # settings. + # FIXME: Currently this is hard coded because the schema requires + # it, but in future, it is almost certainly necessary that + # this will need to be injected by the user or by us to have + # the command execute next to any generated files. A similar + # problem exists for the other settings. # TODO: Find a way to inject this path path=os.getcwd(), env=env, - current_env=dict(os.environ), + # TODO: Not sure how this info is injected + name=None, + output_file=None, + error_file=None, **self._launch_args, ) diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index 013ae6319..c731d0940 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -30,6 +30,7 @@ import typing as t import uuid +from smartsim._core.utils import helpers from smartsim.types import LaunchedJobID if t.TYPE_CHECKING: @@ -139,7 +140,7 @@ def get_launcher_for( def create_job_id() -> LaunchedJobID: - return LaunchedJobID(uuid.uuid4()) + return LaunchedJobID(str(uuid.uuid4())) class LauncherLike(t.Protocol[_T_contra]): @@ -160,7 +161,8 @@ def __init__(self) -> None: def start(self, launchable: t.Sequence[str]) -> LaunchedJobID: id_ = create_job_id() - self._launched[id_] = sp.Popen(launchable) + exe, *rest = launchable + self._launched[id_] = sp.Popen((helpers.expand_exe_path(exe), *rest)) return id_ @classmethod diff --git a/smartsim/types.py b/smartsim/types.py index 84eb31a85..e806d8130 100644 --- a/smartsim/types.py +++ b/smartsim/types.py @@ -27,6 +27,5 @@ from __future__ import annotations import typing as t -import uuid -LaunchedJobID = t.NewType("LaunchedJobID", uuid.UUID) +LaunchedJobID = t.NewType("LaunchedJobID", str) From f4ebada3981bdc6cc1fbcab588f5da483cb893e4 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Tue, 2 Jul 2024 12:13:19 -0500 Subject: [PATCH 14/34] Import sort --- smartsim/_core/launcher/dragon/dragonLauncher.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/smartsim/_core/launcher/dragon/dragonLauncher.py b/smartsim/_core/launcher/dragon/dragonLauncher.py index 829828224..020f7682e 100644 --- a/smartsim/_core/launcher/dragon/dragonLauncher.py +++ b/smartsim/_core/launcher/dragon/dragonLauncher.py @@ -340,8 +340,8 @@ def _assert_schema_type(obj: object, typ: t.Type[_SchemaT], /) -> _SchemaT: # TODO: Remove this registry and move back to builder file after fixing # circular import # ----------------------------------------------------------------------------- -from smartsim.settings.dispatch import default_dispatcher from smartsim.settings.builders.launch.dragon import DragonArgBuilder +from smartsim.settings.dispatch import default_dispatcher default_dispatcher.dispatch(DragonArgBuilder, to_launcher=DragonLauncher) # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< From 9d2901f1e100f36e57a88b2f75be70558442fe46 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Tue, 2 Jul 2024 14:14:40 -0500 Subject: [PATCH 15/34] Remove old TODOs --- smartsim/_core/launcher/dragon/dragonConnector.py | 10 ---------- smartsim/_core/launcher/dragon/dragonLauncher.py | 2 +- smartsim/experiment.py | 2 -- smartsim/settings/dispatch.py | 6 +----- smartsim/settings/launchSettings.py | 2 +- 5 files changed, 3 insertions(+), 19 deletions(-) diff --git a/smartsim/_core/launcher/dragon/dragonConnector.py b/smartsim/_core/launcher/dragon/dragonConnector.py index 6ce0e257f..ca721eeaa 100644 --- a/smartsim/_core/launcher/dragon/dragonConnector.py +++ b/smartsim/_core/launcher/dragon/dragonConnector.py @@ -96,16 +96,6 @@ def __init__(self, path: str | os.PathLike[str]) -> None: logger.debug(f"Dragon Server path was set to {self._dragon_server_path}") self._env_vars: t.Dict[str, str] = {} - # TODO: Remove! in theory this is unreachable - if self._dragon_server_path is None: - raise SmartSimError( - "DragonConnector could not find the dragon server path. " - "This should not happen if the Connector was started by an " - "experiment.\nIf the DragonConnector was started manually, " - "then the environment variable SMARTSIM_DRAGON_SERVER_PATH " - "should be set to an existing directory." - ) - @property def is_connected(self) -> bool: """Whether the Connector established a connection to the server diff --git a/smartsim/_core/launcher/dragon/dragonLauncher.py b/smartsim/_core/launcher/dragon/dragonLauncher.py index 020f7682e..ce11ed91f 100644 --- a/smartsim/_core/launcher/dragon/dragonLauncher.py +++ b/smartsim/_core/launcher/dragon/dragonLauncher.py @@ -130,7 +130,7 @@ def add_step_to_mapping_table(self, name: str, step_map: StepMap) -> None: @classmethod def create(cls, exp: Experiment) -> Self: self = cls(exp.exp_path) - self._connector.connect_to_dragon() # TODO: protected access + self._connector.connect_to_dragon() # pylint: disable=protected-access return self def start(self, req_args: DragonRunRequestView) -> LaunchedJobID: diff --git a/smartsim/experiment.py b/smartsim/experiment.py index 07095461a..182d4663a 100644 --- a/smartsim/experiment.py +++ b/smartsim/experiment.py @@ -573,8 +573,6 @@ def _launch_summary(self, manifest: Manifest) -> None: :param manifest: Manifest of deployables. """ launcher_list = "\n".join(str(launcher) for launcher in self._active_launchers) - # ^^^^^^^^^^^^^ - # TODO: make this a nicer string summary = textwrap.dedent(f"""\ diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index c731d0940..9928bad91 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -150,11 +150,7 @@ def create(cls, exp: Experiment) -> Self: ... class ShellLauncher: - """Mock launcher for launching/tracking simple shell commands - - TODO: this is probably all we need for a "local" launcher, but probably - best to move this to a `smartsim._core.launcher` module/submodule - """ + """Mock launcher for launching/tracking simple shell commands""" def __init__(self) -> None: self._launched: dict[LaunchedJobID, sp.Popen[bytes]] = {} diff --git a/smartsim/settings/launchSettings.py b/smartsim/settings/launchSettings.py index 17dbfa7a2..9718523d2 100644 --- a/smartsim/settings/launchSettings.py +++ b/smartsim/settings/launchSettings.py @@ -75,7 +75,7 @@ def launch_args(self) -> LaunchArgBuilder[t.Any]: # subclasses. Otherwise we have no way of showing what methods # are available at intellisense/static analysis/compile time. # This whole object basically resolves to being one step removed - # from `Any` typed!! + # from `Any` typed (worse even, as type checkers will error)!! return self._arg_builder @launch_args.setter From c413b8f1c25ba425a6f6beee50809d7483fa2047 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Wed, 3 Jul 2024 12:35:41 -0500 Subject: [PATCH 16/34] textwrap.dedent fix --- smartsim/experiment.py | 6 +----- smartsim/settings/__init__.py | 2 -- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/smartsim/experiment.py b/smartsim/experiment.py index 182d4663a..2ce3c287b 100644 --- a/smartsim/experiment.py +++ b/smartsim/experiment.py @@ -574,8 +574,6 @@ def _launch_summary(self, manifest: Manifest) -> None: """ launcher_list = "\n".join(str(launcher) for launcher in self._active_launchers) summary = textwrap.dedent(f"""\ - - === Launch Summary === Experiment: {self.name} Experiment Path: {self.exp_path} @@ -593,9 +591,7 @@ def _launch_summary(self, manifest: Manifest) -> None: else: summary += "Feature Store Status: inactive\n" - summary += f"\n{str(manifest)}" - - logger.info(summary) + logger.info(f"\n\n{summary}\n{manifest}") def _create_entity_dir(self, start_manifest: Manifest) -> None: def create_entity_dir( diff --git a/smartsim/settings/__init__.py b/smartsim/settings/__init__.py index e904240af..e0313f341 100644 --- a/smartsim/settings/__init__.py +++ b/smartsim/settings/__init__.py @@ -24,8 +24,6 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -import typing as t - from .baseSettings import BaseSettings from .batchSettings import BatchSettings from .launchSettings import LaunchSettings From 35e686c6ea554b963deedb23114bfdb740dbe410 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Wed, 3 Jul 2024 20:48:17 -0500 Subject: [PATCH 17/34] Add doc strs --- smartsim/_core/utils/helpers.py | 18 ++++++++++++++++++ smartsim/experiment.py | 5 ++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/smartsim/_core/utils/helpers.py b/smartsim/_core/utils/helpers.py index 646b086dd..d193b6604 100644 --- a/smartsim/_core/utils/helpers.py +++ b/smartsim/_core/utils/helpers.py @@ -415,6 +415,24 @@ def is_crayex_platform() -> bool: def first(predicate: t.Callable[[_T], bool], iterable: t.Iterable[_T]) -> _T | None: + """Return the first instance of an iterable that meets some precondition. + Any elements of the iterable that do not meet the precondition will be + forgotten. If no item in the iterable is found that meets the predicate, + `None` is returned. This is roughly equivalent to + + .. highlight:: python + .. code-block:: python + + next(filter(predicate, iterable), None) + + but does not require the predicate to be a type guard to type check. + + :param predicate: A function that returns `True` or `False` given a element + of the iterable + :param iterable: An iterable that yields elements to evealuate + :returns: The first element of the iterable to make the the `predicate` + return `True` + """ return next((item for item in iterable if predicate(item)), None) diff --git a/smartsim/experiment.py b/smartsim/experiment.py index 2ce3c287b..3b5690f55 100644 --- a/smartsim/experiment.py +++ b/smartsim/experiment.py @@ -113,7 +113,7 @@ def __init__( self, name: str, exp_path: str | None = None, - *, + *, # Keyword arguments only settings_dispatcher: Dispatcher = default_dispatcher, ): """Initialize an Experiment instance. @@ -155,6 +155,9 @@ def __init__( :param name: name for the ``Experiment`` :param exp_path: path to location of ``Experiment`` directory + :param settings_dispatcher: The dispatcher the experiment will use to + figure determine how to launch a job. If none is provided, the + experiment will use the default dispatcher. """ self.name = name if exp_path: From 9b2c2b75bbfe6a730faacd357fd04c6714290514 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Wed, 3 Jul 2024 20:49:18 -0500 Subject: [PATCH 18/34] Make dispatching settings more concise --- smartsim/_core/launcher/dragon/dragonLauncher.py | 4 ++-- smartsim/settings/builders/launch/alps.py | 4 ++-- smartsim/settings/builders/launch/local.py | 4 ++-- smartsim/settings/builders/launch/lsf.py | 4 ++-- smartsim/settings/builders/launch/mpi.py | 8 ++++---- smartsim/settings/builders/launch/pals.py | 4 ++-- smartsim/settings/builders/launch/slurm.py | 4 ++-- smartsim/settings/dispatch.py | 1 + 8 files changed, 17 insertions(+), 16 deletions(-) diff --git a/smartsim/_core/launcher/dragon/dragonLauncher.py b/smartsim/_core/launcher/dragon/dragonLauncher.py index ce11ed91f..e137037c1 100644 --- a/smartsim/_core/launcher/dragon/dragonLauncher.py +++ b/smartsim/_core/launcher/dragon/dragonLauncher.py @@ -341,7 +341,7 @@ def _assert_schema_type(obj: object, typ: t.Type[_SchemaT], /) -> _SchemaT: # circular import # ----------------------------------------------------------------------------- from smartsim.settings.builders.launch.dragon import DragonArgBuilder -from smartsim.settings.dispatch import default_dispatcher +from smartsim.settings.dispatch import dispatch -default_dispatcher.dispatch(DragonArgBuilder, to_launcher=DragonLauncher) +dispatch(DragonArgBuilder, to_launcher=DragonLauncher) # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< diff --git a/smartsim/settings/builders/launch/alps.py b/smartsim/settings/builders/launch/alps.py index f1a196e7c..de2d7b91d 100644 --- a/smartsim/settings/builders/launch/alps.py +++ b/smartsim/settings/builders/launch/alps.py @@ -29,7 +29,7 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, default_dispatcher +from smartsim.settings.dispatch import ShellLauncher, dispatch from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType @@ -41,7 +41,7 @@ logger = get_logger(__name__) -@default_dispatcher.dispatch(to_launcher=ShellLauncher) +@dispatch(to_launcher=ShellLauncher) class AprunArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def _reserved_launch_args(self) -> set[str]: """Return reserved launch arguments.""" diff --git a/smartsim/settings/builders/launch/local.py b/smartsim/settings/builders/launch/local.py index 64770f696..aeb018bc4 100644 --- a/smartsim/settings/builders/launch/local.py +++ b/smartsim/settings/builders/launch/local.py @@ -29,7 +29,7 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, default_dispatcher +from smartsim.settings.dispatch import ShellLauncher, dispatch from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType @@ -41,7 +41,7 @@ logger = get_logger(__name__) -@default_dispatcher.dispatch(to_launcher=ShellLauncher) +@dispatch(to_launcher=ShellLauncher) class LocalArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/lsf.py b/smartsim/settings/builders/launch/lsf.py index e1a03ef3b..bec63a802 100644 --- a/smartsim/settings/builders/launch/lsf.py +++ b/smartsim/settings/builders/launch/lsf.py @@ -29,7 +29,7 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, default_dispatcher +from smartsim.settings.dispatch import ShellLauncher, dispatch from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType @@ -41,7 +41,7 @@ logger = get_logger(__name__) -@default_dispatcher.dispatch(to_launcher=ShellLauncher) +@dispatch(to_launcher=ShellLauncher) class JsrunArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/mpi.py b/smartsim/settings/builders/launch/mpi.py index 6eac12a24..7ce79fbc3 100644 --- a/smartsim/settings/builders/launch/mpi.py +++ b/smartsim/settings/builders/launch/mpi.py @@ -29,7 +29,7 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, default_dispatcher +from smartsim.settings.dispatch import ShellLauncher, dispatch from ...common import set_check_input from ...launchCommand import LauncherType @@ -218,7 +218,7 @@ def set(self, key: str, value: str | None) -> None: self._launch_args[key] = value -@default_dispatcher.dispatch(to_launcher=ShellLauncher) +@dispatch(to_launcher=ShellLauncher) class MpiArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" @@ -230,7 +230,7 @@ def finalize( return ("mpirun", *self.format_launch_args(), "--", *exe.as_program_arguments()) -@default_dispatcher.dispatch(to_launcher=ShellLauncher) +@dispatch(to_launcher=ShellLauncher) class MpiexecArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" @@ -247,7 +247,7 @@ def finalize( ) -@default_dispatcher.dispatch(to_launcher=ShellLauncher) +@dispatch(to_launcher=ShellLauncher) class OrteArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/pals.py b/smartsim/settings/builders/launch/pals.py index d21edc8bd..1b2ed17bf 100644 --- a/smartsim/settings/builders/launch/pals.py +++ b/smartsim/settings/builders/launch/pals.py @@ -29,7 +29,7 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, default_dispatcher +from smartsim.settings.dispatch import ShellLauncher, dispatch from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType @@ -41,7 +41,7 @@ logger = get_logger(__name__) -@default_dispatcher.dispatch(to_launcher=ShellLauncher) +@dispatch(to_launcher=ShellLauncher) class PalsMpiexecArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/slurm.py b/smartsim/settings/builders/launch/slurm.py index 1125c2611..db2d673cb 100644 --- a/smartsim/settings/builders/launch/slurm.py +++ b/smartsim/settings/builders/launch/slurm.py @@ -31,7 +31,7 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, default_dispatcher +from smartsim.settings.dispatch import ShellLauncher, dispatch from ...common import set_check_input from ...launchCommand import LauncherType @@ -43,7 +43,7 @@ logger = get_logger(__name__) -@default_dispatcher.dispatch(to_launcher=ShellLauncher) +@dispatch(to_launcher=ShellLauncher) class SlurmArgBuilder(LaunchArgBuilder[t.Sequence[str]]): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index 9928bad91..f2c945f3a 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -132,6 +132,7 @@ def get_launcher_for( default_dispatcher: t.Final = Dispatcher() +dispatch: t.Final = default_dispatcher.dispatch # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> From 219d481bb71514e0b46e54db9fb4c07922d5d1d8 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Mon, 8 Jul 2024 19:09:28 -0500 Subject: [PATCH 19/34] Address reviewer feedback --- smartsim/experiment.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/smartsim/experiment.py b/smartsim/experiment.py index c6e86e366..cd38454a2 100644 --- a/smartsim/experiment.py +++ b/smartsim/experiment.py @@ -177,7 +177,9 @@ def __init__( # TODO: Remove this! The contoller is becoming obsolete self._control = Controller(launcher="local") self._dispatcher = settings_dispatcher + self._active_launchers: set[LauncherLike[t.Any]] = set() + """The active launchers created, used, and reused by the experiment""" self.fs_identifiers: t.Set[str] = set() self._telemetry_cfg = ExperimentTelemetryConfiguration() @@ -192,7 +194,7 @@ def start_jobs(self, *jobs: Job) -> tuple[LaunchedJobID, ...]: ) def _start(job: Job) -> LaunchedJobID: - builder = job.launch_settings.launch_args + builder: LaunchArgBuilder[t.Any] = job.launch_settings.launch_args launcher_type = self._dispatcher.get_launcher_for(builder) launcher = first( lambda launcher: type(launcher) is launcher_type, From 979744266eea868472d5d77a523025e98bdfee62 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Tue, 16 Jul 2024 12:18:11 -0500 Subject: [PATCH 20/34] Remove stale FIXME comment --- smartsim/settings/launchSettings.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/smartsim/settings/launchSettings.py b/smartsim/settings/launchSettings.py index 24ee8993c..9078a04d9 100644 --- a/smartsim/settings/launchSettings.py +++ b/smartsim/settings/launchSettings.py @@ -69,13 +69,6 @@ def launcher(self) -> str: @property def launch_args(self) -> LaunchArgBuilder[t.Any]: """Return the launch argument translator.""" - # FIXME: We _REALLY_ need to make the `LaunchSettings` class generic at - # on `_arg_builder` if we are expecting users to call specific - # `set_*` methods defined specificially on each of the - # subclasses. Otherwise we have no way of showing what methods - # are available at intellisense/static analysis/compile time. - # This whole object basically resolves to being one step removed - # from `Any` typed (worse even, as type checkers will error)!! return self._arg_builder @launch_args.setter From 4e9f5b28b662494920c60eec66a52bc860791426 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Thu, 18 Jul 2024 02:25:55 -0500 Subject: [PATCH 21/34] Dispatcher takes a format function with launcher like --- setup.py | 4 +- smartsim/error/errors.py | 4 + smartsim/experiment.py | 49 ++++++------- smartsim/settings/dispatch.py | 134 ++++++++++++++++++++++++++++------ 4 files changed, 140 insertions(+), 51 deletions(-) diff --git a/setup.py b/setup.py index 05b6ef70b..d820563d9 100644 --- a/setup.py +++ b/setup.py @@ -179,7 +179,8 @@ def has_ext_modules(_placeholder): "pydantic==1.10.14", "pyzmq>=25.1.2", "pygithub>=2.3.0", - "numpy<2" + "numpy<2", + "typing_extensions>=4.1.0", ] # Add SmartRedis at specific version @@ -203,7 +204,6 @@ def has_ext_modules(_placeholder): "types-tqdm", "types-tensorflow==2.12.0.9", "types-setuptools", - "typing_extensions>=4.1.0", ], # see smartsim/_core/_install/buildenv.py for more details **versions.ml_extras_required(), diff --git a/smartsim/error/errors.py b/smartsim/error/errors.py index 8500e4947..3f32bd3f0 100644 --- a/smartsim/error/errors.py +++ b/smartsim/error/errors.py @@ -112,6 +112,10 @@ class LauncherUnsupportedFeature(LauncherError): """Raised when the launcher does not support a given method""" +class LauncherNotFoundError(LauncherError): + """A requested launcher could not be found""" + + class AllocationError(LauncherError): """Raised when there is a problem with the user WLM allocation""" diff --git a/smartsim/experiment.py b/smartsim/experiment.py index 8a4ed42f6..c2e56d19a 100644 --- a/smartsim/experiment.py +++ b/smartsim/experiment.py @@ -38,8 +38,7 @@ from tabulate import tabulate from smartsim._core.config import CONFIG -from smartsim._core.utils.helpers import first -from smartsim.error.errors import SSUnsupportedError +from smartsim.error.errors import LauncherNotFoundError, SSUnsupportedError from smartsim.settings.dispatch import default_dispatcher from smartsim.status import SmartSimStatus @@ -58,11 +57,8 @@ if t.TYPE_CHECKING: from smartsim.launchable.job import Job - from smartsim.settings.builders.launchArgBuilder import ( - ExecutableLike, - LaunchArgBuilder, - ) - from smartsim.settings.dispatch import Dispatcher, LauncherLike + from smartsim.settings.builders.launchArgBuilder import LaunchArgBuilder + from smartsim.settings.dispatch import Dispatcher, ExecutableLike, LauncherLike from smartsim.types import LaunchedJobID logger = get_logger(__name__) @@ -194,27 +190,30 @@ def start_jobs(self, *jobs: Job) -> tuple[LaunchedJobID, ...]: ) def _start(job: Job) -> LaunchedJobID: - builder: LaunchArgBuilder[t.Any] = job.launch_settings.launch_args - launcher_type = self._dispatcher.get_launcher_for(builder) - launcher = first( - lambda launcher: type(launcher) is launcher_type, - self._active_launchers, - ) - if launcher is None: - launcher = launcher_type.create(self) - self._active_launchers.add(launcher) + args = job.launch_settings.launch_args + env = job.launch_settings.env_vars # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - # FIXME: Opting out of type check here. Fix this later!! - # TODO: Very much dislike that we have to pass in attrs off of `job` - # into `builder`, which is itself an attr of an attr of `job`. - # Why is `Job` not generic based on launch arg builder? - # FIXME: Remove this dangerous cast after `SmartSimEntity` conforms - # to protocol + # FIXME: Remove this cast after `SmartSimEntity` conforms to + # protocol. For now, live with the "dangerous" type cast # --------------------------------------------------------------------- - exe_like = t.cast("ExecutableLike", job.entity) - finalized = builder.finalize(exe_like, job.launch_settings.env_vars) + exe = t.cast("ExecutableLike", job.entity) # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - return launcher.start(finalized) + dispatch = self._dispatcher.get_dispatch(args) + try: + launch_config = dispatch.configure_first_compatible_launcher( + from_available_launchers=self._active_launchers, + with_settings=args, + ) + except LauncherNotFoundError: + launch_config = dispatch.create_new_launcher_configuration( + for_experiment=self, with_settings=args + ) + # Save the underlying launcher instance. That way we do not need to + # spin up a launcher instance for each individual job, and it makes + # it easier to monitor job statuses + # pylint: disable-next=protected-access + self._active_launchers.add(launch_config._adapted_launcher) + return launch_config.start(exe, env) return tuple(map(_start, jobs)) diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index f2c945f3a..71343438e 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -26,21 +26,33 @@ from __future__ import annotations +import dataclasses import subprocess as sp import typing as t import uuid +from typing_extensions import Self, TypeVarTuple, Unpack + from smartsim._core.utils import helpers +from smartsim.error import errors from smartsim.types import LaunchedJobID if t.TYPE_CHECKING: - from typing_extensions import Self from smartsim.experiment import Experiment from smartsim.settings.builders import LaunchArgBuilder + _T = t.TypeVar("_T") +_Ts = TypeVarTuple("_Ts") _T_contra = t.TypeVar("_T_contra", contravariant=True) +_TDispatchable = t.TypeVar("_TDispatchable", bound="LaunchArgBuilder") +_EnvironMappingType: t.TypeAlias = t.Mapping[str, str | None] +_FormatterType: t.TypeAlias = t.Callable[ + [_TDispatchable, "ExecutableLike", _EnvironMappingType], _T +] +_LaunchConfigType: t.TypeAlias = "_LauncherAdapter[ExecutableLike, _EnvironMappingType]" +_UnkownType: t.TypeAlias = t.NoReturn @t.final @@ -53,7 +65,8 @@ def __init__( self, *, dispatch_registry: ( - t.Mapping[type[LaunchArgBuilder[t.Any]], type[LauncherLike[t.Any]]] | None + t.Mapping[type[LaunchArgBuilder], _DispatchRegistration[t.Any, t.Any]] + | None ) = None, ) -> None: self._dispatch_registry = ( @@ -69,38 +82,48 @@ def dispatch( self, args: None = ..., *, + with_format: _FormatterType[_TDispatchable, _T], to_launcher: type[LauncherLike[_T]], allow_overwrite: bool = ..., - ) -> t.Callable[[type[LaunchArgBuilder[_T]]], type[LaunchArgBuilder[_T]]]: ... + ) -> t.Callable[[type[_TDispatchable]], type[_TDispatchable]]: ... @t.overload def dispatch( self, - args: type[LaunchArgBuilder[_T]], + args: type[_TDispatchable], *, + with_format: _FormatterType[_TDispatchable, _T], to_launcher: type[LauncherLike[_T]], allow_overwrite: bool = ..., ) -> None: ... def dispatch( self, - args: type[LaunchArgBuilder[_T]] | None = None, + args: type[_TDispatchable] | None = None, *, + with_format: _FormatterType[_TDispatchable, _T], to_launcher: type[LauncherLike[_T]], allow_overwrite: bool = False, - ) -> t.Callable[[type[LaunchArgBuilder[_T]]], type[LaunchArgBuilder[_T]]] | None: + ) -> t.Callable[[type[_TDispatchable]], type[_TDispatchable]] | None: """A type safe way to add a mapping of settings builder to launcher to handle the settings at launch time. """ + err_msg: str | None = None + if getattr(to_launcher, "_is_protocol", False): + err_msg = f"Cannot dispatch to protocol class `{to_launcher.__name__}`" + elif getattr(to_launcher, "__abstractmethods__", frozenset()): + err_msg = f"Cannot dispatch to abstract class `{to_launcher.__name__}`" + if err_msg is not None: + raise TypeError(err_msg) - def register( - args_: type[LaunchArgBuilder[_T]], / - ) -> type[LaunchArgBuilder[_T]]: + def register(args_: type[_TDispatchable], /) -> type[_TDispatchable]: if args_ in self._dispatch_registry and not allow_overwrite: - launcher_type = self._dispatch_registry[args_] + launcher_type = self._dispatch_registry[args_].launcher_type raise TypeError( f"{args_.__name__} has already been registered to be " f"launched with {launcher_type}" ) - self._dispatch_registry[args_] = to_launcher + self._dispatch_registry[args_] = _DispatchRegistration( + with_format, to_launcher + ) return args_ if args is not None: @@ -108,27 +131,90 @@ def register( return None return register - def get_launcher_for( - self, args: LaunchArgBuilder[_T] | type[LaunchArgBuilder[_T]], / - ) -> type[LauncherLike[_T]]: + def get_dispatch( + self, args: _TDispatchable | type[_TDispatchable] + ) -> _DispatchRegistration[_TDispatchable, _UnkownType]: """Find a type of launcher that is registered as being able to launch the output of the provided builder """ if not isinstance(args, type): args = type(args) - launcher_type = self._dispatch_registry.get(args, None) - if launcher_type is None: + dispatch = self._dispatch_registry.get(args, None) + if dispatch is None: raise TypeError( - f"{type(self).__name__} {self} has no launcher type to " - f"dispatch to for argument builder of type {args}" + f"No dispatch for `{type(args).__name__}` has been registered " + f"has been registered with {type(self).__name__} `{self}`" ) # Note the sleight-of-hand here: we are secretly casting a type of - # `LauncherLike[Any]` to `LauncherLike[_T]`. This is safe to do if all - # entries in the mapping were added using a type safe method (e.g. - # `Dispatcher.dispatch`), but if a user were to supply a custom - # dispatch registry or otherwise modify the registry THIS IS NOT - # NECESSARILY 100% TYPE SAFE!! - return launcher_type + # `_DispatchRegistration[Any, Any]` -> + # `_DispatchRegistration[_TDispatchable, _T]`. + # where `_T` is unbound! + # + # This is safe to do if all entries in the mapping were added using a + # type safe method (e.g. `Dispatcher.dispatch`), but if a user were to + # supply a custom dispatch registry or otherwise modify the registry + # this is not necessarily 100% type safe!! + return dispatch + + +@t.final +@dataclasses.dataclass(frozen=True) +class _DispatchRegistration(t.Generic[_TDispatchable, _T]): + formatter: _FormatterType[_TDispatchable, _T] + launcher_type: type[LauncherLike[_T]] + + def _is_compatible_launcher(self, launcher) -> bool: + return type(launcher) is self.launcher_type + + def create_new_launcher_configuration( + self, for_experiment: Experiment, with_settings: _TDispatchable + ) -> _LaunchConfigType: + launcher = self.launcher_type.create(for_experiment) + return self.create_adapter_from_launcher(launcher, with_settings) + + def create_adapter_from_launcher( + self, launcher: LauncherLike[_T], settings: _TDispatchable + ) -> _LaunchConfigType: + if not self._is_compatible_launcher(launcher): + raise TypeError( + f"Cannot create launcher adapter from launcher `{launcher}` " + f"of type `{type(launcher)}`; expected launcher of type " + f"exactly `{self.launcher_type}`" + ) + + def format_(exe: ExecutableLike, env: _EnvironMappingType) -> _T: + return self.formatter(settings, exe, env) + + return _LauncherAdapter(launcher, format_) + + def configure_first_compatible_launcher( + self, + with_settings: _TDispatchable, + from_available_launchers: t.Iterable[LauncherLike[t.Any]], + ) -> _LaunchConfigType: + launcher = helpers.first(self._is_compatible_launcher, from_available_launchers) + if launcher is None: + raise errors.LauncherNotFoundError( + f"No launcher of exactly type `{self.launcher_type.__name__}` " + "could be found from provided launchers" + ) + return self.create_adapter_from_launcher(launcher, with_settings) + + +@t.final +class _LauncherAdapter(t.Generic[Unpack[_Ts]]): + def __init__( + self, launcher: LauncherLike[_T], map_: t.Callable[[Unpack[_Ts]], _T] + ) -> None: + # NOTE: We need to cast off the `_T` -> `Any` in the `__init__` + # signature to hide the transform from users of this class. If + # possible, try not to expose outside of protected methods! + self._adapt: t.Callable[[Unpack[_Ts]], t.Any] = map_ + self._adapted_launcher: LauncherLike[t.Any] = launcher + + def start(self, *args: Unpack[_Ts]) -> LaunchedJobID: + payload = self._adapt(*args) + return self._adapted_launcher.start(payload) default_dispatcher: t.Final = Dispatcher() From 0b9ec1a46c07a5d752652974d9250e2cf1f7f3dd Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Thu, 18 Jul 2024 02:27:32 -0500 Subject: [PATCH 22/34] Re-wrire up default dispatcher --- .../_core/launcher/dragon/dragonLauncher.py | 30 +++++++++++++-- smartsim/settings/builders/launch/alps.py | 19 ++-------- smartsim/settings/builders/launch/dragon.py | 27 +------------ smartsim/settings/builders/launch/local.py | 14 ++----- smartsim/settings/builders/launch/lsf.py | 19 ++-------- smartsim/settings/builders/launch/mpi.py | 38 +++---------------- smartsim/settings/builders/launch/pals.py | 19 ++-------- smartsim/settings/builders/launch/slurm.py | 19 ++-------- .../settings/builders/launchArgBuilder.py | 11 +----- smartsim/settings/dispatch.py | 28 ++++++++++++++ smartsim/settings/launchSettings.py | 6 +-- 11 files changed, 80 insertions(+), 150 deletions(-) diff --git a/smartsim/_core/launcher/dragon/dragonLauncher.py b/smartsim/_core/launcher/dragon/dragonLauncher.py index e137037c1..74e22e160 100644 --- a/smartsim/_core/launcher/dragon/dragonLauncher.py +++ b/smartsim/_core/launcher/dragon/dragonLauncher.py @@ -341,7 +341,31 @@ def _assert_schema_type(obj: object, typ: t.Type[_SchemaT], /) -> _SchemaT: # circular import # ----------------------------------------------------------------------------- from smartsim.settings.builders.launch.dragon import DragonArgBuilder -from smartsim.settings.dispatch import dispatch - -dispatch(DragonArgBuilder, to_launcher=DragonLauncher) +from smartsim.settings.dispatch import ExecutableLike, dispatch + + +def _as_run_request_view( + run_req_args: DragonArgBuilder, exe: ExecutableLike, env: t.Mapping[str, str | None] +) -> DragonRunRequestView: + exe_, *args = exe.as_program_arguments() + return DragonRunRequestView( + exe=exe_, + exe_args=args, + # FIXME: Currently this is hard coded because the schema requires + # it, but in future, it is almost certainly necessary that + # this will need to be injected by the user or by us to have + # the command execute next to any generated files. A similar + # problem exists for the other settings. + # TODO: Find a way to inject this path + path=os.getcwd(), + env=env, + # TODO: Not sure how this info is injected + name=None, + output_file=None, + error_file=None, + **run_req_args._launch_args, + ) + + +dispatch(DragonArgBuilder, with_format=_as_run_request_view, to_launcher=DragonLauncher) # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< diff --git a/smartsim/settings/builders/launch/alps.py b/smartsim/settings/builders/launch/alps.py index de2d7b91d..f325777e4 100644 --- a/smartsim/settings/builders/launch/alps.py +++ b/smartsim/settings/builders/launch/alps.py @@ -29,20 +29,17 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, dispatch +from smartsim.settings.dispatch import ShellLauncher, dispatch, shell_format from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder -if t.TYPE_CHECKING: - from smartsim.settings.builders.launchArgBuilder import ExecutableLike - logger = get_logger(__name__) -@dispatch(to_launcher=ShellLauncher) -class AprunArgBuilder(LaunchArgBuilder[t.Sequence[str]]): +@dispatch(with_format=shell_format(run_command="aprun"), to_launcher=ShellLauncher) +class AprunArgBuilder(LaunchArgBuilder): def _reserved_launch_args(self) -> set[str]: """Return reserved launch arguments.""" return {"wdir"} @@ -218,13 +215,3 @@ def set(self, key: str, value: str | None) -> None: if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") self._launch_args[key] = value - - def finalize( - self, exe: ExecutableLike, env: t.Mapping[str, str | None] - ) -> t.Sequence[str]: - return ( - "aprun", - *(self.format_launch_args() or ()), - "--", - *exe.as_program_arguments(), - ) diff --git a/smartsim/settings/builders/launch/dragon.py b/smartsim/settings/builders/launch/dragon.py index 8fb2ebc48..3cbb1c6d0 100644 --- a/smartsim/settings/builders/launch/dragon.py +++ b/smartsim/settings/builders/launch/dragon.py @@ -29,7 +29,6 @@ import os import typing as t -from smartsim._core.schemas.dragonRequests import DragonRunRequestView from smartsim.log import get_logger from ...common import StringArgument, set_check_input @@ -37,12 +36,12 @@ from ..launchArgBuilder import LaunchArgBuilder if t.TYPE_CHECKING: - from smartsim.settings.builders.launchArgBuilder import ExecutableLike + from smartsim.settings.dispatch import ExecutableLike logger = get_logger(__name__) -class DragonArgBuilder(LaunchArgBuilder[DragonRunRequestView]): +class DragonArgBuilder(LaunchArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Dragon.value @@ -67,25 +66,3 @@ def set(self, key: str, value: str | None) -> None: if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") self._launch_args[key] = value - - def finalize( - self, exe: ExecutableLike, env: t.Mapping[str, str | None] - ) -> DragonRunRequestView: - exe_, *args = exe.as_program_arguments() - return DragonRunRequestView( - exe=exe_, - exe_args=args, - # FIXME: Currently this is hard coded because the schema requires - # it, but in future, it is almost certainly necessary that - # this will need to be injected by the user or by us to have - # the command execute next to any generated files. A similar - # problem exists for the other settings. - # TODO: Find a way to inject this path - path=os.getcwd(), - env=env, - # TODO: Not sure how this info is injected - name=None, - output_file=None, - error_file=None, - **self._launch_args, - ) diff --git a/smartsim/settings/builders/launch/local.py b/smartsim/settings/builders/launch/local.py index aeb018bc4..21fb71c8a 100644 --- a/smartsim/settings/builders/launch/local.py +++ b/smartsim/settings/builders/launch/local.py @@ -29,20 +29,17 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, dispatch +from smartsim.settings.dispatch import ShellLauncher, dispatch, shell_format from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder -if t.TYPE_CHECKING: - from smartsim.settings.builders.launchArgBuilder import ExecutableLike - logger = get_logger(__name__) -@dispatch(to_launcher=ShellLauncher) -class LocalArgBuilder(LaunchArgBuilder[t.Sequence[str]]): +@dispatch(with_format=shell_format(run_command=None), to_launcher=ShellLauncher) +class LocalArgBuilder(LaunchArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Local.value @@ -77,8 +74,3 @@ def set(self, key: str, value: str | None) -> None: if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") self._launch_args[key] = value - - def finalize( - self, exe: ExecutableLike, env: t.Mapping[str, str | None] - ) -> t.Sequence[str]: - return exe.as_program_arguments() diff --git a/smartsim/settings/builders/launch/lsf.py b/smartsim/settings/builders/launch/lsf.py index bec63a802..13a32fd73 100644 --- a/smartsim/settings/builders/launch/lsf.py +++ b/smartsim/settings/builders/launch/lsf.py @@ -29,20 +29,17 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, dispatch +from smartsim.settings.dispatch import ShellLauncher, dispatch, shell_format from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder -if t.TYPE_CHECKING: - from smartsim.settings.builders.launchArgBuilder import ExecutableLike - logger = get_logger(__name__) -@dispatch(to_launcher=ShellLauncher) -class JsrunArgBuilder(LaunchArgBuilder[t.Sequence[str]]): +@dispatch(with_format=shell_format(run_command="jsrun"), to_launcher=ShellLauncher) +class JsrunArgBuilder(LaunchArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Lsf.value @@ -120,13 +117,3 @@ def set(self, key: str, value: str | None) -> None: if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") self._launch_args[key] = value - - def finalize( - self, exe: ExecutableLike, env: t.Mapping[str, str | None] - ) -> t.Sequence[str]: - return ( - "jsrun", - *(self.format_launch_args() or ()), - "--", - *exe.as_program_arguments(), - ) diff --git a/smartsim/settings/builders/launch/mpi.py b/smartsim/settings/builders/launch/mpi.py index 7ce79fbc3..ea24564da 100644 --- a/smartsim/settings/builders/launch/mpi.py +++ b/smartsim/settings/builders/launch/mpi.py @@ -29,19 +29,16 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, dispatch +from smartsim.settings.dispatch import ShellLauncher, dispatch, shell_format from ...common import set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder -if t.TYPE_CHECKING: - from smartsim.settings.builders.launchArgBuilder import ExecutableLike - logger = get_logger(__name__) -class _BaseMPIArgBuilder(LaunchArgBuilder[t.Sequence[str]]): +class _BaseMPIArgBuilder(LaunchArgBuilder): def _reserved_launch_args(self) -> set[str]: """Return reserved launch arguments.""" return {"wd", "wdir"} @@ -218,47 +215,22 @@ def set(self, key: str, value: str | None) -> None: self._launch_args[key] = value -@dispatch(to_launcher=ShellLauncher) +@dispatch(with_format=shell_format(run_command="mpirun"), to_launcher=ShellLauncher) class MpiArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Mpirun.value - def finalize( - self, exe: ExecutableLike, env: t.Mapping[str, str | None] - ) -> t.Sequence[str]: - return ("mpirun", *self.format_launch_args(), "--", *exe.as_program_arguments()) - -@dispatch(to_launcher=ShellLauncher) +@dispatch(with_format=shell_format(run_command="mpiexec"), to_launcher=ShellLauncher) class MpiexecArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Mpiexec.value - def finalize( - self, exe: ExecutableLike, env: t.Mapping[str, str | None] - ) -> t.Sequence[str]: - return ( - "mpiexec", - *self.format_launch_args(), - "--", - *exe.as_program_arguments(), - ) - -@dispatch(to_launcher=ShellLauncher) +@dispatch(with_format=shell_format(run_command="orterun"), to_launcher=ShellLauncher) class OrteArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Orterun.value - - def finalize( - self, exe: ExecutableLike, env: t.Mapping[str, str | None] - ) -> t.Sequence[str]: - return ( - "orterun", - *self.format_launch_args(), - "--", - *exe.as_program_arguments(), - ) diff --git a/smartsim/settings/builders/launch/pals.py b/smartsim/settings/builders/launch/pals.py index 1b2ed17bf..4f2155c1f 100644 --- a/smartsim/settings/builders/launch/pals.py +++ b/smartsim/settings/builders/launch/pals.py @@ -29,20 +29,17 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, dispatch +from smartsim.settings.dispatch import ShellLauncher, dispatch, shell_format from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder -if t.TYPE_CHECKING: - from smartsim.settings.builders.launchArgBuilder import ExecutableLike - logger = get_logger(__name__) -@dispatch(to_launcher=ShellLauncher) -class PalsMpiexecArgBuilder(LaunchArgBuilder[t.Sequence[str]]): +@dispatch(with_format=shell_format(run_command="mpiexec"), to_launcher=ShellLauncher) +class PalsMpiexecArgBuilder(LaunchArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Pals.value @@ -154,13 +151,3 @@ def set(self, key: str, value: str | None) -> None: if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") self._launch_args[key] = value - - def finalize( - self, exe: ExecutableLike, env: t.Mapping[str, str | None] - ) -> t.Sequence[str]: - return ( - "mpiexec", - *(self.format_launch_args() or ()), - "--", - *exe.as_program_arguments(), - ) diff --git a/smartsim/settings/builders/launch/slurm.py b/smartsim/settings/builders/launch/slurm.py index db2d673cb..907a6da6c 100644 --- a/smartsim/settings/builders/launch/slurm.py +++ b/smartsim/settings/builders/launch/slurm.py @@ -31,20 +31,17 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, dispatch +from smartsim.settings.dispatch import ShellLauncher, dispatch, shell_format from ...common import set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder -if t.TYPE_CHECKING: - from smartsim.settings.builders.launchArgBuilder import ExecutableLike - logger = get_logger(__name__) -@dispatch(to_launcher=ShellLauncher) -class SlurmArgBuilder(LaunchArgBuilder[t.Sequence[str]]): +@dispatch(with_format=shell_format(run_command="srun"), to_launcher=ShellLauncher) +class SlurmArgBuilder(LaunchArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Slurm.value @@ -320,13 +317,3 @@ def set(self, key: str, value: str | None) -> None: if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") self._launch_args[key] = value - - def finalize( - self, exe: ExecutableLike, env: t.Mapping[str, str | None] - ) -> t.Sequence[str]: - return ( - "srun", - *(self.format_launch_args() or ()), - "--", - *exe.as_program_arguments(), - ) diff --git a/smartsim/settings/builders/launchArgBuilder.py b/smartsim/settings/builders/launchArgBuilder.py index b125046cd..2c09dd2e8 100644 --- a/smartsim/settings/builders/launchArgBuilder.py +++ b/smartsim/settings/builders/launchArgBuilder.py @@ -40,7 +40,7 @@ _T = t.TypeVar("_T") -class LaunchArgBuilder(ABC, t.Generic[_T]): +class LaunchArgBuilder(ABC): """Abstract base class that defines all generic launcher argument methods that are not supported. It is the responsibility of child classes for each launcher to translate @@ -58,10 +58,6 @@ def launcher_str(self) -> str: def set(self, arg: str, val: str | None) -> None: """Set the launch arguments""" - @abstractmethod - def finalize(self, exe: ExecutableLike, env: t.Mapping[str, str | None]) -> _T: - """Prepare an entity for launch using the built options""" - def format_launch_args(self) -> t.Union[t.List[str], None]: """Build formatted launch arguments""" logger.warning( @@ -95,8 +91,3 @@ def format_env_vars( def __str__(self) -> str: # pragma: no-cover string = f"\nLaunch Arguments:\n{fmt_dict(self._launch_args)}" return string - - -class ExecutableLike(t.Protocol): - @abstractmethod - def as_program_arguments(self) -> t.Sequence[str]: ... diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index 71343438e..a69a6587a 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -230,12 +230,40 @@ def create_job_id() -> LaunchedJobID: return LaunchedJobID(str(uuid.uuid4())) +class ExecutableLike(t.Protocol): + def as_program_arguments(self) -> t.Sequence[str]: ... + + class LauncherLike(t.Protocol[_T_contra]): def start(self, launchable: _T_contra) -> LaunchedJobID: ... @classmethod def create(cls, exp: Experiment) -> Self: ... +# TODO: This is just a nice helper function that I am using for the time being +# to wire everything up! In reality it might be a bit too confusing and +# meta-program-y for production code. Check with the core team to see +# what they think!! +def shell_format( + run_command: str | None, +) -> _FormatterType[LaunchArgBuilder, t.Sequence[str]]: + def impl( + args: LaunchArgBuilder, exe: ExecutableLike, env: _EnvironMappingType + ) -> t.Sequence[str]: + return ( + ( + run_command, + *(args.format_launch_args() or ()), + "--", + *exe.as_program_arguments(), + ) + if run_command is not None + else exe.as_program_arguments() + ) + + return impl + + class ShellLauncher: """Mock launcher for launching/tracking simple shell commands""" diff --git a/smartsim/settings/launchSettings.py b/smartsim/settings/launchSettings.py index 9078a04d9..dec6034d8 100644 --- a/smartsim/settings/launchSettings.py +++ b/smartsim/settings/launchSettings.py @@ -67,7 +67,7 @@ def launcher(self) -> str: return self._launcher.value @property - def launch_args(self) -> LaunchArgBuilder[t.Any]: + def launch_args(self) -> LaunchArgBuilder: """Return the launch argument translator.""" return self._arg_builder @@ -88,9 +88,7 @@ def env_vars(self, value: dict[str, str | None]) -> None: """Set the environment variables.""" self._env_vars = copy.deepcopy(value) - def _get_arg_builder( - self, launch_args: StringArgument | None - ) -> LaunchArgBuilder[t.Any]: + def _get_arg_builder(self, launch_args: StringArgument | None) -> LaunchArgBuilder: """Map the Launcher to the LaunchArgBuilder""" if self._launcher == LauncherType.Slurm: return SlurmArgBuilder(launch_args) From 6d5a3c46707e4441670b71369147cae104167df4 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Thu, 18 Jul 2024 22:43:22 -0500 Subject: [PATCH 23/34] Add tests for new dispatch API, make old tests pass --- smartsim/settings/builders/launch/alps.py | 3 +- smartsim/settings/builders/launch/local.py | 3 +- smartsim/settings/builders/launch/lsf.py | 3 +- smartsim/settings/builders/launch/mpi.py | 9 +- smartsim/settings/builders/launch/pals.py | 3 +- smartsim/settings/builders/launch/slurm.py | 3 +- tests/temp_tests/test_settings/conftest.py | 9 +- .../test_settings/test_alpsLauncher.py | 7 +- .../temp_tests/test_settings/test_dispatch.py | 333 ++++++++++++++++-- .../test_settings/test_dragonLauncher.py | 9 +- .../test_settings/test_localLauncher.py | 7 +- .../test_settings/test_lsfLauncher.py | 4 +- .../test_settings/test_mpiLauncher.py | 21 +- .../test_settings/test_palsLauncher.py | 7 +- .../test_settings/test_slurmLauncher.py | 7 +- .../test_settings/test_slurmScheduler.py | 1 - 16 files changed, 357 insertions(+), 72 deletions(-) diff --git a/smartsim/settings/builders/launch/alps.py b/smartsim/settings/builders/launch/alps.py index f325777e4..09d5931ac 100644 --- a/smartsim/settings/builders/launch/alps.py +++ b/smartsim/settings/builders/launch/alps.py @@ -36,9 +36,10 @@ from ..launchArgBuilder import LaunchArgBuilder logger = get_logger(__name__) +_format_aprun_command = shell_format(run_command="aprun") -@dispatch(with_format=shell_format(run_command="aprun"), to_launcher=ShellLauncher) +@dispatch(with_format=_format_aprun_command, to_launcher=ShellLauncher) class AprunArgBuilder(LaunchArgBuilder): def _reserved_launch_args(self) -> set[str]: """Return reserved launch arguments.""" diff --git a/smartsim/settings/builders/launch/local.py b/smartsim/settings/builders/launch/local.py index 21fb71c8a..7002a6831 100644 --- a/smartsim/settings/builders/launch/local.py +++ b/smartsim/settings/builders/launch/local.py @@ -36,9 +36,10 @@ from ..launchArgBuilder import LaunchArgBuilder logger = get_logger(__name__) +_format_local_command = shell_format(run_command=None) -@dispatch(with_format=shell_format(run_command=None), to_launcher=ShellLauncher) +@dispatch(with_format=_format_local_command, to_launcher=ShellLauncher) class LocalArgBuilder(LaunchArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/lsf.py b/smartsim/settings/builders/launch/lsf.py index 13a32fd73..ec99d51b9 100644 --- a/smartsim/settings/builders/launch/lsf.py +++ b/smartsim/settings/builders/launch/lsf.py @@ -36,9 +36,10 @@ from ..launchArgBuilder import LaunchArgBuilder logger = get_logger(__name__) +_format_jsrun_command = shell_format(run_command="jsrun") -@dispatch(with_format=shell_format(run_command="jsrun"), to_launcher=ShellLauncher) +@dispatch(with_format=_format_jsrun_command, to_launcher=ShellLauncher) class JsrunArgBuilder(LaunchArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/mpi.py b/smartsim/settings/builders/launch/mpi.py index ea24564da..139096010 100644 --- a/smartsim/settings/builders/launch/mpi.py +++ b/smartsim/settings/builders/launch/mpi.py @@ -36,6 +36,9 @@ from ..launchArgBuilder import LaunchArgBuilder logger = get_logger(__name__) +_format_mpirun_command = shell_format("mpirun") +_format_mpiexec_command = shell_format("mpiexec") +_format_orterun_command = shell_format("orterun") class _BaseMPIArgBuilder(LaunchArgBuilder): @@ -215,21 +218,21 @@ def set(self, key: str, value: str | None) -> None: self._launch_args[key] = value -@dispatch(with_format=shell_format(run_command="mpirun"), to_launcher=ShellLauncher) +@dispatch(with_format=_format_mpirun_command, to_launcher=ShellLauncher) class MpiArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Mpirun.value -@dispatch(with_format=shell_format(run_command="mpiexec"), to_launcher=ShellLauncher) +@dispatch(with_format=_format_mpiexec_command, to_launcher=ShellLauncher) class MpiexecArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Mpiexec.value -@dispatch(with_format=shell_format(run_command="orterun"), to_launcher=ShellLauncher) +@dispatch(with_format=_format_orterun_command, to_launcher=ShellLauncher) class OrteArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/pals.py b/smartsim/settings/builders/launch/pals.py index 4f2155c1f..1e7ed814e 100644 --- a/smartsim/settings/builders/launch/pals.py +++ b/smartsim/settings/builders/launch/pals.py @@ -36,9 +36,10 @@ from ..launchArgBuilder import LaunchArgBuilder logger = get_logger(__name__) +_format_mpiexec_command = shell_format(run_command="mpiexec") -@dispatch(with_format=shell_format(run_command="mpiexec"), to_launcher=ShellLauncher) +@dispatch(with_format=_format_mpiexec_command, to_launcher=ShellLauncher) class PalsMpiexecArgBuilder(LaunchArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/slurm.py b/smartsim/settings/builders/launch/slurm.py index 907a6da6c..72058f983 100644 --- a/smartsim/settings/builders/launch/slurm.py +++ b/smartsim/settings/builders/launch/slurm.py @@ -38,9 +38,10 @@ from ..launchArgBuilder import LaunchArgBuilder logger = get_logger(__name__) +_format_srun_command = shell_format(run_command="srun") -@dispatch(with_format=shell_format(run_command="srun"), to_launcher=ShellLauncher) +@dispatch(with_format=_format_srun_command, to_launcher=ShellLauncher) class SlurmArgBuilder(LaunchArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/tests/temp_tests/test_settings/conftest.py b/tests/temp_tests/test_settings/conftest.py index ebf361e97..72061264f 100644 --- a/tests/temp_tests/test_settings/conftest.py +++ b/tests/temp_tests/test_settings/conftest.py @@ -24,8 +24,6 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -from unittest.mock import Mock - import pytest from smartsim.settings import dispatch @@ -34,7 +32,7 @@ @pytest.fixture def echo_executable_like(): - class _ExeLike(launch.ExecutableLike): + class _ExeLike(dispatch.ExecutableLike): def as_program_arguments(self): return ("echo", "hello", "world") @@ -44,13 +42,10 @@ def as_program_arguments(self): @pytest.fixture def settings_builder(): class _SettingsBuilder(launch.LaunchArgBuilder): + def set(self, arg, val): ... def launcher_str(self): return "Mock Settings Builder" - def set(self, arg, val): ... - def finalize(self, exe, env): - return Mock() - yield _SettingsBuilder({}) diff --git a/tests/temp_tests/test_settings/test_alpsLauncher.py b/tests/temp_tests/test_settings/test_alpsLauncher.py index 7fa95cb6d..5ac2f8e11 100644 --- a/tests/temp_tests/test_settings/test_alpsLauncher.py +++ b/tests/temp_tests/test_settings/test_alpsLauncher.py @@ -1,7 +1,10 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.alps import AprunArgBuilder +from smartsim.settings.builders.launch.alps import ( + AprunArgBuilder, + _format_aprun_command, +) from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -183,5 +186,5 @@ def test_invalid_exclude_hostlist_format(): ), ) def test_formatting_launch_args(echo_executable_like, args, expected): - cmd = AprunArgBuilder(args).finalize(echo_executable_like, {}) + cmd = _format_aprun_command(AprunArgBuilder(args), echo_executable_like, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_dispatch.py b/tests/temp_tests/test_settings/test_dispatch.py index ccd1e81cd..78c44ad54 100644 --- a/tests/temp_tests/test_settings/test_dispatch.py +++ b/tests/temp_tests/test_settings/test_dispatch.py @@ -24,31 +24,56 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +import abc import contextlib +import dataclasses +import io import pytest +from smartsim.error import errors from smartsim.settings import dispatch pytestmark = pytest.mark.group_a +FORMATTED = object() -def test_declaritive_form_dispatch_declaration(launcher_like, settings_builder): + +def format_fn(args, exe, env): + return FORMATTED + + +@pytest.fixture +def expected_dispatch_registry(launcher_like, settings_builder): + yield { + type(settings_builder): dispatch._DispatchRegistration( + format_fn, type(launcher_like) + ) + } + + +def test_declaritive_form_dispatch_declaration( + launcher_like, settings_builder, expected_dispatch_registry +): d = dispatch.Dispatcher() - assert type(settings_builder) == d.dispatch(to_launcher=type(launcher_like))( - type(settings_builder) - ) - assert d._dispatch_registry == {type(settings_builder): type(launcher_like)} + assert type(settings_builder) == d.dispatch( + with_format=format_fn, to_launcher=type(launcher_like) + )(type(settings_builder)) + assert d._dispatch_registry == expected_dispatch_registry -def test_imperative_form_dispatch_declaration(launcher_like, settings_builder): +def test_imperative_form_dispatch_declaration( + launcher_like, settings_builder, expected_dispatch_registry +): d = dispatch.Dispatcher() - assert None == d.dispatch(type(settings_builder), to_launcher=type(launcher_like)) - assert d._dispatch_registry == {type(settings_builder): type(launcher_like)} + assert None == d.dispatch( + type(settings_builder), to_launcher=type(launcher_like), with_format=format_fn + ) + assert d._dispatch_registry == expected_dispatch_registry def test_dispatchers_from_same_registry_do_not_cross_polute( - launcher_like, settings_builder + launcher_like, settings_builder, expected_dispatch_registry ): some_starting_registry = {} d1 = dispatch.Dispatcher(dispatch_registry=some_starting_registry) @@ -60,12 +85,16 @@ def test_dispatchers_from_same_registry_do_not_cross_polute( d1._dispatch_registry is not d2._dispatch_registry is not some_starting_registry ) - d2.dispatch(type(settings_builder), to_launcher=type(launcher_like)) + d2.dispatch( + type(settings_builder), with_format=format_fn, to_launcher=type(launcher_like) + ) assert d1._dispatch_registry == {} - assert d2._dispatch_registry == {type(settings_builder): type(launcher_like)} + assert d2._dispatch_registry == expected_dispatch_registry -def test_copied_dispatchers_do_not_cross_pollute(launcher_like, settings_builder): +def test_copied_dispatchers_do_not_cross_pollute( + launcher_like, settings_builder, expected_dispatch_registry +): some_starting_registry = {} d1 = dispatch.Dispatcher(dispatch_registry=some_starting_registry) d2 = d1.copy() @@ -76,70 +105,304 @@ def test_copied_dispatchers_do_not_cross_pollute(launcher_like, settings_builder d1._dispatch_registry is not d2._dispatch_registry is not some_starting_registry ) - d2.dispatch(type(settings_builder), to_launcher=type(launcher_like)) + d2.dispatch( + type(settings_builder), to_launcher=type(launcher_like), with_format=format_fn + ) assert d1._dispatch_registry == {} - assert d2._dispatch_registry == {type(settings_builder): type(launcher_like)} + assert d2._dispatch_registry == expected_dispatch_registry @pytest.mark.parametrize( "add_dispatch, expected_ctx", ( pytest.param( - lambda d, s, l: d.dispatch(s, to_launcher=l), + lambda d, s, l: d.dispatch(s, to_launcher=l, with_format=format_fn), pytest.raises(TypeError, match="has already been registered"), id="Imperative -- Disallowed implicitly", ), pytest.param( - lambda d, s, l: d.dispatch(s, to_launcher=l, allow_overwrite=True), + lambda d, s, l: d.dispatch( + s, to_launcher=l, with_format=format_fn, allow_overwrite=True + ), contextlib.nullcontext(), id="Imperative -- Allowed with flag", ), pytest.param( - lambda d, s, l: d.dispatch(to_launcher=l)(s), + lambda d, s, l: d.dispatch(to_launcher=l, with_format=format_fn)(s), pytest.raises(TypeError, match="has already been registered"), id="Declarative -- Disallowed implicitly", ), pytest.param( - lambda d, s, l: d.dispatch(to_launcher=l, allow_overwrite=True)(s), + lambda d, s, l: d.dispatch( + to_launcher=l, with_format=format_fn, allow_overwrite=True + )(s), contextlib.nullcontext(), id="Declarative -- Allowed with flag", ), ), ) def test_dispatch_overwriting( - add_dispatch, expected_ctx, launcher_like, settings_builder + add_dispatch, + expected_ctx, + launcher_like, + settings_builder, + expected_dispatch_registry, ): - registry = {type(settings_builder): type(launcher_like)} - d = dispatch.Dispatcher(dispatch_registry=registry) + d = dispatch.Dispatcher(dispatch_registry=expected_dispatch_registry) with expected_ctx: add_dispatch(d, type(settings_builder), type(launcher_like)) @pytest.mark.parametrize( - "map_settings", + "type_or_instance", ( - pytest.param(type, id="From settings type"), - pytest.param(lambda s: s, id="From settings instance"), + pytest.param(type, id="type"), + pytest.param(lambda x: x, id="instance"), ), ) -def test_dispatch_can_retrieve_launcher_to_dispatch_to( - map_settings, launcher_like, settings_builder +def test_dispatch_can_retrieve_dispatch_info_from_dispatch_registry( + expected_dispatch_registry, launcher_like, settings_builder, type_or_instance ): - registry = {type(settings_builder): type(launcher_like)} - d = dispatch.Dispatcher(dispatch_registry=registry) - assert type(launcher_like) == d.get_launcher_for(map_settings(settings_builder)) + d = dispatch.Dispatcher(dispatch_registry=expected_dispatch_registry) + assert dispatch._DispatchRegistration( + format_fn, type(launcher_like) + ) == d.get_dispatch(type_or_instance(settings_builder)) @pytest.mark.parametrize( - "map_settings", + "type_or_instance", ( - pytest.param(type, id="From settings type"), - pytest.param(lambda s: s, id="From settings instance"), + pytest.param(type, id="type"), + pytest.param(lambda x: x, id="instance"), ), ) def test_dispatch_raises_if_settings_type_not_registered( - map_settings, launcher_like, settings_builder + settings_builder, type_or_instance ): d = dispatch.Dispatcher(dispatch_registry={}) - with pytest.raises(TypeError, match="no launcher type to dispatch to"): - d.get_launcher_for(map_settings(settings_builder)) + with pytest.raises( + TypeError, match="No dispatch for `.+?(?=`)` has been registered" + ): + d.get_dispatch(type_or_instance(settings_builder)) + + +class LauncherABC(abc.ABC): + @abc.abstractmethod + def start(self, launchable): ... + @classmethod + @abc.abstractmethod + def create(cls, exp): ... + + +class PartImplLauncherABC(LauncherABC): + def start(self, launchable): + return dispatch.create_job_id() + + +class FullImplLauncherABC(PartImplLauncherABC): + @classmethod + def create(cls, exp): + return cls() + + +@pytest.mark.parametrize( + "cls, ctx", + ( + pytest.param( + dispatch.LauncherLike, + pytest.raises(TypeError, match="Cannot dispatch to protocol"), + id="Cannot dispatch to protocol class", + ), + pytest.param( + "launcher_like", + contextlib.nullcontext(None), + id="Can dispatch to protocol implementation", + ), + pytest.param( + LauncherABC, + pytest.raises(TypeError, match="Cannot dispatch to abstract class"), + id="Cannot dispatch to abstract class", + ), + pytest.param( + PartImplLauncherABC, + pytest.raises(TypeError, match="Cannot dispatch to abstract class"), + id="Cannot dispatch to partially implemented abstract class", + ), + pytest.param( + FullImplLauncherABC, + contextlib.nullcontext(None), + id="Can dispatch to fully implemented abstract class", + ), + ), +) +def test_register_dispatch_to_launcher_types(request, cls, ctx): + if isinstance(cls, str): + cls = request.getfixturevalue(cls) + d = dispatch.Dispatcher() + with ctx: + d.dispatch(to_launcher=cls, with_format=format_fn) + + +@dataclasses.dataclass +class BufferWriterLauncher(dispatch.LauncherLike[list[str]]): + buf: io.StringIO + + @classmethod + def create(cls, exp): + return cls(io.StringIO()) + + def start(self, strs): + self.buf.writelines(f"{s}\n" for s in strs) + return dispatch.create_job_id() + + +class BufferWriterLauncherSubclass(BufferWriterLauncher): ... + + +@pytest.fixture +def buffer_writer_dispatch(): + stub_format_fn = lambda *a, **kw: ["some", "strings"] + return dispatch._DispatchRegistration(stub_format_fn, BufferWriterLauncher) + + +@pytest.mark.parametrize( + "input_, map_, expected", + ( + pytest.param( + ["list", "of", "strings"], + lambda xs: xs, + ["list\n", "of\n", "strings\n"], + id="[str] -> [str]", + ), + pytest.param( + "words on new lines", + lambda x: x.split(), + ["words\n", "on\n", "new\n", "lines\n"], + id="str -> [str]", + ), + pytest.param( + range(1, 4), + lambda xs: [str(x) for x in xs], + ["1\n", "2\n", "3\n"], + id="[int] -> [str]", + ), + ), +) +def test_launcher_adapter_correctly_adapts_input_to_launcher(input_, map_, expected): + buf = io.StringIO() + adapter = dispatch._LauncherAdapter(BufferWriterLauncher(buf), map_) + adapter.start(input_) + buf.seek(0) + assert buf.readlines() == expected + + +@pytest.mark.parametrize( + "launcher_instance, ctx", + ( + pytest.param( + BufferWriterLauncher(io.StringIO()), + contextlib.nullcontext(None), + id="Correctly configures expected launcher", + ), + pytest.param( + BufferWriterLauncherSubclass(io.StringIO()), + pytest.raises( + TypeError, + match="^Cannot create launcher adapter.*expected launcher of type .+$", + ), + id="Errors if launcher types are disparate", + ), + pytest.param( + "launcher_like", + pytest.raises( + TypeError, + match="^Cannot create launcher adapter.*expected launcher of type .+$", + ), + id="Errors if types are not an exact match", + ), + ), +) +def test_dispatch_registration_can_configure_adapter_for_existing_launcher_instance( + request, settings_builder, buffer_writer_dispatch, launcher_instance, ctx +): + if isinstance(launcher_instance, str): + launcher_instance = request.getfixturevalue(launcher_instance) + with ctx: + adapter = buffer_writer_dispatch.create_adapter_from_launcher( + launcher_instance, settings_builder + ) + assert adapter._adapted_launcher is launcher_instance + + +@pytest.mark.parametrize( + "launcher_instances, ctx", + ( + pytest.param( + (BufferWriterLauncher(io.StringIO()),), + contextlib.nullcontext(None), + id="Correctly configures expected launcher", + ), + pytest.param( + ( + "launcher_like", + "launcher_like", + BufferWriterLauncher(io.StringIO()), + "launcher_like", + ), + contextlib.nullcontext(None), + id="Correctly ignores incompatible launchers instances", + ), + pytest.param( + (), + pytest.raises( + errors.LauncherNotFoundError, + match="^No launcher of exactly type.+could be found from provided launchers$", + ), + id="Errors if no launcher could be found", + ), + pytest.param( + ( + "launcher_like", + BufferWriterLauncherSubclass(io.StringIO), + "launcher_like", + ), + pytest.raises( + errors.LauncherNotFoundError, + match="^No launcher of exactly type.+could be found from provided launchers$", + ), + id="Errors if no launcher matches expected type exactly", + ), + ), +) +def test_dispatch_registration_configures_first_compatible_launcher_from_sequence_of_launchers( + request, settings_builder, buffer_writer_dispatch, launcher_instances, ctx +): + def resolve_instance(inst): + return request.getfixturevalue(inst) if isinstance(inst, str) else inst + + launcher_instances = tuple(map(resolve_instance, launcher_instances)) + + with ctx: + adapter = buffer_writer_dispatch.configure_first_compatible_launcher( + with_settings=settings_builder, from_available_launchers=launcher_instances + ) + + +def test_dispatch_registration_can_create_a_laucher_for_an_experiment_and_can_reconfigure_it_later( + settings_builder, buffer_writer_dispatch +): + class MockExperiment: ... + + exp = MockExperiment() + adapter_1 = buffer_writer_dispatch.create_new_launcher_configuration( + for_experiment=exp, with_settings=settings_builder + ) + assert type(adapter_1._adapted_launcher) == buffer_writer_dispatch.launcher_type + existing_launcher = adapter_1._adapted_launcher + + adapter_2 = buffer_writer_dispatch.create_adapter_from_launcher( + existing_launcher, settings_builder + ) + assert type(adapter_2._adapted_launcher) == buffer_writer_dispatch.launcher_type + assert adapter_1._adapted_launcher is adapter_2._adapted_launcher + assert adapter_1 is not adapter_2 diff --git a/tests/temp_tests/test_settings/test_dragonLauncher.py b/tests/temp_tests/test_settings/test_dragonLauncher.py index 004090eef..57ae67d68 100644 --- a/tests/temp_tests/test_settings/test_dragonLauncher.py +++ b/tests/temp_tests/test_settings/test_dragonLauncher.py @@ -1,5 +1,6 @@ import pytest +from smartsim._core.launcher.dragon.dragonLauncher import _as_run_request_view from smartsim._core.schemas.dragonRequests import DragonRunRequest from smartsim.settings import LaunchSettings from smartsim.settings.builders.launch.dragon import DragonArgBuilder @@ -38,12 +39,12 @@ def test_dragon_class_methods(function, value, flag, result): def test_formatting_launch_args_into_request( echo_executable_like, nodes, tasks_per_node ): - builder = DragonArgBuilder({}) + args = DragonArgBuilder({}) if nodes is not NOT_SET: - builder.set_nodes(nodes) + args.set_nodes(nodes) if tasks_per_node is not NOT_SET: - builder.set_tasks_per_node(tasks_per_node) - req = builder.finalize(echo_executable_like, {}) + args.set_tasks_per_node(tasks_per_node) + req = _as_run_request_view(args, echo_executable_like, {}) args = dict( (k, v) diff --git a/tests/temp_tests/test_settings/test_localLauncher.py b/tests/temp_tests/test_settings/test_localLauncher.py index 4eb314a8b..d69657f23 100644 --- a/tests/temp_tests/test_settings/test_localLauncher.py +++ b/tests/temp_tests/test_settings/test_localLauncher.py @@ -1,7 +1,10 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.local import LocalArgBuilder +from smartsim.settings.builders.launch.local import ( + LocalArgBuilder, + _format_local_command, +) from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -115,5 +118,5 @@ def test_format_env_vars(): def test_formatting_returns_original_exe(echo_executable_like): - cmd = LocalArgBuilder({}).finalize(echo_executable_like, {}) + cmd = _format_local_command(LocalArgBuilder({}), echo_executable_like, {}) assert tuple(cmd) == ("echo", "hello", "world") diff --git a/tests/temp_tests/test_settings/test_lsfLauncher.py b/tests/temp_tests/test_settings/test_lsfLauncher.py index 592c80ce7..91f65821b 100644 --- a/tests/temp_tests/test_settings/test_lsfLauncher.py +++ b/tests/temp_tests/test_settings/test_lsfLauncher.py @@ -1,7 +1,7 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.lsf import JsrunArgBuilder +from smartsim.settings.builders.launch.lsf import JsrunArgBuilder, _format_jsrun_command from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -92,5 +92,5 @@ def test_launch_args(): ), ) def test_formatting_launch_args(echo_executable_like, args, expected): - cmd = JsrunArgBuilder(args).finalize(echo_executable_like, {}) + cmd = _format_jsrun_command(JsrunArgBuilder(args), echo_executable_like, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_mpiLauncher.py b/tests/temp_tests/test_settings/test_mpiLauncher.py index 9b651c220..54fed657e 100644 --- a/tests/temp_tests/test_settings/test_mpiLauncher.py +++ b/tests/temp_tests/test_settings/test_mpiLauncher.py @@ -7,6 +7,9 @@ MpiArgBuilder, MpiexecArgBuilder, OrteArgBuilder, + _format_mpiexec_command, + _format_mpirun_command, + _format_orterun_command, ) from smartsim.settings.launchCommand import LauncherType @@ -210,11 +213,15 @@ def test_invalid_hostlist_format(launcher): @pytest.mark.parametrize( - "cls, cmd", + "cls, fmt, cmd", ( - pytest.param(MpiArgBuilder, "mpirun", id="w/ mpirun"), - pytest.param(MpiexecArgBuilder, "mpiexec", id="w/ mpiexec"), - pytest.param(OrteArgBuilder, "orterun", id="w/ orterun"), + pytest.param(MpiArgBuilder, _format_mpirun_command, "mpirun", id="w/ mpirun"), + pytest.param( + MpiexecArgBuilder, _format_mpiexec_command, "mpiexec", id="w/ mpiexec" + ), + pytest.param( + OrteArgBuilder, _format_orterun_command, "orterun", id="w/ orterun" + ), ), ) @pytest.mark.parametrize( @@ -248,6 +255,6 @@ def test_invalid_hostlist_format(launcher): ), ), ) -def test_formatting_launch_args(echo_executable_like, cls, cmd, args, expected): - fmt = cls(args).finalize(echo_executable_like, {}) - assert tuple(fmt) == (cmd,) + expected +def test_formatting_launch_args(echo_executable_like, cls, fmt, cmd, args, expected): + fmt_cmd = fmt(cls(args), echo_executable_like, {}) + assert tuple(fmt_cmd) == (cmd,) + expected diff --git a/tests/temp_tests/test_settings/test_palsLauncher.py b/tests/temp_tests/test_settings/test_palsLauncher.py index a0bc7821c..5b74e2d0c 100644 --- a/tests/temp_tests/test_settings/test_palsLauncher.py +++ b/tests/temp_tests/test_settings/test_palsLauncher.py @@ -1,7 +1,10 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.pals import PalsMpiexecArgBuilder +from smartsim.settings.builders.launch.pals import ( + PalsMpiexecArgBuilder, + _format_mpiexec_command, +) from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -103,5 +106,5 @@ def test_invalid_hostlist_format(): ), ) def test_formatting_launch_args(echo_executable_like, args, expected): - cmd = PalsMpiexecArgBuilder(args).finalize(echo_executable_like, {}) + cmd = _format_mpiexec_command(PalsMpiexecArgBuilder(args), echo_executable_like, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_slurmLauncher.py b/tests/temp_tests/test_settings/test_slurmLauncher.py index bfa7dd9e1..2a84c831e 100644 --- a/tests/temp_tests/test_settings/test_slurmLauncher.py +++ b/tests/temp_tests/test_settings/test_slurmLauncher.py @@ -1,7 +1,10 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.slurm import SlurmArgBuilder +from smartsim.settings.builders.launch.slurm import ( + SlurmArgBuilder, + _format_srun_command, +) from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -289,5 +292,5 @@ def test_set_het_groups(monkeypatch): ), ) def test_formatting_launch_args(echo_executable_like, args, expected): - cmd = SlurmArgBuilder(args).finalize(echo_executable_like, {}) + cmd = _format_srun_command(SlurmArgBuilder(args), echo_executable_like, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_slurmScheduler.py b/tests/temp_tests/test_settings/test_slurmScheduler.py index 0a34b6473..5c65d367a 100644 --- a/tests/temp_tests/test_settings/test_slurmScheduler.py +++ b/tests/temp_tests/test_settings/test_slurmScheduler.py @@ -105,6 +105,5 @@ def test_sbatch_manual(): slurmScheduler.scheduler_args.set_account("A3531") slurmScheduler.scheduler_args.set_walltime("10:00:00") formatted = slurmScheduler.format_batch_args() - print(f"here: {formatted}") result = ["--nodes=5", "--account=A3531", "--time=10:00:00"] assert formatted == result From 6b7943a6f7b47dcc7bb0449e1587c6b0d7156a30 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Fri, 19 Jul 2024 22:37:01 -0500 Subject: [PATCH 24/34] Upper pin typing extensions (thanks TF) --- setup.py | 2 +- smartsim/settings/dispatch.py | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/setup.py b/setup.py index d820563d9..05e1c6436 100644 --- a/setup.py +++ b/setup.py @@ -180,7 +180,7 @@ def has_ext_modules(_placeholder): "pyzmq>=25.1.2", "pygithub>=2.3.0", "numpy<2", - "typing_extensions>=4.1.0", + "typing_extensions>=4.1.0,<4.6", ] # Add SmartRedis at specific version diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index a69a6587a..814797824 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -38,7 +38,6 @@ from smartsim.types import LaunchedJobID if t.TYPE_CHECKING: - from smartsim.experiment import Experiment from smartsim.settings.builders import LaunchArgBuilder @@ -47,7 +46,7 @@ _Ts = TypeVarTuple("_Ts") _T_contra = t.TypeVar("_T_contra", contravariant=True) _TDispatchable = t.TypeVar("_TDispatchable", bound="LaunchArgBuilder") -_EnvironMappingType: t.TypeAlias = t.Mapping[str, str | None] +_EnvironMappingType: t.TypeAlias = t.Mapping[str, "str | None"] _FormatterType: t.TypeAlias = t.Callable[ [_TDispatchable, "ExecutableLike", _EnvironMappingType], _T ] @@ -163,7 +162,7 @@ class _DispatchRegistration(t.Generic[_TDispatchable, _T]): formatter: _FormatterType[_TDispatchable, _T] launcher_type: type[LauncherLike[_T]] - def _is_compatible_launcher(self, launcher) -> bool: + def _is_compatible_launcher(self, launcher: LauncherLike[t.Any]) -> bool: return type(launcher) is self.launcher_type def create_new_launcher_configuration( From 82ee19fd59ca12be5c4c04e18c4fae967d19c5a6 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Sat, 20 Jul 2024 18:44:01 -0500 Subject: [PATCH 25/34] Make 3.9 compatable, lint, better names --- pyproject.toml | 12 ++- .../_core/launcher/dragon/dragonLauncher.py | 2 +- smartsim/experiment.py | 10 +- smartsim/settings/builders/launch/alps.py | 8 +- smartsim/settings/builders/launch/dragon.py | 8 +- smartsim/settings/builders/launch/local.py | 6 +- smartsim/settings/builders/launch/lsf.py | 8 +- smartsim/settings/builders/launch/mpi.py | 14 +-- smartsim/settings/builders/launch/pals.py | 8 +- smartsim/settings/builders/launch/slurm.py | 6 +- smartsim/settings/dispatch.py | 102 +++++++++--------- .../test_settings/test_alpsLauncher.py | 7 +- .../test_settings/test_dragonLauncher.py | 19 ++-- .../test_settings/test_localLauncher.py | 7 +- .../test_settings/test_lsfLauncher.py | 4 +- .../test_settings/test_mpiLauncher.py | 14 ++- .../test_settings/test_palsLauncher.py | 4 +- .../test_settings/test_slurmLauncher.py | 7 +- 18 files changed, 121 insertions(+), 125 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bda99459d..5df64aa97 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -151,7 +151,17 @@ module = [ "smartsim._core.control.controller", "smartsim._core.control.manifest", "smartsim._core.entrypoints.dragon_client", - "smartsim._core.launcher.*", + "smartsim._core.launcher.colocated", + "smartsim._core.launcher.launcher", + "smartsim._core.launcher.local.*", + "smartsim._core.launcher.lsf.*", + "smartsim._core.launcher.pbs.*", + "smartsim._core.launcher.sge.*", + "smartsim._core.launcher.slurm.*", + "smartsim._core.launcher.step.*", + "smartsim._core.launcher.stepInfo", + "smartsim._core.launcher.stepMapping", + "smartsim._core.launcher.taskManager", "smartsim._core.utils.serialize", "smartsim._core.utils.telemetry.*", "smartsim.database.*", diff --git a/smartsim/_core/launcher/dragon/dragonLauncher.py b/smartsim/_core/launcher/dragon/dragonLauncher.py index 74e22e160..949fcb044 100644 --- a/smartsim/_core/launcher/dragon/dragonLauncher.py +++ b/smartsim/_core/launcher/dragon/dragonLauncher.py @@ -338,7 +338,7 @@ def _assert_schema_type(obj: object, typ: t.Type[_SchemaT], /) -> _SchemaT: # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> # TODO: Remove this registry and move back to builder file after fixing -# circular import +# circular import caused by `DragonLauncher.supported_rs` # ----------------------------------------------------------------------------- from smartsim.settings.builders.launch.dragon import DragonArgBuilder from smartsim.settings.dispatch import ExecutableLike, dispatch diff --git a/smartsim/experiment.py b/smartsim/experiment.py index c2e56d19a..ca29191df 100644 --- a/smartsim/experiment.py +++ b/smartsim/experiment.py @@ -28,7 +28,6 @@ from __future__ import annotations -import itertools import os import os.path as osp import textwrap @@ -38,8 +37,8 @@ from tabulate import tabulate from smartsim._core.config import CONFIG -from smartsim.error.errors import LauncherNotFoundError, SSUnsupportedError -from smartsim.settings.dispatch import default_dispatcher +from smartsim.error import errors +from smartsim.settings.dispatch import DEFAULT_DISPATCHER from smartsim.status import SmartSimStatus from ._core import Controller, Generator, Manifest, previewrenderer @@ -53,7 +52,6 @@ ) from .error import SmartSimError from .log import ctx_exp_path, get_logger, method_contextualizer -from .settings import BatchSettings, Container, RunSettings if t.TYPE_CHECKING: from smartsim.launchable.job import Job @@ -113,7 +111,7 @@ def __init__( name: str, exp_path: str | None = None, *, # Keyword arguments only - settings_dispatcher: Dispatcher = default_dispatcher, + settings_dispatcher: Dispatcher = DEFAULT_DISPATCHER, ): """Initialize an Experiment instance. @@ -204,7 +202,7 @@ def _start(job: Job) -> LaunchedJobID: from_available_launchers=self._active_launchers, with_settings=args, ) - except LauncherNotFoundError: + except errors.LauncherNotFoundError: launch_config = dispatch.create_new_launcher_configuration( for_experiment=self, with_settings=args ) diff --git a/smartsim/settings/builders/launch/alps.py b/smartsim/settings/builders/launch/alps.py index 09d5931ac..5826a2de4 100644 --- a/smartsim/settings/builders/launch/alps.py +++ b/smartsim/settings/builders/launch/alps.py @@ -29,17 +29,17 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, dispatch, shell_format +from smartsim.settings.dispatch import ShellLauncher, dispatch, make_shell_format_fn -from ...common import StringArgument, set_check_input +from ...common import set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder logger = get_logger(__name__) -_format_aprun_command = shell_format(run_command="aprun") +_as_aprun_command = make_shell_format_fn(run_command="aprun") -@dispatch(with_format=_format_aprun_command, to_launcher=ShellLauncher) +@dispatch(with_format=_as_aprun_command, to_launcher=ShellLauncher) class AprunArgBuilder(LaunchArgBuilder): def _reserved_launch_args(self) -> set[str]: """Return reserved launch arguments.""" diff --git a/smartsim/settings/builders/launch/dragon.py b/smartsim/settings/builders/launch/dragon.py index 3cbb1c6d0..4ba793bf7 100644 --- a/smartsim/settings/builders/launch/dragon.py +++ b/smartsim/settings/builders/launch/dragon.py @@ -26,18 +26,12 @@ from __future__ import annotations -import os -import typing as t - from smartsim.log import get_logger -from ...common import StringArgument, set_check_input +from ...common import set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder -if t.TYPE_CHECKING: - from smartsim.settings.dispatch import ExecutableLike - logger = get_logger(__name__) diff --git a/smartsim/settings/builders/launch/local.py b/smartsim/settings/builders/launch/local.py index 7002a6831..49ef3ad92 100644 --- a/smartsim/settings/builders/launch/local.py +++ b/smartsim/settings/builders/launch/local.py @@ -29,17 +29,17 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, dispatch, shell_format +from smartsim.settings.dispatch import ShellLauncher, dispatch, make_shell_format_fn from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder logger = get_logger(__name__) -_format_local_command = shell_format(run_command=None) +_as_local_command = make_shell_format_fn(run_command=None) -@dispatch(with_format=_format_local_command, to_launcher=ShellLauncher) +@dispatch(with_format=_as_local_command, to_launcher=ShellLauncher) class LocalArgBuilder(LaunchArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/lsf.py b/smartsim/settings/builders/launch/lsf.py index ec99d51b9..1fe8ea30b 100644 --- a/smartsim/settings/builders/launch/lsf.py +++ b/smartsim/settings/builders/launch/lsf.py @@ -29,17 +29,17 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, dispatch, shell_format +from smartsim.settings.dispatch import ShellLauncher, dispatch, make_shell_format_fn -from ...common import StringArgument, set_check_input +from ...common import set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder logger = get_logger(__name__) -_format_jsrun_command = shell_format(run_command="jsrun") +_as_jsrun_command = make_shell_format_fn(run_command="jsrun") -@dispatch(with_format=_format_jsrun_command, to_launcher=ShellLauncher) +@dispatch(with_format=_as_jsrun_command, to_launcher=ShellLauncher) class JsrunArgBuilder(LaunchArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/mpi.py b/smartsim/settings/builders/launch/mpi.py index 139096010..a0d947418 100644 --- a/smartsim/settings/builders/launch/mpi.py +++ b/smartsim/settings/builders/launch/mpi.py @@ -29,16 +29,16 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, dispatch, shell_format +from smartsim.settings.dispatch import ShellLauncher, dispatch, make_shell_format_fn from ...common import set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder logger = get_logger(__name__) -_format_mpirun_command = shell_format("mpirun") -_format_mpiexec_command = shell_format("mpiexec") -_format_orterun_command = shell_format("orterun") +_as_mpirun_command = make_shell_format_fn("mpirun") +_as_mpiexec_command = make_shell_format_fn("mpiexec") +_as_orterun_command = make_shell_format_fn("orterun") class _BaseMPIArgBuilder(LaunchArgBuilder): @@ -218,21 +218,21 @@ def set(self, key: str, value: str | None) -> None: self._launch_args[key] = value -@dispatch(with_format=_format_mpirun_command, to_launcher=ShellLauncher) +@dispatch(with_format=_as_mpirun_command, to_launcher=ShellLauncher) class MpiArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Mpirun.value -@dispatch(with_format=_format_mpiexec_command, to_launcher=ShellLauncher) +@dispatch(with_format=_as_mpiexec_command, to_launcher=ShellLauncher) class MpiexecArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Mpiexec.value -@dispatch(with_format=_format_orterun_command, to_launcher=ShellLauncher) +@dispatch(with_format=_as_orterun_command, to_launcher=ShellLauncher) class OrteArgBuilder(_BaseMPIArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/pals.py b/smartsim/settings/builders/launch/pals.py index 1e7ed814e..eeb438455 100644 --- a/smartsim/settings/builders/launch/pals.py +++ b/smartsim/settings/builders/launch/pals.py @@ -29,17 +29,17 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, dispatch, shell_format +from smartsim.settings.dispatch import ShellLauncher, dispatch, make_shell_format_fn -from ...common import StringArgument, set_check_input +from ...common import set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder logger = get_logger(__name__) -_format_mpiexec_command = shell_format(run_command="mpiexec") +_as_pals_command = make_shell_format_fn(run_command="mpiexec") -@dispatch(with_format=_format_mpiexec_command, to_launcher=ShellLauncher) +@dispatch(with_format=_as_pals_command, to_launcher=ShellLauncher) class PalsMpiexecArgBuilder(LaunchArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/builders/launch/slurm.py b/smartsim/settings/builders/launch/slurm.py index 72058f983..f80a5b8b9 100644 --- a/smartsim/settings/builders/launch/slurm.py +++ b/smartsim/settings/builders/launch/slurm.py @@ -31,17 +31,17 @@ import typing as t from smartsim.log import get_logger -from smartsim.settings.dispatch import ShellLauncher, dispatch, shell_format +from smartsim.settings.dispatch import ShellLauncher, dispatch, make_shell_format_fn from ...common import set_check_input from ...launchCommand import LauncherType from ..launchArgBuilder import LaunchArgBuilder logger = get_logger(__name__) -_format_srun_command = shell_format(run_command="srun") +_as_srun_command = make_shell_format_fn(run_command="srun") -@dispatch(with_format=_format_srun_command, to_launcher=ShellLauncher) +@dispatch(with_format=_as_srun_command, to_launcher=ShellLauncher) class SlurmArgBuilder(LaunchArgBuilder): def launcher_str(self) -> str: """Get the string representation of the launcher""" diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index 814797824..cda8dd9ba 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -31,7 +31,7 @@ import typing as t import uuid -from typing_extensions import Self, TypeVarTuple, Unpack +from typing_extensions import Self, TypeAlias, TypeVarTuple, Unpack from smartsim._core.utils import helpers from smartsim.error import errors @@ -41,17 +41,17 @@ from smartsim.experiment import Experiment from smartsim.settings.builders import LaunchArgBuilder - -_T = t.TypeVar("_T") _Ts = TypeVarTuple("_Ts") _T_contra = t.TypeVar("_T_contra", contravariant=True) -_TDispatchable = t.TypeVar("_TDispatchable", bound="LaunchArgBuilder") -_EnvironMappingType: t.TypeAlias = t.Mapping[str, "str | None"] -_FormatterType: t.TypeAlias = t.Callable[ - [_TDispatchable, "ExecutableLike", _EnvironMappingType], _T +_DispatchableT = t.TypeVar("_DispatchableT", bound="LaunchArgBuilder") +_LaunchableT = t.TypeVar("_LaunchableT") + +_EnvironMappingType: TypeAlias = t.Mapping[str, "str | None"] +_FormatterType: TypeAlias = t.Callable[ + [_DispatchableT, "ExecutableLike", _EnvironMappingType], _LaunchableT ] -_LaunchConfigType: t.TypeAlias = "_LauncherAdapter[ExecutableLike, _EnvironMappingType]" -_UnkownType: t.TypeAlias = t.NoReturn +_LaunchConfigType: TypeAlias = "_LauncherAdapter[ExecutableLike, _EnvironMappingType]" +_UnkownType: TypeAlias = t.NoReturn @t.final @@ -81,27 +81,27 @@ def dispatch( self, args: None = ..., *, - with_format: _FormatterType[_TDispatchable, _T], - to_launcher: type[LauncherLike[_T]], + with_format: _FormatterType[_DispatchableT, _LaunchableT], + to_launcher: type[LauncherLike[_LaunchableT]], allow_overwrite: bool = ..., - ) -> t.Callable[[type[_TDispatchable]], type[_TDispatchable]]: ... + ) -> t.Callable[[type[_DispatchableT]], type[_DispatchableT]]: ... @t.overload def dispatch( self, - args: type[_TDispatchable], + args: type[_DispatchableT], *, - with_format: _FormatterType[_TDispatchable, _T], - to_launcher: type[LauncherLike[_T]], + with_format: _FormatterType[_DispatchableT, _LaunchableT], + to_launcher: type[LauncherLike[_LaunchableT]], allow_overwrite: bool = ..., ) -> None: ... def dispatch( self, - args: type[_TDispatchable] | None = None, + args: type[_DispatchableT] | None = None, *, - with_format: _FormatterType[_TDispatchable, _T], - to_launcher: type[LauncherLike[_T]], + with_format: _FormatterType[_DispatchableT, _LaunchableT], + to_launcher: type[LauncherLike[_LaunchableT]], allow_overwrite: bool = False, - ) -> t.Callable[[type[_TDispatchable]], type[_TDispatchable]] | None: + ) -> t.Callable[[type[_DispatchableT]], type[_DispatchableT]] | None: """A type safe way to add a mapping of settings builder to launcher to handle the settings at launch time. """ @@ -113,7 +113,7 @@ def dispatch( if err_msg is not None: raise TypeError(err_msg) - def register(args_: type[_TDispatchable], /) -> type[_TDispatchable]: + def register(args_: type[_DispatchableT], /) -> type[_DispatchableT]: if args_ in self._dispatch_registry and not allow_overwrite: launcher_type = self._dispatch_registry[args_].launcher_type raise TypeError( @@ -131,48 +131,51 @@ def register(args_: type[_TDispatchable], /) -> type[_TDispatchable]: return register def get_dispatch( - self, args: _TDispatchable | type[_TDispatchable] - ) -> _DispatchRegistration[_TDispatchable, _UnkownType]: + self, args: _DispatchableT | type[_DispatchableT] + ) -> _DispatchRegistration[_DispatchableT, _UnkownType]: """Find a type of launcher that is registered as being able to launch the output of the provided builder """ if not isinstance(args, type): args = type(args) - dispatch = self._dispatch_registry.get(args, None) - if dispatch is None: + dispatch_ = self._dispatch_registry.get(args, None) + if dispatch_ is None: raise TypeError( f"No dispatch for `{type(args).__name__}` has been registered " f"has been registered with {type(self).__name__} `{self}`" ) # Note the sleight-of-hand here: we are secretly casting a type of # `_DispatchRegistration[Any, Any]` -> - # `_DispatchRegistration[_TDispatchable, _T]`. - # where `_T` is unbound! + # `_DispatchRegistration[_DispatchableT, _LaunchableT]`. + # where `_LaunchableT` is unbound! # # This is safe to do if all entries in the mapping were added using a # type safe method (e.g. `Dispatcher.dispatch`), but if a user were to # supply a custom dispatch registry or otherwise modify the registry # this is not necessarily 100% type safe!! - return dispatch + return dispatch_ @t.final @dataclasses.dataclass(frozen=True) -class _DispatchRegistration(t.Generic[_TDispatchable, _T]): - formatter: _FormatterType[_TDispatchable, _T] - launcher_type: type[LauncherLike[_T]] +class _DispatchRegistration(t.Generic[_DispatchableT, _LaunchableT]): + formatter: _FormatterType[_DispatchableT, _LaunchableT] + launcher_type: type[LauncherLike[_LaunchableT]] def _is_compatible_launcher(self, launcher: LauncherLike[t.Any]) -> bool: + # Disabling because we want to match the the type of the dispatch + # *exactly* as specified by the user + # pylint: disable-next=unidiomatic-typecheck return type(launcher) is self.launcher_type def create_new_launcher_configuration( - self, for_experiment: Experiment, with_settings: _TDispatchable + self, for_experiment: Experiment, with_settings: _DispatchableT ) -> _LaunchConfigType: launcher = self.launcher_type.create(for_experiment) return self.create_adapter_from_launcher(launcher, with_settings) def create_adapter_from_launcher( - self, launcher: LauncherLike[_T], settings: _TDispatchable + self, launcher: LauncherLike[_LaunchableT], settings: _DispatchableT ) -> _LaunchConfigType: if not self._is_compatible_launcher(launcher): raise TypeError( @@ -181,14 +184,14 @@ def create_adapter_from_launcher( f"exactly `{self.launcher_type}`" ) - def format_(exe: ExecutableLike, env: _EnvironMappingType) -> _T: + def format_(exe: ExecutableLike, env: _EnvironMappingType) -> _LaunchableT: return self.formatter(settings, exe, env) return _LauncherAdapter(launcher, format_) def configure_first_compatible_launcher( self, - with_settings: _TDispatchable, + with_settings: _DispatchableT, from_available_launchers: t.Iterable[LauncherLike[t.Any]], ) -> _LaunchConfigType: launcher = helpers.first(self._is_compatible_launcher, from_available_launchers) @@ -203,11 +206,14 @@ def configure_first_compatible_launcher( @t.final class _LauncherAdapter(t.Generic[Unpack[_Ts]]): def __init__( - self, launcher: LauncherLike[_T], map_: t.Callable[[Unpack[_Ts]], _T] + self, + launcher: LauncherLike[_LaunchableT], + map_: t.Callable[[Unpack[_Ts]], _LaunchableT], ) -> None: - # NOTE: We need to cast off the `_T` -> `Any` in the `__init__` - # signature to hide the transform from users of this class. If - # possible, try not to expose outside of protected methods! + # NOTE: We need to cast off the `_LaunchableT` -> `Any` in the + # `__init__` method signature to hide the transform from users of + # this class. If possible, this type should not be exposed to + # users of this class! self._adapt: t.Callable[[Unpack[_Ts]], t.Any] = map_ self._adapted_launcher: LauncherLike[t.Any] = launcher @@ -216,8 +222,11 @@ def start(self, *args: Unpack[_Ts]) -> LaunchedJobID: return self._adapted_launcher.start(payload) -default_dispatcher: t.Final = Dispatcher() -dispatch: t.Final = default_dispatcher.dispatch +DEFAULT_DISPATCHER: t.Final = Dispatcher() +# Disabling because we want this to look and feel like a top level function, +# but don't want to have a second copy of the nasty overloads +# pylint: disable-next=invalid-name +dispatch: t.Final = DEFAULT_DISPATCHER.dispatch # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> @@ -236,18 +245,14 @@ def as_program_arguments(self) -> t.Sequence[str]: ... class LauncherLike(t.Protocol[_T_contra]): def start(self, launchable: _T_contra) -> LaunchedJobID: ... @classmethod - def create(cls, exp: Experiment) -> Self: ... + def create(cls, exp: Experiment, /) -> Self: ... -# TODO: This is just a nice helper function that I am using for the time being -# to wire everything up! In reality it might be a bit too confusing and -# meta-program-y for production code. Check with the core team to see -# what they think!! -def shell_format( +def make_shell_format_fn( run_command: str | None, ) -> _FormatterType[LaunchArgBuilder, t.Sequence[str]]: def impl( - args: LaunchArgBuilder, exe: ExecutableLike, env: _EnvironMappingType + args: LaunchArgBuilder, exe: ExecutableLike, _env: _EnvironMappingType ) -> t.Sequence[str]: return ( ( @@ -272,11 +277,12 @@ def __init__(self) -> None: def start(self, launchable: t.Sequence[str]) -> LaunchedJobID: id_ = create_job_id() exe, *rest = launchable + # pylint: disable-next=consider-using-with self._launched[id_] = sp.Popen((helpers.expand_exe_path(exe), *rest)) return id_ @classmethod - def create(cls, exp: Experiment) -> Self: + def create(cls, _: Experiment) -> Self: return cls() diff --git a/tests/temp_tests/test_settings/test_alpsLauncher.py b/tests/temp_tests/test_settings/test_alpsLauncher.py index 5ac2f8e11..cb6e829eb 100644 --- a/tests/temp_tests/test_settings/test_alpsLauncher.py +++ b/tests/temp_tests/test_settings/test_alpsLauncher.py @@ -1,10 +1,7 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.alps import ( - AprunArgBuilder, - _format_aprun_command, -) +from smartsim.settings.builders.launch.alps import AprunArgBuilder, _as_aprun_command from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -186,5 +183,5 @@ def test_invalid_exclude_hostlist_format(): ), ) def test_formatting_launch_args(echo_executable_like, args, expected): - cmd = _format_aprun_command(AprunArgBuilder(args), echo_executable_like, {}) + cmd = _as_aprun_command(AprunArgBuilder(args), echo_executable_like, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_dragonLauncher.py b/tests/temp_tests/test_settings/test_dragonLauncher.py index 57ae67d68..cc57e329f 100644 --- a/tests/temp_tests/test_settings/test_dragonLauncher.py +++ b/tests/temp_tests/test_settings/test_dragonLauncher.py @@ -1,7 +1,7 @@ import pytest from smartsim._core.launcher.dragon.dragonLauncher import _as_run_request_view -from smartsim._core.schemas.dragonRequests import DragonRunRequest +from smartsim._core.schemas.dragonRequests import DragonRunRequestView from smartsim.settings import LaunchSettings from smartsim.settings.builders.launch.dragon import DragonArgBuilder from smartsim.settings.launchCommand import LauncherType @@ -46,18 +46,17 @@ def test_formatting_launch_args_into_request( args.set_tasks_per_node(tasks_per_node) req = _as_run_request_view(args, echo_executable_like, {}) - args = dict( - (k, v) - for k, v in ( - ("nodes", nodes), - ("tasks_per_node", tasks_per_node), - ) + args = { + k: v + for k, v in { + "nodes": nodes, + "tasks_per_node": tasks_per_node, + }.items() if v is not NOT_SET - ) - expected = DragonRunRequest( + } + expected = DragonRunRequestView( exe="echo", exe_args=["hello", "world"], path="/tmp", env={}, **args ) - assert req.nodes == expected.nodes assert req.tasks_per_node == expected.tasks_per_node assert req.hostlist == expected.hostlist diff --git a/tests/temp_tests/test_settings/test_localLauncher.py b/tests/temp_tests/test_settings/test_localLauncher.py index d69657f23..07adb231b 100644 --- a/tests/temp_tests/test_settings/test_localLauncher.py +++ b/tests/temp_tests/test_settings/test_localLauncher.py @@ -1,10 +1,7 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.local import ( - LocalArgBuilder, - _format_local_command, -) +from smartsim.settings.builders.launch.local import LocalArgBuilder, _as_local_command from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -118,5 +115,5 @@ def test_format_env_vars(): def test_formatting_returns_original_exe(echo_executable_like): - cmd = _format_local_command(LocalArgBuilder({}), echo_executable_like, {}) + cmd = _as_local_command(LocalArgBuilder({}), echo_executable_like, {}) assert tuple(cmd) == ("echo", "hello", "world") diff --git a/tests/temp_tests/test_settings/test_lsfLauncher.py b/tests/temp_tests/test_settings/test_lsfLauncher.py index 91f65821b..9d276c0e7 100644 --- a/tests/temp_tests/test_settings/test_lsfLauncher.py +++ b/tests/temp_tests/test_settings/test_lsfLauncher.py @@ -1,7 +1,7 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.lsf import JsrunArgBuilder, _format_jsrun_command +from smartsim.settings.builders.launch.lsf import JsrunArgBuilder, _as_jsrun_command from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -92,5 +92,5 @@ def test_launch_args(): ), ) def test_formatting_launch_args(echo_executable_like, args, expected): - cmd = _format_jsrun_command(JsrunArgBuilder(args), echo_executable_like, {}) + cmd = _as_jsrun_command(JsrunArgBuilder(args), echo_executable_like, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_mpiLauncher.py b/tests/temp_tests/test_settings/test_mpiLauncher.py index 54fed657e..16947590c 100644 --- a/tests/temp_tests/test_settings/test_mpiLauncher.py +++ b/tests/temp_tests/test_settings/test_mpiLauncher.py @@ -7,9 +7,9 @@ MpiArgBuilder, MpiexecArgBuilder, OrteArgBuilder, - _format_mpiexec_command, - _format_mpirun_command, - _format_orterun_command, + _as_mpiexec_command, + _as_mpirun_command, + _as_orterun_command, ) from smartsim.settings.launchCommand import LauncherType @@ -215,13 +215,11 @@ def test_invalid_hostlist_format(launcher): @pytest.mark.parametrize( "cls, fmt, cmd", ( - pytest.param(MpiArgBuilder, _format_mpirun_command, "mpirun", id="w/ mpirun"), + pytest.param(MpiArgBuilder, _as_mpirun_command, "mpirun", id="w/ mpirun"), pytest.param( - MpiexecArgBuilder, _format_mpiexec_command, "mpiexec", id="w/ mpiexec" - ), - pytest.param( - OrteArgBuilder, _format_orterun_command, "orterun", id="w/ orterun" + MpiexecArgBuilder, _as_mpiexec_command, "mpiexec", id="w/ mpiexec" ), + pytest.param(OrteArgBuilder, _as_orterun_command, "orterun", id="w/ orterun"), ), ) @pytest.mark.parametrize( diff --git a/tests/temp_tests/test_settings/test_palsLauncher.py b/tests/temp_tests/test_settings/test_palsLauncher.py index 5b74e2d0c..47fa64713 100644 --- a/tests/temp_tests/test_settings/test_palsLauncher.py +++ b/tests/temp_tests/test_settings/test_palsLauncher.py @@ -3,7 +3,7 @@ from smartsim.settings import LaunchSettings from smartsim.settings.builders.launch.pals import ( PalsMpiexecArgBuilder, - _format_mpiexec_command, + _as_pals_command, ) from smartsim.settings.launchCommand import LauncherType @@ -106,5 +106,5 @@ def test_invalid_hostlist_format(): ), ) def test_formatting_launch_args(echo_executable_like, args, expected): - cmd = _format_mpiexec_command(PalsMpiexecArgBuilder(args), echo_executable_like, {}) + cmd = _as_pals_command(PalsMpiexecArgBuilder(args), echo_executable_like, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_slurmLauncher.py b/tests/temp_tests/test_settings/test_slurmLauncher.py index 2a84c831e..5935ec690 100644 --- a/tests/temp_tests/test_settings/test_slurmLauncher.py +++ b/tests/temp_tests/test_settings/test_slurmLauncher.py @@ -1,10 +1,7 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.slurm import ( - SlurmArgBuilder, - _format_srun_command, -) +from smartsim.settings.builders.launch.slurm import SlurmArgBuilder, _as_srun_command from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -292,5 +289,5 @@ def test_set_het_groups(monkeypatch): ), ) def test_formatting_launch_args(echo_executable_like, args, expected): - cmd = _format_srun_command(SlurmArgBuilder(args), echo_executable_like, {}) + cmd = _as_srun_command(SlurmArgBuilder(args), echo_executable_like, {}) assert tuple(cmd) == expected From 36bcfbb39eb35999b2e5de9439c1366b358ff7b1 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Wed, 17 Jul 2024 18:19:43 -0500 Subject: [PATCH 26/34] Address reviewer comments --- .../_core/launcher/dragon/dragonConnector.py | 16 +++++ .../_core/launcher/dragon/dragonLauncher.py | 6 +- smartsim/experiment.py | 59 +++++++++++-------- .../settings/builders/launchArgBuilder.py | 15 +++-- smartsim/settings/dispatch.py | 34 ++++++----- smartsim/settings/launchSettings.py | 10 +++- tests/temp_tests/test_settings/conftest.py | 12 ++-- .../test_settings/test_alpsLauncher.py | 4 +- .../temp_tests/test_settings/test_dispatch.py | 46 +++++++-------- .../test_settings/test_dragonLauncher.py | 4 +- .../test_settings/test_localLauncher.py | 4 +- .../test_settings/test_lsfLauncher.py | 4 +- .../test_settings/test_mpiLauncher.py | 4 +- .../test_settings/test_palsLauncher.py | 4 +- .../test_settings/test_slurmLauncher.py | 4 +- 15 files changed, 134 insertions(+), 92 deletions(-) diff --git a/smartsim/_core/launcher/dragon/dragonConnector.py b/smartsim/_core/launcher/dragon/dragonConnector.py index ca721eeaa..60fbf3ce7 100644 --- a/smartsim/_core/launcher/dragon/dragonConnector.py +++ b/smartsim/_core/launcher/dragon/dragonConnector.py @@ -522,6 +522,22 @@ def _dragon_cleanup( def _resolve_dragon_path(fallback: t.Union[str, "os.PathLike[str]"]) -> Path: + """Return the path at which a user should set up a dragon server. + + The order of path resolution is: + 1) If the the user has set a global dragon path via + `Config.dragon_server_path` use that without alteration. + 2) Use the `fallback` path which should be the path to an existing + directory. Append the default dragon server subdirectory defined by + `Config.dragon_default_subdir` + + Currently this function will raise if a user attempts to specify multiple + dragon server paths via `:` seperation. + + :param fallback: The path to an existing directory on the file system to + use if the global dragon directory is not set. + :returns: The path to directory in which the dragon server should run. + """ config = get_config() dragon_server_path = config.dragon_server_path or os.path.join( fallback, config.dragon_default_subdir diff --git a/smartsim/_core/launcher/dragon/dragonLauncher.py b/smartsim/_core/launcher/dragon/dragonLauncher.py index 949fcb044..078a89278 100644 --- a/smartsim/_core/launcher/dragon/dragonLauncher.py +++ b/smartsim/_core/launcher/dragon/dragonLauncher.py @@ -341,11 +341,13 @@ def _assert_schema_type(obj: object, typ: t.Type[_SchemaT], /) -> _SchemaT: # circular import caused by `DragonLauncher.supported_rs` # ----------------------------------------------------------------------------- from smartsim.settings.builders.launch.dragon import DragonArgBuilder -from smartsim.settings.dispatch import ExecutableLike, dispatch +from smartsim.settings.dispatch import ExecutableProtocol, dispatch def _as_run_request_view( - run_req_args: DragonArgBuilder, exe: ExecutableLike, env: t.Mapping[str, str | None] + run_req_args: DragonArgBuilder, + exe: ExecutableProtocol, + env: t.Mapping[str, str | None], ) -> DragonRunRequestView: exe_, *args = exe.as_program_arguments() return DragonRunRequestView( diff --git a/smartsim/experiment.py b/smartsim/experiment.py index ca29191df..6f1a5744a 100644 --- a/smartsim/experiment.py +++ b/smartsim/experiment.py @@ -55,8 +55,11 @@ if t.TYPE_CHECKING: from smartsim.launchable.job import Job - from smartsim.settings.builders.launchArgBuilder import LaunchArgBuilder - from smartsim.settings.dispatch import Dispatcher, ExecutableLike, LauncherLike + from smartsim.settings.dispatch import ( + Dispatcher, + ExecutableProtocol, + LauncherProtocol, + ) from smartsim.types import LaunchedJobID logger = get_logger(__name__) @@ -106,13 +109,7 @@ class Experiment: and utilized throughout runtime. """ - def __init__( - self, - name: str, - exp_path: str | None = None, - *, # Keyword arguments only - settings_dispatcher: Dispatcher = DEFAULT_DISPATCHER, - ): + def __init__(self, name: str, exp_path: str | None = None): """Initialize an Experiment instance. With the default settings, the Experiment will use the @@ -152,9 +149,6 @@ def __init__( :param name: name for the ``Experiment`` :param exp_path: path to location of ``Experiment`` directory - :param settings_dispatcher: The dispatcher the experiment will use to - figure determine how to launch a job. If none is provided, the - experiment will use the default dispatcher. """ self.name = name if exp_path: @@ -167,24 +161,39 @@ def __init__( exp_path = osp.join(getcwd(), name) self.exp_path = exp_path + """The path under which the experiment operate""" - # TODO: Remove this! The contoller is becoming obsolete + # TODO: Remove this! The controller is becoming obsolete self._control = Controller(launcher="local") - self._dispatcher = settings_dispatcher - self._active_launchers: set[LauncherLike[t.Any]] = set() + self._active_launchers: set[LauncherProtocol[t.Any]] = set() """The active launchers created, used, and reused by the experiment""" - self.fs_identifiers: t.Set[str] = set() + self._fs_identifiers: t.Set[str] = set() + """Set of feature store identifiers currently in use by this + experiment""" self._telemetry_cfg = ExperimentTelemetryConfiguration() - - def start_jobs(self, *jobs: Job) -> tuple[LaunchedJobID, ...]: - """WIP: replacemnt method to launch jobs using the new API""" + """Switch to specify if telemetry data should be produced for this + experiment""" + + def start_jobs( + self, *jobs: Job, dispatcher: Dispatcher = DEFAULT_DISPATCHER + ) -> tuple[LaunchedJobID, ...]: + """Execute a collection of `Job` instances. + + :param jobs: The collection of jobs instances to start + :param dispatcher: The dispatcher that should be used to determine how + to start a job based on its settings. If not specified it will + default to a dispatcher pre-configured by SmartSim. + :returns: A sequence of ids with order corresponding to the sequence of + jobs that can be used to query or alter the status of that + particular execution of the job. + """ if not jobs: raise TypeError( f"{type(self).__name__}.start_jobs() missing at least 1 required " - "positional argument" + "positional argument of type `Job`" ) def _start(job: Job) -> LaunchedJobID: @@ -194,9 +203,9 @@ def _start(job: Job) -> LaunchedJobID: # FIXME: Remove this cast after `SmartSimEntity` conforms to # protocol. For now, live with the "dangerous" type cast # --------------------------------------------------------------------- - exe = t.cast("ExecutableLike", job.entity) + exe = t.cast("ExecutableProtocol", job.entity) # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - dispatch = self._dispatcher.get_dispatch(args) + dispatch = dispatcher.get_dispatch(args) try: launch_config = dispatch.configure_first_compatible_launcher( from_available_launchers=self._active_launchers, @@ -603,7 +612,7 @@ def _create_entity_dir(self, start_manifest: Manifest) -> None: def create_entity_dir( entity: t.Union[FeatureStore, Application, Ensemble] ) -> None: - if not os.path.isdir(entity.path): + if not osp.isdir(entity.path): os.makedirs(entity.path) for application in start_manifest.applications: @@ -620,11 +629,11 @@ def __str__(self) -> str: def _append_to_fs_identifier_list(self, fs_identifier: str) -> None: """Check if fs_identifier already exists when calling create_feature_store""" - if fs_identifier in self.fs_identifiers: + if fs_identifier in self._fs_identifiers: logger.warning( f"A feature store with the identifier {fs_identifier} has already been made " "An error will be raised if multiple Feature Stores are started " "with the same identifier" ) # Otherwise, add - self.fs_identifiers.add(fs_identifier) + self._fs_identifiers.add(fs_identifier) diff --git a/smartsim/settings/builders/launchArgBuilder.py b/smartsim/settings/builders/launchArgBuilder.py index 2c09dd2e8..8ebfdb0f0 100644 --- a/smartsim/settings/builders/launchArgBuilder.py +++ b/smartsim/settings/builders/launchArgBuilder.py @@ -37,9 +37,6 @@ logger = get_logger(__name__) -_T = t.TypeVar("_T") - - class LaunchArgBuilder(ABC): """Abstract base class that defines all generic launcher argument methods that are not supported. It is the @@ -48,6 +45,11 @@ class LaunchArgBuilder(ABC): """ def __init__(self, launch_args: t.Dict[str, str | None] | None) -> None: + """Initialize a new `LaunchArgBuilder` instance. + + :param launch_args: A mapping of argument to be used to initialize the + argument builder. + """ self._launch_args = copy.deepcopy(launch_args) or {} @abstractmethod @@ -56,7 +58,12 @@ def launcher_str(self) -> str: @abstractmethod def set(self, arg: str, val: str | None) -> None: - """Set the launch arguments""" + """Set a launch argument + + :param arg: The argument name to set + :param val: The value to set the argument to as a `str` (if + applicable). Otherwise `None` + """ def format_launch_args(self) -> t.Union[t.List[str], None]: """Build formatted launch arguments""" diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index cda8dd9ba..aceead700 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -48,9 +48,11 @@ _EnvironMappingType: TypeAlias = t.Mapping[str, "str | None"] _FormatterType: TypeAlias = t.Callable[ - [_DispatchableT, "ExecutableLike", _EnvironMappingType], _LaunchableT + [_DispatchableT, "ExecutableProtocol", _EnvironMappingType], _LaunchableT ] -_LaunchConfigType: TypeAlias = "_LauncherAdapter[ExecutableLike, _EnvironMappingType]" +_LaunchConfigType: TypeAlias = ( + "_LauncherAdapter[ExecutableProtocol, _EnvironMappingType]" +) _UnkownType: TypeAlias = t.NoReturn @@ -82,7 +84,7 @@ def dispatch( args: None = ..., *, with_format: _FormatterType[_DispatchableT, _LaunchableT], - to_launcher: type[LauncherLike[_LaunchableT]], + to_launcher: type[LauncherProtocol[_LaunchableT]], allow_overwrite: bool = ..., ) -> t.Callable[[type[_DispatchableT]], type[_DispatchableT]]: ... @t.overload @@ -91,7 +93,7 @@ def dispatch( args: type[_DispatchableT], *, with_format: _FormatterType[_DispatchableT, _LaunchableT], - to_launcher: type[LauncherLike[_LaunchableT]], + to_launcher: type[LauncherProtocol[_LaunchableT]], allow_overwrite: bool = ..., ) -> None: ... def dispatch( @@ -99,7 +101,7 @@ def dispatch( args: type[_DispatchableT] | None = None, *, with_format: _FormatterType[_DispatchableT, _LaunchableT], - to_launcher: type[LauncherLike[_LaunchableT]], + to_launcher: type[LauncherProtocol[_LaunchableT]], allow_overwrite: bool = False, ) -> t.Callable[[type[_DispatchableT]], type[_DispatchableT]] | None: """A type safe way to add a mapping of settings builder to launcher to @@ -150,7 +152,7 @@ def get_dispatch( # where `_LaunchableT` is unbound! # # This is safe to do if all entries in the mapping were added using a - # type safe method (e.g. `Dispatcher.dispatch`), but if a user were to + # type safe method (e.g. `Dispatcher.dispatch`), but if a user were to # supply a custom dispatch registry or otherwise modify the registry # this is not necessarily 100% type safe!! return dispatch_ @@ -160,9 +162,9 @@ def get_dispatch( @dataclasses.dataclass(frozen=True) class _DispatchRegistration(t.Generic[_DispatchableT, _LaunchableT]): formatter: _FormatterType[_DispatchableT, _LaunchableT] - launcher_type: type[LauncherLike[_LaunchableT]] + launcher_type: type[LauncherProtocol[_LaunchableT]] - def _is_compatible_launcher(self, launcher: LauncherLike[t.Any]) -> bool: + def _is_compatible_launcher(self, launcher: LauncherProtocol[t.Any]) -> bool: # Disabling because we want to match the the type of the dispatch # *exactly* as specified by the user # pylint: disable-next=unidiomatic-typecheck @@ -175,7 +177,7 @@ def create_new_launcher_configuration( return self.create_adapter_from_launcher(launcher, with_settings) def create_adapter_from_launcher( - self, launcher: LauncherLike[_LaunchableT], settings: _DispatchableT + self, launcher: LauncherProtocol[_LaunchableT], settings: _DispatchableT ) -> _LaunchConfigType: if not self._is_compatible_launcher(launcher): raise TypeError( @@ -184,7 +186,7 @@ def create_adapter_from_launcher( f"exactly `{self.launcher_type}`" ) - def format_(exe: ExecutableLike, env: _EnvironMappingType) -> _LaunchableT: + def format_(exe: ExecutableProtocol, env: _EnvironMappingType) -> _LaunchableT: return self.formatter(settings, exe, env) return _LauncherAdapter(launcher, format_) @@ -192,7 +194,7 @@ def format_(exe: ExecutableLike, env: _EnvironMappingType) -> _LaunchableT: def configure_first_compatible_launcher( self, with_settings: _DispatchableT, - from_available_launchers: t.Iterable[LauncherLike[t.Any]], + from_available_launchers: t.Iterable[LauncherProtocol[t.Any]], ) -> _LaunchConfigType: launcher = helpers.first(self._is_compatible_launcher, from_available_launchers) if launcher is None: @@ -207,7 +209,7 @@ def configure_first_compatible_launcher( class _LauncherAdapter(t.Generic[Unpack[_Ts]]): def __init__( self, - launcher: LauncherLike[_LaunchableT], + launcher: LauncherProtocol[_LaunchableT], map_: t.Callable[[Unpack[_Ts]], _LaunchableT], ) -> None: # NOTE: We need to cast off the `_LaunchableT` -> `Any` in the @@ -215,7 +217,7 @@ def __init__( # this class. If possible, this type should not be exposed to # users of this class! self._adapt: t.Callable[[Unpack[_Ts]], t.Any] = map_ - self._adapted_launcher: LauncherLike[t.Any] = launcher + self._adapted_launcher: LauncherProtocol[t.Any] = launcher def start(self, *args: Unpack[_Ts]) -> LaunchedJobID: payload = self._adapt(*args) @@ -238,11 +240,11 @@ def create_job_id() -> LaunchedJobID: return LaunchedJobID(str(uuid.uuid4())) -class ExecutableLike(t.Protocol): +class ExecutableProtocol(t.Protocol): def as_program_arguments(self) -> t.Sequence[str]: ... -class LauncherLike(t.Protocol[_T_contra]): +class LauncherProtocol(t.Protocol[_T_contra]): def start(self, launchable: _T_contra) -> LaunchedJobID: ... @classmethod def create(cls, exp: Experiment, /) -> Self: ... @@ -252,7 +254,7 @@ def make_shell_format_fn( run_command: str | None, ) -> _FormatterType[LaunchArgBuilder, t.Sequence[str]]: def impl( - args: LaunchArgBuilder, exe: ExecutableLike, _env: _EnvironMappingType + args: LaunchArgBuilder, exe: ExecutableProtocol, _env: _EnvironMappingType ) -> t.Sequence[str]: return ( ( diff --git a/smartsim/settings/launchSettings.py b/smartsim/settings/launchSettings.py index dec6034d8..0990e2e82 100644 --- a/smartsim/settings/launchSettings.py +++ b/smartsim/settings/launchSettings.py @@ -63,7 +63,7 @@ def __init__( @property def launcher(self) -> str: - """Return the launcher name.""" + """The launcher type""" return self._launcher.value @property @@ -89,7 +89,13 @@ def env_vars(self, value: dict[str, str | None]) -> None: self._env_vars = copy.deepcopy(value) def _get_arg_builder(self, launch_args: StringArgument | None) -> LaunchArgBuilder: - """Map the Launcher to the LaunchArgBuilder""" + """Map the Launcher to the LaunchArgBuilder. This method should only be + called once during construction. + + :param launch_args: A mapping of argument to be used to initialize the + argument builder. + :returns: The appropriate argument builder for the settings instance. + """ if self._launcher == LauncherType.Slurm: return SlurmArgBuilder(launch_args) elif self._launcher == LauncherType.Mpiexec: diff --git a/tests/temp_tests/test_settings/conftest.py b/tests/temp_tests/test_settings/conftest.py index 72061264f..1368006a9 100644 --- a/tests/temp_tests/test_settings/conftest.py +++ b/tests/temp_tests/test_settings/conftest.py @@ -31,12 +31,12 @@ @pytest.fixture -def echo_executable_like(): - class _ExeLike(dispatch.ExecutableLike): +def mock_echo_executable(): + class _MockExe(dispatch.ExecutableProtocol): def as_program_arguments(self): return ("echo", "hello", "world") - yield _ExeLike() + yield _MockExe() @pytest.fixture @@ -50,8 +50,8 @@ def launcher_str(self): @pytest.fixture -def launcher_like(): - class _LuancherLike(dispatch.LauncherLike): +def mock_launcher(): + class _MockLauncher(dispatch.LauncherProtocol): def start(self, launchable): return dispatch.create_job_id() @@ -59,4 +59,4 @@ def start(self, launchable): def create(cls, exp): return cls() - yield _LuancherLike() + yield _MockLauncher() diff --git a/tests/temp_tests/test_settings/test_alpsLauncher.py b/tests/temp_tests/test_settings/test_alpsLauncher.py index cb6e829eb..367b30c7f 100644 --- a/tests/temp_tests/test_settings/test_alpsLauncher.py +++ b/tests/temp_tests/test_settings/test_alpsLauncher.py @@ -182,6 +182,6 @@ def test_invalid_exclude_hostlist_format(): ), ), ) -def test_formatting_launch_args(echo_executable_like, args, expected): - cmd = _as_aprun_command(AprunArgBuilder(args), echo_executable_like, {}) +def test_formatting_launch_args(mock_echo_executable, args, expected): + cmd = _as_aprun_command(AprunArgBuilder(args), mock_echo_executable, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_dispatch.py b/tests/temp_tests/test_settings/test_dispatch.py index 78c44ad54..673c8998a 100644 --- a/tests/temp_tests/test_settings/test_dispatch.py +++ b/tests/temp_tests/test_settings/test_dispatch.py @@ -44,36 +44,36 @@ def format_fn(args, exe, env): @pytest.fixture -def expected_dispatch_registry(launcher_like, settings_builder): +def expected_dispatch_registry(mock_launcher, settings_builder): yield { type(settings_builder): dispatch._DispatchRegistration( - format_fn, type(launcher_like) + format_fn, type(mock_launcher) ) } def test_declaritive_form_dispatch_declaration( - launcher_like, settings_builder, expected_dispatch_registry + mock_launcher, settings_builder, expected_dispatch_registry ): d = dispatch.Dispatcher() assert type(settings_builder) == d.dispatch( - with_format=format_fn, to_launcher=type(launcher_like) + with_format=format_fn, to_launcher=type(mock_launcher) )(type(settings_builder)) assert d._dispatch_registry == expected_dispatch_registry def test_imperative_form_dispatch_declaration( - launcher_like, settings_builder, expected_dispatch_registry + mock_launcher, settings_builder, expected_dispatch_registry ): d = dispatch.Dispatcher() assert None == d.dispatch( - type(settings_builder), to_launcher=type(launcher_like), with_format=format_fn + type(settings_builder), to_launcher=type(mock_launcher), with_format=format_fn ) assert d._dispatch_registry == expected_dispatch_registry def test_dispatchers_from_same_registry_do_not_cross_polute( - launcher_like, settings_builder, expected_dispatch_registry + mock_launcher, settings_builder, expected_dispatch_registry ): some_starting_registry = {} d1 = dispatch.Dispatcher(dispatch_registry=some_starting_registry) @@ -86,14 +86,14 @@ def test_dispatchers_from_same_registry_do_not_cross_polute( ) d2.dispatch( - type(settings_builder), with_format=format_fn, to_launcher=type(launcher_like) + type(settings_builder), with_format=format_fn, to_launcher=type(mock_launcher) ) assert d1._dispatch_registry == {} assert d2._dispatch_registry == expected_dispatch_registry def test_copied_dispatchers_do_not_cross_pollute( - launcher_like, settings_builder, expected_dispatch_registry + mock_launcher, settings_builder, expected_dispatch_registry ): some_starting_registry = {} d1 = dispatch.Dispatcher(dispatch_registry=some_starting_registry) @@ -106,7 +106,7 @@ def test_copied_dispatchers_do_not_cross_pollute( ) d2.dispatch( - type(settings_builder), to_launcher=type(launcher_like), with_format=format_fn + type(settings_builder), to_launcher=type(mock_launcher), with_format=format_fn ) assert d1._dispatch_registry == {} assert d2._dispatch_registry == expected_dispatch_registry @@ -144,13 +144,13 @@ def test_copied_dispatchers_do_not_cross_pollute( def test_dispatch_overwriting( add_dispatch, expected_ctx, - launcher_like, + mock_launcher, settings_builder, expected_dispatch_registry, ): d = dispatch.Dispatcher(dispatch_registry=expected_dispatch_registry) with expected_ctx: - add_dispatch(d, type(settings_builder), type(launcher_like)) + add_dispatch(d, type(settings_builder), type(mock_launcher)) @pytest.mark.parametrize( @@ -161,11 +161,11 @@ def test_dispatch_overwriting( ), ) def test_dispatch_can_retrieve_dispatch_info_from_dispatch_registry( - expected_dispatch_registry, launcher_like, settings_builder, type_or_instance + expected_dispatch_registry, mock_launcher, settings_builder, type_or_instance ): d = dispatch.Dispatcher(dispatch_registry=expected_dispatch_registry) assert dispatch._DispatchRegistration( - format_fn, type(launcher_like) + format_fn, type(mock_launcher) ) == d.get_dispatch(type_or_instance(settings_builder)) @@ -209,12 +209,12 @@ def create(cls, exp): "cls, ctx", ( pytest.param( - dispatch.LauncherLike, + dispatch.LauncherProtocol, pytest.raises(TypeError, match="Cannot dispatch to protocol"), id="Cannot dispatch to protocol class", ), pytest.param( - "launcher_like", + "mock_launcher", contextlib.nullcontext(None), id="Can dispatch to protocol implementation", ), @@ -244,7 +244,7 @@ def test_register_dispatch_to_launcher_types(request, cls, ctx): @dataclasses.dataclass -class BufferWriterLauncher(dispatch.LauncherLike[list[str]]): +class BufferWriterLauncher(dispatch.LauncherProtocol[list[str]]): buf: io.StringIO @classmethod @@ -313,7 +313,7 @@ def test_launcher_adapter_correctly_adapts_input_to_launcher(input_, map_, expec id="Errors if launcher types are disparate", ), pytest.param( - "launcher_like", + "mock_launcher", pytest.raises( TypeError, match="^Cannot create launcher adapter.*expected launcher of type .+$", @@ -344,10 +344,10 @@ def test_dispatch_registration_can_configure_adapter_for_existing_launcher_insta ), pytest.param( ( - "launcher_like", - "launcher_like", + "mock_launcher", + "mock_launcher", BufferWriterLauncher(io.StringIO()), - "launcher_like", + "mock_launcher", ), contextlib.nullcontext(None), id="Correctly ignores incompatible launchers instances", @@ -362,9 +362,9 @@ def test_dispatch_registration_can_configure_adapter_for_existing_launcher_insta ), pytest.param( ( - "launcher_like", + "mock_launcher", BufferWriterLauncherSubclass(io.StringIO), - "launcher_like", + "mock_launcher", ), pytest.raises( errors.LauncherNotFoundError, diff --git a/tests/temp_tests/test_settings/test_dragonLauncher.py b/tests/temp_tests/test_settings/test_dragonLauncher.py index cc57e329f..c71b1d548 100644 --- a/tests/temp_tests/test_settings/test_dragonLauncher.py +++ b/tests/temp_tests/test_settings/test_dragonLauncher.py @@ -37,14 +37,14 @@ def test_dragon_class_methods(function, value, flag, result): @pytest.mark.parametrize("nodes", (NOT_SET, 20, 40)) @pytest.mark.parametrize("tasks_per_node", (NOT_SET, 1, 20)) def test_formatting_launch_args_into_request( - echo_executable_like, nodes, tasks_per_node + mock_echo_executable, nodes, tasks_per_node ): args = DragonArgBuilder({}) if nodes is not NOT_SET: args.set_nodes(nodes) if tasks_per_node is not NOT_SET: args.set_tasks_per_node(tasks_per_node) - req = _as_run_request_view(args, echo_executable_like, {}) + req = _as_run_request_view(args, mock_echo_executable, {}) args = { k: v diff --git a/tests/temp_tests/test_settings/test_localLauncher.py b/tests/temp_tests/test_settings/test_localLauncher.py index 07adb231b..b3cb4108f 100644 --- a/tests/temp_tests/test_settings/test_localLauncher.py +++ b/tests/temp_tests/test_settings/test_localLauncher.py @@ -114,6 +114,6 @@ def test_format_env_vars(): assert localLauncher.format_env_vars() == ["A=a", "B=", "C=", "D=12"] -def test_formatting_returns_original_exe(echo_executable_like): - cmd = _as_local_command(LocalArgBuilder({}), echo_executable_like, {}) +def test_formatting_returns_original_exe(mock_echo_executable): + cmd = _as_local_command(LocalArgBuilder({}), mock_echo_executable, {}) assert tuple(cmd) == ("echo", "hello", "world") diff --git a/tests/temp_tests/test_settings/test_lsfLauncher.py b/tests/temp_tests/test_settings/test_lsfLauncher.py index 9d276c0e7..636d2896f 100644 --- a/tests/temp_tests/test_settings/test_lsfLauncher.py +++ b/tests/temp_tests/test_settings/test_lsfLauncher.py @@ -91,6 +91,6 @@ def test_launch_args(): ), ), ) -def test_formatting_launch_args(echo_executable_like, args, expected): - cmd = _as_jsrun_command(JsrunArgBuilder(args), echo_executable_like, {}) +def test_formatting_launch_args(mock_echo_executable, args, expected): + cmd = _as_jsrun_command(JsrunArgBuilder(args), mock_echo_executable, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_mpiLauncher.py b/tests/temp_tests/test_settings/test_mpiLauncher.py index 16947590c..23df78c92 100644 --- a/tests/temp_tests/test_settings/test_mpiLauncher.py +++ b/tests/temp_tests/test_settings/test_mpiLauncher.py @@ -253,6 +253,6 @@ def test_invalid_hostlist_format(launcher): ), ), ) -def test_formatting_launch_args(echo_executable_like, cls, fmt, cmd, args, expected): - fmt_cmd = fmt(cls(args), echo_executable_like, {}) +def test_formatting_launch_args(mock_echo_executable, cls, fmt, cmd, args, expected): + fmt_cmd = fmt(cls(args), mock_echo_executable, {}) assert tuple(fmt_cmd) == (cmd,) + expected diff --git a/tests/temp_tests/test_settings/test_palsLauncher.py b/tests/temp_tests/test_settings/test_palsLauncher.py index 47fa64713..18d85a778 100644 --- a/tests/temp_tests/test_settings/test_palsLauncher.py +++ b/tests/temp_tests/test_settings/test_palsLauncher.py @@ -105,6 +105,6 @@ def test_invalid_hostlist_format(): ), ), ) -def test_formatting_launch_args(echo_executable_like, args, expected): - cmd = _as_pals_command(PalsMpiexecArgBuilder(args), echo_executable_like, {}) +def test_formatting_launch_args(mock_echo_executable, args, expected): + cmd = _as_pals_command(PalsMpiexecArgBuilder(args), mock_echo_executable, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_slurmLauncher.py b/tests/temp_tests/test_settings/test_slurmLauncher.py index 5935ec690..e3b73aee7 100644 --- a/tests/temp_tests/test_settings/test_slurmLauncher.py +++ b/tests/temp_tests/test_settings/test_slurmLauncher.py @@ -288,6 +288,6 @@ def test_set_het_groups(monkeypatch): ), ), ) -def test_formatting_launch_args(echo_executable_like, args, expected): - cmd = _as_srun_command(SlurmArgBuilder(args), echo_executable_like, {}) +def test_formatting_launch_args(mock_echo_executable, args, expected): + cmd = _as_srun_command(SlurmArgBuilder(args), mock_echo_executable, {}) assert tuple(cmd) == expected From 0a1ebba97ac56709c089aa862e7e2afdc8eb2f28 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Sun, 21 Jul 2024 05:50:35 -0500 Subject: [PATCH 27/34] Add tests for `Experiment.start_jobs` --- smartsim/experiment.py | 17 ++- smartsim/settings/dispatch.py | 6 +- tests/test_experiment.py | 188 ++++++++++++++++++++++++++++++++++ 3 files changed, 197 insertions(+), 14 deletions(-) create mode 100644 tests/test_experiment.py diff --git a/smartsim/experiment.py b/smartsim/experiment.py index 6f1a5744a..8915a620d 100644 --- a/smartsim/experiment.py +++ b/smartsim/experiment.py @@ -177,11 +177,12 @@ def __init__(self, name: str, exp_path: str | None = None): experiment""" def start_jobs( - self, *jobs: Job, dispatcher: Dispatcher = DEFAULT_DISPATCHER + self, job: Job, *jobs: Job, dispatcher: Dispatcher = DEFAULT_DISPATCHER ) -> tuple[LaunchedJobID, ...]: """Execute a collection of `Job` instances. - :param jobs: The collection of jobs instances to start + :param job: The job instance to start + :param jobs: A collection of other job instances to start :param dispatcher: The dispatcher that should be used to determine how to start a job based on its settings. If not specified it will default to a dispatcher pre-configured by SmartSim. @@ -190,12 +191,6 @@ def start_jobs( particular execution of the job. """ - if not jobs: - raise TypeError( - f"{type(self).__name__}.start_jobs() missing at least 1 required " - "positional argument of type `Job`" - ) - def _start(job: Job) -> LaunchedJobID: args = job.launch_settings.launch_args env = job.launch_settings.env_vars @@ -222,7 +217,7 @@ def _start(job: Job) -> LaunchedJobID: self._active_launchers.add(launch_config._adapted_launcher) return launch_config.start(exe, env) - return tuple(map(_start, jobs)) + return _start(job), *map(_start, jobs) @_contextualize def start( @@ -592,8 +587,8 @@ def _launch_summary(self, manifest: Manifest) -> None: === Launch Summary === Experiment: {self.name} Experiment Path: {self.exp_path} - Launchers: - {textwrap.indent(" - ", launcher_list)} + Launcher(s): + {textwrap.indent(" - ", launcher_list) if launcher_list else " "} """) if manifest.applications: diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index aceead700..9cb6b7f37 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -245,7 +245,7 @@ def as_program_arguments(self) -> t.Sequence[str]: ... class LauncherProtocol(t.Protocol[_T_contra]): - def start(self, launchable: _T_contra) -> LaunchedJobID: ... + def start(self, launchable: _T_contra, /) -> LaunchedJobID: ... @classmethod def create(cls, exp: Experiment, /) -> Self: ... @@ -276,9 +276,9 @@ class ShellLauncher: def __init__(self) -> None: self._launched: dict[LaunchedJobID, sp.Popen[bytes]] = {} - def start(self, launchable: t.Sequence[str]) -> LaunchedJobID: + def start(self, command: t.Sequence[str]) -> LaunchedJobID: id_ = create_job_id() - exe, *rest = launchable + exe, *rest = command # pylint: disable-next=consider-using-with self._launched[id_] = sp.Popen((helpers.expand_exe_path(exe), *rest)) return id_ diff --git a/tests/test_experiment.py b/tests/test_experiment.py new file mode 100644 index 000000000..1cd1ee5c3 --- /dev/null +++ b/tests/test_experiment.py @@ -0,0 +1,188 @@ +# BSD 2-Clause License +# +# Copyright (c) 2021-2024, Hewlett Packard Enterprise +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import dataclasses +import itertools +import tempfile +import typing as t +import uuid +import weakref + +import pytest + +from smartsim.entity import _mock, entity +from smartsim.experiment import Experiment +from smartsim.launchable import job +from smartsim.settings import dispatch, launchSettings +from smartsim.settings.builders import launchArgBuilder + +pytestmark = pytest.mark.group_a + + +@pytest.fixture +def experiment(test_dir): + yield Experiment(f"test-exp-{uuid.uuid4()}", test_dir) + + +@pytest.fixture +def dispatcher(): + d = dispatch.Dispatcher() + to_record = lambda *a: LaunchRecord(*a) + d.dispatch(MockLaunchArgs, with_format=to_record, to_launcher=NoOpRecordLauncher) + yield d + + +@pytest.fixture +def job_maker(monkeypatch): + def iter_jobs(): + for i in itertools.count(): + settings = launchSettings.LaunchSettings("local") + monkeypatch.setattr(settings, "_arg_builder", MockLaunchArgs(i)) + yield job.Job(EchoHelloWorldEntity(), settings) + + jobs = iter_jobs() + return lambda: next(jobs) + + +@dataclasses.dataclass(frozen=True) +class NoOpRecordLauncher(dispatch.LauncherProtocol): + created_by_experiment: Experiment + launched_order: list[LaunchRecord] = dataclasses.field(default_factory=list) + ids_to_launched: dict[dispatch.LaunchedJobID, LaunchRecord] = dataclasses.field( + default_factory=dict + ) + + __hash__ = object.__hash__ + + @classmethod + def create(cls, exp): + return cls(exp) + + def start(self, record: LaunchRecord): + id_ = dispatch.create_job_id() + self.launched_order.append(record) + self.ids_to_launched[id_] = record + return id_ + + +@dataclasses.dataclass(frozen=True) +class LaunchRecord: + launch_args: launchArgBuilder.LaunchArgBuilder + entity: entity.SmartSimEntity + env: t.Mapping[str, str | None] + + @classmethod + def from_job(cls, job): + args = job._launch_settings.launch_args + entity = job._entity + env = job._launch_settings.env_vars + return cls(args, entity, env) + + +class MockLaunchArgs(launchArgBuilder.LaunchArgBuilder): + def __init__(self, count): + super().__init__({}) + self.count = count + + def __eq__(self, other): + if type(self) is not type(other): + return NotImplemented + return other.count == self.count + + def launcher_str(self): + return "test-launch-args" + + def set(self, arg, val): ... + + +class EchoHelloWorldEntity(entity.SmartSimEntity): + def __init__(self): + path = tempfile.TemporaryDirectory() + self._finalizer = weakref.finalize(self, path.cleanup) + super().__init__("test-entity", path, _mock.Mock()) + + def __eq__(self, other): + if type(self) is not type(other): + return NotImplemented + return self.as_program_arguments() == other.as_program_arguments() + + def as_program_arguments(self): + return ("echo", "Hello", "World!") + + +def test_start_raises_if_no_args_supplied(experiment): + with pytest.raises(TypeError, match="missing 1 required positional argument"): + experiment.start_jobs() + + +# fmt: off +@pytest.mark.parametrize( + "num_jobs", [pytest.param(i, id=f"{i} job(s)") for i in (1, 2, 3, 5, 10, 100, 1_000)] +) +@pytest.mark.parametrize( + "make_jobs", ( + pytest.param(lambda maker, n: tuple(maker() for _ in range(n)), id="many job instances"), + pytest.param(lambda maker, n: (maker(),) * n , id="same job instance many times"), + ), +) +# fmt: on +def test_start_can_launch_jobs(experiment, job_maker, dispatcher, make_jobs, num_jobs): + jobs = make_jobs(job_maker, num_jobs) + assert len(experiment._active_launchers) == 0, "Initialized w/ launchers" + launched_ids = experiment.start_jobs(*jobs, dispatcher=dispatcher) + assert len(experiment._active_launchers) == 1, "Unexpected number of launchers" + (launcher,) = experiment._active_launchers + assert isinstance(launcher, NoOpRecordLauncher), "Unexpected launcher type" + assert launcher.created_by_experiment is experiment, "Not created by experiment" + assert ( + len(jobs) == len(launcher.launched_order) == len(launched_ids) == num_jobs + ), "Inconsistent number of jobs/launched jobs/launched ids/expected number of jobs" + expected_launched = [LaunchRecord.from_job(job) for job in jobs] + assert expected_launched == list(launcher.launched_order), "Unexpected launch order" + expected_id_map = dict(zip(launched_ids, expected_launched)) + assert expected_id_map == launcher.ids_to_launched, "IDs returned in wrong order" + + +@pytest.mark.parametrize( + "num_starts", + [pytest.param(i, id=f"{i} start(s)") for i in (1, 2, 3, 5, 10, 100, 1_000)], +) +def test_start_can_start_a_job_multiple_times_accross_multiple_calls( + experiment, job_maker, dispatcher, num_starts +): + assert len(experiment._active_launchers) == 0, "Initialized w/ launchers" + job = job_maker() + ids_to_launches = { + experiment.start_jobs(job, dispatcher=dispatcher)[0]: LaunchRecord.from_job(job) + for _ in range(num_starts) + } + assert len(experiment._active_launchers) == 1, "Did not reuse the launcher" + (launcher,) = experiment._active_launchers + assert isinstance(launcher, NoOpRecordLauncher), "Unexpected launcher type" + assert len(launcher.launched_order) == num_starts, "Unexpected number launches" + assert ids_to_launches == launcher.ids_to_launched, "Job was not re-launched" From b729e9152cfc0c7435c606f151337e2525d86e51 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Mon, 22 Jul 2024 12:59:43 -0500 Subject: [PATCH 28/34] Rename `Builder` -> `Arguments` --- .../_core/launcher/dragon/dragonLauncher.py | 8 ++- .../{builders => arguments}/__init__.py | 6 +- .../{builders => arguments}/batch/__init__.py | 12 ++-- .../{builders => arguments}/batch/lsf.py | 4 +- .../{builders => arguments}/batch/pbs.py | 4 +- .../{builders => arguments}/batch/slurm.py | 4 +- .../batchArguments.py} | 2 +- .../settings/arguments/launch/__init__.py | 19 ++++++ .../{builders => arguments}/launch/alps.py | 4 +- .../{builders => arguments}/launch/dragon.py | 4 +- .../{builders => arguments}/launch/local.py | 4 +- .../{builders => arguments}/launch/lsf.py | 4 +- .../{builders => arguments}/launch/mpi.py | 10 ++-- .../{builders => arguments}/launch/pals.py | 4 +- .../{builders => arguments}/launch/slurm.py | 4 +- .../launchArguments.py} | 7 +-- smartsim/settings/batchSettings.py | 34 ++++++----- smartsim/settings/builders/launch/__init__.py | 19 ------ smartsim/settings/dispatch.py | 21 ++++--- smartsim/settings/launchSettings.py | 60 ++++++++++--------- tests/temp_tests/test_settings/conftest.py | 10 ++-- .../test_settings/test_alpsLauncher.py | 15 +++-- .../temp_tests/test_settings/test_dispatch.py | 48 +++++++-------- .../test_settings/test_dragonLauncher.py | 6 +- .../test_settings/test_localLauncher.py | 9 ++- .../test_settings/test_lsfLauncher.py | 13 ++-- .../test_settings/test_mpiLauncher.py | 26 ++++---- .../test_settings/test_palsLauncher.py | 8 +-- .../test_settings/test_pbsScheduler.py | 4 +- .../test_settings/test_slurmLauncher.py | 13 ++-- .../test_settings/test_slurmScheduler.py | 4 +- tests/test_experiment.py | 8 +-- 32 files changed, 211 insertions(+), 187 deletions(-) rename smartsim/settings/{builders => arguments}/__init__.py (90%) rename smartsim/settings/{builders => arguments}/batch/__init__.py (87%) rename smartsim/settings/{builders => arguments}/batch/lsf.py (98%) rename smartsim/settings/{builders => arguments}/batch/pbs.py (98%) rename smartsim/settings/{builders => arguments}/batch/slurm.py (98%) rename smartsim/settings/{builders/batchArgBuilder.py => arguments/batchArguments.py} (99%) create mode 100644 smartsim/settings/arguments/launch/__init__.py rename smartsim/settings/{builders => arguments}/launch/alps.py (98%) rename smartsim/settings/{builders => arguments}/launch/dragon.py (96%) rename smartsim/settings/{builders => arguments}/launch/local.py (97%) rename smartsim/settings/{builders => arguments}/launch/lsf.py (97%) rename smartsim/settings/{builders => arguments}/launch/mpi.py (96%) rename smartsim/settings/{builders => arguments}/launch/pals.py (98%) rename smartsim/settings/{builders => arguments}/launch/slurm.py (99%) rename smartsim/settings/{builders/launchArgBuilder.py => arguments/launchArguments.py} (95%) delete mode 100644 smartsim/settings/builders/launch/__init__.py diff --git a/smartsim/_core/launcher/dragon/dragonLauncher.py b/smartsim/_core/launcher/dragon/dragonLauncher.py index 078a89278..2a7182eea 100644 --- a/smartsim/_core/launcher/dragon/dragonLauncher.py +++ b/smartsim/_core/launcher/dragon/dragonLauncher.py @@ -340,12 +340,12 @@ def _assert_schema_type(obj: object, typ: t.Type[_SchemaT], /) -> _SchemaT: # TODO: Remove this registry and move back to builder file after fixing # circular import caused by `DragonLauncher.supported_rs` # ----------------------------------------------------------------------------- -from smartsim.settings.builders.launch.dragon import DragonArgBuilder +from smartsim.settings.arguments.launch.dragon import DragonLaunchArguments from smartsim.settings.dispatch import ExecutableProtocol, dispatch def _as_run_request_view( - run_req_args: DragonArgBuilder, + run_req_args: DragonLaunchArguments, exe: ExecutableProtocol, env: t.Mapping[str, str | None], ) -> DragonRunRequestView: @@ -369,5 +369,7 @@ def _as_run_request_view( ) -dispatch(DragonArgBuilder, with_format=_as_run_request_view, to_launcher=DragonLauncher) +dispatch( + DragonLaunchArguments, with_format=_as_run_request_view, to_launcher=DragonLauncher +) # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< diff --git a/smartsim/settings/builders/__init__.py b/smartsim/settings/arguments/__init__.py similarity index 90% rename from smartsim/settings/builders/__init__.py rename to smartsim/settings/arguments/__init__.py index 9cfdd5f9c..cd216526c 100644 --- a/smartsim/settings/builders/__init__.py +++ b/smartsim/settings/arguments/__init__.py @@ -24,7 +24,7 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -from .batchArgBuilder import BatchArgBuilder -from .launchArgBuilder import LaunchArgBuilder +from .batchArguments import BatchArguments +from .launchArguments import LaunchArguments -__all__ = ["LaunchArgBuilder", "BatchArgBuilder"] +__all__ = ["LaunchArguments", "BatchArguments"] diff --git a/smartsim/settings/builders/batch/__init__.py b/smartsim/settings/arguments/batch/__init__.py similarity index 87% rename from smartsim/settings/builders/batch/__init__.py rename to smartsim/settings/arguments/batch/__init__.py index 41dcbbfc2..e6dc055ea 100644 --- a/smartsim/settings/builders/batch/__init__.py +++ b/smartsim/settings/arguments/batch/__init__.py @@ -24,12 +24,12 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -from .lsf import BsubBatchArgBuilder -from .pbs import QsubBatchArgBuilder -from .slurm import SlurmBatchArgBuilder +from .lsf import BsubBatchArguments +from .pbs import QsubBatchArguments +from .slurm import SlurmBatchArguments __all__ = [ - "BsubBatchArgBuilder", - "QsubBatchArgBuilder", - "SlurmBatchArgBuilder", + "BsubBatchArguments", + "QsubBatchArguments", + "SlurmBatchArguments", ] diff --git a/smartsim/settings/builders/batch/lsf.py b/smartsim/settings/arguments/batch/lsf.py similarity index 98% rename from smartsim/settings/builders/batch/lsf.py rename to smartsim/settings/arguments/batch/lsf.py index 4bb7bbd27..4f6e80a70 100644 --- a/smartsim/settings/builders/batch/lsf.py +++ b/smartsim/settings/arguments/batch/lsf.py @@ -32,12 +32,12 @@ from ...batchCommand import SchedulerType from ...common import StringArgument -from ..batchArgBuilder import BatchArgBuilder +from ..batchArguments import BatchArguments logger = get_logger(__name__) -class BsubBatchArgBuilder(BatchArgBuilder): +class BsubBatchArguments(BatchArguments): def scheduler_str(self) -> str: """Get the string representation of the scheduler""" return SchedulerType.Lsf.value diff --git a/smartsim/settings/builders/batch/pbs.py b/smartsim/settings/arguments/batch/pbs.py similarity index 98% rename from smartsim/settings/builders/batch/pbs.py rename to smartsim/settings/arguments/batch/pbs.py index d04b4beba..d67f1be7b 100644 --- a/smartsim/settings/builders/batch/pbs.py +++ b/smartsim/settings/arguments/batch/pbs.py @@ -34,12 +34,12 @@ from ....error import SSConfigError from ...batchCommand import SchedulerType from ...common import StringArgument -from ..batchArgBuilder import BatchArgBuilder +from ..batchArguments import BatchArguments logger = get_logger(__name__) -class QsubBatchArgBuilder(BatchArgBuilder): +class QsubBatchArguments(BatchArguments): def scheduler_str(self) -> str: """Get the string representation of the scheduler""" return SchedulerType.Pbs.value diff --git a/smartsim/settings/builders/batch/slurm.py b/smartsim/settings/arguments/batch/slurm.py similarity index 98% rename from smartsim/settings/builders/batch/slurm.py rename to smartsim/settings/arguments/batch/slurm.py index 5a03f5acd..eca26176d 100644 --- a/smartsim/settings/builders/batch/slurm.py +++ b/smartsim/settings/arguments/batch/slurm.py @@ -33,12 +33,12 @@ from ...batchCommand import SchedulerType from ...common import StringArgument -from ..batchArgBuilder import BatchArgBuilder +from ..batchArguments import BatchArguments logger = get_logger(__name__) -class SlurmBatchArgBuilder(BatchArgBuilder): +class SlurmBatchArguments(BatchArguments): def scheduler_str(self) -> str: """Get the string representation of the scheduler""" return SchedulerType.Slurm.value diff --git a/smartsim/settings/builders/batchArgBuilder.py b/smartsim/settings/arguments/batchArguments.py similarity index 99% rename from smartsim/settings/builders/batchArgBuilder.py rename to smartsim/settings/arguments/batchArguments.py index ad466f254..a85148697 100644 --- a/smartsim/settings/builders/batchArgBuilder.py +++ b/smartsim/settings/arguments/batchArguments.py @@ -37,7 +37,7 @@ logger = get_logger(__name__) -class BatchArgBuilder(ABC): +class BatchArguments(ABC): """Abstract base class that defines all generic scheduler argument methods that are not supported. It is the responsibility of child classes for each launcher to translate diff --git a/smartsim/settings/arguments/launch/__init__.py b/smartsim/settings/arguments/launch/__init__.py new file mode 100644 index 000000000..30502394b --- /dev/null +++ b/smartsim/settings/arguments/launch/__init__.py @@ -0,0 +1,19 @@ +from .alps import AprunLaunchArguments +from .dragon import DragonLaunchArguments +from .local import LocalLaunchArguments +from .lsf import JsrunLaunchArguments +from .mpi import MpiexecLaunchArguments, MpirunLaunchArguments, OrterunLaunchArguments +from .pals import PalsMpiexecLaunchArguments +from .slurm import SlurmLaunchArguments + +__all__ = [ + "AprunLaunchArguments", + "DragonLaunchArguments", + "LocalLaunchArguments", + "JsrunLaunchArguments", + "MpiLaunchArguments", + "MpiexecLaunchArguments", + "OrteLaunchArguments", + "PalsMpiexecLaunchArguments", + "SlurmLaunchArguments", +] diff --git a/smartsim/settings/builders/launch/alps.py b/smartsim/settings/arguments/launch/alps.py similarity index 98% rename from smartsim/settings/builders/launch/alps.py rename to smartsim/settings/arguments/launch/alps.py index 5826a2de4..e92bc7b85 100644 --- a/smartsim/settings/builders/launch/alps.py +++ b/smartsim/settings/arguments/launch/alps.py @@ -33,14 +33,14 @@ from ...common import set_check_input from ...launchCommand import LauncherType -from ..launchArgBuilder import LaunchArgBuilder +from ..launchArguments import LaunchArguments logger = get_logger(__name__) _as_aprun_command = make_shell_format_fn(run_command="aprun") @dispatch(with_format=_as_aprun_command, to_launcher=ShellLauncher) -class AprunArgBuilder(LaunchArgBuilder): +class AprunLaunchArguments(LaunchArguments): def _reserved_launch_args(self) -> set[str]: """Return reserved launch arguments.""" return {"wdir"} diff --git a/smartsim/settings/builders/launch/dragon.py b/smartsim/settings/arguments/launch/dragon.py similarity index 96% rename from smartsim/settings/builders/launch/dragon.py rename to smartsim/settings/arguments/launch/dragon.py index 4ba793bf7..7a6a5aab4 100644 --- a/smartsim/settings/builders/launch/dragon.py +++ b/smartsim/settings/arguments/launch/dragon.py @@ -30,12 +30,12 @@ from ...common import set_check_input from ...launchCommand import LauncherType -from ..launchArgBuilder import LaunchArgBuilder +from ..launchArguments import LaunchArguments logger = get_logger(__name__) -class DragonArgBuilder(LaunchArgBuilder): +class DragonLaunchArguments(LaunchArguments): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Dragon.value diff --git a/smartsim/settings/builders/launch/local.py b/smartsim/settings/arguments/launch/local.py similarity index 97% rename from smartsim/settings/builders/launch/local.py rename to smartsim/settings/arguments/launch/local.py index 49ef3ad92..f89299500 100644 --- a/smartsim/settings/builders/launch/local.py +++ b/smartsim/settings/arguments/launch/local.py @@ -33,14 +33,14 @@ from ...common import StringArgument, set_check_input from ...launchCommand import LauncherType -from ..launchArgBuilder import LaunchArgBuilder +from ..launchArguments import LaunchArguments logger = get_logger(__name__) _as_local_command = make_shell_format_fn(run_command=None) @dispatch(with_format=_as_local_command, to_launcher=ShellLauncher) -class LocalArgBuilder(LaunchArgBuilder): +class LocalLaunchArguments(LaunchArguments): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Local.value diff --git a/smartsim/settings/builders/launch/lsf.py b/smartsim/settings/arguments/launch/lsf.py similarity index 97% rename from smartsim/settings/builders/launch/lsf.py rename to smartsim/settings/arguments/launch/lsf.py index 1fe8ea30b..83e5cdc94 100644 --- a/smartsim/settings/builders/launch/lsf.py +++ b/smartsim/settings/arguments/launch/lsf.py @@ -33,14 +33,14 @@ from ...common import set_check_input from ...launchCommand import LauncherType -from ..launchArgBuilder import LaunchArgBuilder +from ..launchArguments import LaunchArguments logger = get_logger(__name__) _as_jsrun_command = make_shell_format_fn(run_command="jsrun") @dispatch(with_format=_as_jsrun_command, to_launcher=ShellLauncher) -class JsrunArgBuilder(LaunchArgBuilder): +class JsrunLaunchArguments(LaunchArguments): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Lsf.value diff --git a/smartsim/settings/builders/launch/mpi.py b/smartsim/settings/arguments/launch/mpi.py similarity index 96% rename from smartsim/settings/builders/launch/mpi.py rename to smartsim/settings/arguments/launch/mpi.py index a0d947418..edd93b4ac 100644 --- a/smartsim/settings/builders/launch/mpi.py +++ b/smartsim/settings/arguments/launch/mpi.py @@ -33,7 +33,7 @@ from ...common import set_check_input from ...launchCommand import LauncherType -from ..launchArgBuilder import LaunchArgBuilder +from ..launchArguments import LaunchArguments logger = get_logger(__name__) _as_mpirun_command = make_shell_format_fn("mpirun") @@ -41,7 +41,7 @@ _as_orterun_command = make_shell_format_fn("orterun") -class _BaseMPIArgBuilder(LaunchArgBuilder): +class _BaseMPILaunchArguments(LaunchArguments): def _reserved_launch_args(self) -> set[str]: """Return reserved launch arguments.""" return {"wd", "wdir"} @@ -219,21 +219,21 @@ def set(self, key: str, value: str | None) -> None: @dispatch(with_format=_as_mpirun_command, to_launcher=ShellLauncher) -class MpiArgBuilder(_BaseMPIArgBuilder): +class MpirunLaunchArguments(_BaseMPILaunchArguments): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Mpirun.value @dispatch(with_format=_as_mpiexec_command, to_launcher=ShellLauncher) -class MpiexecArgBuilder(_BaseMPIArgBuilder): +class MpiexecLaunchArguments(_BaseMPILaunchArguments): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Mpiexec.value @dispatch(with_format=_as_orterun_command, to_launcher=ShellLauncher) -class OrteArgBuilder(_BaseMPIArgBuilder): +class OrterunLaunchArguments(_BaseMPILaunchArguments): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Orterun.value diff --git a/smartsim/settings/builders/launch/pals.py b/smartsim/settings/arguments/launch/pals.py similarity index 98% rename from smartsim/settings/builders/launch/pals.py rename to smartsim/settings/arguments/launch/pals.py index eeb438455..ed4de2131 100644 --- a/smartsim/settings/builders/launch/pals.py +++ b/smartsim/settings/arguments/launch/pals.py @@ -33,14 +33,14 @@ from ...common import set_check_input from ...launchCommand import LauncherType -from ..launchArgBuilder import LaunchArgBuilder +from ..launchArguments import LaunchArguments logger = get_logger(__name__) _as_pals_command = make_shell_format_fn(run_command="mpiexec") @dispatch(with_format=_as_pals_command, to_launcher=ShellLauncher) -class PalsMpiexecArgBuilder(LaunchArgBuilder): +class PalsMpiexecLaunchArguments(LaunchArguments): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Pals.value diff --git a/smartsim/settings/builders/launch/slurm.py b/smartsim/settings/arguments/launch/slurm.py similarity index 99% rename from smartsim/settings/builders/launch/slurm.py rename to smartsim/settings/arguments/launch/slurm.py index f80a5b8b9..bba79b969 100644 --- a/smartsim/settings/builders/launch/slurm.py +++ b/smartsim/settings/arguments/launch/slurm.py @@ -35,14 +35,14 @@ from ...common import set_check_input from ...launchCommand import LauncherType -from ..launchArgBuilder import LaunchArgBuilder +from ..launchArguments import LaunchArguments logger = get_logger(__name__) _as_srun_command = make_shell_format_fn(run_command="srun") @dispatch(with_format=_as_srun_command, to_launcher=ShellLauncher) -class SlurmArgBuilder(LaunchArgBuilder): +class SlurmLaunchArguments(LaunchArguments): def launcher_str(self) -> str: """Get the string representation of the launcher""" return LauncherType.Slurm.value diff --git a/smartsim/settings/builders/launchArgBuilder.py b/smartsim/settings/arguments/launchArguments.py similarity index 95% rename from smartsim/settings/builders/launchArgBuilder.py rename to smartsim/settings/arguments/launchArguments.py index 8ebfdb0f0..1a407a252 100644 --- a/smartsim/settings/builders/launchArgBuilder.py +++ b/smartsim/settings/arguments/launchArguments.py @@ -37,7 +37,7 @@ logger = get_logger(__name__) -class LaunchArgBuilder(ABC): +class LaunchArguments(ABC): """Abstract base class that defines all generic launcher argument methods that are not supported. It is the responsibility of child classes for each launcher to translate @@ -45,10 +45,9 @@ class LaunchArgBuilder(ABC): """ def __init__(self, launch_args: t.Dict[str, str | None] | None) -> None: - """Initialize a new `LaunchArgBuilder` instance. + """Initialize a new `LaunchArguments` instance. - :param launch_args: A mapping of argument to be used to initialize the - argument builder. + :param launch_args: A mapping of arguments to values to pre-initialize """ self._launch_args = copy.deepcopy(launch_args) or {} diff --git a/smartsim/settings/batchSettings.py b/smartsim/settings/batchSettings.py index 79a559ecb..6649fa5f7 100644 --- a/smartsim/settings/batchSettings.py +++ b/smartsim/settings/batchSettings.py @@ -32,12 +32,12 @@ from smartsim.log import get_logger from .._core.utils.helpers import fmt_dict +from .arguments import BatchArguments +from .arguments.batch.lsf import BsubBatchArguments +from .arguments.batch.pbs import QsubBatchArguments +from .arguments.batch.slurm import SlurmBatchArguments from .baseSettings import BaseSettings from .batchCommand import SchedulerType -from .builders import BatchArgBuilder -from .builders.batch.lsf import BsubBatchArgBuilder -from .builders.batch.pbs import QsubBatchArgBuilder -from .builders.batch.slurm import SlurmBatchArgBuilder from .common import StringArgument logger = get_logger(__name__) @@ -54,7 +54,7 @@ def __init__( self._batch_scheduler = SchedulerType(batch_scheduler) except ValueError: raise ValueError(f"Invalid scheduler type: {batch_scheduler}") from None - self._arg_builder = self._get_arg_builder(scheduler_args) + self._arguments = self._get_arguments(scheduler_args) self.env_vars = env_vars or {} @property @@ -68,9 +68,9 @@ def batch_scheduler(self) -> str: return self._batch_scheduler.value @property - def scheduler_args(self) -> BatchArgBuilder: + def scheduler_args(self) -> BatchArguments: """Return the batch argument translator.""" - return self._arg_builder + return self._arguments @property def env_vars(self) -> StringArgument: @@ -82,16 +82,20 @@ def env_vars(self, value: t.Dict[str, str | None]) -> None: """Set the environment variables.""" self._env_vars = copy.deepcopy(value) - def _get_arg_builder( - self, scheduler_args: StringArgument | None - ) -> BatchArgBuilder: - """Map the Scheduler to the BatchArgBuilder""" + def _get_arguments(self, scheduler_args: StringArgument | None) -> BatchArguments: + """Map the Scheduler to the BatchArguments. This method should only be + called once during construction. + + :param scheduler_args: A mapping of arguments names to values to be + used to initialize the arguments + :returns: The appropriate type for the settings instance. + """ if self._batch_scheduler == SchedulerType.Slurm: - return SlurmBatchArgBuilder(scheduler_args) + return SlurmBatchArguments(scheduler_args) elif self._batch_scheduler == SchedulerType.Lsf: - return BsubBatchArgBuilder(scheduler_args) + return BsubBatchArguments(scheduler_args) elif self._batch_scheduler == SchedulerType.Pbs: - return QsubBatchArgBuilder(scheduler_args) + return QsubBatchArguments(scheduler_args) else: raise ValueError(f"Invalid scheduler type: {self._batch_scheduler}") @@ -100,7 +104,7 @@ def format_batch_args(self) -> t.List[str]: :return: batch arguments for Sbatch """ - return self._arg_builder.format_batch_args() + return self._arguments.format_batch_args() def __str__(self) -> str: # pragma: no-cover string = f"\nScheduler: {self.scheduler}{self.scheduler_args}" diff --git a/smartsim/settings/builders/launch/__init__.py b/smartsim/settings/builders/launch/__init__.py deleted file mode 100644 index d593c59f7..000000000 --- a/smartsim/settings/builders/launch/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -from .alps import AprunArgBuilder -from .dragon import DragonArgBuilder -from .local import LocalArgBuilder -from .lsf import JsrunArgBuilder -from .mpi import MpiArgBuilder, MpiexecArgBuilder, OrteArgBuilder -from .pals import PalsMpiexecArgBuilder -from .slurm import SlurmArgBuilder - -__all__ = [ - "AprunArgBuilder", - "DragonArgBuilder", - "LocalArgBuilder", - "JsrunArgBuilder", - "MpiArgBuilder", - "MpiexecArgBuilder", - "OrteArgBuilder", - "PalsMpiexecArgBuilder", - "SlurmArgBuilder", -] diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index 9cb6b7f37..6fa6fb864 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -39,11 +39,11 @@ if t.TYPE_CHECKING: from smartsim.experiment import Experiment - from smartsim.settings.builders import LaunchArgBuilder + from smartsim.settings.arguments import LaunchArguments _Ts = TypeVarTuple("_Ts") _T_contra = t.TypeVar("_T_contra", contravariant=True) -_DispatchableT = t.TypeVar("_DispatchableT", bound="LaunchArgBuilder") +_DispatchableT = t.TypeVar("_DispatchableT", bound="LaunchArguments") _LaunchableT = t.TypeVar("_LaunchableT") _EnvironMappingType: TypeAlias = t.Mapping[str, "str | None"] @@ -59,15 +59,14 @@ @t.final class Dispatcher: """A class capable of deciding which launcher type should be used to launch - a given settings builder type. + a given settings type. """ def __init__( self, *, dispatch_registry: ( - t.Mapping[type[LaunchArgBuilder], _DispatchRegistration[t.Any, t.Any]] - | None + t.Mapping[type[LaunchArguments], _DispatchRegistration[t.Any, t.Any]] | None ) = None, ) -> None: self._dispatch_registry = ( @@ -104,8 +103,8 @@ def dispatch( to_launcher: type[LauncherProtocol[_LaunchableT]], allow_overwrite: bool = False, ) -> t.Callable[[type[_DispatchableT]], type[_DispatchableT]] | None: - """A type safe way to add a mapping of settings builder to launcher to - handle the settings at launch time. + """A type safe way to add a mapping of settings type to launcher type + to handle a settings instance at launch time. """ err_msg: str | None = None if getattr(to_launcher, "_is_protocol", False): @@ -135,8 +134,8 @@ def register(args_: type[_DispatchableT], /) -> type[_DispatchableT]: def get_dispatch( self, args: _DispatchableT | type[_DispatchableT] ) -> _DispatchRegistration[_DispatchableT, _UnkownType]: - """Find a type of launcher that is registered as being able to launch - the output of the provided builder + """Find a type of launcher that is registered as being able to launch a + settings instance of the provided type """ if not isinstance(args, type): args = type(args) @@ -252,9 +251,9 @@ def create(cls, exp: Experiment, /) -> Self: ... def make_shell_format_fn( run_command: str | None, -) -> _FormatterType[LaunchArgBuilder, t.Sequence[str]]: +) -> _FormatterType[LaunchArguments, t.Sequence[str]]: def impl( - args: LaunchArgBuilder, exe: ExecutableProtocol, _env: _EnvironMappingType + args: LaunchArguments, exe: ExecutableProtocol, _env: _EnvironMappingType ) -> t.Sequence[str]: return ( ( diff --git a/smartsim/settings/launchSettings.py b/smartsim/settings/launchSettings.py index 0990e2e82..066309d19 100644 --- a/smartsim/settings/launchSettings.py +++ b/smartsim/settings/launchSettings.py @@ -32,15 +32,19 @@ from smartsim.log import get_logger from .._core.utils.helpers import fmt_dict +from .arguments import LaunchArguments +from .arguments.launch.alps import AprunLaunchArguments +from .arguments.launch.dragon import DragonLaunchArguments +from .arguments.launch.local import LocalLaunchArguments +from .arguments.launch.lsf import JsrunLaunchArguments +from .arguments.launch.mpi import ( + MpiexecLaunchArguments, + MpirunLaunchArguments, + OrterunLaunchArguments, +) +from .arguments.launch.pals import PalsMpiexecLaunchArguments +from .arguments.launch.slurm import SlurmLaunchArguments from .baseSettings import BaseSettings -from .builders import LaunchArgBuilder -from .builders.launch.alps import AprunArgBuilder -from .builders.launch.dragon import DragonArgBuilder -from .builders.launch.local import LocalArgBuilder -from .builders.launch.lsf import JsrunArgBuilder -from .builders.launch.mpi import MpiArgBuilder, MpiexecArgBuilder, OrteArgBuilder -from .builders.launch.pals import PalsMpiexecArgBuilder -from .builders.launch.slurm import SlurmArgBuilder from .common import StringArgument from .launchCommand import LauncherType @@ -58,7 +62,7 @@ def __init__( self._launcher = LauncherType(launcher) except ValueError: raise ValueError(f"Invalid launcher type: {launcher}") - self._arg_builder = self._get_arg_builder(launch_args) + self._arguments = self._get_arguments(launch_args) self.env_vars = env_vars or {} @property @@ -67,9 +71,9 @@ def launcher(self) -> str: return self._launcher.value @property - def launch_args(self) -> LaunchArgBuilder: + def launch_args(self) -> LaunchArguments: """Return the launch argument translator.""" - return self._arg_builder + return self._arguments @launch_args.setter def launch_args(self, args: t.Mapping[str, str]) -> None: @@ -88,32 +92,32 @@ def env_vars(self, value: dict[str, str | None]) -> None: """Set the environment variables.""" self._env_vars = copy.deepcopy(value) - def _get_arg_builder(self, launch_args: StringArgument | None) -> LaunchArgBuilder: - """Map the Launcher to the LaunchArgBuilder. This method should only be + def _get_arguments(self, launch_args: StringArgument | None) -> LaunchArguments: + """Map the Launcher to the LaunchArguments. This method should only be called once during construction. - :param launch_args: A mapping of argument to be used to initialize the - argument builder. - :returns: The appropriate argument builder for the settings instance. + :param launch_args: A mapping of arguments names to values to be used + to initialize the arguments + :returns: The appropriate type for the settings instance. """ if self._launcher == LauncherType.Slurm: - return SlurmArgBuilder(launch_args) + return SlurmLaunchArguments(launch_args) elif self._launcher == LauncherType.Mpiexec: - return MpiexecArgBuilder(launch_args) + return MpiexecLaunchArguments(launch_args) elif self._launcher == LauncherType.Mpirun: - return MpiArgBuilder(launch_args) + return MpirunLaunchArguments(launch_args) elif self._launcher == LauncherType.Orterun: - return OrteArgBuilder(launch_args) + return OrterunLaunchArguments(launch_args) elif self._launcher == LauncherType.Alps: - return AprunArgBuilder(launch_args) + return AprunLaunchArguments(launch_args) elif self._launcher == LauncherType.Lsf: - return JsrunArgBuilder(launch_args) + return JsrunLaunchArguments(launch_args) elif self._launcher == LauncherType.Pals: - return PalsMpiexecArgBuilder(launch_args) + return PalsMpiexecLaunchArguments(launch_args) elif self._launcher == LauncherType.Dragon: - return DragonArgBuilder(launch_args) + return DragonLaunchArguments(launch_args) elif self._launcher == LauncherType.Local: - return LocalArgBuilder(launch_args) + return LocalLaunchArguments(launch_args) else: raise ValueError(f"Invalid launcher type: {self._launcher}") @@ -143,7 +147,7 @@ def format_env_vars(self) -> t.Union[t.List[str], None]: """Build bash compatible environment variable string for Slurm :returns: the formatted string of environment variables """ - return self._arg_builder.format_env_vars(self._env_vars) + return self._arguments.format_env_vars(self._env_vars) def format_comma_sep_env_vars(self) -> t.Union[t.Tuple[str, t.List[str]], None]: """Build environment variable string for Slurm @@ -152,7 +156,7 @@ def format_comma_sep_env_vars(self) -> t.Union[t.Tuple[str, t.List[str]], None]: for more information on this, see the slurm documentation for srun :returns: the formatted string of environment variables """ - return self._arg_builder.format_comma_sep_env_vars(self._env_vars) + return self._arguments.format_comma_sep_env_vars(self._env_vars) def format_launch_args(self) -> t.Union[t.List[str], None]: """Return formatted launch arguments @@ -160,7 +164,7 @@ def format_launch_args(self) -> t.Union[t.List[str], None]: literally with no formatting. :return: list run arguments for these settings """ - return self._arg_builder.format_launch_args() + return self._arguments.format_launch_args() def __str__(self) -> str: # pragma: no-cover string = f"\nLauncher: {self.launcher}{self.launch_args}" diff --git a/tests/temp_tests/test_settings/conftest.py b/tests/temp_tests/test_settings/conftest.py index 1368006a9..3edf5af6b 100644 --- a/tests/temp_tests/test_settings/conftest.py +++ b/tests/temp_tests/test_settings/conftest.py @@ -27,7 +27,7 @@ import pytest from smartsim.settings import dispatch -from smartsim.settings.builders import launchArgBuilder as launch +from smartsim.settings.arguments import launchArguments as launch @pytest.fixture @@ -40,13 +40,13 @@ def as_program_arguments(self): @pytest.fixture -def settings_builder(): - class _SettingsBuilder(launch.LaunchArgBuilder): +def mock_launch_args(): + class _MockLaunchArgs(launch.LaunchArguments): def set(self, arg, val): ... def launcher_str(self): - return "Mock Settings Builder" + return "mock-laucnh-args" - yield _SettingsBuilder({}) + yield _MockLaunchArgs({}) @pytest.fixture diff --git a/tests/temp_tests/test_settings/test_alpsLauncher.py b/tests/temp_tests/test_settings/test_alpsLauncher.py index 367b30c7f..3b3084c45 100644 --- a/tests/temp_tests/test_settings/test_alpsLauncher.py +++ b/tests/temp_tests/test_settings/test_alpsLauncher.py @@ -1,7 +1,10 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.alps import AprunArgBuilder, _as_aprun_command +from smartsim.settings.arguments.launch.alps import ( + AprunLaunchArguments, + _as_aprun_command, +) from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -86,14 +89,14 @@ def test_launcher_str(): ) def test_alps_class_methods(function, value, flag, result): alpsLauncher = LaunchSettings(launcher=LauncherType.Alps) - assert isinstance(alpsLauncher._arg_builder, AprunArgBuilder) + assert isinstance(alpsLauncher._arguments, AprunLaunchArguments) getattr(alpsLauncher.launch_args, function)(*value) assert alpsLauncher.launch_args._launch_args[flag] == result def test_set_verbose_launch(): alpsLauncher = LaunchSettings(launcher=LauncherType.Alps) - assert isinstance(alpsLauncher._arg_builder, AprunArgBuilder) + assert isinstance(alpsLauncher._arguments, AprunLaunchArguments) alpsLauncher.launch_args.set_verbose_launch(True) assert alpsLauncher.launch_args._launch_args == {"debug": "7"} alpsLauncher.launch_args.set_verbose_launch(False) @@ -102,7 +105,7 @@ def test_set_verbose_launch(): def test_set_quiet_launch(): aprunLauncher = LaunchSettings(launcher=LauncherType.Alps) - assert isinstance(aprunLauncher._arg_builder, AprunArgBuilder) + assert isinstance(aprunLauncher._arguments, AprunLaunchArguments) aprunLauncher.launch_args.set_quiet_launch(True) assert aprunLauncher.launch_args._launch_args == {"quiet": None} aprunLauncher.launch_args.set_quiet_launch(False) @@ -112,7 +115,7 @@ def test_set_quiet_launch(): def test_format_env_vars(): env_vars = {"OMP_NUM_THREADS": "20", "LOGGING": "verbose"} aprunLauncher = LaunchSettings(launcher=LauncherType.Alps, env_vars=env_vars) - assert isinstance(aprunLauncher._arg_builder, AprunArgBuilder) + assert isinstance(aprunLauncher._arguments, AprunLaunchArguments) aprunLauncher.update_env({"OMP_NUM_THREADS": "10"}) formatted = aprunLauncher.format_env_vars() result = ["-e", "OMP_NUM_THREADS=10", "-e", "LOGGING=verbose"] @@ -183,5 +186,5 @@ def test_invalid_exclude_hostlist_format(): ), ) def test_formatting_launch_args(mock_echo_executable, args, expected): - cmd = _as_aprun_command(AprunArgBuilder(args), mock_echo_executable, {}) + cmd = _as_aprun_command(AprunLaunchArguments(args), mock_echo_executable, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_dispatch.py b/tests/temp_tests/test_settings/test_dispatch.py index 673c8998a..78e4ec349 100644 --- a/tests/temp_tests/test_settings/test_dispatch.py +++ b/tests/temp_tests/test_settings/test_dispatch.py @@ -44,36 +44,36 @@ def format_fn(args, exe, env): @pytest.fixture -def expected_dispatch_registry(mock_launcher, settings_builder): +def expected_dispatch_registry(mock_launcher, mock_launch_args): yield { - type(settings_builder): dispatch._DispatchRegistration( + type(mock_launch_args): dispatch._DispatchRegistration( format_fn, type(mock_launcher) ) } def test_declaritive_form_dispatch_declaration( - mock_launcher, settings_builder, expected_dispatch_registry + mock_launcher, mock_launch_args, expected_dispatch_registry ): d = dispatch.Dispatcher() - assert type(settings_builder) == d.dispatch( + assert type(mock_launch_args) == d.dispatch( with_format=format_fn, to_launcher=type(mock_launcher) - )(type(settings_builder)) + )(type(mock_launch_args)) assert d._dispatch_registry == expected_dispatch_registry def test_imperative_form_dispatch_declaration( - mock_launcher, settings_builder, expected_dispatch_registry + mock_launcher, mock_launch_args, expected_dispatch_registry ): d = dispatch.Dispatcher() assert None == d.dispatch( - type(settings_builder), to_launcher=type(mock_launcher), with_format=format_fn + type(mock_launch_args), to_launcher=type(mock_launcher), with_format=format_fn ) assert d._dispatch_registry == expected_dispatch_registry def test_dispatchers_from_same_registry_do_not_cross_polute( - mock_launcher, settings_builder, expected_dispatch_registry + mock_launcher, mock_launch_args, expected_dispatch_registry ): some_starting_registry = {} d1 = dispatch.Dispatcher(dispatch_registry=some_starting_registry) @@ -86,14 +86,14 @@ def test_dispatchers_from_same_registry_do_not_cross_polute( ) d2.dispatch( - type(settings_builder), with_format=format_fn, to_launcher=type(mock_launcher) + type(mock_launch_args), with_format=format_fn, to_launcher=type(mock_launcher) ) assert d1._dispatch_registry == {} assert d2._dispatch_registry == expected_dispatch_registry def test_copied_dispatchers_do_not_cross_pollute( - mock_launcher, settings_builder, expected_dispatch_registry + mock_launcher, mock_launch_args, expected_dispatch_registry ): some_starting_registry = {} d1 = dispatch.Dispatcher(dispatch_registry=some_starting_registry) @@ -106,7 +106,7 @@ def test_copied_dispatchers_do_not_cross_pollute( ) d2.dispatch( - type(settings_builder), to_launcher=type(mock_launcher), with_format=format_fn + type(mock_launch_args), to_launcher=type(mock_launcher), with_format=format_fn ) assert d1._dispatch_registry == {} assert d2._dispatch_registry == expected_dispatch_registry @@ -145,12 +145,12 @@ def test_dispatch_overwriting( add_dispatch, expected_ctx, mock_launcher, - settings_builder, + mock_launch_args, expected_dispatch_registry, ): d = dispatch.Dispatcher(dispatch_registry=expected_dispatch_registry) with expected_ctx: - add_dispatch(d, type(settings_builder), type(mock_launcher)) + add_dispatch(d, type(mock_launch_args), type(mock_launcher)) @pytest.mark.parametrize( @@ -161,12 +161,12 @@ def test_dispatch_overwriting( ), ) def test_dispatch_can_retrieve_dispatch_info_from_dispatch_registry( - expected_dispatch_registry, mock_launcher, settings_builder, type_or_instance + expected_dispatch_registry, mock_launcher, mock_launch_args, type_or_instance ): d = dispatch.Dispatcher(dispatch_registry=expected_dispatch_registry) assert dispatch._DispatchRegistration( format_fn, type(mock_launcher) - ) == d.get_dispatch(type_or_instance(settings_builder)) + ) == d.get_dispatch(type_or_instance(mock_launch_args)) @pytest.mark.parametrize( @@ -177,13 +177,13 @@ def test_dispatch_can_retrieve_dispatch_info_from_dispatch_registry( ), ) def test_dispatch_raises_if_settings_type_not_registered( - settings_builder, type_or_instance + mock_launch_args, type_or_instance ): d = dispatch.Dispatcher(dispatch_registry={}) with pytest.raises( TypeError, match="No dispatch for `.+?(?=`)` has been registered" ): - d.get_dispatch(type_or_instance(settings_builder)) + d.get_dispatch(type_or_instance(mock_launch_args)) class LauncherABC(abc.ABC): @@ -323,13 +323,13 @@ def test_launcher_adapter_correctly_adapts_input_to_launcher(input_, map_, expec ), ) def test_dispatch_registration_can_configure_adapter_for_existing_launcher_instance( - request, settings_builder, buffer_writer_dispatch, launcher_instance, ctx + request, mock_launch_args, buffer_writer_dispatch, launcher_instance, ctx ): if isinstance(launcher_instance, str): launcher_instance = request.getfixturevalue(launcher_instance) with ctx: adapter = buffer_writer_dispatch.create_adapter_from_launcher( - launcher_instance, settings_builder + launcher_instance, mock_launch_args ) assert adapter._adapted_launcher is launcher_instance @@ -375,7 +375,7 @@ def test_dispatch_registration_can_configure_adapter_for_existing_launcher_insta ), ) def test_dispatch_registration_configures_first_compatible_launcher_from_sequence_of_launchers( - request, settings_builder, buffer_writer_dispatch, launcher_instances, ctx + request, mock_launch_args, buffer_writer_dispatch, launcher_instances, ctx ): def resolve_instance(inst): return request.getfixturevalue(inst) if isinstance(inst, str) else inst @@ -384,24 +384,24 @@ def resolve_instance(inst): with ctx: adapter = buffer_writer_dispatch.configure_first_compatible_launcher( - with_settings=settings_builder, from_available_launchers=launcher_instances + with_settings=mock_launch_args, from_available_launchers=launcher_instances ) def test_dispatch_registration_can_create_a_laucher_for_an_experiment_and_can_reconfigure_it_later( - settings_builder, buffer_writer_dispatch + mock_launch_args, buffer_writer_dispatch ): class MockExperiment: ... exp = MockExperiment() adapter_1 = buffer_writer_dispatch.create_new_launcher_configuration( - for_experiment=exp, with_settings=settings_builder + for_experiment=exp, with_settings=mock_launch_args ) assert type(adapter_1._adapted_launcher) == buffer_writer_dispatch.launcher_type existing_launcher = adapter_1._adapted_launcher adapter_2 = buffer_writer_dispatch.create_adapter_from_launcher( - existing_launcher, settings_builder + existing_launcher, mock_launch_args ) assert type(adapter_2._adapted_launcher) == buffer_writer_dispatch.launcher_type assert adapter_1._adapted_launcher is adapter_2._adapted_launcher diff --git a/tests/temp_tests/test_settings/test_dragonLauncher.py b/tests/temp_tests/test_settings/test_dragonLauncher.py index c71b1d548..0c6fb80ac 100644 --- a/tests/temp_tests/test_settings/test_dragonLauncher.py +++ b/tests/temp_tests/test_settings/test_dragonLauncher.py @@ -3,7 +3,7 @@ from smartsim._core.launcher.dragon.dragonLauncher import _as_run_request_view from smartsim._core.schemas.dragonRequests import DragonRunRequestView from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.dragon import DragonArgBuilder +from smartsim.settings.arguments.launch.dragon import DragonLaunchArguments from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -26,7 +26,7 @@ def test_launcher_str(): ) def test_dragon_class_methods(function, value, flag, result): dragonLauncher = LaunchSettings(launcher=LauncherType.Dragon) - assert isinstance(dragonLauncher._arg_builder, DragonArgBuilder) + assert isinstance(dragonLauncher._arguments, DragonLaunchArguments) getattr(dragonLauncher.launch_args, function)(*value) assert dragonLauncher.launch_args._launch_args[flag] == result @@ -39,7 +39,7 @@ def test_dragon_class_methods(function, value, flag, result): def test_formatting_launch_args_into_request( mock_echo_executable, nodes, tasks_per_node ): - args = DragonArgBuilder({}) + args = DragonLaunchArguments({}) if nodes is not NOT_SET: args.set_nodes(nodes) if tasks_per_node is not NOT_SET: diff --git a/tests/temp_tests/test_settings/test_localLauncher.py b/tests/temp_tests/test_settings/test_localLauncher.py index b3cb4108f..580e53d36 100644 --- a/tests/temp_tests/test_settings/test_localLauncher.py +++ b/tests/temp_tests/test_settings/test_localLauncher.py @@ -1,7 +1,10 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.local import LocalArgBuilder, _as_local_command +from smartsim.settings.arguments.launch.local import ( + LocalLaunchArguments, + _as_local_command, +) from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -110,10 +113,10 @@ def test_format_env_vars(): "D": "12", } localLauncher = LaunchSettings(launcher=LauncherType.Local, env_vars=env_vars) - assert isinstance(localLauncher._arg_builder, LocalArgBuilder) + assert isinstance(localLauncher._arguments, LocalLaunchArguments) assert localLauncher.format_env_vars() == ["A=a", "B=", "C=", "D=12"] def test_formatting_returns_original_exe(mock_echo_executable): - cmd = _as_local_command(LocalArgBuilder({}), mock_echo_executable, {}) + cmd = _as_local_command(LocalLaunchArguments({}), mock_echo_executable, {}) assert tuple(cmd) == ("echo", "hello", "world") diff --git a/tests/temp_tests/test_settings/test_lsfLauncher.py b/tests/temp_tests/test_settings/test_lsfLauncher.py index 636d2896f..c73edb6a9 100644 --- a/tests/temp_tests/test_settings/test_lsfLauncher.py +++ b/tests/temp_tests/test_settings/test_lsfLauncher.py @@ -1,7 +1,10 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.lsf import JsrunArgBuilder, _as_jsrun_command +from smartsim.settings.arguments.launch.lsf import ( + JsrunLaunchArguments, + _as_jsrun_command, +) from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -24,7 +27,7 @@ def test_launcher_str(): ) def test_lsf_class_methods(function, value, flag, result): lsfLauncher = LaunchSettings(launcher=LauncherType.Lsf) - assert isinstance(lsfLauncher._arg_builder, JsrunArgBuilder) + assert isinstance(lsfLauncher._arguments, JsrunLaunchArguments) getattr(lsfLauncher.launch_args, function)(*value) assert lsfLauncher.launch_args._launch_args[flag] == result @@ -32,7 +35,7 @@ def test_lsf_class_methods(function, value, flag, result): def test_format_env_vars(): env_vars = {"OMP_NUM_THREADS": None, "LOGGING": "verbose"} lsfLauncher = LaunchSettings(launcher=LauncherType.Lsf, env_vars=env_vars) - assert isinstance(lsfLauncher._arg_builder, JsrunArgBuilder) + assert isinstance(lsfLauncher._arguments, JsrunLaunchArguments) formatted = lsfLauncher.format_env_vars() assert formatted == ["-E", "OMP_NUM_THREADS", "-E", "LOGGING=verbose"] @@ -47,7 +50,7 @@ def test_launch_args(): "np": 100, } lsfLauncher = LaunchSettings(launcher=LauncherType.Lsf, launch_args=launch_args) - assert isinstance(lsfLauncher._arg_builder, JsrunArgBuilder) + assert isinstance(lsfLauncher._arguments, JsrunLaunchArguments) formatted = lsfLauncher.format_launch_args() result = [ "--latency_priority=gpu-gpu", @@ -92,5 +95,5 @@ def test_launch_args(): ), ) def test_formatting_launch_args(mock_echo_executable, args, expected): - cmd = _as_jsrun_command(JsrunArgBuilder(args), mock_echo_executable, {}) + cmd = _as_jsrun_command(JsrunLaunchArguments(args), mock_echo_executable, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_mpiLauncher.py b/tests/temp_tests/test_settings/test_mpiLauncher.py index 23df78c92..350555ae7 100644 --- a/tests/temp_tests/test_settings/test_mpiLauncher.py +++ b/tests/temp_tests/test_settings/test_mpiLauncher.py @@ -3,10 +3,10 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.mpi import ( - MpiArgBuilder, - MpiexecArgBuilder, - OrteArgBuilder, +from smartsim.settings.arguments.launch.mpi import ( + MpiexecLaunchArguments, + MpirunLaunchArguments, + OrterunLaunchArguments, _as_mpiexec_command, _as_mpirun_command, _as_orterun_command, @@ -107,9 +107,9 @@ def test_launcher_str(launcher): ), ) for l in ( - [LauncherType.Mpirun, MpiArgBuilder], - [LauncherType.Mpiexec, MpiexecArgBuilder], - [LauncherType.Orterun, OrteArgBuilder], + [LauncherType.Mpirun, MpirunLaunchArguments], + [LauncherType.Mpiexec, MpiexecLaunchArguments], + [LauncherType.Orterun, OrterunLaunchArguments], ) ) ) @@ -117,7 +117,7 @@ def test_launcher_str(launcher): ) def test_mpi_class_methods(l, function, value, flag, result): mpiSettings = LaunchSettings(launcher=l[0]) - assert isinstance(mpiSettings._arg_builder, l[1]) + assert isinstance(mpiSettings._arguments, l[1]) getattr(mpiSettings.launch_args, function)(*value) assert mpiSettings.launch_args._launch_args[flag] == result @@ -215,11 +215,15 @@ def test_invalid_hostlist_format(launcher): @pytest.mark.parametrize( "cls, fmt, cmd", ( - pytest.param(MpiArgBuilder, _as_mpirun_command, "mpirun", id="w/ mpirun"), pytest.param( - MpiexecArgBuilder, _as_mpiexec_command, "mpiexec", id="w/ mpiexec" + MpirunLaunchArguments, _as_mpirun_command, "mpirun", id="w/ mpirun" + ), + pytest.param( + MpiexecLaunchArguments, _as_mpiexec_command, "mpiexec", id="w/ mpiexec" + ), + pytest.param( + OrterunLaunchArguments, _as_orterun_command, "orterun", id="w/ orterun" ), - pytest.param(OrteArgBuilder, _as_orterun_command, "orterun", id="w/ orterun"), ), ) @pytest.mark.parametrize( diff --git a/tests/temp_tests/test_settings/test_palsLauncher.py b/tests/temp_tests/test_settings/test_palsLauncher.py index 18d85a778..c348fe96f 100644 --- a/tests/temp_tests/test_settings/test_palsLauncher.py +++ b/tests/temp_tests/test_settings/test_palsLauncher.py @@ -1,8 +1,8 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.pals import ( - PalsMpiexecArgBuilder, +from smartsim.settings.arguments.launch.pals import ( + PalsMpiexecLaunchArguments, _as_pals_command, ) from smartsim.settings.launchCommand import LauncherType @@ -49,7 +49,7 @@ def test_launcher_str(): ) def test_pals_class_methods(function, value, flag, result): palsLauncher = LaunchSettings(launcher=LauncherType.Pals) - assert isinstance(palsLauncher.launch_args, PalsMpiexecArgBuilder) + assert isinstance(palsLauncher.launch_args, PalsMpiexecLaunchArguments) getattr(palsLauncher.launch_args, function)(*value) assert palsLauncher.launch_args._launch_args[flag] == result assert palsLauncher.format_launch_args() == ["--" + flag, str(result)] @@ -106,5 +106,5 @@ def test_invalid_hostlist_format(): ), ) def test_formatting_launch_args(mock_echo_executable, args, expected): - cmd = _as_pals_command(PalsMpiexecArgBuilder(args), mock_echo_executable, {}) + cmd = _as_pals_command(PalsMpiexecLaunchArguments(args), mock_echo_executable, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_pbsScheduler.py b/tests/temp_tests/test_settings/test_pbsScheduler.py index ab3435df5..94da0411a 100644 --- a/tests/temp_tests/test_settings/test_pbsScheduler.py +++ b/tests/temp_tests/test_settings/test_pbsScheduler.py @@ -1,8 +1,8 @@ import pytest from smartsim.settings import BatchSettings +from smartsim.settings.arguments.batch.pbs import QsubBatchArguments from smartsim.settings.batchCommand import SchedulerType -from smartsim.settings.builders.batch.pbs import QsubBatchArgBuilder def test_scheduler_str(): @@ -35,7 +35,7 @@ def test_scheduler_str(): ) def test_create_pbs_batch(function, value, flag, result): pbsScheduler = BatchSettings(batch_scheduler=SchedulerType.Pbs) - assert isinstance(pbsScheduler.scheduler_args, QsubBatchArgBuilder) + assert isinstance(pbsScheduler.scheduler_args, QsubBatchArguments) getattr(pbsScheduler.scheduler_args, function)(*value) assert pbsScheduler.scheduler_args._scheduler_args[flag] == result diff --git a/tests/temp_tests/test_settings/test_slurmLauncher.py b/tests/temp_tests/test_settings/test_slurmLauncher.py index e3b73aee7..61c0d55c4 100644 --- a/tests/temp_tests/test_settings/test_slurmLauncher.py +++ b/tests/temp_tests/test_settings/test_slurmLauncher.py @@ -1,7 +1,10 @@ import pytest from smartsim.settings import LaunchSettings -from smartsim.settings.builders.launch.slurm import SlurmArgBuilder, _as_srun_command +from smartsim.settings.arguments.launch.slurm import ( + SlurmLaunchArguments, + _as_srun_command, +) from smartsim.settings.launchCommand import LauncherType pytestmark = pytest.mark.group_a @@ -83,7 +86,7 @@ def test_launcher_str(): ) def test_slurm_class_methods(function, value, flag, result): slurmLauncher = LaunchSettings(launcher=LauncherType.Slurm) - assert isinstance(slurmLauncher.launch_args, SlurmArgBuilder) + assert isinstance(slurmLauncher.launch_args, SlurmLaunchArguments) getattr(slurmLauncher.launch_args, function)(*value) assert slurmLauncher.launch_args._launch_args[flag] == result @@ -250,9 +253,9 @@ def test_set_het_groups(monkeypatch): monkeypatch.setenv("SLURM_HET_SIZE", "4") slurmLauncher = LaunchSettings(launcher=LauncherType.Slurm) slurmLauncher.launch_args.set_het_group([1]) - assert slurmLauncher._arg_builder._launch_args["het-group"] == "1" + assert slurmLauncher._arguments._launch_args["het-group"] == "1" slurmLauncher.launch_args.set_het_group([3, 2]) - assert slurmLauncher._arg_builder._launch_args["het-group"] == "3,2" + assert slurmLauncher._arguments._launch_args["het-group"] == "3,2" with pytest.raises(ValueError): slurmLauncher.launch_args.set_het_group([4]) @@ -289,5 +292,5 @@ def test_set_het_groups(monkeypatch): ), ) def test_formatting_launch_args(mock_echo_executable, args, expected): - cmd = _as_srun_command(SlurmArgBuilder(args), mock_echo_executable, {}) + cmd = _as_srun_command(SlurmLaunchArguments(args), mock_echo_executable, {}) assert tuple(cmd) == expected diff --git a/tests/temp_tests/test_settings/test_slurmScheduler.py b/tests/temp_tests/test_settings/test_slurmScheduler.py index 5c65d367a..38c98b171 100644 --- a/tests/temp_tests/test_settings/test_slurmScheduler.py +++ b/tests/temp_tests/test_settings/test_slurmScheduler.py @@ -1,8 +1,8 @@ import pytest from smartsim.settings import BatchSettings +from smartsim.settings.arguments.batch.slurm import SlurmBatchArguments from smartsim.settings.batchCommand import SchedulerType -from smartsim.settings.builders.batch.slurm import SlurmBatchArgBuilder def test_scheduler_str(): @@ -57,7 +57,7 @@ def test_create_sbatch(): slurmScheduler = BatchSettings( batch_scheduler=SchedulerType.Slurm, scheduler_args=batch_args ) - assert isinstance(slurmScheduler._arg_builder, SlurmBatchArgBuilder) + assert isinstance(slurmScheduler._arguments, SlurmBatchArguments) args = slurmScheduler.format_batch_args() assert args == ["--exclusive", "--oversubscribe"] diff --git a/tests/test_experiment.py b/tests/test_experiment.py index 1cd1ee5c3..132fb92f3 100644 --- a/tests/test_experiment.py +++ b/tests/test_experiment.py @@ -39,7 +39,7 @@ from smartsim.experiment import Experiment from smartsim.launchable import job from smartsim.settings import dispatch, launchSettings -from smartsim.settings.builders import launchArgBuilder +from smartsim.settings.arguments import launchArguments pytestmark = pytest.mark.group_a @@ -62,7 +62,7 @@ def job_maker(monkeypatch): def iter_jobs(): for i in itertools.count(): settings = launchSettings.LaunchSettings("local") - monkeypatch.setattr(settings, "_arg_builder", MockLaunchArgs(i)) + monkeypatch.setattr(settings, "_arguments", MockLaunchArgs(i)) yield job.Job(EchoHelloWorldEntity(), settings) jobs = iter_jobs() @@ -92,7 +92,7 @@ def start(self, record: LaunchRecord): @dataclasses.dataclass(frozen=True) class LaunchRecord: - launch_args: launchArgBuilder.LaunchArgBuilder + launch_args: launchArguments.LaunchArguments entity: entity.SmartSimEntity env: t.Mapping[str, str | None] @@ -104,7 +104,7 @@ def from_job(cls, job): return cls(args, entity, env) -class MockLaunchArgs(launchArgBuilder.LaunchArgBuilder): +class MockLaunchArgs(launchArguments.LaunchArguments): def __init__(self, count): super().__init__({}) self.count = count From f0bf80ff611736e635c33f813530028616b83ed7 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Tue, 23 Jul 2024 21:20:55 -0500 Subject: [PATCH 29/34] Remove `Experiment._control` --- smartsim/experiment.py | 280 +-------------------------------------- tests/test_experiment.py | 6 +- 2 files changed, 5 insertions(+), 281 deletions(-) diff --git a/smartsim/experiment.py b/smartsim/experiment.py index 8915a620d..7bab5146e 100644 --- a/smartsim/experiment.py +++ b/smartsim/experiment.py @@ -163,9 +163,6 @@ def __init__(self, name: str, exp_path: str | None = None): self.exp_path = exp_path """The path under which the experiment operate""" - # TODO: Remove this! The controller is becoming obsolete - self._control = Controller(launcher="local") - self._active_launchers: set[LauncherProtocol[t.Any]] = set() """The active launchers created, used, and reused by the experiment""" @@ -176,7 +173,7 @@ def __init__(self, name: str, exp_path: str | None = None): """Switch to specify if telemetry data should be produced for this experiment""" - def start_jobs( + def start( self, job: Job, *jobs: Job, dispatcher: Dispatcher = DEFAULT_DISPATCHER ) -> tuple[LaunchedJobID, ...]: """Execute a collection of `Job` instances. @@ -219,113 +216,6 @@ def _start(job: Job) -> LaunchedJobID: return _start(job), *map(_start, jobs) - @_contextualize - def start( - self, - *args: t.Union[SmartSimEntity, EntitySequence[SmartSimEntity]], - block: bool = True, - summary: bool = False, - kill_on_interrupt: bool = True, - ) -> None: - """Start passed instances using Experiment launcher - - Any instance ``Application``, ``Ensemble`` or ``FeatureStore`` - instance created by the Experiment can be passed as - an argument to the start method. - - .. highlight:: python - .. code-block:: python - - exp = Experiment(name="my_exp", launcher="slurm") - settings = exp.create_run_settings(exe="./path/to/binary") - application = exp.create_application("my_application", settings) - exp.start(application) - - Multiple entity instances can also be passed to the start method - at once no matter which type of instance they are. These will - all be launched together. - - .. highlight:: python - .. code-block:: python - - exp.start(application_1, application_2, fs, ensemble, block=True) - # alternatively - stage_1 = [application_1, application_2, fs, ensemble] - exp.start(*stage_1, block=True) - - - If `block==True` the Experiment will poll the launched instances - at runtime until all non-feature store jobs have completed. Feature store - jobs *must* be killed by the user by passing them to - ``Experiment.stop``. This allows for multiple stages of a workflow - to produce to and consume from the same FeatureStore feature store. - - If `kill_on_interrupt=True`, then all jobs launched by this - experiment are guaranteed to be killed when ^C (SIGINT) signal is - received. If `kill_on_interrupt=False`, then it is not guaranteed - that all jobs launched by this experiment will be killed, and the - zombie processes will need to be manually killed. - - :param block: block execution until all non-feature store - jobs are finished - :param summary: print a launch summary prior to launch - :param kill_on_interrupt: flag for killing jobs when ^C (SIGINT) - signal is received. - """ - start_manifest = Manifest(*args) - self._create_entity_dir(start_manifest) - try: - if summary: - self._launch_summary(start_manifest) - self._control.start( - exp_name=self.name, - exp_path=self.exp_path, - manifest=start_manifest, - block=block, - kill_on_interrupt=kill_on_interrupt, - ) - except SmartSimError as e: - logger.error(e) - raise - - @_contextualize - def stop( - self, *args: t.Union[SmartSimEntity, EntitySequence[SmartSimEntity]] - ) -> None: - """Stop specific instances launched by this ``Experiment`` - - Instances of ``Application``, ``Ensemble`` and ``FeatureStore`` - can all be passed as arguments to the stop method. - - Whichever launcher was specified at Experiment initialization - will be used to stop the instance. For example, which using - the slurm launcher, this equates to running `scancel` on the - instance. - - Example - - .. highlight:: python - .. code-block:: python - - exp.stop(application) - # multiple - exp.stop(application_1, application_2, fs, ensemble) - - :param args: One or more SmartSimEntity or EntitySequence objects. - :raises TypeError: if wrong type - :raises SmartSimError: if stop request fails - """ - stop_manifest = Manifest(*args) - try: - for entity in stop_manifest.applications: - self._control.stop_entity(entity) - fss = stop_manifest.fss - for fs in fss: - self._control.stop_fs(fs) - except SmartSimError as e: - logger.error(e) - raise - @_contextualize def generate( self, @@ -360,128 +250,6 @@ def generate( logger.error(e) raise - @_contextualize - def poll( - self, interval: int = 10, verbose: bool = True, kill_on_interrupt: bool = True - ) -> None: - """Monitor jobs through logging to stdout. - - This method should only be used if jobs were launched - with ``Experiment.start(block=False)`` - - The internal specified will control how often the - logging is performed, not how often the polling occurs. - By default, internal polling is set to every second for - local launcher jobs and every 10 seconds for all other - launchers. - - If internal polling needs to be slower or faster based on - system or site standards, set the ``SMARTSIM_JM_INTERNAL`` - environment variable to control the internal polling interval - for SmartSim. - - For more verbose logging output, the ``SMARTSIM_LOG_LEVEL`` - environment variable can be set to `debug` - - If `kill_on_interrupt=True`, then all jobs launched by this - experiment are guaranteed to be killed when ^C (SIGINT) signal is - received. If `kill_on_interrupt=False`, then it is not guaranteed - that all jobs launched by this experiment will be killed, and the - zombie processes will need to be manually killed. - - :param interval: frequency (in seconds) of logging to stdout - :param verbose: set verbosity - :param kill_on_interrupt: flag for killing jobs when SIGINT is received - :raises SmartSimError: if poll request fails - """ - try: - self._control.poll(interval, verbose, kill_on_interrupt=kill_on_interrupt) - except SmartSimError as e: - logger.error(e) - raise - - @_contextualize - def finished(self, entity: SmartSimEntity) -> bool: - """Query if a job has completed. - - An instance of ``application`` or ``Ensemble`` can be passed - as an argument. - - Passing ``FeatureStore`` will return an error as a - feature store deployment is never finished until stopped - by the user. - - :param entity: object launched by this ``Experiment`` - :returns: True if the job has finished, False otherwise - :raises SmartSimError: if entity has not been launched - by this ``Experiment`` - """ - try: - return self._control.finished(entity) - except SmartSimError as e: - logger.error(e) - raise - - @_contextualize - def get_status( - self, *args: t.Union[SmartSimEntity, EntitySequence[SmartSimEntity]] - ) -> t.List[SmartSimStatus]: - """Query the status of launched entity instances - - Return a smartsim.status string representing - the status of the launched instance. - - .. highlight:: python - .. code-block:: python - - exp.get_status(application) - - As with an Experiment method, multiple instance of - varying types can be passed to and all statuses will - be returned at once. - - .. highlight:: python - .. code-block:: python - - statuses = exp.get_status(application, ensemble, featurestore) - complete = [s == smartsim.status.STATUS_COMPLETED for s in statuses] - assert all(complete) - - :returns: status of the instances passed as arguments - :raises SmartSimError: if status retrieval fails - """ - try: - manifest = Manifest(*args) - statuses: t.List[SmartSimStatus] = [] - for entity in manifest.applications: - statuses.append(self._control.get_entity_status(entity)) - for entity_list in manifest.all_entity_lists: - statuses.extend(self._control.get_entity_list_status(entity_list)) - return statuses - except SmartSimError as e: - logger.error(e) - raise - - @_contextualize - def reconnect_feature_store(self, checkpoint: str) -> FeatureStore: - """Reconnect to a running ``FeatureStore`` - - This method can be used to connect to a ``FeatureStore`` deployment - that was launched by a previous ``Experiment``. This can be - helpful in the case where separate runs of an ``Experiment`` - wish to use the same ``FeatureStore`` instance currently - running on a system. - - :param checkpoint: the `smartsim_db.dat` file created - when an ``FeatureStore`` is launched - """ - try: - feature_store = self._control.reload_saved_fs(checkpoint) - return feature_store - except SmartSimError as e: - logger.error(e) - raise - def preview( self, *args: t.Any, @@ -511,9 +279,6 @@ def preview( output to stdout. Defaults to None. """ - # Retrieve any active feature store jobs - active_fsjobs = self._control.active_feature_store_jobs - preview_manifest = Manifest(*args) previewrenderer.render( @@ -522,7 +287,6 @@ def preview( verbosity_level, output_format, output_filename, - active_fsjobs, ) @_contextualize @@ -537,7 +301,6 @@ def summary(self, style: str = "github") -> str: https://github.com/astanin/python-tabulate :return: tabulate string of ``Experiment`` history """ - values = [] headers = [ "Name", "Entity-Type", @@ -547,21 +310,8 @@ def summary(self, style: str = "github") -> str: "Status", "Returncode", ] - for job in self._control.get_jobs().values(): - for run in range(job.history.runs + 1): - values.append( - [ - job.entity.name, - job.entity.type, - job.history.jids[run], - run, - f"{job.history.job_times[run]:.4f}", - job.history.statuses[run], - job.history.returns[run], - ] - ) return tabulate( - values, + [], headers, showindex=True, tablefmt=style, @@ -577,32 +327,6 @@ def telemetry(self) -> TelemetryConfiguration: """ return self._telemetry_cfg - def _launch_summary(self, manifest: Manifest) -> None: - """Experiment pre-launch summary of entities that will be launched - - :param manifest: Manifest of deployables. - """ - launcher_list = "\n".join(str(launcher) for launcher in self._active_launchers) - summary = textwrap.dedent(f"""\ - === Launch Summary === - Experiment: {self.name} - Experiment Path: {self.exp_path} - Launcher(s): - {textwrap.indent(" - ", launcher_list) if launcher_list else " "} - """) - - if manifest.applications: - summary += f"Applications: {len(manifest.applications)}\n" - - if self._control.feature_store_active: - summary += "Feature Store Status: active\n" - elif manifest.fss: - summary += "Feature Store Status: launching\n" - else: - summary += "Feature Store Status: inactive\n" - - logger.info(f"\n\n{summary}\n{manifest}") - def _create_entity_dir(self, start_manifest: Manifest) -> None: def create_entity_dir( entity: t.Union[FeatureStore, Application, Ensemble] diff --git a/tests/test_experiment.py b/tests/test_experiment.py index 132fb92f3..b7f4b1efa 100644 --- a/tests/test_experiment.py +++ b/tests/test_experiment.py @@ -137,7 +137,7 @@ def as_program_arguments(self): def test_start_raises_if_no_args_supplied(experiment): with pytest.raises(TypeError, match="missing 1 required positional argument"): - experiment.start_jobs() + experiment.start() # fmt: off @@ -154,7 +154,7 @@ def test_start_raises_if_no_args_supplied(experiment): def test_start_can_launch_jobs(experiment, job_maker, dispatcher, make_jobs, num_jobs): jobs = make_jobs(job_maker, num_jobs) assert len(experiment._active_launchers) == 0, "Initialized w/ launchers" - launched_ids = experiment.start_jobs(*jobs, dispatcher=dispatcher) + launched_ids = experiment.start(*jobs, dispatcher=dispatcher) assert len(experiment._active_launchers) == 1, "Unexpected number of launchers" (launcher,) = experiment._active_launchers assert isinstance(launcher, NoOpRecordLauncher), "Unexpected launcher type" @@ -178,7 +178,7 @@ def test_start_can_start_a_job_multiple_times_accross_multiple_calls( assert len(experiment._active_launchers) == 0, "Initialized w/ launchers" job = job_maker() ids_to_launches = { - experiment.start_jobs(job, dispatcher=dispatcher)[0]: LaunchRecord.from_job(job) + experiment.start(job, dispatcher=dispatcher)[0]: LaunchRecord.from_job(job) for _ in range(num_starts) } assert len(experiment._active_launchers) == 1, "Did not reuse the launcher" From 7a0b8938fa72672ef74a1a7bfab4add428493bd5 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Tue, 23 Jul 2024 22:17:01 -0500 Subject: [PATCH 30/34] Address some reviewer feedback --- smartsim/experiment.py | 4 +++ .../settings/arguments/launchArguments.py | 19 ++++++---- smartsim/settings/dispatch.py | 35 +++++++++++++++++-- 3 files changed, 49 insertions(+), 9 deletions(-) diff --git a/smartsim/experiment.py b/smartsim/experiment.py index 7bab5146e..d0ad321d5 100644 --- a/smartsim/experiment.py +++ b/smartsim/experiment.py @@ -199,11 +199,15 @@ def _start(job: Job) -> LaunchedJobID: # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< dispatch = dispatcher.get_dispatch(args) try: + # Check to see if one of the existing launchers can be + # configured to handle the launch arguments ... launch_config = dispatch.configure_first_compatible_launcher( from_available_launchers=self._active_launchers, with_settings=args, ) except errors.LauncherNotFoundError: + # ... otherwise create a new launcher that _can_ handle the + # launch arguments and configure _that_ one launch_config = dispatch.create_new_launcher_configuration( for_experiment=self, with_settings=args ) diff --git a/smartsim/settings/arguments/launchArguments.py b/smartsim/settings/arguments/launchArguments.py index 1a407a252..da27487f3 100644 --- a/smartsim/settings/arguments/launchArguments.py +++ b/smartsim/settings/arguments/launchArguments.py @@ -27,6 +27,7 @@ from __future__ import annotations import copy +import textwrap import typing as t from abc import ABC, abstractmethod @@ -38,16 +39,15 @@ class LaunchArguments(ABC): - """Abstract base class that defines all generic launcher - argument methods that are not supported. It is the - responsibility of child classes for each launcher to translate - the input parameter to a properly formatted launcher argument. + """Abstract base class for launcher arguments. It is the responsibility of + child classes for each launcher to add methods to set input parameters and + to maintain valid state between parameters set by a user. """ def __init__(self, launch_args: t.Dict[str, str | None] | None) -> None: """Initialize a new `LaunchArguments` instance. - :param launch_args: A mapping of arguments to values to pre-initialize + :param launch_args: A mapping of arguments to (optional) values """ self._launch_args = copy.deepcopy(launch_args) or {} @@ -95,5 +95,10 @@ def format_env_vars( return None def __str__(self) -> str: # pragma: no-cover - string = f"\nLaunch Arguments:\n{fmt_dict(self._launch_args)}" - return string + return textwrap.dedent(f"""\ + Launch Arguments: + Launcher: {self.launcher_str()} + Name: {type(self).__name__} + Arguments: + {fmt_dict(self._launch_args)} + """) diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index 6fa6fb864..7b6b9a87a 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -142,7 +142,7 @@ def get_dispatch( dispatch_ = self._dispatch_registry.get(args, None) if dispatch_ is None: raise TypeError( - f"No dispatch for `{type(args).__name__}` has been registered " + f"No dispatch for `{args.__name__}` has been registered " f"has been registered with {type(self).__name__} `{self}`" ) # Note the sleight-of-hand here: we are secretly casting a type of @@ -160,11 +160,16 @@ def get_dispatch( @t.final @dataclasses.dataclass(frozen=True) class _DispatchRegistration(t.Generic[_DispatchableT, _LaunchableT]): + """An entry into the `Dispatcher`'s dispatch registry. This class is simply + a wrapper around a launcher and how to format a `_DispatchableT` instance + to be launched by the afore mentioned launcher. + """ + formatter: _FormatterType[_DispatchableT, _LaunchableT] launcher_type: type[LauncherProtocol[_LaunchableT]] def _is_compatible_launcher(self, launcher: LauncherProtocol[t.Any]) -> bool: - # Disabling because we want to match the the type of the dispatch + # Disabling because we want to match the type of the dispatch # *exactly* as specified by the user # pylint: disable-next=unidiomatic-typecheck return type(launcher) is self.launcher_type @@ -172,12 +177,28 @@ def _is_compatible_launcher(self, launcher: LauncherProtocol[t.Any]) -> bool: def create_new_launcher_configuration( self, for_experiment: Experiment, with_settings: _DispatchableT ) -> _LaunchConfigType: + """Create a new instance of a launcher for an experiment that the + provided settings where set to dispatch to, and configure it with the + provided launch settings. + + :param for_experiment: The experiment responsible creating the launcher + :param with_settings: The settings with which to configure the newly + created launcher + :returns: A configured launcher + """ launcher = self.launcher_type.create(for_experiment) return self.create_adapter_from_launcher(launcher, with_settings) def create_adapter_from_launcher( self, launcher: LauncherProtocol[_LaunchableT], settings: _DispatchableT ) -> _LaunchConfigType: + """Creates configured launcher from an existing launcher using the provided settings + + :param launcher: A launcher that the type of `settings` has been + configured to dispatch to. + :param settings: A settings with which to configure the launcher. + :returns: A configured launcher. + """ if not self._is_compatible_launcher(launcher): raise TypeError( f"Cannot create launcher adapter from launcher `{launcher}` " @@ -195,6 +216,16 @@ def configure_first_compatible_launcher( with_settings: _DispatchableT, from_available_launchers: t.Iterable[LauncherProtocol[t.Any]], ) -> _LaunchConfigType: + """Configure the first compatible adapter launch to launch with the + provided settings. Launchers are iterated and discarded from the + iterator until the iterator is exhausted. + + :param with_settings: The settings with which to configure the launcher + :param from_available_launchers: An iterable that yields launcher instances + :raises errors.LauncherNotFoundError: No compatible launcher was + yielded from the provided iterator. + :returns: A launcher configured with the provided settings. + """ launcher = helpers.first(self._is_compatible_launcher, from_available_launchers) if launcher is None: raise errors.LauncherNotFoundError( From 9984d18d6286e9cb82f35224d711fc367426c7c7 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Wed, 24 Jul 2024 12:50:00 -0500 Subject: [PATCH 31/34] Moar reviewer feedback --- smartsim/settings/launchSettings.py | 7 --- tests/test_experiment.py | 98 ++++++++++++++++++++++++++--- 2 files changed, 88 insertions(+), 17 deletions(-) diff --git a/smartsim/settings/launchSettings.py b/smartsim/settings/launchSettings.py index 066309d19..1a606a1bc 100644 --- a/smartsim/settings/launchSettings.py +++ b/smartsim/settings/launchSettings.py @@ -75,13 +75,6 @@ def launch_args(self) -> LaunchArguments: """Return the launch argument translator.""" return self._arguments - @launch_args.setter - def launch_args(self, args: t.Mapping[str, str]) -> None: - """Update the launch arguments.""" - self.launch_args._launch_args.clear() - for k, v in args.items(): - self.launch_args.set(k, v) - @property def env_vars(self) -> dict[str, str | None]: """Return an immutable list of attached environment variables.""" diff --git a/tests/test_experiment.py b/tests/test_experiment.py index b7f4b1efa..b2d18a25f 100644 --- a/tests/test_experiment.py +++ b/tests/test_experiment.py @@ -46,19 +46,32 @@ @pytest.fixture def experiment(test_dir): + """A simple experiment instance with a unique name anda unique name and its + own directory to be used by tests + """ yield Experiment(f"test-exp-{uuid.uuid4()}", test_dir) @pytest.fixture def dispatcher(): + """A pre-configured dispatcher to be used by experiments that simply + dispatches any jobs with `MockLaunchArgs` to a `NoOpRecordLauncher` + """ d = dispatch.Dispatcher() - to_record = lambda *a: LaunchRecord(*a) + to_record: dispatch._FormatterType[MockLaunchArgs, LaunchRecord] = ( + lambda settings, exe, env: LaunchRecord(settings, exe, env) + ) d.dispatch(MockLaunchArgs, with_format=to_record, to_launcher=NoOpRecordLauncher) yield d @pytest.fixture def job_maker(monkeypatch): + """A fixture to generate a never ending stream of `Job` instances each + configured with a unique `MockLaunchArgs` instance, but identical + executable. + """ + def iter_jobs(): for i in itertools.count(): settings = launchSettings.LaunchSettings("local") @@ -66,11 +79,35 @@ def iter_jobs(): yield job.Job(EchoHelloWorldEntity(), settings) jobs = iter_jobs() - return lambda: next(jobs) + yield lambda: next(jobs) -@dataclasses.dataclass(frozen=True) +JobMakerType: t.TypeAlias = t.Callable[[], job.Job] + + +@dataclasses.dataclass(frozen=True, eq=False) class NoOpRecordLauncher(dispatch.LauncherProtocol): + """Simple launcher to track the order of and mapping of ids to `start` + method calls. It has exactly three attrs: + + - `created_by_experiment`: + A back ref to the experiment used when calling + `NoOpRecordLauncher.create`. + + - `launched_order`: + An append-only list of `LaunchRecord`s that it has "started". Notice + that this launcher will not actually open any subprocesses/run any + threads/otherwise execute the contents of the record on the system + + - `ids_to_launched`: + A mapping where keys are the generated launched id returned from + a `NoOpRecordLauncher.start` call and the values are the + `LaunchRecord` that was passed into `NoOpRecordLauncher.start` to + cause the id to be generated. + + This is helpful for testing that launchers are handling the expected input + """ + created_by_experiment: Experiment launched_order: list[LaunchRecord] = dataclasses.field(default_factory=list) ids_to_launched: dict[dispatch.LaunchedJobID, LaunchRecord] = dataclasses.field( @@ -97,7 +134,15 @@ class LaunchRecord: env: t.Mapping[str, str | None] @classmethod - def from_job(cls, job): + def from_job(cls, job: job.Job): + """Create a launch record for what we would expect a launch record to + look like having gone through the launching process + + :param job: A job that has or will be launched through an experiment + and dispatched to a `NoOpRecordLauncher` + :returns: A `LaunchRecord` that should evaluate to being equivilient to + that of the one stored in the `NoOpRecordLauncher` + """ args = job._launch_settings.launch_args entity = job._entity env = job._launch_settings.env_vars @@ -105,14 +150,19 @@ def from_job(cls, job): class MockLaunchArgs(launchArguments.LaunchArguments): - def __init__(self, count): + """A `LaunchArguments` subclass that will evaluate as true with another if + and only if they were initialized with the same id. In practice this class + has no arguments to set. + """ + + def __init__(self, id_: int): super().__init__({}) - self.count = count + self.id = id_ def __eq__(self, other): if type(self) is not type(other): return NotImplemented - return other.count == self.count + return other.id == self.id def launcher_str(self): return "test-launch-args" @@ -121,6 +171,8 @@ def set(self, arg, val): ... class EchoHelloWorldEntity(entity.SmartSimEntity): + """A simple smartsim entity that meets the `ExecutableProtocol` protocol""" + def __init__(self): path = tempfile.TemporaryDirectory() self._finalizer = weakref.finalize(self, path.cleanup) @@ -151,7 +203,13 @@ def test_start_raises_if_no_args_supplied(experiment): ), ) # fmt: on -def test_start_can_launch_jobs(experiment, job_maker, dispatcher, make_jobs, num_jobs): +def test_start_can_launch_jobs( + experiment: Experiment, + job_maker: JobMakerType, + dispatcher: dispatch.Dispatcher, + make_jobs: t.Callable[[JobMakerType, int], tuple[job.Job, ...]], + num_jobs: int, +) -> None: jobs = make_jobs(job_maker, num_jobs) assert len(experiment._active_launchers) == 0, "Initialized w/ launchers" launched_ids = experiment.start(*jobs, dispatcher=dispatcher) @@ -163,7 +221,14 @@ def test_start_can_launch_jobs(experiment, job_maker, dispatcher, make_jobs, num len(jobs) == len(launcher.launched_order) == len(launched_ids) == num_jobs ), "Inconsistent number of jobs/launched jobs/launched ids/expected number of jobs" expected_launched = [LaunchRecord.from_job(job) for job in jobs] + + # Check that `job_a, job_b, job_c, ...` are started in that order when + # calling `experiemnt.start(job_a, job_b, job_c, ...)` assert expected_launched == list(launcher.launched_order), "Unexpected launch order" + + # Similarly, check that `id_a, id_b, id_c, ...` corresponds to + # `job_a, job_b, job_c, ...` when calling + # `id_a, id_b, id_c, ... = experiemnt.start(job_a, job_b, job_c, ...)` expected_id_map = dict(zip(launched_ids, expected_launched)) assert expected_id_map == launcher.ids_to_launched, "IDs returned in wrong order" @@ -173,8 +238,11 @@ def test_start_can_launch_jobs(experiment, job_maker, dispatcher, make_jobs, num [pytest.param(i, id=f"{i} start(s)") for i in (1, 2, 3, 5, 10, 100, 1_000)], ) def test_start_can_start_a_job_multiple_times_accross_multiple_calls( - experiment, job_maker, dispatcher, num_starts -): + experiment: Experiment, + job_maker: JobMakerType, + dispatcher: dispatch.Dispatcher, + num_starts: int, +) -> None: assert len(experiment._active_launchers) == 0, "Initialized w/ launchers" job = job_maker() ids_to_launches = { @@ -185,4 +253,14 @@ def test_start_can_start_a_job_multiple_times_accross_multiple_calls( (launcher,) = experiment._active_launchers assert isinstance(launcher, NoOpRecordLauncher), "Unexpected launcher type" assert len(launcher.launched_order) == num_starts, "Unexpected number launches" + + # Check that a single `job` instance can be launched and re-launcherd and + # that `id_a, id_b, id_c, ...` corresponds to + # `"start_a", "start_b", "start_c", ...` when calling + # ```py + # id_a = experiment.start(job) # "start_a" + # id_b = experiment.start(job) # "start_b" + # id_c = experiment.start(job) # "start_c" + # ... + # ``` assert ids_to_launches == launcher.ids_to_launched, "Job was not re-launched" From dbe581d55737d7844757851e6154db5ba1bc8bb1 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Wed, 24 Jul 2024 17:33:03 -0500 Subject: [PATCH 32/34] Moar reviewer requested docs --- smartsim/settings/arguments/batch/lsf.py | 13 +++++- smartsim/settings/arguments/batch/pbs.py | 14 ++++-- smartsim/settings/arguments/batch/slurm.py | 13 +++++- smartsim/settings/arguments/launch/alps.py | 17 +++++-- smartsim/settings/arguments/launch/dragon.py | 12 ++++- smartsim/settings/arguments/launch/local.py | 17 +++++-- smartsim/settings/arguments/launch/lsf.py | 17 +++++-- smartsim/settings/arguments/launch/mpi.py | 27 +++++++++--- smartsim/settings/arguments/launch/pals.py | 23 +++++++--- smartsim/settings/arguments/launch/slurm.py | 17 +++++-- .../settings/arguments/launchArguments.py | 23 +++++++++- smartsim/settings/dispatch.py | 44 +++++++++++++++++++ smartsim/settings/launchSettings.py | 26 ++++++++--- 13 files changed, 223 insertions(+), 40 deletions(-) diff --git a/smartsim/settings/arguments/batch/lsf.py b/smartsim/settings/arguments/batch/lsf.py index 4f6e80a70..10dc85763 100644 --- a/smartsim/settings/arguments/batch/lsf.py +++ b/smartsim/settings/arguments/batch/lsf.py @@ -39,7 +39,10 @@ class BsubBatchArguments(BatchArguments): def scheduler_str(self) -> str: - """Get the string representation of the scheduler""" + """Get the string representation of the scheduler + + :returns: The string representation of the scheduler + """ return SchedulerType.Lsf.value def set_walltime(self, walltime: str) -> None: @@ -130,7 +133,7 @@ def set_queue(self, queue: str) -> None: def format_batch_args(self) -> t.List[str]: """Get the formatted batch arguments for a preview - :return: list of batch arguments for Qsub + :return: list of batch arguments for `bsub` """ opts = [] @@ -146,5 +149,11 @@ def format_batch_args(self) -> t.List[str]: return opts def set(self, key: str, value: str | None) -> None: + """Set an arbitrary scheduler argument + + :param key: The launch argument + :param value: A string representation of the value for the launch + argument (if applicable), otherwise `None` + """ # Store custom arguments in the launcher_args self._scheduler_args[key] = value diff --git a/smartsim/settings/arguments/batch/pbs.py b/smartsim/settings/arguments/batch/pbs.py index d67f1be7b..192874c16 100644 --- a/smartsim/settings/arguments/batch/pbs.py +++ b/smartsim/settings/arguments/batch/pbs.py @@ -41,7 +41,10 @@ class QsubBatchArguments(BatchArguments): def scheduler_str(self) -> str: - """Get the string representation of the scheduler""" + """Get the string representation of the scheduler + + :returns: The string representation of the scheduler + """ return SchedulerType.Pbs.value def set_nodes(self, num_nodes: int) -> None: @@ -113,7 +116,7 @@ def set_account(self, account: str) -> None: def format_batch_args(self) -> t.List[str]: """Get the formatted batch arguments for a preview - :return: batch arguments for Qsub + :return: batch arguments for `qsub` :raises ValueError: if options are supplied without values """ opts, batch_arg_copy = self._create_resource_list(self._scheduler_args) @@ -170,5 +173,10 @@ def _create_resource_list( return res, batch_arg_copy def set(self, key: str, value: str | None) -> None: - # Store custom arguments in the launcher_args + """Set an arbitrary launch argument + + :param key: The launch argument + :param value: A string representation of the value for the launch + argument (if applicable), otherwise `None` + """ self._scheduler_args[key] = value diff --git a/smartsim/settings/arguments/batch/slurm.py b/smartsim/settings/arguments/batch/slurm.py index eca26176d..f4725a117 100644 --- a/smartsim/settings/arguments/batch/slurm.py +++ b/smartsim/settings/arguments/batch/slurm.py @@ -40,7 +40,10 @@ class SlurmBatchArguments(BatchArguments): def scheduler_str(self) -> str: - """Get the string representation of the scheduler""" + """Get the string representation of the scheduler + + :returns: The string representation of the scheduler + """ return SchedulerType.Slurm.value def set_walltime(self, walltime: str) -> None: @@ -120,7 +123,7 @@ def set_hostlist(self, host_list: t.Union[str, t.List[str]]) -> None: def format_batch_args(self) -> t.List[str]: """Get the formatted batch arguments for a preview - :return: batch arguments for Sbatch + :return: batch arguments for `sbatch` """ opts = [] # TODO add restricted here @@ -139,5 +142,11 @@ def format_batch_args(self) -> t.List[str]: return opts def set(self, key: str, value: str | None) -> None: + """Set an arbitrary scheduler argument + + :param key: The launch argument + :param value: A string representation of the value for the launch + argument (if applicable), otherwise `None` + """ # Store custom arguments in the launcher_args self._scheduler_args[key] = value diff --git a/smartsim/settings/arguments/launch/alps.py b/smartsim/settings/arguments/launch/alps.py index e92bc7b85..1879dd102 100644 --- a/smartsim/settings/arguments/launch/alps.py +++ b/smartsim/settings/arguments/launch/alps.py @@ -42,11 +42,17 @@ @dispatch(with_format=_as_aprun_command, to_launcher=ShellLauncher) class AprunLaunchArguments(LaunchArguments): def _reserved_launch_args(self) -> set[str]: - """Return reserved launch arguments.""" + """Return reserved launch arguments. + + :returns: The set of reserved launcher arguments + """ return {"wdir"} def launcher_str(self) -> str: - """Get the string representation of the launcher""" + """Get the string representation of the launcher + + :returns: The string representation of the launcher + """ return LauncherType.Alps.value def set_cpus_per_task(self, cpus_per_task: int) -> None: @@ -203,7 +209,12 @@ def format_launch_args(self) -> t.Union[t.List[str], None]: return args def set(self, key: str, value: str | None) -> None: - """Set the launch arguments""" + """Set an arbitrary launch argument + + :param key: The launch argument + :param value: A string representation of the value for the launch + argument (if applicable), otherwise `None` + """ set_check_input(key, value) if key in self._reserved_launch_args(): logger.warning( diff --git a/smartsim/settings/arguments/launch/dragon.py b/smartsim/settings/arguments/launch/dragon.py index 7a6a5aab4..ca383ad91 100644 --- a/smartsim/settings/arguments/launch/dragon.py +++ b/smartsim/settings/arguments/launch/dragon.py @@ -37,7 +37,10 @@ class DragonLaunchArguments(LaunchArguments): def launcher_str(self) -> str: - """Get the string representation of the launcher""" + """Get the string representation of the launcher + + :returns: The string representation of the launcher + """ return LauncherType.Dragon.value def set_nodes(self, nodes: int) -> None: @@ -55,7 +58,12 @@ def set_tasks_per_node(self, tasks_per_node: int) -> None: self.set("tasks_per_node", str(tasks_per_node)) def set(self, key: str, value: str | None) -> None: - """Set the launch arguments""" + """Set an arbitrary launch argument + + :param key: The launch argument + :param value: A string representation of the value for the launch + argument (if applicable), otherwise `None` + """ set_check_input(key, value) if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") diff --git a/smartsim/settings/arguments/launch/local.py b/smartsim/settings/arguments/launch/local.py index f89299500..5b598b9e4 100644 --- a/smartsim/settings/arguments/launch/local.py +++ b/smartsim/settings/arguments/launch/local.py @@ -42,13 +42,17 @@ @dispatch(with_format=_as_local_command, to_launcher=ShellLauncher) class LocalLaunchArguments(LaunchArguments): def launcher_str(self) -> str: - """Get the string representation of the launcher""" + """Get the string representation of the launcher + + :returns: The string representation of the launcher + """ return LauncherType.Local.value def format_env_vars(self, env_vars: StringArgument) -> t.Union[t.List[str], None]: - """Build environment variable string + """Build bash compatible sequence of strins to specify and environment - :returns: formatted list of strings to export variables + :param env_vars: An environment mapping + :returns: the formatted string of environment variables """ formatted = [] for key, val in env_vars.items(): @@ -70,7 +74,12 @@ def format_launch_args(self) -> t.Union[t.List[str], None]: return formatted def set(self, key: str, value: str | None) -> None: - """Set the launch arguments""" + """Set an arbitrary launch argument + + :param key: The launch argument + :param value: A string representation of the value for the launch + argument (if applicable), otherwise `None` + """ set_check_input(key, value) if key in self._launch_args and key != self._launch_args[key]: logger.warning(f"Overwritting argument '{key}' with value '{value}'") diff --git a/smartsim/settings/arguments/launch/lsf.py b/smartsim/settings/arguments/launch/lsf.py index 83e5cdc94..80cd748f1 100644 --- a/smartsim/settings/arguments/launch/lsf.py +++ b/smartsim/settings/arguments/launch/lsf.py @@ -42,11 +42,17 @@ @dispatch(with_format=_as_jsrun_command, to_launcher=ShellLauncher) class JsrunLaunchArguments(LaunchArguments): def launcher_str(self) -> str: - """Get the string representation of the launcher""" + """Get the string representation of the launcher + + :returns: The string representation of the launcher + """ return LauncherType.Lsf.value def _reserved_launch_args(self) -> set[str]: - """Return reserved launch arguments.""" + """Return reserved launch arguments. + + :returns: The set of reserved launcher arguments + """ return {"chdir", "h", "stdio_stdout", "o", "stdio_stderr", "k"} def set_tasks(self, tasks: int) -> None: @@ -105,7 +111,12 @@ def format_launch_args(self) -> t.Union[t.List[str], None]: return args def set(self, key: str, value: str | None) -> None: - """Set the launch arguments""" + """Set an arbitrary launch argument + + :param key: The launch argument + :param value: A string representation of the value for the launch + argument (if applicable), otherwise `None` + """ set_check_input(key, value) if key in self._reserved_launch_args(): logger.warning( diff --git a/smartsim/settings/arguments/launch/mpi.py b/smartsim/settings/arguments/launch/mpi.py index edd93b4ac..85fd38145 100644 --- a/smartsim/settings/arguments/launch/mpi.py +++ b/smartsim/settings/arguments/launch/mpi.py @@ -43,7 +43,10 @@ class _BaseMPILaunchArguments(LaunchArguments): def _reserved_launch_args(self) -> set[str]: - """Return reserved launch arguments.""" + """Return reserved launch arguments. + + :returns: The set of reserved launcher arguments + """ return {"wd", "wdir"} def set_task_map(self, task_mapping: str) -> None: @@ -203,7 +206,12 @@ def format_launch_args(self) -> t.List[str]: return args def set(self, key: str, value: str | None) -> None: - """Set the launch arguments""" + """Set an arbitrary launch argument + + :param key: The launch argument + :param value: A string representation of the value for the launch + argument (if applicable), otherwise `None` + """ set_check_input(key, value) if key in self._reserved_launch_args(): logger.warning( @@ -221,19 +229,28 @@ def set(self, key: str, value: str | None) -> None: @dispatch(with_format=_as_mpirun_command, to_launcher=ShellLauncher) class MpirunLaunchArguments(_BaseMPILaunchArguments): def launcher_str(self) -> str: - """Get the string representation of the launcher""" + """Get the string representation of the launcher + + :returns: The string representation of the launcher + """ return LauncherType.Mpirun.value @dispatch(with_format=_as_mpiexec_command, to_launcher=ShellLauncher) class MpiexecLaunchArguments(_BaseMPILaunchArguments): def launcher_str(self) -> str: - """Get the string representation of the launcher""" + """Get the string representation of the launcher + + :returns: The string representation of the launcher + """ return LauncherType.Mpiexec.value @dispatch(with_format=_as_orterun_command, to_launcher=ShellLauncher) class OrterunLaunchArguments(_BaseMPILaunchArguments): def launcher_str(self) -> str: - """Get the string representation of the launcher""" + """Get the string representation of the launcher + + :returns: The string representation of the launcher + """ return LauncherType.Orterun.value diff --git a/smartsim/settings/arguments/launch/pals.py b/smartsim/settings/arguments/launch/pals.py index ed4de2131..3132f1b02 100644 --- a/smartsim/settings/arguments/launch/pals.py +++ b/smartsim/settings/arguments/launch/pals.py @@ -42,11 +42,17 @@ @dispatch(with_format=_as_pals_command, to_launcher=ShellLauncher) class PalsMpiexecLaunchArguments(LaunchArguments): def launcher_str(self) -> str: - """Get the string representation of the launcher""" + """Get the string representation of the launcher + + :returns: The string representation of the launcher + """ return LauncherType.Pals.value def _reserved_launch_args(self) -> set[str]: - """Return reserved launch arguments.""" + """Return reserved launch arguments. + + :returns: The set of reserved launcher arguments + """ return {"wdir", "wd"} def set_cpu_binding_type(self, bind_type: str) -> None: @@ -139,14 +145,17 @@ def format_launch_args(self) -> t.List[str]: return args def set(self, key: str, value: str | None) -> None: - """Set the launch arguments""" + """Set an arbitrary launch argument + + :param key: The launch argument + :param value: A string representation of the value for the launch + argument (if applicable), otherwise `None` + """ set_check_input(key, value) if key in self._reserved_launch_args(): logger.warning( - ( - f"Could not set argument '{key}': " - f"it is a reserved argument of '{type(self).__name__}'" - ) + f"Could not set argument '{key}': " + f"it is a reserved argument of '{type(self).__name__}'" ) return if key in self._launch_args and key != self._launch_args[key]: diff --git a/smartsim/settings/arguments/launch/slurm.py b/smartsim/settings/arguments/launch/slurm.py index bba79b969..ac485b7c8 100644 --- a/smartsim/settings/arguments/launch/slurm.py +++ b/smartsim/settings/arguments/launch/slurm.py @@ -44,11 +44,17 @@ @dispatch(with_format=_as_srun_command, to_launcher=ShellLauncher) class SlurmLaunchArguments(LaunchArguments): def launcher_str(self) -> str: - """Get the string representation of the launcher""" + """Get the string representation of the launcher + + :returns: The string representation of the launcher + """ return LauncherType.Slurm.value def _reserved_launch_args(self) -> set[str]: - """Return reserved launch arguments.""" + """Return reserved launch arguments. + + :returns: The set of reserved launcher arguments + """ return {"chdir", "D"} def set_nodes(self, nodes: int) -> None: @@ -305,7 +311,12 @@ def _check_env_vars(self, env_vars: t.Dict[str, t.Optional[str]]) -> None: logger.warning(msg) def set(self, key: str, value: str | None) -> None: - """Set the launch arguments""" + """Set an arbitrary launch argument + + :param key: The launch argument + :param value: A string representation of the value for the launch + argument (if applicable), otherwise `None` + """ set_check_input(key, value) if key in self._reserved_launch_args(): logger.warning( diff --git a/smartsim/settings/arguments/launchArguments.py b/smartsim/settings/arguments/launchArguments.py index da27487f3..61f837d98 100644 --- a/smartsim/settings/arguments/launchArguments.py +++ b/smartsim/settings/arguments/launchArguments.py @@ -65,7 +65,14 @@ def set(self, arg: str, val: str | None) -> None: """ def format_launch_args(self) -> t.Union[t.List[str], None]: - """Build formatted launch arguments""" + """Build formatted launch arguments + + .. warning:: + This method will be removed from this class in a future ticket + + :returns: The launch arguments formatted as a list or `None` if the + arguments cannot be formatted. + """ logger.warning( f"format_launcher_args() not supported for {self.launcher_str()}." ) @@ -78,6 +85,15 @@ def format_comma_sep_env_vars( Slurm takes exports in comma separated lists the list starts with all as to not disturb the rest of the environment for more information on this, see the slurm documentation for srun + + .. warning:: + The return value described in this docstring does not match the + type hint, but I have no idea how this is supposed to be used or + how to resolve the descrepency. I'm not going to try and fix it and + the point is moot as this method is almost certainly going to be + removed in a later ticket. + + :param env_vars: An environment mapping :returns: the formatted string of environment variables """ logger.warning( @@ -89,6 +105,11 @@ def format_env_vars( self, env_vars: t.Dict[str, t.Optional[str]] ) -> t.Union[t.List[str], None]: """Build bash compatible environment variable string for Slurm + + .. warning:: + This method will be removed from this class in a future ticket + + :param env_vars: An environment mapping :returns: the formatted string of environment variables """ logger.warning(f"format_env_vars() not supported for {self.launcher_str()}.") diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index 7b6b9a87a..62a318aed 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -43,17 +43,36 @@ _Ts = TypeVarTuple("_Ts") _T_contra = t.TypeVar("_T_contra", contravariant=True) + _DispatchableT = t.TypeVar("_DispatchableT", bound="LaunchArguments") +"""Any type of luanch arguments, typically used when the type bound by the type +argument is a key a `Dispatcher` dispatch registry +""" _LaunchableT = t.TypeVar("_LaunchableT") +"""Any type, typically used to bind to a type accepted as the input parameter +to the to the `LauncherProtocol.start` method +""" _EnvironMappingType: TypeAlias = t.Mapping[str, "str | None"] +"""A mapping of user provided mapping of environment variables in which to run +a job +""" _FormatterType: TypeAlias = t.Callable[ [_DispatchableT, "ExecutableProtocol", _EnvironMappingType], _LaunchableT ] +"""A callable that is capable of formatting the components of a job into a type +capable of being launched by a launcher. +""" _LaunchConfigType: TypeAlias = ( "_LauncherAdapter[ExecutableProtocol, _EnvironMappingType]" ) +"""A launcher adapater that has configured a launcher to launch the components +of a job with some pre-determined launch settings +""" _UnkownType: TypeAlias = t.NoReturn +"""A type alias for a bottom type. Use this to inform a user that the parameter +a parameter should never be set or a callable will never return +""" @t.final @@ -69,6 +88,14 @@ def __init__( t.Mapping[type[LaunchArguments], _DispatchRegistration[t.Any, t.Any]] | None ) = None, ) -> None: + """Initialize a new `Dispatcher` + + :param dispatch_registry: A pre-configured dispatch registry that the + dispatcher should use. This registry is not type checked and is + used blindly. This registry is shallow copied, meaning that adding + into the original registry after construction will not mutate the + state of the registry. + """ self._dispatch_registry = ( dict(dispatch_registry) if dispatch_registry is not None else {} ) @@ -237,11 +264,22 @@ def configure_first_compatible_launcher( @t.final class _LauncherAdapter(t.Generic[Unpack[_Ts]]): + """An adapter class that will wrap a launcher and allow for a unique + signature of for its `start` method + """ + def __init__( self, launcher: LauncherProtocol[_LaunchableT], map_: t.Callable[[Unpack[_Ts]], _LaunchableT], ) -> None: + """Initialize a launcher adapter + + :param launcher: The launcher instance this class should wrap + :param map_: A callable with arguments for the new `start` method that + can translate them into the expected launching type for the wrapped + launcher. + """ # NOTE: We need to cast off the `_LaunchableT` -> `Any` in the # `__init__` method signature to hide the transform from users of # this class. If possible, this type should not be exposed to @@ -250,6 +288,12 @@ def __init__( self._adapted_launcher: LauncherProtocol[t.Any] = launcher def start(self, *args: Unpack[_Ts]) -> LaunchedJobID: + """Start a new job through the wrapped launcher using the custom + `start` signature + + :param args: The custom start arguments + :returns: The launched job id provided by the wrapped launcher + """ payload = self._adapt(*args) return self._adapted_launcher.start(payload) diff --git a/smartsim/settings/launchSettings.py b/smartsim/settings/launchSettings.py index 1a606a1bc..98c199b83 100644 --- a/smartsim/settings/launchSettings.py +++ b/smartsim/settings/launchSettings.py @@ -67,22 +67,38 @@ def __init__( @property def launcher(self) -> str: - """The launcher type""" + """The launcher type + + :returns: The launcher type's string representation + """ return self._launcher.value @property def launch_args(self) -> LaunchArguments: - """Return the launch argument translator.""" + """The launch argument + + :returns: The launch arguments + """ return self._arguments @property - def env_vars(self) -> dict[str, str | None]: - """Return an immutable list of attached environment variables.""" + def env_vars(self) -> t.Mapping[str, str | None]: + """A mapping of environment variables to set or remove. This mapping is + a deep copy of the mapping used by the settings and as such altering + will not mutate the settings. + + :returns: An environment mapping + """ return copy.deepcopy(self._env_vars) @env_vars.setter def env_vars(self, value: dict[str, str | None]) -> None: - """Set the environment variables.""" + """Set the environment variables to a new mapping. This setter will + make a copy of the mapping and as such altering the original mapping + will not mutate the settings. + + :param value: The new environment mapping + """ self._env_vars = copy.deepcopy(value) def _get_arguments(self, launch_args: StringArgument | None) -> LaunchArguments: From 0b4e65d2ca8b79efadbe2ff3b6fffbe33025436a Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Wed, 24 Jul 2024 17:44:07 -0500 Subject: [PATCH 33/34] typo --- smartsim/settings/arguments/launch/local.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/smartsim/settings/arguments/launch/local.py b/smartsim/settings/arguments/launch/local.py index 5b598b9e4..05fd63a8e 100644 --- a/smartsim/settings/arguments/launch/local.py +++ b/smartsim/settings/arguments/launch/local.py @@ -49,7 +49,7 @@ def launcher_str(self) -> str: return LauncherType.Local.value def format_env_vars(self, env_vars: StringArgument) -> t.Union[t.List[str], None]: - """Build bash compatible sequence of strins to specify and environment + """Build bash compatible sequence of strings to specify and environment :param env_vars: An environment mapping :returns: the formatted string of environment variables From 7bcc8fd084f7a414ca17b6fd9de7b7767fb723c9 Mon Sep 17 00:00:00 2001 From: Matt Drozt Date: Fri, 26 Jul 2024 14:21:48 -0500 Subject: [PATCH 34/34] Address (online and offline) reviewer comments --- smartsim/experiment.py | 45 ++++--- smartsim/settings/arguments/launch/local.py | 2 +- smartsim/settings/dispatch.py | 113 ++++++++++++++++-- .../temp_tests/test_settings/test_dispatch.py | 4 +- tests/test_experiment.py | 17 ++- 5 files changed, 137 insertions(+), 44 deletions(-) diff --git a/smartsim/experiment.py b/smartsim/experiment.py index d0ad321d5..35d1a5eb1 100644 --- a/smartsim/experiment.py +++ b/smartsim/experiment.py @@ -38,7 +38,7 @@ from smartsim._core.config import CONFIG from smartsim.error import errors -from smartsim.settings.dispatch import DEFAULT_DISPATCHER +from smartsim.settings import dispatch from smartsim.status import SmartSimStatus from ._core import Controller, Generator, Manifest, previewrenderer @@ -55,11 +55,7 @@ if t.TYPE_CHECKING: from smartsim.launchable.job import Job - from smartsim.settings.dispatch import ( - Dispatcher, - ExecutableProtocol, - LauncherProtocol, - ) + from smartsim.settings.dispatch import ExecutableProtocol, LauncherProtocol from smartsim.types import LaunchedJobID logger = get_logger(__name__) @@ -168,27 +164,38 @@ def __init__(self, name: str, exp_path: str | None = None): self._fs_identifiers: t.Set[str] = set() """Set of feature store identifiers currently in use by this - experiment""" + experiment + """ self._telemetry_cfg = ExperimentTelemetryConfiguration() """Switch to specify if telemetry data should be produced for this - experiment""" + experiment + """ - def start( - self, job: Job, *jobs: Job, dispatcher: Dispatcher = DEFAULT_DISPATCHER - ) -> tuple[LaunchedJobID, ...]: + def start(self, *jobs: Job) -> tuple[LaunchedJobID, ...]: """Execute a collection of `Job` instances. - :param job: The job instance to start :param jobs: A collection of other job instances to start - :param dispatcher: The dispatcher that should be used to determine how - to start a job based on its settings. If not specified it will - default to a dispatcher pre-configured by SmartSim. :returns: A sequence of ids with order corresponding to the sequence of jobs that can be used to query or alter the status of that particular execution of the job. """ + return self._dispatch(dispatch.DEFAULT_DISPATCHER, *jobs) + + def _dispatch( + self, dispatcher: dispatch.Dispatcher, job: Job, *jobs: Job + ) -> tuple[LaunchedJobID, ...]: + """Dispatch a series of jobs with a particular dispatcher + + :param dispatcher: The dispatcher that should be used to determine how + to start a job based on its launch settings. + :param job: The first job instance to dispatch + :param jobs: A collection of other job instances to dispatch + :returns: A sequence of ids with order corresponding to the sequence of + jobs that can be used to query or alter the status of that + particular dispatch of the job. + """ - def _start(job: Job) -> LaunchedJobID: + def execute_dispatch(job: Job) -> LaunchedJobID: args = job.launch_settings.launch_args env = job.launch_settings.env_vars # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> @@ -203,13 +210,13 @@ def _start(job: Job) -> LaunchedJobID: # configured to handle the launch arguments ... launch_config = dispatch.configure_first_compatible_launcher( from_available_launchers=self._active_launchers, - with_settings=args, + with_arguments=args, ) except errors.LauncherNotFoundError: # ... otherwise create a new launcher that _can_ handle the # launch arguments and configure _that_ one launch_config = dispatch.create_new_launcher_configuration( - for_experiment=self, with_settings=args + for_experiment=self, with_arguments=args ) # Save the underlying launcher instance. That way we do not need to # spin up a launcher instance for each individual job, and it makes @@ -218,7 +225,7 @@ def _start(job: Job) -> LaunchedJobID: self._active_launchers.add(launch_config._adapted_launcher) return launch_config.start(exe, env) - return _start(job), *map(_start, jobs) + return execute_dispatch(job), *map(execute_dispatch, jobs) @_contextualize def generate( diff --git a/smartsim/settings/arguments/launch/local.py b/smartsim/settings/arguments/launch/local.py index 05fd63a8e..0bbba2584 100644 --- a/smartsim/settings/arguments/launch/local.py +++ b/smartsim/settings/arguments/launch/local.py @@ -49,7 +49,7 @@ def launcher_str(self) -> str: return LauncherType.Local.value def format_env_vars(self, env_vars: StringArgument) -> t.Union[t.List[str], None]: - """Build bash compatible sequence of strings to specify and environment + """Build bash compatible sequence of strings to specify an environment :param env_vars: An environment mapping :returns: the formatted string of environment variables diff --git a/smartsim/settings/dispatch.py b/smartsim/settings/dispatch.py index 62a318aed..53c6be04d 100644 --- a/smartsim/settings/dispatch.py +++ b/smartsim/settings/dispatch.py @@ -79,6 +79,24 @@ class Dispatcher: """A class capable of deciding which launcher type should be used to launch a given settings type. + + The `Dispatcher` class maintains a type safe API for adding and retrieving + a settings type into the underlying mapping. It does this through two main + methods: `Dispatcher.dispatch` and `Dispatcher.get_dispatch`. + + `Dispatcher.dispatch` takes in a dispatchable type, a launcher type that is + capable of launching a launchable type and formatting function that maps an + instance of the dispatchable type to an instance of the launchable type. + The dispatcher will then take these components and then enter them into its + dispatch registry. `Dispatcher.dispatch` can also be used as a decorator, + to automatically add a dispatchable type dispatch to a dispatcher at type + creation time. + + `Dispatcher.get_dispatch` takes a dispatchable type or instance as a + parameter, and will attempt to look up, in its dispatch registry, how to + dispatch that type. It will then return an object that can configure a + launcher of the expected launcher type. If the dispatchable type was never + registered a `TypeError` will be raised. """ def __init__( @@ -105,7 +123,7 @@ def copy(self) -> Self: return type(self)(dispatch_registry=self._dispatch_registry) @t.overload - def dispatch( + def dispatch( # Signature when used as a decorator self, args: None = ..., *, @@ -114,7 +132,7 @@ def dispatch( allow_overwrite: bool = ..., ) -> t.Callable[[type[_DispatchableT]], type[_DispatchableT]]: ... @t.overload - def dispatch( + def dispatch( # Signature when used as a method self, args: type[_DispatchableT], *, @@ -122,7 +140,7 @@ def dispatch( to_launcher: type[LauncherProtocol[_LaunchableT]], allow_overwrite: bool = ..., ) -> None: ... - def dispatch( + def dispatch( # Actual implementation self, args: type[_DispatchableT] | None = None, *, @@ -202,7 +220,7 @@ def _is_compatible_launcher(self, launcher: LauncherProtocol[t.Any]) -> bool: return type(launcher) is self.launcher_type def create_new_launcher_configuration( - self, for_experiment: Experiment, with_settings: _DispatchableT + self, for_experiment: Experiment, with_arguments: _DispatchableT ) -> _LaunchConfigType: """Create a new instance of a launcher for an experiment that the provided settings where set to dispatch to, and configure it with the @@ -214,12 +232,13 @@ def create_new_launcher_configuration( :returns: A configured launcher """ launcher = self.launcher_type.create(for_experiment) - return self.create_adapter_from_launcher(launcher, with_settings) + return self.create_adapter_from_launcher(launcher, with_arguments) def create_adapter_from_launcher( - self, launcher: LauncherProtocol[_LaunchableT], settings: _DispatchableT + self, launcher: LauncherProtocol[_LaunchableT], arguments: _DispatchableT ) -> _LaunchConfigType: - """Creates configured launcher from an existing launcher using the provided settings + """Creates configured launcher from an existing launcher using the + provided settings. :param launcher: A launcher that the type of `settings` has been configured to dispatch to. @@ -234,13 +253,13 @@ def create_adapter_from_launcher( ) def format_(exe: ExecutableProtocol, env: _EnvironMappingType) -> _LaunchableT: - return self.formatter(settings, exe, env) + return self.formatter(arguments, exe, env) return _LauncherAdapter(launcher, format_) def configure_first_compatible_launcher( self, - with_settings: _DispatchableT, + with_arguments: _DispatchableT, from_available_launchers: t.Iterable[LauncherProtocol[t.Any]], ) -> _LaunchConfigType: """Configure the first compatible adapter launch to launch with the @@ -259,13 +278,47 @@ def configure_first_compatible_launcher( f"No launcher of exactly type `{self.launcher_type.__name__}` " "could be found from provided launchers" ) - return self.create_adapter_from_launcher(launcher, with_settings) + return self.create_adapter_from_launcher(launcher, with_arguments) @t.final class _LauncherAdapter(t.Generic[Unpack[_Ts]]): - """An adapter class that will wrap a launcher and allow for a unique - signature of for its `start` method + """The launcher adapter is an adapter class takes a launcher that is + capable of launching some type `LaunchableT` and a function with a generic + argument list that returns a `LaunchableT`. The launcher adapter will then + provide `start` method that will have the same argument list as the + provided function and launch the output through the provided launcher. + + For example, the launcher adapter could be used like so: + + .. highlight:: python + .. code-block:: python + + class SayHelloLauncher(LauncherProtocol[str]): + ... + def start(self, title: str): + ... + print(f"Hello, {title}") + ... + ... + + @dataclasses.dataclass + class Person: + name: str + honorific: str + + def full_title(self) -> str: + return f"{honorific}. {self.name}" + + mark = Person("Jim", "Mr") + sally = Person("Sally", "Ms") + matt = Person("Matt", "Dr") + hello_person_launcher = _LauncherAdapter(SayHelloLauncher, + Person.full_title) + + hello_person_launcher.start(mark) # prints: "Hello, Mr. Mark" + hello_person_launcher.start(sally) # prints: "Hello, Ms. Sally" + hello_person_launcher.start(matt) # prints: "Hello, Dr. Matt" """ def __init__( @@ -299,10 +352,17 @@ def start(self, *args: Unpack[_Ts]) -> LaunchedJobID: DEFAULT_DISPATCHER: t.Final = Dispatcher() +"""A global `Dispatcher` instance that SmartSim automatically configures to +launch its built in launchables +""" + # Disabling because we want this to look and feel like a top level function, # but don't want to have a second copy of the nasty overloads # pylint: disable-next=invalid-name dispatch: t.Final = DEFAULT_DISPATCHER.dispatch +"""Function that can be used as a decorator to add a dispatch registration into +`DEFAULT_DISPATCHER`. +""" # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> @@ -327,6 +387,35 @@ def create(cls, exp: Experiment, /) -> Self: ... def make_shell_format_fn( run_command: str | None, ) -> _FormatterType[LaunchArguments, t.Sequence[str]]: + """A function that builds a function that formats a `LaunchArguments` as a + shell executable sequence of strings for a given launching utility. + + Example usage: + + .. highlight:: python + .. code-block:: python + + echo_hello_world: ExecutableProtocol = ... + env = {} + slurm_args: SlurmLaunchArguments = ... + slurm_args.set_nodes(3) + + as_srun_command = make_shell_format_fn("srun") + fmt_cmd = as_srun_command(slurm_args, echo_hello_world, env) + print(list(fmt_cmd)) + # prints: "['srun', '--nodes=3', '--', 'echo', 'Hello World!']" + + .. note:: + This function was/is a kind of slap-dash implementation, and is likely + to change or be removed entierely as more functionality is added to the + shell launcher. Use with caution and at your own risk! + + :param run_command: Name or path of the launching utility to invoke with + the arguments. + :returns: A function to format an arguments, an executable, and an + environment as a shell launchable sequence for strings. + """ + def impl( args: LaunchArguments, exe: ExecutableProtocol, _env: _EnvironMappingType ) -> t.Sequence[str]: diff --git a/tests/temp_tests/test_settings/test_dispatch.py b/tests/temp_tests/test_settings/test_dispatch.py index 78e4ec349..9c99cb7d0 100644 --- a/tests/temp_tests/test_settings/test_dispatch.py +++ b/tests/temp_tests/test_settings/test_dispatch.py @@ -384,7 +384,7 @@ def resolve_instance(inst): with ctx: adapter = buffer_writer_dispatch.configure_first_compatible_launcher( - with_settings=mock_launch_args, from_available_launchers=launcher_instances + with_arguments=mock_launch_args, from_available_launchers=launcher_instances ) @@ -395,7 +395,7 @@ class MockExperiment: ... exp = MockExperiment() adapter_1 = buffer_writer_dispatch.create_new_launcher_configuration( - for_experiment=exp, with_settings=mock_launch_args + for_experiment=exp, with_arguments=mock_launch_args ) assert type(adapter_1._adapted_launcher) == buffer_writer_dispatch.launcher_type existing_launcher = adapter_1._adapted_launcher diff --git a/tests/test_experiment.py b/tests/test_experiment.py index b2d18a25f..6571763d7 100644 --- a/tests/test_experiment.py +++ b/tests/test_experiment.py @@ -45,11 +45,13 @@ @pytest.fixture -def experiment(test_dir): +def experiment(monkeypatch, test_dir, dispatcher): """A simple experiment instance with a unique name anda unique name and its own directory to be used by tests """ - yield Experiment(f"test-exp-{uuid.uuid4()}", test_dir) + exp = Experiment(f"test-exp-{uuid.uuid4()}", test_dir) + monkeypatch.setattr(dispatch, "DEFAULT_DISPATCHER", dispatcher) + yield exp @pytest.fixture @@ -206,13 +208,12 @@ def test_start_raises_if_no_args_supplied(experiment): def test_start_can_launch_jobs( experiment: Experiment, job_maker: JobMakerType, - dispatcher: dispatch.Dispatcher, make_jobs: t.Callable[[JobMakerType, int], tuple[job.Job, ...]], num_jobs: int, ) -> None: jobs = make_jobs(job_maker, num_jobs) assert len(experiment._active_launchers) == 0, "Initialized w/ launchers" - launched_ids = experiment.start(*jobs, dispatcher=dispatcher) + launched_ids = experiment.start(*jobs) assert len(experiment._active_launchers) == 1, "Unexpected number of launchers" (launcher,) = experiment._active_launchers assert isinstance(launcher, NoOpRecordLauncher), "Unexpected launcher type" @@ -238,16 +239,12 @@ def test_start_can_launch_jobs( [pytest.param(i, id=f"{i} start(s)") for i in (1, 2, 3, 5, 10, 100, 1_000)], ) def test_start_can_start_a_job_multiple_times_accross_multiple_calls( - experiment: Experiment, - job_maker: JobMakerType, - dispatcher: dispatch.Dispatcher, - num_starts: int, + experiment: Experiment, job_maker: JobMakerType, num_starts: int ) -> None: assert len(experiment._active_launchers) == 0, "Initialized w/ launchers" job = job_maker() ids_to_launches = { - experiment.start(job, dispatcher=dispatcher)[0]: LaunchRecord.from_job(job) - for _ in range(num_starts) + experiment.start(job)[0]: LaunchRecord.from_job(job) for _ in range(num_starts) } assert len(experiment._active_launchers) == 1, "Did not reuse the launcher" (launcher,) = experiment._active_launchers