Removed Slurm specific methods from LaunchArguments

smartsim/_core/arguments/shell.py

@@ -0,0 +1,42 @@
+# 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
+from abc import abstractmethod
+
+from smartsim.log import get_logger
+from smartsim.settings.arguments.launchArguments import LaunchArguments
+
+logger = get_logger(__name__)
+
+
+class ShellLaunchArguments(LaunchArguments):
+    @abstractmethod
+    def format_env_vars(self, env_vars: t.Mapping[str, str | None]) -> list[str]: ...
+    @abstractmethod
+    def format_launch_args(self) -> list[str]: ...

smartsim/_core/dispatch.py

@@ -37,6 +37,7 @@
 from smartsim.types import LaunchedJobID
 
 if t.TYPE_CHECKING:
+    from smartsim._core.arguments.shell import ShellLaunchArguments
     from smartsim._core.utils.launcher import ExecutableProtocol, LauncherProtocol
     from smartsim.experiment import Experiment
     from smartsim.settings.arguments import LaunchArguments

smartsim/_core/shell/shellLauncher.py

@@ -33,6 +33,7 @@
 
 import psutil
 
+from smartsim._core.arguments.shell import ShellLaunchArguments
 from smartsim._core.dispatch import _EnvironMappingType, _FormatterType, dispatch
 from smartsim._core.utils import helpers
 from smartsim._core.utils.launcher import ExecutableProtocol, create_job_id
