Improved handling of SIGTERM

CHANGELOG.md

@@ -39,6 +39,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 - Properly record cmake variables in the cmake build directory instead of the source tree.
 - Added default "none" for option shuffle in BatchGenerator, so that it works in executables where shuffle is not an option.
 - Added a few missing header files in shortlist.h and beam_search.h.
+- Improved handling for graceful shutdown upon receiving SIGTERM. SIGTERM now also interrupts batch prefetching, which runs in a separate thread. Graceful shutdown can be disabled with --sigterm 'immediate'.
 
 ### Changed
 - Move Simple-WebSocket-Server to submodule

src/CMakeLists.txt

@@ -25,6 +25,7 @@ add_library(marian STATIC
   common/filesystem.cpp
   common/file_stream.cpp
   common/file_utils.cpp
+  common/signal_handling.cpp
   common/types.cpp
 
   data/alignment.cpp
@@ -99,7 +100,6 @@ add_library(marian STATIC
   training/graph_group_singleton.cpp
   training/validator.cpp
   training/communicator.cpp
-  training/scheduler.cpp
 
   # this is only compiled to catch build errors, but not linked
   microsoft/quicksand.cpp

src/command/marian_train.cpp

@@ -1,6 +1,7 @@
 #include <signal.h>
 #include "marian.h"
 
+#include "common/signal_handling.h"
 #include "training/graph_group_async.h"
 #include "training/graph_group_singleton.h"
 #include "training/graph_group_sync.h"
@@ -42,14 +43,12 @@ int main(int argc, char** argv) {
       New<Train<AsyncGraphGroup>>(options)->run();
     }
   }
-
-  // If we exit due to SIGTERM, exit with 128 + the signal number, as suggested
-  // for bash in http://tldp.org/LDP/abs/html/exitcodes.html. This allows parent
+  // If we exit due to a graceful exit request via SIGTERM, exit with 128 + SIGTERM,
+  // as suggested for bash in http://tldp.org/LDP/abs/html/exitcodes.html. This allows parent
   // scripts to determine if training terminated naturally or via SIGTERM.
-  // Whith this approach we can accommodate additional signals in the future.
-  // An alternative would be to return 124, which is what the timeout command
+  // An alternative would be to exit with code 124, which is what the timeout command
   // returns for timeout -s SIGTERM <seconds> ...., because exiting after SIGTERM
   // is not technically a fatal error (which is what the 128+x convention usually
   // stands for).
-  return getSigtermFlag() ? (128 + SIGTERM) : 0;
+  exit(getSignalFlag(SIGTERM) ? 128 + SIGTERM : EXIT_SUCCESS);
 }

src/common/config_parser.cpp

@@ -324,13 +324,23 @@ void ConfigParser::addOptionsModel(cli::CLIWrapper& cli) {
         "Dropout for transformer attention (0 = no dropout)");
     cli.add<float>("--transformer-dropout-ffn",
         "Dropout for transformer filter (0 = no dropout)");
