Re-wrire up default dispatcher

CrayLabs · Jul 20, 2024 · 0b9ec1a · 0b9ec1a
1 parent 4e9f5b2
commit 0b9ec1a
Show file tree

Hide file tree

Showing 11 changed files with 80 additions and 150 deletions.
diff --git 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
@@ -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
@@ -29,20 +29,19 @@
 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
 from ...launchCommand import LauncherType
 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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]: ...