From 0569c05d4ea798af583fbd33f1c90685eccaa9cd Mon Sep 17 00:00:00 2001 From: Julia Putko <81587103+juliaputko@users.noreply.github.com> Date: Wed, 21 Aug 2024 15:42:44 -0700 Subject: [PATCH] Move `ShellLauncher` (#658) Move Shelllauncher to _core/shell. [ committed by @juliaputko ] [ reviewed by @MattToast , @mellis13 ] --- smartsim/_core/control/launch_history.py | 2 +- smartsim/_core/dispatch.py | 172 +----------------- .../_core/launcher/dragon/dragonLauncher.py | 4 +- smartsim/_core/shell/__init__.py | 25 +++ smartsim/_core/shell/shellLauncher.py | 153 ++++++++++++++++ smartsim/_core/utils/launcher.py | 91 +++++++++ smartsim/experiment.py | 2 +- smartsim/settings/arguments/launch/alps.py | 3 +- smartsim/settings/arguments/launch/local.py | 3 +- smartsim/settings/arguments/launch/lsf.py | 3 +- smartsim/settings/arguments/launch/mpi.py | 3 +- smartsim/settings/arguments/launch/pals.py | 3 +- smartsim/settings/arguments/launch/slurm.py | 3 +- tests/temp_tests/test_settings/conftest.py | 12 +- .../temp_tests/test_settings/test_dispatch.py | 9 +- tests/test_experiment.py | 11 +- tests/test_launch_history.py | 2 +- 17 files changed, 310 insertions(+), 191 deletions(-) create mode 100644 smartsim/_core/shell/__init__.py create mode 100644 smartsim/_core/shell/shellLauncher.py create mode 100644 smartsim/_core/utils/launcher.py diff --git a/smartsim/_core/control/launch_history.py b/smartsim/_core/control/launch_history.py index b8b9f4c7e..e7f04a4ff 100644 --- a/smartsim/_core/control/launch_history.py +++ b/smartsim/_core/control/launch_history.py @@ -32,7 +32,7 @@ from smartsim._core.utils import helpers as _helpers if t.TYPE_CHECKING: - from smartsim._core.dispatch import LauncherProtocol + from smartsim._core.utils.launcher import LauncherProtocol from smartsim.types import LaunchedJobID diff --git a/smartsim/_core/dispatch.py b/smartsim/_core/dispatch.py index 95e80b121..de767d6a4 100644 --- a/smartsim/_core/dispatch.py +++ b/smartsim/_core/dispatch.py @@ -26,28 +26,23 @@ from __future__ import annotations -import abc -import collections.abc import dataclasses import os -import subprocess as sp import typing as t -import uuid -import psutil from typing_extensions import Self, TypeAlias, TypeVarTuple, Unpack from smartsim._core.utils import helpers from smartsim.error import errors -from smartsim.status import JobStatus from smartsim.types import LaunchedJobID if t.TYPE_CHECKING: + from smartsim._core.utils.launcher import ExecutableProtocol, LauncherProtocol from smartsim.experiment import Experiment from smartsim.settings.arguments import LaunchArguments _Ts = TypeVarTuple("_Ts") -_T_contra = t.TypeVar("_T_contra", contravariant=True) + _WorkingDirectory: TypeAlias = t.Union[str, os.PathLike[str]] """A working directory represented as a string or PathLike object""" @@ -75,6 +70,7 @@ _LaunchConfigType: TypeAlias = ( "_LauncherAdapter[ExecutableProtocol, _WorkingDirectory, _EnvironMappingType]" ) + """A launcher adapater that has configured a launcher to launch the components of a job with some pre-determined launch settings """ @@ -232,7 +228,7 @@ def create_new_launcher_configuration( 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 + provided settings were set to dispatch, and configure it with the provided launch settings. :param for_experiment: The experiment responsible creating the launcher @@ -376,163 +372,3 @@ def start(self, *args: Unpack[_Ts]) -> LaunchedJobID: """Function that can be used as a decorator to add a dispatch registration into `DEFAULT_DISPATCHER`. """ - - -# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> -# TODO: move these to a common module under `smartsim._core.launcher` -# ----------------------------------------------------------------------------- - - -def create_job_id() -> LaunchedJobID: - return LaunchedJobID(str(uuid.uuid4())) - - -class ExecutableProtocol(t.Protocol): - def as_program_arguments(self) -> t.Sequence[str]: ... - - -class LauncherProtocol(collections.abc.Hashable, t.Protocol[_T_contra]): - """The protocol defining a launcher that can be used by a SmartSim - experiment - """ - - @classmethod - @abc.abstractmethod - def create(cls, exp: Experiment, /) -> Self: - """Create an new launcher instance from and to be used by the passed in - experiment instance - - :param: An experiment to use the newly created launcher instance - :returns: The newly constructed launcher instance - """ - - @abc.abstractmethod - def start(self, launchable: _T_contra, /) -> LaunchedJobID: - """Given input that this launcher understands, create a new process and - issue a launched job id to query the status of the job in future. - - :param launchable: The input to start a new process - :returns: The id to query the status of the process in future - """ - - @abc.abstractmethod - def get_status( - self, *launched_ids: LaunchedJobID - ) -> t.Mapping[LaunchedJobID, JobStatus]: - """Given a collection of launched job ids, return a mapping of id to - current status of the launched job. If a job id is no recognized by the - launcher, a `smartsim.error.errors.LauncherJobNotFound` error should be - raised. - - :param launched_ids: The collection of ids of launched jobs to query - for current status - :raises smartsim.error.errors.LauncherJobNotFound: If at least one of - the ids of the `launched_ids` collection is not recognized. - :returns: A mapping of launched id to current status - """ - - -def make_shell_format_fn( - run_command: str | None, -) -> _FormatterType[LaunchArguments, 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. - - 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, - path: str | os.PathLike[str], - _env: _EnvironMappingType, - ) -> t.Tuple[str | os.PathLike[str], t.Sequence[str]]: - return path, ( - ( - 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""" - - def __init__(self) -> None: - self._launched: dict[LaunchedJobID, sp.Popen[bytes]] = {} - - def start( - self, command: tuple[str | os.PathLike[str], t.Sequence[str]] - ) -> LaunchedJobID: - id_ = create_job_id() - path, args = command - exe, *rest = args - # pylint: disable-next=consider-using-with - self._launched[id_] = sp.Popen((helpers.expand_exe_path(exe), *rest), cwd=path) - return id_ - - def get_status( - self, *launched_ids: LaunchedJobID - ) -> t.Mapping[LaunchedJobID, JobStatus]: - return {id_: self._get_status(id_) for id_ in launched_ids} - - def _get_status(self, id_: LaunchedJobID, /) -> JobStatus: - if (proc := self._launched.get(id_)) is None: - msg = f"Launcher `{self}` has not launched a job with id `{id_}`" - raise errors.LauncherJobNotFound(msg) - ret_code = proc.poll() - if ret_code is None: - status = psutil.Process(proc.pid).status() - return { - psutil.STATUS_RUNNING: JobStatus.RUNNING, - psutil.STATUS_SLEEPING: JobStatus.RUNNING, - psutil.STATUS_WAKING: JobStatus.RUNNING, - psutil.STATUS_DISK_SLEEP: JobStatus.RUNNING, - psutil.STATUS_DEAD: JobStatus.FAILED, - psutil.STATUS_TRACING_STOP: JobStatus.PAUSED, - psutil.STATUS_WAITING: JobStatus.PAUSED, - psutil.STATUS_STOPPED: JobStatus.PAUSED, - psutil.STATUS_LOCKED: JobStatus.PAUSED, - psutil.STATUS_PARKED: JobStatus.PAUSED, - psutil.STATUS_IDLE: JobStatus.PAUSED, - psutil.STATUS_ZOMBIE: JobStatus.COMPLETED, - }.get(status, JobStatus.UNKNOWN) - if ret_code == 0: - return JobStatus.COMPLETED - return JobStatus.FAILED - - @classmethod - def create(cls, _: Experiment) -> Self: - return cls() - - -# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< diff --git a/smartsim/_core/launcher/dragon/dragonLauncher.py b/smartsim/_core/launcher/dragon/dragonLauncher.py index 26a3d3daf..39e195881 100644 --- a/smartsim/_core/launcher/dragon/dragonLauncher.py +++ b/smartsim/_core/launcher/dragon/dragonLauncher.py @@ -63,8 +63,10 @@ if t.TYPE_CHECKING: from typing_extensions import Self + from smartsim._core.utils.launcher import ExecutableProtocol from smartsim.experiment import Experiment + logger = get_logger(__name__) @@ -355,7 +357,7 @@ def _assert_schema_type(obj: object, typ: t.Type[_SchemaT], /) -> _SchemaT: return obj -from smartsim._core.dispatch import ExecutableProtocol, dispatch +from smartsim._core.dispatch import dispatch # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> # TODO: Remove this registry and move back to builder file after fixing diff --git a/smartsim/_core/shell/__init__.py b/smartsim/_core/shell/__init__.py new file mode 100644 index 000000000..efe03908e --- /dev/null +++ b/smartsim/_core/shell/__init__.py @@ -0,0 +1,25 @@ +# 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. diff --git a/smartsim/_core/shell/shellLauncher.py b/smartsim/_core/shell/shellLauncher.py new file mode 100644 index 000000000..0d0912ab2 --- /dev/null +++ b/smartsim/_core/shell/shellLauncher.py @@ -0,0 +1,153 @@ +# 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 os +import subprocess as sp +import typing as t + +import psutil + +from smartsim._core.dispatch import _EnvironMappingType, _FormatterType, dispatch +from smartsim._core.utils import helpers +from smartsim._core.utils.launcher import ExecutableProtocol, create_job_id +from smartsim.error import errors +from smartsim.log import get_logger +from smartsim.settings.arguments.launchArguments import LaunchArguments +from smartsim.status import JobStatus +from smartsim.types import LaunchedJobID + +if t.TYPE_CHECKING: + from typing_extensions import Self + + from smartsim.experiment import Experiment + +logger = get_logger(__name__) + + +class ShellLauncher: + """Mock launcher for launching/tracking simple shell commands""" + + def __init__(self) -> None: + self._launched: dict[LaunchedJobID, sp.Popen[bytes]] = {} + + def start( + self, command: tuple[str | os.PathLike[str], t.Sequence[str]] + ) -> LaunchedJobID: + id_ = create_job_id() + path, args = command + exe, *rest = args + # pylint: disable-next=consider-using-with + self._launched[id_] = sp.Popen((helpers.expand_exe_path(exe), *rest), cwd=path) + return id_ + + def get_status( + self, *launched_ids: LaunchedJobID + ) -> t.Mapping[LaunchedJobID, JobStatus]: + return {id_: self._get_status(id_) for id_ in launched_ids} + + def _get_status(self, id_: LaunchedJobID, /) -> JobStatus: + if (proc := self._launched.get(id_)) is None: + msg = f"Launcher `{self}` has not launched a job with id `{id_}`" + raise errors.LauncherJobNotFound(msg) + ret_code = proc.poll() + if ret_code is None: + status = psutil.Process(proc.pid).status() + return { + psutil.STATUS_RUNNING: JobStatus.RUNNING, + psutil.STATUS_SLEEPING: JobStatus.RUNNING, + psutil.STATUS_WAKING: JobStatus.RUNNING, + psutil.STATUS_DISK_SLEEP: JobStatus.RUNNING, + psutil.STATUS_DEAD: JobStatus.FAILED, + psutil.STATUS_TRACING_STOP: JobStatus.PAUSED, + psutil.STATUS_WAITING: JobStatus.PAUSED, + psutil.STATUS_STOPPED: JobStatus.PAUSED, + psutil.STATUS_LOCKED: JobStatus.PAUSED, + psutil.STATUS_PARKED: JobStatus.PAUSED, + psutil.STATUS_IDLE: JobStatus.PAUSED, + psutil.STATUS_ZOMBIE: JobStatus.COMPLETED, + }.get(status, JobStatus.UNKNOWN) + if ret_code == 0: + return JobStatus.COMPLETED + return JobStatus.FAILED + + @classmethod + def create(cls, _: Experiment) -> Self: + return cls() + + +def make_shell_format_fn( + run_command: str | None, +) -> _FormatterType[LaunchArguments, 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. + + 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, + path: str | os.PathLike[str], + _env: _EnvironMappingType, + ) -> t.Tuple[str | os.PathLike[str], t.Sequence[str]]: + return path, ( + ( + run_command, + *(args.format_launch_args() or ()), + "--", + *exe.as_program_arguments(), + ) + if run_command is not None + else exe.as_program_arguments() + ) + + return impl diff --git a/smartsim/_core/utils/launcher.py b/smartsim/_core/utils/launcher.py new file mode 100644 index 000000000..32ca3b2e5 --- /dev/null +++ b/smartsim/_core/utils/launcher.py @@ -0,0 +1,91 @@ +# 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 abc +import collections.abc +import typing as t +import uuid + +from typing_extensions import Self + +from smartsim.status import JobStatus +from smartsim.types import LaunchedJobID + +if t.TYPE_CHECKING: + from smartsim.experiment import Experiment + +_T_contra = t.TypeVar("_T_contra", contravariant=True) + + +def create_job_id() -> LaunchedJobID: + return LaunchedJobID(str(uuid.uuid4())) + + +class ExecutableProtocol(t.Protocol): + def as_program_arguments(self) -> t.Sequence[str]: ... + + +class LauncherProtocol(collections.abc.Hashable, t.Protocol[_T_contra]): + """The protocol defining a launcher that can be used by a SmartSim + experiment + """ + + @classmethod + @abc.abstractmethod + def create(cls, exp: Experiment, /) -> Self: + """Create an new launcher instance from and to be used by the passed in + experiment instance + + :param: An experiment to use the newly created launcher instance + :returns: The newly constructed launcher instance + """ + + @abc.abstractmethod + def start(self, launchable: _T_contra, /) -> LaunchedJobID: + """Given input that this launcher understands, create a new process and + issue a launched job id to query the status of the job in future. + + :param launchable: The input to start a new process + :returns: The id to query the status of the process in future + """ + + @abc.abstractmethod + def get_status( + self, *launched_ids: LaunchedJobID + ) -> t.Mapping[LaunchedJobID, JobStatus]: + """Given a collection of launched job ids, return a mapping of id to + current status of the launched job. If a job id is no recognized by the + launcher, a `smartsim.error.errors.LauncherJobNotFound` error should be + raised. + + :param launched_ids: The collection of ids of launched jobs to query + for current status + :raises smartsim.error.errors.LauncherJobNotFound: If at least one of + the ids of the `launched_ids` collection is not recognized. + :returns: A mapping of launched id to current status + """ diff --git a/smartsim/experiment.py b/smartsim/experiment.py index 55ccea7b5..8cb4dad24 100644 --- a/smartsim/experiment.py +++ b/smartsim/experiment.py @@ -59,7 +59,7 @@ from .log import ctx_exp_path, get_logger, method_contextualizer if t.TYPE_CHECKING: - from smartsim._core.dispatch import ExecutableProtocol + from smartsim._core.utils.launcher import ExecutableProtocol from smartsim.launchable.job import Job from smartsim.types import LaunchedJobID diff --git a/smartsim/settings/arguments/launch/alps.py b/smartsim/settings/arguments/launch/alps.py index 6375a4141..51af8ee1b 100644 --- a/smartsim/settings/arguments/launch/alps.py +++ b/smartsim/settings/arguments/launch/alps.py @@ -28,7 +28,8 @@ import typing as t -from smartsim._core.dispatch import ShellLauncher, dispatch, make_shell_format_fn +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 diff --git a/smartsim/settings/arguments/launch/local.py b/smartsim/settings/arguments/launch/local.py index 97b300bce..2ed57861a 100644 --- a/smartsim/settings/arguments/launch/local.py +++ b/smartsim/settings/arguments/launch/local.py @@ -28,7 +28,8 @@ import typing as t -from smartsim._core.dispatch import ShellLauncher, dispatch, make_shell_format_fn +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 diff --git a/smartsim/settings/arguments/launch/lsf.py b/smartsim/settings/arguments/launch/lsf.py index 34db91ff2..00a2c1bbc 100644 --- a/smartsim/settings/arguments/launch/lsf.py +++ b/smartsim/settings/arguments/launch/lsf.py @@ -28,7 +28,8 @@ import typing as t -from smartsim._core.dispatch import ShellLauncher, dispatch, make_shell_format_fn +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 diff --git a/smartsim/settings/arguments/launch/mpi.py b/smartsim/settings/arguments/launch/mpi.py index 04ae55b57..72605c1b3 100644 --- a/smartsim/settings/arguments/launch/mpi.py +++ b/smartsim/settings/arguments/launch/mpi.py @@ -28,7 +28,8 @@ import typing as t -from smartsim._core.dispatch import ShellLauncher, dispatch, make_shell_format_fn +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 diff --git a/smartsim/settings/arguments/launch/pals.py b/smartsim/settings/arguments/launch/pals.py index 2727e47d5..7ebe65dea 100644 --- a/smartsim/settings/arguments/launch/pals.py +++ b/smartsim/settings/arguments/launch/pals.py @@ -28,7 +28,8 @@ import typing as t -from smartsim._core.dispatch import ShellLauncher, dispatch, make_shell_format_fn +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 diff --git a/smartsim/settings/arguments/launch/slurm.py b/smartsim/settings/arguments/launch/slurm.py index 0e057e386..63e4c134b 100644 --- a/smartsim/settings/arguments/launch/slurm.py +++ b/smartsim/settings/arguments/launch/slurm.py @@ -30,7 +30,8 @@ import re import typing as t -from smartsim._core.dispatch import ShellLauncher, dispatch, make_shell_format_fn +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 diff --git a/tests/temp_tests/test_settings/conftest.py b/tests/temp_tests/test_settings/conftest.py index d5e66e94d..70b03630a 100644 --- a/tests/temp_tests/test_settings/conftest.py +++ b/tests/temp_tests/test_settings/conftest.py @@ -26,13 +26,17 @@ import pytest -from smartsim._core import dispatch +from smartsim._core.utils.launcher import ( + ExecutableProtocol, + LauncherProtocol, + create_job_id, +) from smartsim.settings.arguments import launchArguments as launch @pytest.fixture def mock_echo_executable(): - class _MockExe(dispatch.ExecutableProtocol): + class _MockExe(ExecutableProtocol): def as_program_arguments(self): return ("echo", "hello", "world") @@ -51,11 +55,11 @@ def launcher_str(self): @pytest.fixture def mock_launcher(): - class _MockLauncher(dispatch.LauncherProtocol): + class _MockLauncher(LauncherProtocol): __hash__ = object.__hash__ def start(self, launchable): - return dispatch.create_job_id() + return create_job_id() @classmethod def create(cls, exp): diff --git a/tests/temp_tests/test_settings/test_dispatch.py b/tests/temp_tests/test_settings/test_dispatch.py index db346ab98..f1545f58e 100644 --- a/tests/temp_tests/test_settings/test_dispatch.py +++ b/tests/temp_tests/test_settings/test_dispatch.py @@ -33,6 +33,7 @@ import pytest from smartsim._core import dispatch +from smartsim._core.utils.launcher import LauncherProtocol, create_job_id from smartsim.error import errors pytestmark = pytest.mark.group_a @@ -197,7 +198,7 @@ def create(cls, exp): ... class PartImplLauncherABC(LauncherABC): def start(self, launchable): - return dispatch.create_job_id() + return create_job_id() class FullImplLauncherABC(PartImplLauncherABC): @@ -210,7 +211,7 @@ def create(cls, exp): "cls, ctx", ( pytest.param( - dispatch.LauncherProtocol, + LauncherProtocol, pytest.raises(TypeError, match="Cannot dispatch to protocol"), id="Cannot dispatch to protocol class", ), @@ -245,7 +246,7 @@ def test_register_dispatch_to_launcher_types(request, cls, ctx): @dataclasses.dataclass(frozen=True) -class BufferWriterLauncher(dispatch.LauncherProtocol[list[str]]): +class BufferWriterLauncher(LauncherProtocol[list[str]]): buf: io.StringIO if sys.version_info < (3, 10): @@ -257,7 +258,7 @@ def create(cls, exp): def start(self, strs): self.buf.writelines(f"{s}\n" for s in strs) - return dispatch.create_job_id() + return create_job_id() def get_status(self, *ids): raise NotImplementedError diff --git a/tests/test_experiment.py b/tests/test_experiment.py index 8671bfedb..2af864ab8 100644 --- a/tests/test_experiment.py +++ b/tests/test_experiment.py @@ -37,6 +37,7 @@ from smartsim._core import dispatch from smartsim._core.control.launch_history import LaunchHistory +from smartsim._core.utils.launcher import LauncherProtocol, create_job_id from smartsim.entity import _mock, entity from smartsim.experiment import Experiment from smartsim.launchable import job @@ -92,7 +93,7 @@ def iter_jobs(): @dataclasses.dataclass(frozen=True, eq=False) -class NoOpRecordLauncher(dispatch.LauncherProtocol): +class NoOpRecordLauncher(LauncherProtocol): """Simple launcher to track the order of and mapping of ids to `start` method calls. It has exactly three attrs: @@ -127,7 +128,7 @@ def create(cls, exp): return cls(exp) def start(self, record: LaunchRecord): - id_ = dispatch.create_job_id() + id_ = create_job_id() self.launched_order.append(record) self.ids_to_launched[id_] = record return id_ @@ -284,9 +285,9 @@ def test_start_can_start_a_job_multiple_times_accross_multiple_calls( assert sorted(ids_to_launches) == sorted(exp_cached_ids), "Exp did not cache ids" -class GetStatusLauncher(dispatch.LauncherProtocol): +class GetStatusLauncher(LauncherProtocol): def __init__(self): - self.id_to_status = {dispatch.create_job_id(): stat for stat in JobStatus} + self.id_to_status = {create_job_id(): stat for stat in JobStatus} __hash__ = object.__hash__ @@ -356,7 +357,7 @@ def test_get_status_returns_not_started_for_unrecognized_ids( monkeypatch, make_populated_experment ): exp = make_populated_experment(num_active_launchers=1) - brand_new_id = dispatch.create_job_id() + brand_new_id = create_job_id() ((launcher, (id_not_known_by_exp, *rest)),) = ( exp._launch_history.group_by_launcher().items() ) diff --git a/tests/test_launch_history.py b/tests/test_launch_history.py index d076e41a3..9d3bb31ac 100644 --- a/tests/test_launch_history.py +++ b/tests/test_launch_history.py @@ -30,7 +30,7 @@ import pytest from smartsim._core.control.launch_history import LaunchHistory -from smartsim._core.dispatch import LauncherProtocol, create_job_id +from smartsim._core.utils.launcher import LauncherProtocol, create_job_id pytestmark = pytest.mark.group_a