Allow process signals to trigger checkpoint or break behavior

Docs/source/usage/parameters.rst

@@ -64,6 +64,15 @@ Overall simulation parameters
     It is mainly intended for debug purposes, and is best used with
     ``warpx.always_warn_immediately=1``.
 
+* ``warpx.break_signals`` (array of `string`, separated by spaces) optional
+    A list of signal names or numbers that the simulation should
+    handle by cleanly terminating at the next timestep
+
+* ``warpx.checkpoint_signals`` (array of `string`, separated by spaces) optional
+    A list of signal names or numbers that the simulation should
+    handle by outputting a checkpoint at the next timestep. A
+    diagnostic of type `checkpoint` must be configured.
+
 * ``warpx.random_seed`` (`string` or `int` > 0) optional
     If provided ``warpx.random_seed = random``, the random seed will be determined
     using `std::random_device` and `std::clock()`,

Python/pywarpx/picmi.py

@@ -981,6 +981,9 @@ def init(self, kw):
         self.collisions = kw.pop('warpx_collisions', None)
         self.embedded_boundary = kw.pop('warpx_embedded_boundary', None)
 
+        self.break_signals = kw.pop('warpx_break_signals', None)
+        self.checkpoint_signals = kw.pop('warpx_checkpoint_signals', None)
+
         self.inputs_initialized = False
         self.warpx_initialized = False
 
@@ -1019,6 +1022,9 @@ def initialize_inputs(self):
 
         pywarpx.amr.check_input = self.amr_check_input
 
+        pywarpx.warpx.break_signals = self.break_signals
+        pywarpx.warpx.checkpoint_signals = self.checkpoint_signals
+
         particle_shape = self.particle_shape
         for s in self.species:
             if s.particle_shape is not None:

Source/Evolve/WarpXEvolve.cpp

@@ -33,6 +33,8 @@
 #include "Utils/WarpXProfilerWrapper.H"
 #include "Utils/WarpXUtil.H"
 
+#include <ablastr/utils/SignalHandling.H>
+
 #include <AMReX.H>
 #include <AMReX_Array.H>
 #include <AMReX_BLassert.H>
@@ -53,6 +55,7 @@
 #include <vector>
 
 using namespace amrex;
+using ablastr::utils::SignalHandling;
 
 void
 WarpX::Evolve (int numsteps)
@@ -79,6 +82,8 @@ WarpX::Evolve (int numsteps)
         WARPX_PROFILE("WarpX::Evolve::step");
         Real evolve_time_beg_step = amrex::second();
 
+        CheckSignals();
+
         multi_diags->NewIteration();
 
         // Start loop on time steps
@@ -342,6 +347,8 @@ WarpX::Evolve (int numsteps)
         Real evolve_time_end_step = amrex::second();
         evolve_time += evolve_time_end_step - evolve_time_beg_step;
 
+        HandleSignals();
+
         if (verbose) {
             amrex::Print()<< "STEP " << step+1 << " ends." << " TIME = " << cur_time
                         << " DT = " << dt[0] << "\n";
@@ -350,7 +357,7 @@ WarpX::Evolve (int numsteps)
                       << " s; Avg. per step = " << evolve_time/(step-step_begin+1) << " s\n";
         }
 
-        if (cur_time >= stop_time - 1.e-3*dt[0]) {
+        if (cur_time >= stop_time - 1.e-3*dt[0] || SignalHandling::TestAndResetActionRequestFlag(SignalHandling::SIGNAL_REQUESTS_BREAK)) {
             break;
         }
 
@@ -929,3 +936,21 @@ WarpX::applyMirrors(Real time){
         }
     }
 }
+
+void
+WarpX::CheckSignals()
+{
+    SignalHandling::CheckSignals();
+}
+
+void
+WarpX::HandleSignals()
+{
+    SignalHandling::WaitSignals();
+
+    // SIGNAL_REQUESTS_BREAK is handled directly in WarpX::Evolve
+
+    if (SignalHandling::TestAndResetActionRequestFlag(SignalHandling::SIGNAL_REQUESTS_CHECKPOINT)) {
+        multi_diags->FilterComputePackFlushLastTimestep( istep[0] );
+    }
+}

Source/WarpX.H

@@ -951,6 +951,11 @@ private:
     // Singleton is used when the code is run from python
     static WarpX* m_instance;
 
+    //! Check and clear signal flags and asynchronously broadcast them from process 0
+    static void CheckSignals ();
+    //! Complete the asynchronous broadcast of signal flags, and initiate a checkpoint if requested
+    void HandleSignals ();
+
     ///
     /// Advance the simulation by numsteps steps, electromagnetic case.
     ///

Source/WarpX.cpp

@@ -38,6 +38,8 @@
 #include "Utils/WarpXProfilerWrapper.H"
 #include "Utils/WarpXUtil.H"
 
+#include <ablastr/utils/SignalHandling.H>
+
 #ifdef AMREX_USE_SENSEI_INSITU
 #   include <AMReX_AmrMeshInSituBridge.H>
 #endif
@@ -233,6 +235,8 @@ WarpX::WarpX ()
 
     InitEB();
 
+    ablastr::utils::SignalHandling::InitSignalHandling();
+
     // Geometry on all levels has been defined already.
     // No valid BoxArray and DistributionMapping have been defined.
     // But the arrays for them have been resized.
@@ -555,6 +559,50 @@ WarpX::ReadParameters ()
             }
         }
 
