-
Notifications
You must be signed in to change notification settings - Fork 18.7k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #2253 from jyegerlehner/snapshot_on_signal
Snapshot on signal
- Loading branch information
Showing
5 changed files
with
268 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,24 @@ | ||
| #ifndef INCLUDE_CAFFE_UTIL_SIGNAL_HANDLER_H_ | ||
| #define INCLUDE_CAFFE_UTIL_SIGNAL_HANDLER_H_ | ||
|
|
||
| #include "caffe/proto/caffe.pb.h" | ||
| #include "caffe/solver.hpp" | ||
|
|
||
| namespace caffe { | ||
|
|
||
| class SignalHandler { | ||
| public: | ||
| // Contructor. Specify what action to take when a signal is received. | ||
| SignalHandler(SolverAction::Enum SIGINT_action, | ||
| SolverAction::Enum SIGHUP_action); | ||
| ~SignalHandler(); | ||
| ActionCallback GetActionFunction(); | ||
| private: | ||
| SolverAction::Enum CheckForSignals() const; | ||
| SolverAction::Enum SIGINT_action_; | ||
| SolverAction::Enum SIGHUP_action_; | ||
| }; | ||
|
|
||
| } // namespace caffe | ||
|
|
||
| #endif // INCLUDE_CAFFE_UTIL_SIGNAL_HANDLER_H_ |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
Oops, something went wrong.