Merge pull request #206 from facebookresearch/abort_tuning_sigterm

Genetic tuner: std::abort on SIGTERM

Author: prigoyal

include/tc/autotuner/genetic_tuning_harness.h

@@ -32,9 +32,6 @@ namespace tc {
 namespace autotune {
 namespace detail {
 
-extern volatile std::sig_atomic_t signal_;
-extern volatile std::sig_atomic_t killRequested_;
-
 class GeneticTunerHarness {
  public:
   GeneticTunerHarness(
@@ -50,6 +47,7 @@ class GeneticTunerHarness {
       std::vector<MappingOptions> startingPoints,
       const TuningParameterFixer& fixedParams);
   void run(size_t numGenerations);
+  void stopAfterCurrentGeneration();
 
  private:
   void setupTuningParameters();
@@ -110,6 +108,8 @@ class GeneticTunerHarness {
   std::unordered_map<size_t, std::vector<DLTensor*>> outputs_;
   const MappingOptions kBaseMapping_;
   const std::vector<MappingOptions> kStartingPoints_;
+
+  std::atomic_bool stopRequested_{false};
 };
 
 std::vector<size_t> parseGpus();

src/autotuner/genetic_autotuner.cc

@@ -83,6 +83,11 @@ std::vector<MappingOptions> GeneticAutotuner::load(
   return tc::autotune::restoreCandidates(tcName, inputs, outputs);
 }
 
+namespace {
+volatile std::sig_atomic_t sigint_ = 0;
+volatile std::sig_atomic_t sigterm_ = 0;
+} // namespace
+
 llvm::Optional<MappingOptions> GeneticAutotuner::tune(
     const std::string& cacheFileName,
     const std::string& tcName,
@@ -122,13 +127,13 @@ llvm::Optional<MappingOptions> GeneticAutotuner::tune(
       startingPoints,
       fixedParams);
 
-  std::signal(SIGTERM, [](int sig) {
-    signal_ = sig;
-    killRequested_ = 1;
-  });
-  std::signal(SIGINT, [](int sig) {
-    signal_ = sig;
-    killRequested_ = 1;
+  sigterm_ = 0;
+  sigint_ = 0;
+  auto sigtermOrigHandler = std::signal(SIGTERM, [](int) { sigterm_ = 1; });
+  auto sigintOrigHandler = std::signal(SIGINT, [](int) { sigint_ = 1; });
+  ScopeGuard handlersGuard([&]() {
+    std::signal(SIGTERM, sigtermOrigHandler);
+    std::signal(SIGINT, sigintOrigHandler);
   });
 
   std::atomic_bool tunerFinished(false);
@@ -148,15 +153,23 @@ llvm::Optional<MappingOptions> GeneticAutotuner::tune(
   });
   while (not tunerFinished) {
     std::this_thread::sleep_for(std::chrono::milliseconds(100));
-    if (killRequested_) {
+    if (sigint_) {
+      std::cerr
+          << "Autotuning will stop after the current generation has finished."
+          << std::endl;
+      tuner.stopAfterCurrentGeneration();
+      tunerThread.join();
+      storeCaches(cacheFileName);
+    }
+    if (sigterm_) {
       std::cerr << "Autotuning aborted." << std::endl;
       storeCaches(cacheFileName);
-      tunerThread.join();
-      killRequested_ = 0;
-      throw std::runtime_error("Abort requested");
+      std::abort();
     }
   }
+
   tunerThread.join();
+
   if (tunerThreadEx) {
     std::rethrow_exception(tunerThreadEx);
   }

src/autotuner/genetic_tuning_harness.cc

@@ -18,7 +18,6 @@
 
 #include <atomic>
 #include <chrono>
-#include <csignal>
 #include <numeric>
 #include <thread>
 
@@ -41,9 +40,6 @@ namespace tc {
 namespace autotune {
 namespace detail {
 
-volatile std::sig_atomic_t signal_ = 0;
-volatile std::sig_atomic_t killRequested_ = 0;
-
 GeneticTunerHarness::GeneticTunerHarness(
     size_t n,
     uint8_t crossoverRate,
@@ -102,12 +98,16 @@ GeneticTunerHarness::GeneticTunerHarness(
 
 void GeneticTunerHarness::run(size_t numGenerations) {
   for (size_t i = 0; i < numGenerations; ++i) {
-    if (not killRequested_) {
+    if (not stopRequested_) {
       runOneGeneration(i);
     }
   }
 }
 
+void GeneticTunerHarness::stopAfterCurrentGeneration() {
+  stopRequested_ = true;
+}
+
 namespace {
 
 std::vector<size_t> filterHigherThan(