+        using ablastr::utils::SignalHandling;
+        std::vector<std::string> signals_in;
+        pp_warpx.queryarr("break_signals", signals_in);
+
+#if defined(__linux__) || defined(__APPLE__)
+        for (const std::string &str : signals_in) {
+            int sig = SignalHandling::parseSignalNameToNumber(str);
+            SignalHandling::signal_conf_requests[SignalHandling::SIGNAL_REQUESTS_BREAK][sig] = true;
+        }
+        signals_in.clear();
+#else
+        WARPX_ALWAYS_ASSERT_WITH_MESSAGE(signals_in.empty(),
+                                         "Signal handling requested in input, but is not supported on this platform");
+#endif
+
+        bool have_checkpoint_diagnostic = false;
+
+        ParmParse pp("diagnostics");
+        std::vector<std::string> diags_names;
+        pp.queryarr("diags_names", diags_names);
+
+        for (const auto &diag : diags_names) {
+            ParmParse dd(diag);
+            std::string format;
+            dd.query("format", format);
+            if (format == "checkpoint") {
+                have_checkpoint_diagnostic = true;
+                break;
+            }
+        }
+
+        pp_warpx.queryarr("checkpoint_signals", signals_in);
+#if defined(__linux__) || defined(__APPLE__)
+        for (const std::string &str : signals_in) {
+            int sig = SignalHandling::parseSignalNameToNumber(str);
+            SignalHandling::signal_conf_requests[SignalHandling::SIGNAL_REQUESTS_CHECKPOINT][sig] = true;
+            WARPX_ALWAYS_ASSERT_WITH_MESSAGE(have_checkpoint_diagnostic,
+                                             "Signal handling was requested to checkpoint, but no checkpoint diagnostic is configured");
+        }
+#else
+        WARPX_ALWAYS_ASSERT_WITH_MESSAGE(signals_in.empty(),
+                                         "Signal handling requested in input, but is not supported on this platform");
+#endif
+
         // set random seed
         std::string random_seed = "default";
         pp_warpx.query("random_seed", random_seed);

Source/ablastr/utils/CMakeLists.txt

@@ -1,4 +1,5 @@
 target_sources(ablastr
   PRIVATE
     TextMsg.cpp
+    SignalHandling.cpp
 )

Source/ablastr/utils/Make.package

@@ -1,3 +1,3 @@
-CEXE_sources += TextMsg.cpp
+CEXE_sources += TextMsg.cpp SignalHandling.cpp
 
 VPATH_LOCATIONS   += $(WARPX_HOME)/Source/ablastr/utils

Source/ablastr/utils/SignalHandling.H

@@ -0,0 +1,88 @@
+/* Copyright 2022 Philip Miller
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ */
+
+#ifndef ABLASTR_SIGNAL_HANDLING_H_
+#define ABLASTR_SIGNAL_HANDLING_H_
+
+#include <AMReX_Config.H>
+
+#if defined(AMREX_USE_MPI)
+#   include <mpi.h>
+#endif
+
+#include <atomic>
+#include <string>
+
+namespace ablastr::utils {
+
+/**
+ * \brief
+ * Signal handling
+ *
+ * Rank 0 will accept signals and asynchronously broadcast the
+ * configured response; other processes will ignore them and
+ * follow the lead of rank 0 to avoid potential for deadlocks or
+ * timestep-skewed response.
+ *
+ * Variables and functions are static rather than per-instance
+ * because signal handlers are configured at the process level.
+ */
+class SignalHandling
+{
+public:
+    //! The range of signal values to accept
+    static constexpr int NUM_SIGNALS = 32;
+
+    //! Labels for indexed positions in signal_actions_requests
+    enum signal_action_requested_labels {
+        //! Cleanly stop execution, as if the simulation reached its configured end
+        SIGNAL_REQUESTS_BREAK = 0,
+        //! Produce a checkpoint
+        SIGNAL_REQUESTS_CHECKPOINT = 1,
+        SIGNAL_REQUESTS_SIZE = 2 // This should always be 1 greater than the last valid value
+    };
+
+    //! Whether configuration requests the code take a particular action on a particular signal
+    static bool signal_conf_requests[SIGNAL_REQUESTS_SIZE][NUM_SIGNALS];
+
+    //! Take a string and convert it to a corresponding signal number if possible
+    static int parseSignalNameToNumber (const std::string &str);
+
+    //! Set up signal handlers based on input configuration provided in `signal_conf_requests_*`
+    static void InitSignalHandling ();
+
+    //! Check and clear signal flags and asynchronously broadcast them from process 0
+    static void CheckSignals ();
+    //! Complete the asynchronous broadcast of signal flags
+    static void WaitSignals ();
+
+    //! Check whether a given action has been requested, and reset the associated flag
+    static bool TestAndResetActionRequestFlag (int action_to_test);
+
+private:
+    //! On process 0, whether a given signal has been received since the last check
+    static std::atomic<bool> signal_received_flags[NUM_SIGNALS];
+
+#if defined(AMREX_USE_MPI)
+    //! MPI requests for the asynchronous broadcasts of the signal-requested actions
+    static MPI_Request signal_mpi_ibcast_request;
+#endif
+
+    //! Signal handler to set flags on process 0 (other processes ignore configured signals)
+    static void SignalSetFlag (int signal_number);
+
+    //! Boolean flags transmitted between CheckSignals() and
+    //! HandleSignals() to indicate actions requested by signals
+    static bool signal_actions_requested[SIGNAL_REQUESTS_SIZE];
+
+     // Don't allow clients to incorrectly try to construct and use an instance of this type
+    SignalHandling() = delete;
+};
+
+} // namespace ablastr::utils
+
+#endif // ABLASTR_SIGNAL_HANDLING_H_