add queue timeout for 30 seconds

VIDA-NYU · Aug 4, 2023 · 034c29c · 034c29c
1 parent 6ee7d76
commit 034c29c
Showing 1 changed file with 22 additions and 18 deletions.
diff --git a/alpha_automl/automl_manager.py b/alpha_automl/automl_manager.py
@@ -67,31 +67,35 @@ def _search_pipelines(self, automl_hyperparams):
         found_pipelines = 0
 
         while True:
-            result = queue.get()
+            try:
+                result = queue.get(timeout=30)
+            except Exception:
+                logger.info('Queue Timeout after 30s')
+                result = None
 
             if result == 'DONE':
                 search_process.terminate()
                 search_process.join(30)
                 logger.info(f'Found {found_pipelines} pipelines')
                 logger.info('Search done')
                 break
-
-            pipeline = result
-            score = pipeline.get_score()
-            logger.info('Found new pipeline')
-            yield {'pipeline': pipeline, 'message': 'FOUND'}
-
-            if need_rescoring:
-                score, start_time, end_time = score_pipeline(pipeline.get_pipeline(), self.X, self.y, self.scoring,
-                                                             self.splitting_strategy)
-                pipeline.set_score(score)
-                pipeline.set_start_time(start_time)
-                pipeline.set_end_time(end_time)
-
-            if score is not None:
-                logger.info(f'Pipeline scored successfully, score={score}')
-                found_pipelines += 1
-                yield {'pipeline': pipeline, 'message': 'SCORED'}
+            elif result is not None:
+                pipeline = result
+                score = pipeline.get_score()
+                logger.info('Found new pipeline')
+                yield {'pipeline': pipeline, 'message': 'FOUND'}
+
+                if need_rescoring:
+                    score, start_time, end_time = score_pipeline(pipeline.get_pipeline(), self.X, self.y, self.scoring,
+                                                                self.splitting_strategy)
+                    pipeline.set_score(score)
+                    pipeline.set_start_time(start_time)
+                    pipeline.set_end_time(end_time)
+
+                if score is not None:
+                    logger.info(f'Pipeline scored successfully, score={score}')
+                    found_pipelines += 1
+                    yield {'pipeline': pipeline, 'message': 'SCORED'}
 
             if time.time() > search_start_time + self.time_bound:
                 search_process.terminate()