@@ -103,7 +104,9 @@ def create(cls, _: Experiment) -> Self:
 
 def make_shell_format_fn(
     run_command: str | None,
-) -> _FormatterType[LaunchArguments, tuple[str | os.PathLike[str], t.Sequence[str]]]:
+) -> _FormatterType[
+    ShellLaunchArguments, tuple[str | os.PathLike[str], 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.
 
@@ -134,7 +137,7 @@ def make_shell_format_fn(
     """
 
     def impl(
-        args: LaunchArguments,
+        args: ShellLaunchArguments,
         exe: ExecutableProtocol,
         path: str | os.PathLike[str],
         _env: _EnvironMappingType,

smartsim/settings/arguments/launch/alps.py

@@ -28,20 +28,20 @@
 
 import typing as t
 
+from smartsim._core.arguments.shell import ShellLaunchArguments
 from smartsim._core.dispatch import dispatch
 from smartsim._core.shell.shellLauncher import ShellLauncher, make_shell_format_fn
 from smartsim.log import get_logger
 
 from ...common import set_check_input
 from ...launchCommand import LauncherType
-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 AprunLaunchArguments(LaunchArguments):
+class AprunLaunchArguments(ShellLaunchArguments):
     def _reserved_launch_args(self) -> set[str]:
         """Return reserved launch arguments.
 
@@ -177,9 +177,7 @@ def set_quiet_launch(self, quiet: bool) -> None:
         else:
             self._launch_args.pop("quiet", None)
 
-    def format_env_vars(
-        self, env_vars: t.Optional[t.Dict[str, t.Optional[str]]]
-    ) -> t.Union[t.List[str], None]:
+    def format_env_vars(self, env_vars: t.Mapping[str, str | None]) -> list[str]:
         """Format the environment variables for aprun
 
         :return: list of env vars
@@ -190,7 +188,7 @@ def format_env_vars(
                 formatted += ["-e", name + "=" + str(value)]
         return formatted
 
-    def format_launch_args(self) -> t.Union[t.List[str], None]:
+    def format_launch_args(self) -> t.List[str]:
         """Return a list of ALPS formatted run arguments
 
         :return: list of ALPS arguments for these settings

smartsim/settings/arguments/launch/local.py

@@ -28,28 +28,28 @@
 
 import typing as t
 
+from smartsim._core.arguments.shell import ShellLaunchArguments
 from smartsim._core.dispatch import dispatch
 from smartsim._core.shell.shellLauncher import ShellLauncher, make_shell_format_fn
 from smartsim.log import get_logger
 
 from ...common import StringArgument, set_check_input
 from ...launchCommand import LauncherType
-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 LocalLaunchArguments(LaunchArguments):
+class LocalLaunchArguments(ShellLaunchArguments):
     def launcher_str(self) -> str:
         """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]:
+    def format_env_vars(self, env_vars: t.Mapping[str, str | None]) -> list[str]:
         """Build bash compatible sequence of strings to specify an environment
 
         :param env_vars: An environment mapping
@@ -63,7 +63,7 @@ def format_env_vars(self, env_vars: StringArgument) -> t.Union[t.List[str], None
                 formatted.append(f"{key}={val}")
         return formatted
 
-    def format_launch_args(self) -> t.Union[t.List[str], None]:
+    def format_launch_args(self) -> t.List[str]:
         """Build launcher argument string
 
         :returns: formatted list of launcher arguments

smartsim/settings/arguments/launch/lsf.py

@@ -28,20 +28,20 @@
 
 import typing as t
 
+from smartsim._core.arguments.shell import ShellLaunchArguments
 from smartsim._core.dispatch import dispatch
 from smartsim._core.shell.shellLauncher import ShellLauncher, make_shell_format_fn
 from smartsim.log import get_logger
 
 from ...common import set_check_input
 from ...launchCommand import LauncherType
-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 JsrunLaunchArguments(LaunchArguments):
+class JsrunLaunchArguments(ShellLaunchArguments):
     def launcher_str(self) -> str:
         """Get the string representation of the launcher
 
@@ -74,9 +74,7 @@ def set_binding(self, binding: str) -> None:
         """
         self.set("bind", binding)
 
-    def format_env_vars(
-        self, env_vars: t.Dict[str, t.Optional[str]]
-    ) -> t.Union[t.List[str], None]:
+    def format_env_vars(self, env_vars: t.Mapping[str, str | None]) -> list[str]:
         """Format environment variables. Each variable needs
         to be passed with ``--env``. If a variable is set to ``None``,
         its value is propagated from the current environment.
@@ -91,7 +89,7 @@ def format_env_vars(
                 format_str += ["-E", f"{k}"]
         return format_str
 
-    def format_launch_args(self) -> t.Union[t.List[str], None]:
+    def format_launch_args(self) -> t.List[str]:
         """Return a list of LSF formatted run arguments
 
         :return: list of LSF arguments for these settings

smartsim/settings/arguments/launch/mpi.py

@@ -28,21 +28,21 @@
 
 import typing as t
 
+from smartsim._core.arguments.shell import ShellLaunchArguments
 from smartsim._core.dispatch import dispatch
 from smartsim._core.shell.shellLauncher import ShellLauncher, make_shell_format_fn
 from smartsim.log import get_logger
 
 from ...common import set_check_input
 from ...launchCommand import LauncherType
-from ..launchArguments import LaunchArguments
 
 logger = get_logger(__name__)
 _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 _BaseMPILaunchArguments(LaunchArguments):
+class _BaseMPILaunchArguments(ShellLaunchArguments):
     def _reserved_launch_args(self) -> set[str]:
         """Return reserved launch arguments.
 
@@ -172,9 +172,7 @@ def set_quiet_launch(self, quiet: bool) -> None:
         else:
             self._launch_args.pop("quiet", None)
 
-    def format_env_vars(
-        self, env_vars: t.Optional[t.Dict[str, t.Optional[str]]]
-    ) -> t.Union[t.List[str], None]:
+    def format_env_vars(self, env_vars: t.Mapping[str, str | None]) -> list[str]:
         """Format the environment variables for mpirun
 
         :return: list of env vars

smartsim/settings/arguments/launch/pals.py

@@ -28,20 +28,20 @@
 
 import typing as t
 
+from smartsim._core.arguments.shell import ShellLaunchArguments
 from smartsim._core.dispatch import dispatch
 from smartsim._core.shell.shellLauncher import ShellLauncher, make_shell_format_fn
 from smartsim.log import get_logger
 
 from ...common import set_check_input
 from ...launchCommand import LauncherType
-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 PalsMpiexecLaunchArguments(LaunchArguments):
+class PalsMpiexecLaunchArguments(ShellLaunchArguments):
     def launcher_str(self) -> str:
         """Get the string representation of the launcher
 
@@ -106,9 +106,7 @@ def set_hostlist(self, host_list: t.Union[str, t.List[str]]) -> None:
             raise TypeError("host_list argument must be list of strings")
         self.set("hosts", ",".join(host_list))
 
-    def format_env_vars(
-        self, env_vars: t.Optional[t.Dict[str, t.Optional[str]]]
-    ) -> t.Union[t.List[str], None]:
+    def format_env_vars(self, env_vars: t.Mapping[str, str | None]) -> list[str]:
         """Format the environment variables for mpirun
 
         :return: list of env vars

smartsim/settings/arguments/launch/slurm.py

@@ -30,20 +30,20 @@
 import re
 import typing as t
 
+from smartsim._core.arguments.shell import ShellLaunchArguments
 from smartsim._core.dispatch import dispatch
 from smartsim._core.shell.shellLauncher import ShellLauncher, make_shell_format_fn
 from smartsim.log import get_logger
 
 from ...common import set_check_input
 from ...launchCommand import LauncherType
-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 SlurmLaunchArguments(LaunchArguments):
+class SlurmLaunchArguments(ShellLaunchArguments):
     def launcher_str(self) -> str:
         """Get the string representation of the launcher
 
@@ -234,7 +234,7 @@ def set_quiet_launch(self, quiet: bool) -> None:
         else:
             self._launch_args.pop("quiet", None)
 
-    def format_launch_args(self) -> t.Union[t.List[str], None]:
+    def format_launch_args(self) -> t.List[str]:
         """Return a list of slurm formatted launch arguments
 
         :return: list of slurm arguments for these settings
@@ -252,9 +252,7 @@ def format_launch_args(self) -> t.Union[t.List[str], None]:
                     formatted += ["=".join((prefix + key, str(value)))]
         return formatted
 
-    def format_env_vars(
-        self, env_vars: t.Dict[str, t.Optional[str]]
-    ) -> t.Union[t.List[str], None]:
+    def format_env_vars(self, env_vars: t.Mapping[str, str | None]) -> list[str]:
         """Build bash compatible environment variable string for Slurm
 
         :returns: the formatted string of environment variables
@@ -271,6 +269,7 @@ def format_comma_sep_env_vars(
         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
 
+        :param env_vars: An environment mapping
         :returns: the formatted string of environment variables
         """
         self._check_env_vars(env_vars)
@@ -290,7 +289,7 @@ def format_comma_sep_env_vars(
 
         return fmt_exported_env, compound_env
 
-    def _check_env_vars(self, env_vars: t.Dict[str, t.Optional[str]]) -> None:
+    def _check_env_vars(self, env_vars: t.Mapping[str, str | None]) -> None:
         """Warn a user trying to set a variable which is set in the environment
 
         Given Slurm's env var precedence, trying to export a variable which is already

smartsim/settings/arguments/launchArguments.py

@@ -64,57 +64,6 @@ def set(self, arg: str, val: str | None) -> None:
             applicable). Otherwise `None`
         """
 
-    def format_launch_args(self) -> t.Union[t.List[str], None]:
-        """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()}."
-        )
-        return None
-
-    def format_comma_sep_env_vars(
-        self, env_vars: t.Dict[str, t.Optional[str]]
-    ) -> t.Union[t.Tuple[str, t.List[str]], None]:
-        """Build environment variable string for Slurm
-        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(
-            f"format_comma_sep_env_vars() not supported for {self.launcher_str()}."
-        )
-        return None
-
-    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()}.")
-        return None
-
     def __str__(self) -> str:  # pragma: no-cover
         return textwrap.dedent(f"""\
             Launch Arguments: