Improved handling of SIGTERM

CHANGELOG.md

@@ -35,8 +35,11 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 - In concatenation make sure that we do not multiply 0 with nan (which results in nan)
 - Change Approx.epsilon(0.01) to Approx.margin(0.001) in unit tests. Tolerance is now
   absolute and not relative. We assumed incorrectly that epsilon is absolute tolerance.
-- Improved handling for graceful shutdown upon receiving SIGTERM.
-  SIGTERM now also interrupts batch prefetching, which runs in a separate thread.
+- Fixed bug in finding .git/logs/HEAD when Marian is a submodule in another project.
+- 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/command/marian_train.cpp

@@ -43,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 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).
-  exit(getSignalFlag(SIGTERM) ? (128 + SIGTERM) : EXIT_SUCCESS);
+  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

@@ -7,15 +7,27 @@
 
 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).", sig);
+  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

@@ -1,27 +1,39 @@
 #pragma once
 #include <csignal>
+#include <string>
 
 // SIGNAL HANDLING
 
-// The Marian signal handlers set global flags that thread can
-// consider when a signal is received. This can be used for a graceful
-// shutdown instead of a hard abandonment, e.g.  after receiving
-// SIGTERM during training.
-
-// When SIGTERM is received, the global (static member) flag sigterm_
-// (false by default) is set to true by signalHandler(). When sigterm_
-// is true, keepGoing() returns false, and the current state of
-// training models is saved prior to exiting.  This functionality is
-// helpful when training on clusters with time limits on compute
-// slots, e.g., on s clusters managed by slurm. Slurm can be asked to
-// sending a (custom) warning signal to a process at a given point in
-// time prior to the hard "time's up".
+// 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
 //
-// Correspondingly, fetchBatches in the batch generator checks the flag
-// frequently and quits after the overall process receives a SIGTERM.
-
+// - 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 {
-bool getSignalFlag(int sig); // return true if sig was received, false otherwise
-void setSignalFlag(int sig); // set custom handler (set flag) for sig
-}
+
+
+/// 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

@@ -135,14 +135,10 @@ class BatchGenerator : public RNGEngine {
       if(current_ != data_->end())
         ++current_;
     }
-
-    std::deque<BatchPtr> tempBatches;
-
     size_t sets = 0;
     while(current_ != data_->end() && maxiBatch->size() < maxSize) { // loop over data
-      if (getSignalFlag(SIGTERM)) { // received SIGTERM, abandon ship ...
-        return tempBatches;
-      }
+      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
@@ -158,20 +154,18 @@ class BatchGenerator : public RNGEngine {
     size_t currentWords = 0;
     std::vector<size_t> lengths(sets, 0); // records maximum length observed within current batch
 
+    std::deque<BatchPtr> tempBatches;
+
     // process all loaded sentences in order of increasing length
     // @TODO: we could just use a vector and do a sort() here; would make the cost more explicit
     const size_t mbWords = options_->get<size_t>("mini-batch-words", 0);
     const bool useDynamicBatching = options_->has("mini-batch-fit");
     BatchStats::const_iterator cachedStatsIter;
     if (stats_)
       cachedStatsIter = stats_->begin();
-
     while(!maxiBatch->empty()) { // while there are sentences in the queue
-
-      if (getSignalFlag(SIGTERM)) { // received SIGTERM, abandon ship ...
-        return tempBatches;
-      }
-
+      if (gracefulExitRequested()) // stop generating batches
+        return std::deque<BatchPtr>();
       // push item onto batch
       batchVector.push_back(maxiBatch->top());
       maxiBatch->pop(); // fetch next-shortest
@@ -252,18 +246,26 @@ class BatchGenerator : public RNGEngine {
 
   BatchPtr next() {
     if(bufferedBatches_.empty()) {
-      // out of data: need to get next batch from background thread
-      // We only get here if the future has been scheduled to run; it must be valid.
-      ABORT_IF(!futureBufferedBatches_.valid(), "Attempted to wait for futureBufferedBatches_ when none pending.\n"
-          "This error often occurs when Marian tries to restore the training data iterator, but the corpus has been changed or replaced.\n"
-          "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 (bufferedBatches_.empty() // i.e., end of Epoch
-          || getSignalFlag(SIGTERM)) { // process received SIGTERM, abandon ship ...
-        return nullptr;
+      if(runAsync_) { // by default we will run in asynchronous mode
+        // out of data: need to get next batch from background thread
+        // We only get here if the future has been scheduled to run; it must be valid.
+        ABORT_IF(!futureBufferedBatches_.valid(), "Attempted to wait for futureBufferedBatches_ when none pending.\n"
+            "This error often occurs when Marian tries to restore the training data iterator, but the corpus has been changed or replaced.\n"
+            "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() || 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() || gracefulExitRequested()) {
+          return nullptr;
+        }
       }
-      fetchBatchesAsync(); // pre-fetch next slew of batches in separate thread
     }
 
     auto batch = bufferedBatches_.front();

src/training/scheduler.h

@@ -155,7 +155,7 @@ class Scheduler : public TrainingObserver {
   }
 
   bool keepGoing() {
-    if(getSignalFlag(SIGTERM)) // received signal SIGERM => exit gracefully
+    if(gracefulExitRequested()) // via SIGTERM
       return false;
 
     // stop if it reached the maximum number of epochs
@@ -189,13 +189,12 @@ class Scheduler : public TrainingObserver {
 
   void started() { LOG(info, "Training started"); }
   void finished() {
-    if (getSignalFlag(SIGTERM))
-      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);
 
@@ -220,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(getSignalFlag(SIGTERM) // 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,11 +78,8 @@ class Train : public ModelTask {
     bool restored = !options_->get<bool>("no-restore-corpus")
                     && batchGenerator->restore(trainState);
 
-    // Install custom handler for SIGTERM, to allow for a graceful
-    // shutdown that saves the current state of training before exiting.
-    // This signal handler simply sets a flag that can be checked from
-    // everywhere (getSignalFLAG(SIGTERM); #include common/signal_handling.h)
-    signal(SIGTERM,setSignalFlag);
+    // We only want custom behavior once training starts.
+    installCustomSignalHandlers();
 
     // -- main training loop
     scheduler->started();
@@ -113,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