Interactive: Interrupt interface bootup when the executor is shutdown during bootup (#801)

* Interactive: Interrupt interface bootup when the executor is shutdown during bootup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix flux interface

* sync changes

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix type hints

* type fixes

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fixes

* remove shutdown

* fixes

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* type fixes

* add tests

* more tests

* fix minimal

* fix interface

* clean up interface

* fixes

* fix type hints

* fix tests

* fix tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update docstring

* fix

* Add restart_limit to resource_dict

* Add error messages

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use random hash

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: jan-janssen

executorlib/executor/flux.py

@@ -43,6 +43,7 @@ class FluxJobExecutor(BaseExecutor):
                                                   compute notes. Defaults to False.
                               - error_log_file (str): Name of the error log file to use for storing exceptions raised
                                                       by the Python functions submitted to the Executor.
+                              - restart_limit (int): The maximum number of restarting worker processes. Default: 0
         pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
         flux_executor (flux.job.FluxExecutor): Flux Python interface to submit the workers to flux
         flux_executor_nesting (bool): Provide hierarchically nested Flux job scheduler inside the submitted function.

executorlib/executor/single.py

@@ -120,6 +120,7 @@ def __init__(
                                                            only)
                                   - error_log_file (str): Name of the error log file to use for storing exceptions
                                                           raised by the Python functions submitted to the Executor.
+                                  - restart_limit (int): The maximum number of restarting worker processes. Default: 0
             hostname_localhost (boolean): use localhost instead of the hostname to establish the zmq connection. In the
                                       context of an HPC cluster this essential to be able to communicate to an
                                       Executor running on a different compute node within the same allocation. And

executorlib/executor/slurm.py

@@ -43,6 +43,7 @@ class SlurmClusterExecutor(BaseExecutor):
                               - slurm_cmd_args (list): Additional command line arguments for the srun call (SLURM only)
                               - error_log_file (str): Name of the error log file to use for storing exceptions raised
                                                       by the Python functions submitted to the Executor.
+                              - restart_limit (int): The maximum number of restarting worker processes. Default: 0
         pysqa_config_directory (str, optional): path to the pysqa config directory (only for pysqa based backend).
         pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
         hostname_localhost (boolean): use localhost instead of the hostname to establish the zmq connection. In the

executorlib/standalone/interactive/communication.py

@@ -1,7 +1,7 @@
 import logging
 import sys
 from socket import gethostname
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 import cloudpickle
 import zmq
@@ -42,6 +42,17 @@ def __init__(
         if log_obj_size:
             self._logger = logging.getLogger("executorlib")
         self._spawner = spawner
+        self._command_lst: list[str] = []
+        self._booted_sucessfully: bool = False
+        self._stop_function: Optional[Callable] = None
+
+    @property
+    def status(self) -> bool:
+        return self._booted_sucessfully
+
+    @status.setter
+    def status(self, status: bool):
+        self._booted_sucessfully = status
 
     def send_dict(self, input_dict: dict):
         """
@@ -67,7 +78,9 @@ def receive_dict(self) -> dict:
         while len(response_lst) == 0:
             response_lst = self._poller.poll(self._time_out_ms)
             if not self._spawner.poll():
-                raise ExecutorlibSocketError()
+                raise ExecutorlibSocketError(
+                    "SocketInterface crashed during execution."
+                )
         data = self._socket.recv(zmq.NOBLOCK)
         if self._logger is not None:
             self._logger.warning(
@@ -105,20 +118,30 @@ def bind_to_random_port(self) -> int:
 
     def bootup(
         self,
-        command_lst: list[str],
-    ) -> bool:
+        command_lst: Optional[list[str]] = None,
+        stop_function: Optional[Callable] = None,
+    ):
         """
         Boot up the client process to connect to the SocketInterface.
 
         Args:
             command_lst (list): list of strings to start the client process
-
-        Returns:
-            bool: Whether the interface was successfully started.
+            stop_function (Callable): Function to stop the interface.
         """
-        return self._spawner.bootup(
-            command_lst=command_lst,
-        )
+        if command_lst is not None:
+            self._command_lst = command_lst
+        if stop_function is not None:
+            self._stop_function = stop_function
+        if len(self._command_lst) == 0:
+            raise ValueError("No command defined to boot up SocketInterface.")
+        if not self._spawner.bootup(
+            command_lst=self._command_lst,
+            stop_function=self._stop_function,
+        ):
+            self._reset_socket()
+            self._booted_sucessfully = False
+        else:
+            self._booted_sucessfully = True
 
     def shutdown(self, wait: bool = True):
         """
@@ -162,6 +185,7 @@ def interface_bootup(
     hostname_localhost: Optional[bool] = None,
     log_obj_size: bool = False,
     worker_id: Optional[int] = None,
+    stop_function: Optional[Callable] = None,
 ) -> SocketInterface:
     """
     Start interface for ZMQ communication
@@ -180,6 +204,7 @@ def interface_bootup(
         log_obj_size (boolean): Enable debug mode which reports the size of the communicated objects.
         worker_id (int): Communicate the worker which ID was assigned to it for future reference and resource
                          distribution.
+        stop_function (Callable): Function to stop the interface.
 
     Returns:
          executorlib.shared.communication.SocketInterface: socket interface for zmq communication
@@ -203,6 +228,7 @@ def interface_bootup(
     ]
     interface.bootup(
         command_lst=command_lst,
+        stop_function=stop_function,
     )
     return interface
 

executorlib/standalone/interactive/spawner.py

@@ -1,7 +1,7 @@
 import os
 import subprocess
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Callable, Optional
 
 MPI_COMMAND = "mpiexec"
 
@@ -29,12 +29,17 @@ def __init__(
     def bootup(
         self,
         command_lst: list[str],
-    ):
+        stop_function: Optional[Callable] = None,
+    ) -> bool:
         """
         Method to start the interface.
 
         Args:
             command_lst (list[str]): The command list to execute.
+            stop_function (Callable): Function to stop the interface.
+
+        Returns:
+            bool: Whether the interface was successfully started.
         """
         raise NotImplementedError
 
@@ -87,12 +92,17 @@ def __init__(
     def bootup(
         self,
         command_lst: list[str],
-    ):
+        stop_function: Optional[Callable] = None,
+    ) -> bool:
         """
         Method to start the subprocess interface.
 
         Args:
             command_lst (list[str]): The command list to execute.
+            stop_function (Callable): Function to stop the interface.
+
+        Returns:
+            bool: Whether the interface was successfully started.
         """
         if self._cwd is not None:
             os.makedirs(self._cwd, exist_ok=True)

executorlib/task_scheduler/interactive/blockallocation.py

@@ -1,4 +1,5 @@
 import queue
+import random
 from concurrent.futures import Future
 from threading import Thread
 from typing import Callable, Optional
@@ -8,11 +9,21 @@
     check_resource_dict,
     check_resource_dict_is_empty,
 )
-from executorlib.standalone.interactive.communication import interface_bootup
+from executorlib.standalone.interactive.communication import (
+    ExecutorlibSocketError,
+    SocketInterface,
+    interface_bootup,
+)
 from executorlib.standalone.interactive.spawner import BaseSpawner, MpiExecSpawner
 from executorlib.standalone.queue import cancel_items_in_queue
 from executorlib.task_scheduler.base import TaskSchedulerBase
-from executorlib.task_scheduler.interactive.shared import execute_task_dict, task_done
+from executorlib.task_scheduler.interactive.shared import (
+    execute_task_dict,
+    reset_task_dict,
+    task_done,
+)
+
+_interrupt_bootup_dict: dict = {}
 
 
 class BlockAllocationTaskScheduler(TaskSchedulerBase):
@@ -63,11 +74,18 @@ def __init__(
         executor_kwargs["queue_join_on_shutdown"] = False
         self._process_kwargs = executor_kwargs
         self._max_workers = max_workers
+        self_id = random.getrandbits(128)
+        self._self_id = self_id
+        _interrupt_bootup_dict[self._self_id] = False
         self._set_process(
             process=[
                 Thread(
                     target=_execute_multiple_tasks,
-                    kwargs=executor_kwargs | {"worker_id": worker_id},
+                    kwargs=executor_kwargs
+                    | {
+                        "worker_id": worker_id,
+                        "stop_function": lambda: _interrupt_bootup_dict[self_id],
+                    },
                 )
                 for worker_id in range(self._max_workers)
             ],
@@ -157,6 +175,7 @@ def shutdown(self, wait: bool = True, *, cancel_futures: bool = False):
             if cancel_futures:
                 cancel_items_in_queue(que=self._future_queue)
             if isinstance(self._process, list):
+                _interrupt_bootup_dict[self._self_id] = True
                 for _ in range(len(self._process)):
                     self._future_queue.put({"shutdown": True, "wait": wait})
                 if wait:
@@ -190,6 +209,8 @@ def _execute_multiple_tasks(
     log_obj_size: bool = False,
     error_log_file: Optional[str] = None,
     worker_id: Optional[int] = None,
+    stop_function: Optional[Callable] = None,
+    restart_limit: int = 0,
     **kwargs,
 ) -> None:
     """
@@ -216,6 +237,8 @@ def _execute_multiple_tasks(
                               submitted to the Executor.
         worker_id (int): Communicate the worker which ID was assigned to it for future reference and resource
                          distribution.
+        stop_function (Callable): Function to stop the interface.
+        restart_limit (int): The maximum number of restarting worker processes.
     """
     interface = interface_bootup(
         command_lst=get_interactive_execute_command(
@@ -225,34 +248,66 @@ def _execute_multiple_tasks(
         hostname_localhost=hostname_localhost,
         log_obj_size=log_obj_size,
         worker_id=worker_id,
+        stop_function=stop_function,
+    )
+    interface_initialization_exception = _set_init_function(
+        interface=interface,
+        init_function=init_function,
     )
+    restart_counter = 0
+    while True:
+        if not interface.status and restart_counter > restart_limit:
+            interface.status = True  # no more restarts
+            interface_initialization_exception = ExecutorlibSocketError(
+                "SocketInterface crashed during execution."
+            )
+        elif not interface.status:
+            interface.bootup()
+            interface_initialization_exception = _set_init_function(
+                interface=interface,
+                init_function=init_function,
+            )
+            restart_counter += 1
+        else:  # interface.status == True
+            task_dict = future_queue.get()
+            if "shutdown" in task_dict and task_dict["shutdown"]:
+                if interface.status:
+                    interface.shutdown(wait=task_dict["wait"])
+                task_done(future_queue=future_queue)
+                if queue_join_on_shutdown:
+                    future_queue.join()
+                break
+            elif "fn" in task_dict and "future" in task_dict:
+                f = task_dict.pop("future")
+                if interface_initialization_exception is not None:
+                    f.set_exception(exception=interface_initialization_exception)
+                else:
+                    # The interface failed during the execution
+                    interface.status = execute_task_dict(
+                        task_dict=task_dict,
+                        future_obj=f,
+                        interface=interface,
+                        cache_directory=cache_directory,
+                        cache_key=cache_key,
+                        error_log_file=error_log_file,
+                    )
+                    if not interface.status:
+                        reset_task_dict(
+                            future_obj=f, future_queue=future_queue, task_dict=task_dict
+                        )
+                task_done(future_queue=future_queue)
+
+
+def _set_init_function(
+    interface: SocketInterface,
+    init_function: Optional[Callable] = None,
+) -> Optional[Exception]:
     interface_initialization_exception = None
-    if init_function is not None:
+    if init_function is not None and interface.status:
         try:
             _ = interface.send_and_receive_dict(
                 input_dict={"init": True, "fn": init_function, "args": (), "kwargs": {}}
             )
         except Exception as init_exception:
             interface_initialization_exception = init_exception
-    while True:
-        task_dict = future_queue.get()
-        if "shutdown" in task_dict and task_dict["shutdown"]:
-            interface.shutdown(wait=task_dict["wait"])
-            task_done(future_queue=future_queue)
-            if queue_join_on_shutdown:
-                future_queue.join()
-            break
-        elif "fn" in task_dict and "future" in task_dict:
-            f = task_dict.pop("future")
-            if interface_initialization_exception is not None:
-                f.set_exception(exception=interface_initialization_exception)
-            else:
-                execute_task_dict(
-                    task_dict=task_dict,
-                    future_obj=f,
-                    interface=interface,
-                    cache_directory=cache_directory,
-                    cache_key=cache_key,
-                    error_log_file=error_log_file,
-                )
-            task_done(future_queue=future_queue)
+    return interface_initialization_exception

executorlib/task_scheduler/interactive/onetoone.py

@@ -4,7 +4,10 @@
 from typing import Optional
 
 from executorlib.standalone.command import get_interactive_execute_command
-from executorlib.standalone.interactive.communication import interface_bootup
+from executorlib.standalone.interactive.communication import (
+    ExecutorlibSocketError,
+    interface_bootup,
+)
 from executorlib.standalone.interactive.spawner import BaseSpawner, MpiExecSpawner
 from executorlib.task_scheduler.base import TaskSchedulerBase
 from executorlib.task_scheduler.interactive.shared import execute_task_dict
@@ -230,7 +233,7 @@ def _execute_task_in_thread(
     error_log_file: Optional[str] = None,
     worker_id: Optional[int] = None,
     **kwargs,
-) -> None:
+):
     """
     Execute a single tasks in parallel using the message passing interface (MPI).
 
@@ -256,7 +259,7 @@ def _execute_task_in_thread(
         worker_id (int): Communicate the worker which ID was assigned to it for future reference and resource
                          distribution.
     """
-    execute_task_dict(
+    if not execute_task_dict(
         task_dict=task_dict,
         future_obj=future_obj,
         interface=interface_bootup(
@@ -271,4 +274,7 @@ def _execute_task_in_thread(
         cache_directory=cache_directory,
         cache_key=cache_key,
         error_log_file=error_log_file,
-    )
+    ):
+        future_obj.set_exception(
+            ExecutorlibSocketError("SocketInterface crashed during execution.")
+        )