From a29f2b4335946d61b20f435bad802ffde596e56f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jan=20Jan=C3=9Fen?= Date: Fri, 11 Aug 2023 19:47:13 -0600 Subject: [PATCH] Implement a central _set_init_function() --- pympipool/flux/fluxtask.py | 5 +---- pympipool/mpi/mpitask.py | 5 +---- pympipool/shared/executorbase.py | 6 ++++++ 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pympipool/flux/fluxtask.py b/pympipool/flux/fluxtask.py index 89925cea..b7a59331 100644 --- a/pympipool/flux/fluxtask.py +++ b/pympipool/flux/fluxtask.py @@ -73,10 +73,7 @@ def __init__( }, ) self._process.start() - if init_function is not None: - self._future_queue.put( - {"init": True, "fn": init_function, "args": (), "kwargs": {}} - ) + self._set_init_function(init_function=init_function) cloudpickle_register(ind=3) diff --git a/pympipool/mpi/mpitask.py b/pympipool/mpi/mpitask.py index 61834f19..fdcc17b7 100644 --- a/pympipool/mpi/mpitask.py +++ b/pympipool/mpi/mpitask.py @@ -72,10 +72,7 @@ def __init__( }, ) self._process.start() - if init_function is not None: - self._future_queue.put( - {"init": True, "fn": init_function, "args": (), "kwargs": {}} - ) + self._set_init_function(init_function=init_function) cloudpickle_register(ind=3) diff --git a/pympipool/shared/executorbase.py b/pympipool/shared/executorbase.py index ae4b687b..c950eae1 100644 --- a/pympipool/shared/executorbase.py +++ b/pympipool/shared/executorbase.py @@ -58,6 +58,12 @@ def shutdown(self, wait=True, *, cancel_futures=False): def __len__(self): return self._future_queue.qsize() + def _set_init_function(self, init_function): + if init_function is not None: + self._future_queue.put( + {"init": True, "fn": init_function, "args": (), "kwargs": {}} + ) + def cancel_items_in_queue(que): """