From 2f53006db89de2199e87891199985ed709644c40 Mon Sep 17 00:00:00 2001
From: Colm Talbot <talbotcolm@gmail.com>
Date: Tue, 15 Aug 2023 14:25:55 -0400
Subject: [PATCH] DEV: make pool terminate on exit
---
bilby/core/sampler/base_sampler.py | 19 ++++++++++++++++---
1 file changed, 16 insertions(+), 3 deletions(-)
diff --git a/bilby/core/sampler/base_sampler.py b/bilby/core/sampler/base_sampler.py
index 6b4f34619..c1b34e57d 100644
--- a/bilby/core/sampler/base_sampler.py
+++ b/bilby/core/sampler/base_sampler.py
@@ -727,16 +727,29 @@ class Sampler(object):
if self.npool in (1, None) or getattr(self, "pool", None) is not None:
self._log_interruption(signum=signum)
self.write_current_state()
- self._close_pool()
+ self._close_pool(hard=True)
if self.hard_exit:
os._exit(self.exit_code)
else:
sys.exit(self.exit_code)
- def _close_pool(self):
+ def _close_pool(self, hard=False):
+ """
+ Close the worker pool.
+
+ Parameters
+ ==========
+ hard: bool
+ Whether to terminate the pool or just close it. Terminating the
+ pool will kill all running jobs, closing it will wait for them to
+ finish.
+ """
if getattr(self, "pool", None) is not None:
logger.info("Starting to close worker pool.")
- self.pool.close()
+ if hard:
+ self.pool.terminate()
+ else:
+ self.pool.close()
self.pool.join()
self.pool = None
self.kwargs["pool"] = self.pool
--
GitLab