+
   }
   cli.switchGroup(previous_group);
   // clang-format on
 }
 
 void ConfigParser::addOptionsTraining(cli::CLIWrapper& cli) {
-  auto previous_group = cli.switchGroup("Training options");
+  auto previous_group = cli.switchGroup("Signal Handling");
+  // --sigterm is deliberately not a boolean, to allow for a consistent
+  // pattern of specifying custom signal handling in the future.
+  // (e.g., dump model but continue training upon SIGUSR1, or report current
+  // training status upon SIGINFO.)
+  cli.add<std::string>("--sigterm",
+                       "What to do with SIGTERM: 'graceful' => save and exit (default); "
+                       "'immediate' => exit immediately.", "graceful");
+
+  cli.switchGroup("Training options");
   // clang-format off
   cli.add<std::string>("--cost-type", // @TODO: rename to loss-type
       "Optimization criterion: ce-mean, ce-mean-words, ce-sum, perplexity", "ce-mean");

src/common/signal_handling.cpp

@@ -0,0 +1,33 @@
+#include "common/logging.h"
+#include "signal_handling.h"
+
+// We use signal() here instead of the usual strong recommendation for
+// using sigaction, which apparently is not available for Windows (cf.
+// https://stackoverflow.com/questions/231912/what-is-the-difference-between-sigaction-and-signal).
+
+namespace marian{
+volatile std::sig_atomic_t sigflags_{0};
+volatile std::sig_atomic_t gracefulExitRequested_{0};
+
+bool getSignalFlag(const int sig) {
+  // sig_atomic_t has 32 bits. We don't accommodate signals beyond that.
+  ABORT_IF(sig >= 32, "Signal out of range (must be < 32, is {}).", sig);
+  return sigflags_ & (1<<sig);
+}
+
+void requestGracefulExit(int sig) {
+  setSignalFlag(sig);         // keep track of triggering signal
+  gracefulExitRequested_ = 1; // set flag to exit gracefully
+  LOG(debug, "Graceful exit requested via signal {}.", sig);
+}
+
+bool gracefulExitRequested() {
+  return gracefulExitRequested_;
+}
+
+void setSignalFlag(int sig) {
+  // sig_atomic_t has 32 bits. We don't accommodate signals beyond that.
+  if (sig < 32) sigflags_ |= (1<<sig);
+}
+
+}

src/common/signal_handling.h

@@ -0,0 +1,39 @@
+#pragma once
+#include <csignal>
+#include <string>
+
+// SIGNAL HANDLING
+
+// The signal handlers (and checkers) here are implemented in line with with the recommendations
+// for signal handling in the SEI CERT C Coding Standard, specifically
+//
+// - SIG30-C:
+//   https://wiki.sei.cmu.edu/confluence/display/c/SIG30-C.+Call+only+asynchronous-safe+functions+within+signal+handlers
+//
+// - SIG31-C:
+//   https://wiki.sei.cmu.edu/confluence/display/c/SIG31-C.+Do+not+access+shared+objects+in+signal+handlers
+//
+// The exact behavior of 'graceful exit' depends on the application; for training, it means 'save model and exit',
+// for a server (not implemented yet): 'block new requests but serve pending requests and then exit'.
+//
+// Graceful exit for training is useful for training on clusters with time limits on jobs. Slurm, for example, can be
+// set up to send a custom signal at a set time before the end of the time slot, giving Marian time to save its current
+// state before getting killed.
+
+namespace marian {
+
+
+/// Request graceful exit (signal handler)
+void requestGracefulExit(const int sig);
+
+/// Check if graceful exit was requested.
+bool gracefulExitRequested();
+
+/// General purpose signal handler that simply sets a flag when a signal is received.
+//  (only for SIGNAL No. < 32).
+void setSignalFlag(const int sig);  // custom handler (set flag) for sig
+
+/// Check if a setSignalFlag was triggered for this signal
+bool getSignalFlag(const int sig);
+
+} // End of namespace marian

src/data/batch_generator.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "common/options.h"
+#include "common/signal_handling.h"
 #include "data/batch_stats.h"
 #include "data/rng_engine.h"
 #include "training/training_state.h"
@@ -136,6 +137,8 @@ class BatchGenerator : public RNGEngine {
     }
     size_t sets = 0;
     while(current_ != data_->end() && maxiBatch->size() < maxSize) { // loop over data
+      if (gracefulExitRequested()) // stop generating batches
+        return std::deque<BatchPtr>();
       maxiBatch->push(*current_);
       sets = current_->size();
       // do not consume more than required for the maxi batch as this causes
@@ -161,6 +164,8 @@ class BatchGenerator : public RNGEngine {
     if (stats_)
       cachedStatsIter = stats_->begin();
     while(!maxiBatch->empty()) { // while there are sentences in the queue
+      if (gracefulExitRequested()) // stop generating batches
+        return std::deque<BatchPtr>();
       // push item onto batch
       batchVector.push_back(maxiBatch->top());
       maxiBatch->pop(); // fetch next-shortest
@@ -249,15 +254,15 @@ class BatchGenerator : public RNGEngine {
             "If you have changed the training corpus, add --no-restore-corpus to the training command and run it again.");
         bufferedBatches_ = std::move(futureBufferedBatches_.get());
         // if bg thread returns an empty swath, we hit the end of the epoch
-        if (bufferedBatches_.empty()) {
+        if (bufferedBatches_.empty() || gracefulExitRequested()) {
           return nullptr;
         }
         // and kick off the next bg operation
         fetchBatchesAsync();
       } else { // don't spawn any threads, i.e. batch fetching is blocking.
         bufferedBatches_ = fetchBatches();
         // if bufferedBatches is empty we hit the end of the epoch
-        if (bufferedBatches_.empty()) {
+        if (bufferedBatches_.empty() || gracefulExitRequested()) {
           return nullptr;
         }
       }

src/training/scheduler.h

@@ -1,16 +1,14 @@
 #pragma once
 
 #include "common/options.h"
+#include "common/signal_handling.h"
 #include "training/training_state.h"
 #include "training/validator.h"
 #include "training/communicator.h"
 #include "layers/loss.h"
 
 namespace marian {
 
-bool getSigtermFlag();
-void installSignalHandlers();
-
 class Scheduler : public TrainingObserver {
 private:
   Ptr<Options> options_;
@@ -154,11 +152,10 @@ class Scheduler : public TrainingObserver {
       : options_(options), state_(state) {
     ABORT_IF(state_->factor != 1, "state.factor unexpectedly not 1 at this point??");
     updateLearningRate(*state);
-    installSignalHandlers();
   }
 
   bool keepGoing() {
-    if(getSigtermFlag()) // received signal SIGERM => exit gracefully
+    if(gracefulExitRequested()) // via SIGTERM
       return false;
 
     // stop if it reached the maximum number of epochs
@@ -192,13 +189,12 @@ class Scheduler : public TrainingObserver {
 
   void started() { LOG(info, "Training started"); }
   void finished() {
-    if (getSigtermFlag())
-      LOG(info, "Training interrupted (SIGTERM).");
+    if (gracefulExitRequested())
+      LOG(info, "Training interrupted (via signal).");
     else
       LOG(info, "Training finished");
   }
 
-
   void addValidator(Ptr<ValidatorBase> validator) {
     validators_.push_back(validator);
 
@@ -223,9 +219,10 @@ class Scheduler : public TrainingObserver {
 
   void validate(const std::vector<Ptr<ExpressionGraph>>& graphs,
                 bool isFinal = false) {
-    // Do not validate if already validated (for instance, after the model is
-    // loaded) or if validation is scheduled for another update, or when signal SIGTERM was received
-    if(getSigtermFlag() // SIGTERM was received
+    // Do not validate if already validated (for instance, after the model is loaded)
+    // or if validation is scheduled for another update, or when a graceful shutdown
+    // was requested.
+    if(gracefulExitRequested()
        || state_->validated // already validated (in resumed training, for example)
        || (!state_->enteredNewPeriodOf(options_->get<std::string>("valid-freq")) && !isFinal)) // not now
       return;

src/training/training.h

@@ -16,6 +16,7 @@ template <class ModelWrapper>
 class Train : public ModelTask {
 private:
   Ptr<Options> options_;
+  void installCustomSignalHandlers();
 
 public:
   Train(Ptr<Options> options) : options_(options) {}
@@ -77,6 +78,9 @@ class Train : public ModelTask {
     bool restored = !options_->get<bool>("no-restore-corpus")
                     && batchGenerator->restore(trainState);
 
+    // We only want custom behavior once training starts.
+    installCustomSignalHandlers();
+
     // -- main training loop
     scheduler->started();
     while(scheduler->keepGoing()) {
@@ -107,4 +111,17 @@ class Train : public ModelTask {
     finalizeMPI(std::move(mpi));
   }
 };
+
+template <class ModelWrapper>
+void Train<ModelWrapper>::installCustomSignalHandlers()
+{
+  const std::string sigTermAction = options_->get<std::string>("sigterm");
+  if (sigTermAction == "graceful") {
+    LOG(debug, "Enabling graceful shutdown for SIGTERM.");
+    signal(SIGTERM, requestGracefulExit);
+  }
+  else if (sigTermAction != "immediate")
+    ABORT("Unrecognized value '{}' for --sigterm", sigTermAction);
+}
+
 }  // namespace marian