diff --git a/smartsim/_core/shell/shellLauncher.py b/smartsim/_core/shell/shellLauncher.py index 60713d37e..237269d5d 100644 --- a/smartsim/_core/shell/shellLauncher.py +++ b/smartsim/_core/shell/shellLauncher.py @@ -101,55 +101,55 @@ 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() +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 + return impl diff --git a/smartsim/settings/arguments/launch/alps.py b/smartsim/settings/arguments/launch/alps.py index a63fa8b87..51af8ee1b 100644 --- a/smartsim/settings/arguments/launch/alps.py +++ b/smartsim/settings/arguments/launch/alps.py @@ -29,7 +29,7 @@ import typing as t from smartsim._core.dispatch import dispatch -from smartsim._core.shell.shellLauncher import ShellLauncher +from smartsim._core.shell.shellLauncher import ShellLauncher, make_shell_format_fn from smartsim.log import get_logger from ...common import set_check_input @@ -37,7 +37,7 @@ from ..launchArguments import LaunchArguments logger = get_logger(__name__) -_as_aprun_command = ShellLauncher.make_shell_format_fn(run_command="aprun") +_as_aprun_command = make_shell_format_fn(run_command="aprun") @dispatch(with_format=_as_aprun_command, to_launcher=ShellLauncher) diff --git a/smartsim/settings/arguments/launch/local.py b/smartsim/settings/arguments/launch/local.py index 47b8b25b7..2ed57861a 100644 --- a/smartsim/settings/arguments/launch/local.py +++ b/smartsim/settings/arguments/launch/local.py @@ -29,7 +29,7 @@ import typing as t from smartsim._core.dispatch import dispatch -from smartsim._core.shell.shellLauncher import ShellLauncher +from smartsim._core.shell.shellLauncher import ShellLauncher, make_shell_format_fn from smartsim.log import get_logger from ...common import StringArgument, set_check_input @@ -37,7 +37,7 @@ from ..launchArguments import LaunchArguments logger = get_logger(__name__) -_as_local_command = ShellLauncher.make_shell_format_fn(run_command=None) +_as_local_command = make_shell_format_fn(run_command=None) @dispatch(with_format=_as_local_command, to_launcher=ShellLauncher) diff --git a/smartsim/settings/arguments/launch/lsf.py b/smartsim/settings/arguments/launch/lsf.py index 3b9640e25..00a2c1bbc 100644 --- a/smartsim/settings/arguments/launch/lsf.py +++ b/smartsim/settings/arguments/launch/lsf.py @@ -29,7 +29,7 @@ import typing as t from smartsim._core.dispatch import dispatch -from smartsim._core.shell.shellLauncher import ShellLauncher +from smartsim._core.shell.shellLauncher import ShellLauncher, make_shell_format_fn from smartsim.log import get_logger from ...common import set_check_input @@ -37,7 +37,7 @@ from ..launchArguments import LaunchArguments logger = get_logger(__name__) -_as_jsrun_command = ShellLauncher.make_shell_format_fn(run_command="jsrun") +_as_jsrun_command = make_shell_format_fn(run_command="jsrun") @dispatch(with_format=_as_jsrun_command, to_launcher=ShellLauncher) diff --git a/smartsim/settings/arguments/launch/mpi.py b/smartsim/settings/arguments/launch/mpi.py index 894f4f760..72605c1b3 100644 --- a/smartsim/settings/arguments/launch/mpi.py +++ b/smartsim/settings/arguments/launch/mpi.py @@ -29,7 +29,7 @@ import typing as t from smartsim._core.dispatch import dispatch -from smartsim._core.shell.shellLauncher import ShellLauncher +from smartsim._core.shell.shellLauncher import ShellLauncher, make_shell_format_fn from smartsim.log import get_logger from ...common import set_check_input @@ -37,9 +37,9 @@ from ..launchArguments import LaunchArguments logger = get_logger(__name__) -_as_mpirun_command = ShellLauncher.make_shell_format_fn("mpirun") -_as_mpiexec_command = ShellLauncher.make_shell_format_fn("mpiexec") -_as_orterun_command = ShellLauncher.make_shell_format_fn("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 _BaseMPILaunchArguments(LaunchArguments): diff --git a/smartsim/settings/arguments/launch/pals.py b/smartsim/settings/arguments/launch/pals.py index 7e2b2e933..7ebe65dea 100644 --- a/smartsim/settings/arguments/launch/pals.py +++ b/smartsim/settings/arguments/launch/pals.py @@ -29,7 +29,7 @@ import typing as t from smartsim._core.dispatch import dispatch -from smartsim._core.shell.shellLauncher import ShellLauncher +from smartsim._core.shell.shellLauncher import ShellLauncher, make_shell_format_fn from smartsim.log import get_logger from ...common import set_check_input @@ -37,7 +37,7 @@ from ..launchArguments import LaunchArguments logger = get_logger(__name__) -_as_pals_command = ShellLauncher.make_shell_format_fn(run_command="mpiexec") +_as_pals_command = make_shell_format_fn(run_command="mpiexec") @dispatch(with_format=_as_pals_command, to_launcher=ShellLauncher) diff --git a/smartsim/settings/arguments/launch/slurm.py b/smartsim/settings/arguments/launch/slurm.py index 6578d8666..63e4c134b 100644 --- a/smartsim/settings/arguments/launch/slurm.py +++ b/smartsim/settings/arguments/launch/slurm.py @@ -31,7 +31,7 @@ import typing as t from smartsim._core.dispatch import dispatch -from smartsim._core.shell.shellLauncher import ShellLauncher +from smartsim._core.shell.shellLauncher import ShellLauncher, make_shell_format_fn from smartsim.log import get_logger from ...common import set_check_input @@ -39,7 +39,7 @@ from ..launchArguments import LaunchArguments logger = get_logger(__name__) -_as_srun_command = ShellLauncher.make_shell_format_fn(run_command="srun") +_as_srun_command = make_shell_format_fn(run_command="srun") @dispatch(with_format=_as_srun_command, to_launcher=ShellLauncher)