Add signal handler and early exit/snapshot to Solver.

Add signal handler and early exit/snapshot to Solver.

Add signal handler and early exit/snapshot to Solver.

Also check for exit and snapshot when testing.

Skip running test after early exit.

Fix more lint.

Rebase on master.

Finish rebase on master.

Fixups per review comments.

Redress review comments.

Lint.

Correct error message wording.

1 files changed, 115 insertions, 0 deletions

diff --git a/src/caffe/util/signal_handler.cpp b/src/caffe/util/signal_handler.cpp
new file mode 100644
index 00000000..5d764ec5
--- /dev/null
+++ b/src/caffe/util/signal_handler.cpp
@@ -0,0 +1,115 @@
+#include <boost/bind.hpp>
+#include <glog/logging.h>
+
+#include <signal.h>
+#include <csignal>
+
+#include "caffe/util/signal_handler.h"
+
+namespace {
+  static volatile sig_atomic_t got_sigint = false;
+  static volatile sig_atomic_t got_sighup = false;
+  static bool already_hooked_up = false;
+
+  void handle_signal(int signal) {
+    switch (signal) {
+    case SIGHUP:
+      got_sighup = true;
+      break;
+    case SIGINT:
+      got_sigint = true;
+      break;
+    }
+  }
+
+  void HookupHandler() {
+    if (already_hooked_up) {
+      LOG(FATAL) << "Tried to hookup signal handlers more than once.";
+    }
+    already_hooked_up = true;
+
+    struct sigaction sa;
+    // Setup the handler
+    sa.sa_handler = &handle_signal;
+    // Restart the system call, if at all possible
+    sa.sa_flags = SA_RESTART;
+    // Block every signal during the handler
+    sigfillset(&sa.sa_mask);
+    // Intercept SIGHUP and SIGINT
+    if (sigaction(SIGHUP, &sa, NULL) == -1) {
+      LOG(FATAL) << "Cannot install SIGHUP handler.";
+    }
+    if (sigaction(SIGINT, &sa, NULL) == -1) {
+      LOG(FATAL) << "Cannot install SIGINT handler.";
+    }
+  }
+
+  // Set the signal handlers to the default.
+  void UnhookHandler() {
+    if (already_hooked_up) {
+      struct sigaction sa;
+      // Setup the sighub handler
+      sa.sa_handler = SIG_DFL;
+      // Restart the system call, if at all possible
+      sa.sa_flags = SA_RESTART;
+      // Block every signal during the handler
+      sigfillset(&sa.sa_mask);
+      // Intercept SIGHUP and SIGINT
+      if (sigaction(SIGHUP, &sa, NULL) == -1) {
+        LOG(FATAL) << "Cannot uninstall SIGHUP handler.";
+      }
+      if (sigaction(SIGINT, &sa, NULL) == -1) {
+        LOG(FATAL) << "Cannot uninstall SIGINT handler.";
+      }
+
+      already_hooked_up = false;
+    }
+  }
+
+  // Return true iff a SIGINT has been received since the last time this
+  // function was called.
+  bool GotSIGINT() {
+    bool result = got_sigint;
+    got_sigint = false;
+    return result;
+  }
+
+  // Return true iff a SIGHUP has been received since the last time this
+  // function was called.
+  bool GotSIGHUP() {
+    bool result = got_sighup;
+    got_sighup = false;
+    return result;
+  }
+}  // namespace
+
+namespace caffe {
+
+SignalHandler::SignalHandler(SolverAction::Enum SIGINT_action,
+                             SolverAction::Enum SIGHUP_action):
+  SIGINT_action_(SIGINT_action),
+  SIGHUP_action_(SIGHUP_action) {
+  HookupHandler();
+}
+
+SignalHandler::~SignalHandler() {
+  UnhookHandler();
+}
+
+SolverAction::Enum SignalHandler::CheckForSignals() const {
+  if (GotSIGHUP()) {
+    return SIGHUP_action_;
+  }
+  if (GotSIGINT()) {
+    return SIGINT_action_;
+  }
+  return SolverAction::NONE;
+}
+
+// Return the function that the solver can use to find out if a snapshot or
+// early exit is being requested.
+ActionCallback SignalHandler::GetActionFunction() {
+  return boost::bind(&SignalHandler::CheckForSignals, this);
+}
+
+}  // namespace caffe