Dragon Launcher Batch Job Support

README.md

@@ -174,13 +174,17 @@ system with which it has a corresponding `RunSettings` class. If one can be foun
 ## Experiments on HPC Systems
 
 SmartSim integrates with common HPC schedulers providing batch and interactive
-launch capabilities for all applications.
+launch capabilities for all applications:
 
  - Slurm
  - LSF
  - PBSPro
  - Local (for laptops/single node, no batch)
 
+In addition, on Slurm and PBS systems, [Dragon](https://dragonhpc.github.io/dragon/doc/_build/html/index.html)
+can be used as a launcher. Please refer to the documentation for instructions on
+how to insall it on your system and use it in SmartSim.
+
 
 ### Interactive Launch Example
 

conftest.py

@@ -46,7 +46,8 @@
 
 import smartsim
 from smartsim import Experiment
-from smartsim._core.launcher.dragon.dragonLauncher import DragonLauncher, _dragon_cleanup
+from smartsim._core.launcher.dragon.dragonConnector import DragonConnector
+from smartsim._core.launcher.dragon.dragonLauncher import DragonLauncher
 from smartsim._core.config import CONFIG
 from smartsim._core.config.config import Config
 from smartsim._core.utils.telemetry.telemetry import JobEntity
@@ -165,9 +166,11 @@ def pytest_sessionfinish(
                     raise
             else:
                 break
-
-    # kill all spawned processes
-    kill_all_test_spawned_processes()
+    else:
+        # kill all spawned processes
+        if CONFIG.test_launcher == "dragon":
+            time.sleep(5)
+        kill_all_test_spawned_processes()
 
 
 def build_mpi_app() -> t.Optional[pathlib.Path]:
@@ -342,8 +345,9 @@ def get_base_run_settings(
             return settings
         if test_launcher == "dragon":
             run_args = {"nodes": nodes}
+            run_args = {"ntasks": ntasks}
             run_args.update(kwargs)
-            settings = RunSettings(exe, args, run_command="", run_args=run_args)
+            settings = DragonRunSettings(exe, args, run_args=run_args)
             return settings
         if test_launcher == "pbs":
             if shutil.which("aprun"):
@@ -798,17 +802,10 @@ def global_dragon_teardown() -> None:
     """
     if test_launcher != "dragon" or CONFIG.dragon_server_path is None:
         return
-    exp_path = os.path.join(test_output_root, "dragon_teardown")
-    os.makedirs(exp_path, exist_ok=True)
-    exp: Experiment = Experiment("dragon_shutdown", exp_path=exp_path, launcher=test_launcher)
-    rs = exp.create_run_settings("sleep", ["0.1"])
-    model = exp.create_model("dummy", run_settings=rs)
-    exp.generate(model, overwrite=True)
-    exp.start(model, block=True)
-
-    launcher: DragonLauncher = exp._control._launcher
-    launcher.cleanup()
-    time.sleep(5)
+    logger.debug(f"Tearing down Dragon infrastructure, server path: {CONFIG.dragon_server_path}")
+    dragon_connector = DragonConnector()
+    dragon_connector.ensure_connected()
+    dragon_connector.cleanup()
 
 
 @pytest.fixture

doc/api/smartsim_api.rst

@@ -60,6 +60,7 @@ Types of Settings:
     MpiexecSettings
     OrterunSettings
     JsrunSettings
+    DragonRunSettings
     SbatchSettings
     QsubBatchSettings
     BsubBatchSettings
@@ -163,6 +164,28 @@ and within batch launches (e.g., ``QsubBatchSettings``)
     :members:
 
 
+.. _dragonsettings_api:
+
+DragonRunSettings
+-----------------
+
+``DragonRunSettings`` can be used on systems that support Slurm or
+PBS, if Dragon is available in the Python environment (see `_dragon_install`
+for instructions on how to install it through ``smart``).
+
+``DragonRunSettings`` can be used in interactive sessions (on allcation)
+and within batch launches (i.e. ``SbatchSettings`` or ``QsubBatchSettings``,
+for Slurm and PBS sessions, respectively).
+
+.. autosummary::
+    DragonRunSettings.set_nodes
+    DragonRunSettings.set_tasks_per_node
+
+.. autoclass:: DragonRunSettings
+    :inherited-members:
+    :undoc-members:
+    :members:
+
 
 .. _jsrun_api:
 

doc/dragon.rst

@@ -0,0 +1,169 @@
+******
+Dragon
+******
+
+========
+Overview
+========
+
+Dragon is a composable distributed run-time targeting HPC workflows. In SmartSim,
+Dragon can be used as a launcher, within a Slurm or PBS allocation or batch job.
+The SmartSim team collaborates with the Dragon team to develop an efficient
+launcher which will enable fast, interactive, and customized execution of
+complex workflows on large HPC systems. As Dragon is scheduler-agnostic,
+the same SmartSim script using Dragon as a launcher can be run indifferently
+on a Slurm or PBS system. Support for additional schedulers is coming soon.
+
+.. warning::
+    The Dragon launcher is currently in its early development stage and should be treated as
+    a prototype implementation. Your assistance is invaluable in identifying any issues
+    encountered during usage and suggesting missing features for implementation. Please
+    provide feedback in the form of a created issue on the
+    `SmartSim issues GitHub page <https://github.com/CrayLabs/SmartSim/issues>`_.
+    The :ref:`Known Issues section<dragon_known_issues>` is also a good starting
+    point when troubleshooting workflows run using the Dragon launcher.
+
+=====
+Usage
+=====
+To use Dragon, you need to install it in your current Python environment. This can
+be accomplished by providing the ``--dragon`` flag to the ``smart build`` command, as
+detailed in the :ref:`Dragon Install <dragon_install>`. Note that specifying the device
+configuration is also required for a proper build.
+
+After installation, specify Dragon as the launcher when creating an ``Experiment``:
+
+.. code-block:: python
+
+    exp = Experiment(name="dragon-example", launcher="dragon")
+
+Dragon introduces its own run settings class, ``DragonRunSettings``, which allows users to
+specify nodes and tasks per node for a ``Model``. For instance, continuing from the previous
+example:
+
+.. code-block:: python
+
+    # Because "dragon" was specified as the launcher during Experiment initialization,
+    # create_run_settings will return a DragonRunSettings object
+    rs = exp.create_run_settings(exe="mpi_app",
+                                 exe_args=["--option", "value"],
+                                 env_vars={"MYVAR": "VALUE"})
+    # Above we specify the executable (exe), executable arguments (exe_args)
+    # and environment variables (env_vars)
+
+    # Sets the number of nodes for this job
+    rs.set_nodes(4)
+    # Set the tasks per node for this job
+    rs.set_tasks_per_node(3)
+    # Initialize the Model and pass in the DragonRunSettings object
+    mpi_app = exp.create_model("MPI_APP", run_settings=rs)
+    # Start the Model
+    exp.start(mpi_app)
+
+SmartSim supports ``DragonRunSettings`` with ``Model``, ``Ensemble`` and ``Orchestrator`` entities.
+In the next sections, we detail how Dragon is integrated into SmartSim.
+
+For more information on HPC launchers, visit the :ref:`Run Settings<run_settings_hpc_ex>` page.
+
+=================
+The Dragon Server
+=================
+
+Dragon can initiate processes on any available resource within an allocation. To facilitate
+this, SmartSim initializes the Dragon infrastructure whenever a ``Model`` is launched and maintains
+it until the parent ``Experiment`` concludes. To facilitate interaction with processes managed by
+Dragon, SmartSim establishes a command server within the Dragon infrastructure. This server,
+known as the `Dragon Server`, is responsible for executing commands to start or stop processes
+and to query their status.
+
+Sharing the Dragon Server across Experiments
+============================================
+
+Currently, SmartSim supports only one Dragon server per allocation. Consequently,
+if multiple Experiments need to run within the same allocation, the Dragon server
+must be shared among them. By default, the server starts from a subdirectory
+of the ``Experiment`` path, where it creates a configuration file.
+To enable server sharing, users can specify a custom path
+from which the server should be launched. This can be achieved by setting the
+environment variable ``SMARTSIM_DRAGON_SERVER_PATH`` to an existing absolute path.
+Each ``Experiment`` will then search for the configuration file in the specified path
+and initiate a new server instance only if the file is not found.
+
+Dragon's High-Speed Transport Agents
+====================================
+
+On systems equipped with the HPE Slingshot interconnect, Dragon utilizes High-Speed
+Transport Agents (HSTA) by default for internal messaging within the infrastructure
+launched by SmartSim. On systems without the HPE Slingshot interconnect,
+TCP agents are employed. To specify the use of TCP agents, users must set the environment
+variable ``SMARTSIM_DRAGON_TRANSPORT`` to ``tcp`` prior to executing the Experiment.
+To specify HSTA, ``SMARTSIM_DRAGON_TRANSPORT`` can be set to ``hsta`` or left unset.
+
+=============
+Communication
+=============
+
+SmartSim and the Dragon Server communicate using `ZeroMQ <https://zeromq.org/>`_.
+
+Similar to other communication protocols, defining timeouts for send and receive operations
+is crucial in SmartSim. SmartSim configures default timeouts that have been tested on various
+systems, such as Polaris, Perlmutter, and other HPE Cray EX and Apollo systems.
+However, if you encounter failed communication attempts, adjusting the timeouts may
+be necessary. You can adjust these timeouts by setting the corresponding environment variables:
+
+- **Server Start-up Timeout**: This timeout specifies the duration the SmartSim ``Experiment``
+  waits when the server is initially started. It must accommodate the time required for
+  Dragon to set up the infrastructure, which varies based on the system's workload manager
+  response time. The default timeout is `"300000"` milliseconds (i.e., five minutes), and you can override
+  it using the ``SMARTSIM_DRAGON_STARTUP_TIMEOUT`` environment variable.
+
+- **Server Send and Receive Timeout**: This timeout dictates how long SmartSim and the Dragon
+  server wait to send or receive a message. The default timeout is `"30000"` milliseconds (i.e., 30 seconds),
+  and you can modify it using the ``SMARTSIM_DRAGON_TIMEOUT`` environment variable.
+
+Setting any timeout to "-1" will result in an infinite waiting time, causing the execution to
+block until the communication is completed, potentially hanging indefinitely if issues occur.
+
+It's important to note that all communications are secured with `elliptic curve cryptography <http://curvezmq.org/>`_.
+SmartSim generates the necessary key-pairs and stores them in the user's home directory by
+default. However, you can specify an alternative absolute path using the ``SMARTSIM_KEY_PATH``
+environment variable.
+
+.. _dragon_known_issues:
+
+============
+Known issues
+============
+
+As previously noted, the integration of SmartSim with Dragon is still in its early
+development stage, and there are known issues that may result in unexpected behavior
+during runs:
+
+- **Incomplete cleanup of Dragon resources**: When SmartSim exits, it attempts to properly
+  shut down the Dragon infrastructure to clean up associated resources, such as shared memory
+  segments, and terminate all processes. However, in rare cases, if the execution is
+  abruptly interrupted (e.g., by terminating SmartSim with ``SIGKILL``), the cleanup process
+  may be incomplete, leaving processes like the Dragon overlay network active on the node
+  where SmartSim was executed (which could be a login node, particularly on Slurm systems).
+  If this occurs, you can use the following command to address the issue:
+
+  .. code-block::
+
+    smart teardown --dragon
+
+  This command will terminate all Dragon-related processes, release shared memory segments,
+  but also terminate all Python processes associated with your username.
+
+- **Dragon server not starting**: This issue may arise due to two main reasons:
+
+  1. *HSTA not available on the system*: Try setting the environment variable
+     ``SMARTSIM_DRAGON_TRANSPORT`` to ``tcp``.
+  2. *System or Workload Manager too busy*: Attempt to mitigate this by setting the environment
+     variable ``SMARTSIM_DRAGON_STARTUP_TIMEOUT`` to a larger value or ``"-1"``.
+
+- **MPI-based applications hanging**: To run MPI-based applications on Dragon, Cray PMI or
+  Cray PALS must be available on the system. This limitation is currently being addressed.
+
+
+Interested users can learn more about the Dragon project at the external
+`Dragon documentation page <https://dragonhpc.github.io/dragon/doc/_build/html/index.html>`_.

doc/experiment.rst

@@ -52,14 +52,18 @@ SmartSim supports launching AI-enabled workflows on a wide variety of systems, i
 Linux machine or on HPC machines with a job scheduler (e.g. Slurm, PBS Pro, and LSF). When creating a SmartSim
 ``Experiment``, the user has the opportunity to specify the `launcher` type or defer to automatic `launcher` selection.
 `Launcher` selection determines how SmartSim translates entity configurations into system calls to launch,
-manage, and monitor. Currently, SmartSim supports 5 `launchers`:
+manage, and monitor. Currently, SmartSim supports 7 `launcher` options:
 
 1. ``local`` **[default]**: for single-node, workstation, or laptop
 2. ``slurm``: for systems using the Slurm scheduler
 3. ``pbs``: for systems using the PBS Pro scheduler
 4. ``pals``: for systems using the PALS scheduler
 5. ``lsf``: for systems using the LSF scheduler
-6. ``auto``: have SmartSim auto-detect the launcher to use
+6. ``dragon``: if Dragon is installed in the current Python environment, see :ref:`Dragon Install <dragon_install>`
+7. ``auto``: have SmartSim auto-detect the launcher to use (will not detect ``dragon``)
+
+The :ref:`Dragon-based launcher <dragon>` can be run on PBS- or Slurm-based systems
+(MPI applications are supported only when Cray PMI or Cray PALS are available).
 
 If the systems `launcher` cannot be found or no `launcher` argument is provided, the default value of
 `"local"` will be assigned which will start all ``Experiment`` launched entities on the

doc/index.rst

@@ -37,6 +37,7 @@
    orchestrator
    ss_logger
    ml_features
+   dragon
    api/smartsim_api
 
 .. toctree::

doc/installation_instructions/basic.rst

@@ -237,6 +237,28 @@ to building SmartSim with GPU support is to specify a different ``device``
   backends look for the CUDA Toolkit and cuDNN libraries. Please see the
   :ref:`Platform Installation Section <install-notes>` section for guidance.
 
+
+.. _dragon_install:
+
+Dragon Install
+--------------
+
+`Dragon <https://dragonhpc.github.io/dragon/doc/_build/html/index.html>`_ is
+an HPC-native library for distributed computing. SmartSim can use Dragon as a
+launcher on systems with Slurm or PBS as schedulers. To install the correct
+version of Dragon, you can add the ``--dragon`` option to ``smart build``.
+For example, to install dragon alongside the RedisAI CPU backends, you can run
+
+.. code-block:: bash
+
+    # run one of the following
+    smart build --device cpu --dragon           # install Dragon, PT and TF for cpu
+    smart build --device cpu --onnx --dragon    # install Dragon and all backends (PT, TF, ONNX) on cpu
+
+.. note::
+  Dragon is only supported on Linux systems. For further information, you
+  can read :ref:`the dedicated documentation page <dragon>`.
+
 ==========
 SmartRedis
 ==========
@@ -300,7 +322,7 @@ source remains at the site of the clone instead of in site-packages.
   pip install -e .[dev,ml]    # for bash users
   pip install -e .\[dev,ml\]  # for zsh users
 
-Use the now installed ``smart`` cli to install the machine learning runtimes.
+Use the now installed ``smart`` cli to install the machine learning runtimes and dragon.
 
 .. tabs::
 
@@ -309,8 +331,8 @@ Use the now installed ``smart`` cli to install the machine learning runtimes.
     .. code-block:: bash
 
       # run one of the following
-      smart build --device cpu --onnx  # install with cpu-only support
-      smart build --device gpu --onnx  # install with both cpu and gpu support
+      smart build --device cpu --onnx --dragon  # install with cpu-only support
+      smart build --device gpu --onnx --dragon  # install with both cpu and gpu support
 
 
   .. tab:: MacOS (Intel x64)