diff --git a/pyiron_workflow/executors/wrapped_executorlib.py b/pyiron_workflow/executors/wrapped_executorlib.py index c8775df30..8289e94f8 100644 --- a/pyiron_workflow/executors/wrapped_executorlib.py +++ b/pyiron_workflow/executors/wrapped_executorlib.py @@ -1,4 +1,5 @@ import inspect +from threading import Thread from typing import Any, ClassVar from executorlib import BaseExecutor, SingleNodeExecutor, SlurmClusterExecutor @@ -22,6 +23,21 @@ class ProtectedResourceError(ValueError): class CacheOverride(BaseExecutor): override_cache_file_name: ClassVar[str] = "executorlib_cache" + def shutdown(self, wait: bool = True, *, cancel_futures: bool = False): + if ( + self._process is not None + and self._future_queue is not None + and wait + and isinstance(self._process, Thread) + ): + self._future_queue.put( + {"shutdown": True, "wait": wait, "cancel_futures": cancel_futures} + ) + self._process.join() + self._future_queue.join() + else: + super().shutdown(wait, cancel_futures=cancel_futures) + def submit(self, fn, /, *args, **kwargs): """ We demand that `fn` be the bound-method `on_run` of a `Lexical`+`Runnable`