Kill child process using signals

VIDA-NYU · Nov 14, 2023 · 66dea64 · 66dea64
1 parent e65da04
commit 66dea64
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 19 deletions.
diff --git a/alpha_automl/automl_manager.py b/alpha_automl/automl_manager.py
@@ -70,7 +70,7 @@ def _search_pipelines(self, automl_hyperparams):
 
             if result == 'DONE':
                 search_process.terminate()
-                search_process.join(30)
+                search_process.join(10)
                 logger.debug(f'Found {found_pipelines} pipelines')
                 logger.debug('Search done')
                 break
@@ -93,10 +93,10 @@ def _search_pipelines(self, automl_hyperparams):
                 yield {'pipeline': pipeline, 'message': 'SCORED'}
 
             if time.time() > search_start_time + self.time_bound:
+                logger.debug('Reached search timeout')
                 search_process.terminate()
-                search_process.join(30)
+                search_process.join(10)
                 logger.debug(f'Found {found_pipelines} pipelines')
-                logger.debug('Reached search timeout')
                 break
 
     def check_automl_hyperparams(self, automl_hyperparams):

diff --git a/alpha_automl/pipeline_synthesis/setup_search.py b/alpha_automl/pipeline_synthesis/setup_search.py
@@ -1,7 +1,8 @@
-import logging
 import os
+import sys
+import signal
+import logging
 from os.path import join
-
 from alpha_automl.grammar_loader import load_automatic_grammar, load_manual_grammar
 from alpha_automl.pipeline import Pipeline
 from alpha_automl.pipeline_search.Coach import Coach
@@ -41,22 +42,16 @@
 }
 
 
+def signal_handler(queue, signum):
+    logger.debug(f'Receiving signal {signum}, terminating process')
+    queue.put('DONE')
+    # TODO: Should it save the last status of the NN model?
+    sys.exit(0)
 
 
-def search_pipelines(
-    X,
-    y,
-    scoring,
-    splitting_strategy,
-    task_name,
-    time_bound,
-    automl_hyperparams,
-    metadata,
-    output_folder,
-    verbose,
-    queue,
-):
-
+def search_pipelines(X, y, scoring, splitting_strategy, task_name, time_bound, automl_hyperparams, metadata,
+                     output_folder, verbose, queue):
+    signal.signal(signal.SIGTERM, lambda signum, frame: signal_handler(queue, signum))
     hide_logs(verbose)  # Hide logs here too, since multiprocessing has some issues with loggers
 
     builder = BaseBuilder(metadata, automl_hyperparams)