Launch Arg Builders Format Input for Launchers

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.*",

setup.py

@@ -182,7 +182,6 @@ def has_ext_modules(_placeholder):
         "types-tqdm",
         "types-tensorflow==2.12.0.9",
         "types-setuptools",
-        "typing_extensions>=4.1.0,<4.6",
     ],
     "docs": [
         "Sphinx==6.2.1",
@@ -226,6 +225,7 @@ def has_ext_modules(_placeholder):
         "pygithub>=2.3.0",
         "numpy<2",
         "smartredis>=0.5,<0.6",
+        "typing_extensions>=4.1.0,<4.6",
     ],
     cmdclass={
         "build_py": SmartSimBuild,

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:

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,29 +74,27 @@ 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] = {}
-        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:
@@ -293,8 +296,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 +522,25 @@ 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"
+    """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
     )
     dragon_server_paths = dragon_server_path.split(":")
     if len(dragon_server_paths) > 1:

smartsim/_core/launcher/dragon/dragonLauncher.py

@@ -30,6 +30,7 @@
 import typing as t
 
 from smartsim._core.schemas.dragonRequests import DragonRunPolicy
+from smartsim.types import LaunchedJobID
 
 from ...._core.launcher.stepMapping import StepMap
 from ....error import LauncherError, SmartSimError
@@ -44,6 +45,7 @@
 from ....status import SmartSimStatus
 from ...schemas import (
     DragonRunRequest,
+    DragonRunRequestView,
     DragonRunResponse,
     DragonStopRequest,
     DragonStopResponse,
@@ -57,6 +59,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__)
 
 
@@ -74,9 +81,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"""
@@ -121,6 +128,22 @@ 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()  # pylint: disable=protected-access
+        return self
+
+    def start(
+        self, args_and_policy: tuple[DragonRunRequestView, DragonRunPolicy]
+    ) -> LaunchedJobID:
+        req_args, policy = args_and_policy
+        self._connector.load_persisted_env()
+        merged_env = self._connector.merge_persisted_env(os.environ.copy())
+        req = DragonRunRequest(**dict(req_args), current_env=merged_env, policy=policy)
+        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
 
@@ -167,31 +190,25 @@ 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)
-
             policy = DragonRunPolicy.from_run_args(run_args)
-
-            response = _assert_schema_type(
-                self._connector.send_request(
-                    DragonRunRequest(
+            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,
-                        current_env=merged_env,
                         output_file=out,
                         error_file=err,
-                        policy=policy,
-                    )
-                ),
-                DragonRunResponse,
+                    ),
+                    policy,
+                )
             )
-            step_id = str(response.step_id)
         else:
             # pylint: disable-next=consider-using-with
             out_strm = open(out, "w+", encoding="utf-8")
@@ -325,3 +342,53 @@ 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 caused by `DragonLauncher.supported_rs`
+# -----------------------------------------------------------------------------
+from smartsim.settings.arguments.launch.dragon import DragonLaunchArguments
+from smartsim.settings.dispatch import ExecutableProtocol, dispatch
+
+
+def _as_run_request_args_and_policy(
+    run_req_args: DragonLaunchArguments,
+    exe: ExecutableProtocol,
+    env: t.Mapping[str, str | None],
+) -> tuple[DragonRunRequestView, DragonRunPolicy]:
+    # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+    # FIXME: This type is 100% unacceptable, but I don't want to spend too much
+    #        time on fixing the dragon launcher API. Something that we need to
+    #        revisit in the future though.
+    exe_, *args = exe.as_program_arguments()
+    run_args = dict[str, "int | str | float | None"](run_req_args._launch_args)
+    policy = DragonRunPolicy.from_run_args(run_args)
+    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_args,
+        ),
+        policy,
+    )
+
+
+dispatch(
+    DragonLaunchArguments,
+    with_format=_as_run_request_args_and_policy,
+    to_launcher=DragonLauncher,
+)
+# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

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]
 
 
@@ -411,6 +414,28 @@ def is_crayex_platform() -> bool:
     return result.is_cray
 
 
+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)
+
+
 @t.final
 class SignalInterceptionStack(collections.abc.Collection[_TSignalHandlerFn]):
     """Registers a stack of callables to be called when a signal is

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"""