[Fix] Provide restart limit during initialization (#938)

* [Fix] Provide restart limit during initialization

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

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

---------

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

Author: jan-janssen

src/executorlib/executor/flux.py

@@ -9,6 +9,7 @@
     check_plot_dependency_graph,
     check_pmi,
     check_refresh_rate,
+    check_restart_limit,
     check_wait_on_shutdown,
     validate_number_of_cores,
 )
@@ -70,6 +71,7 @@ class FluxJobExecutor(BaseExecutor):
         export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
         wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
+        restart_limit (int): The maximum number of restarting worker processes.
         openmpi_oversubscribe (bool): adds the `--oversubscribe` command flag (OpenMPI and SLURM) - default False
 
     Examples:
@@ -113,6 +115,7 @@ def __init__(
         export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
         wait: bool = True,
+        restart_limit: int = 0,
         openmpi_oversubscribe: bool = False,
     ):
         """
@@ -164,6 +167,7 @@ def __init__(
             export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
             wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
+            restart_limit (int): The maximum number of restarting worker processes.
             openmpi_oversubscribe (bool): adds the `--oversubscribe` command flag (OpenMPI and SLURM) - default False
 
         """
@@ -180,6 +184,9 @@ def __init__(
         resource_dict.update(
             {k: v for k, v in default_resource_dict.items() if k not in resource_dict}
         )
+        check_restart_limit(
+            restart_limit=restart_limit, block_allocation=block_allocation
+        )
         if not disable_dependencies:
             super().__init__(
                 executor=DependencyTaskScheduler(
@@ -197,6 +204,7 @@ def __init__(
                         init_function=init_function,
                         log_obj_size=log_obj_size,
                         wait=wait,
+                        restart_limit=restart_limit,
                     ),
                     max_cores=max_cores,
                     refresh_rate=refresh_rate,
@@ -223,6 +231,7 @@ def __init__(
                     init_function=init_function,
                     log_obj_size=log_obj_size,
                     wait=wait,
+                    restart_limit=restart_limit,
                 )
             )
 
@@ -464,6 +473,7 @@ def create_flux_executor(
     init_function: Optional[Callable] = None,
     log_obj_size: bool = False,
     wait: bool = True,
+    restart_limit: int = 0,
 ) -> Union[OneProcessTaskScheduler, BlockAllocationTaskScheduler]:
     """
     Create a flux executor
@@ -504,6 +514,7 @@ def create_flux_executor(
         init_function (None): optional function to preset arguments for functions which are submitted later
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
         wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
+        restart_limit (int): The maximum number of restarting worker processes.
 
     Returns:
         InteractiveStepExecutor/ InteractiveExecutor
@@ -551,6 +562,7 @@ def create_flux_executor(
             max_workers=max_workers,
             executor_kwargs=resource_dict,
             spawner=FluxPythonSpawner,
+            restart_limit=restart_limit,
         )
     else:
         return OneProcessTaskScheduler(

src/executorlib/executor/single.py

@@ -7,6 +7,7 @@
     check_init_function,
     check_plot_dependency_graph,
     check_refresh_rate,
+    check_restart_limit,
     check_wait_on_shutdown,
     validate_number_of_cores,
 )
@@ -61,6 +62,7 @@ class SingleNodeExecutor(BaseExecutor):
         export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
         wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
+        restart_limit (int): The maximum number of restarting worker processes.
         openmpi_oversubscribe (bool): adds the `--oversubscribe` command flag (OpenMPI and SLURM) - default False
 
     Examples:
@@ -100,6 +102,7 @@ def __init__(
         export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
         wait: bool = True,
+        restart_limit: int = 0,
         openmpi_oversubscribe: bool = False,
     ):
         """
@@ -146,6 +149,7 @@ def __init__(
             export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
             wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
+            restart_limit (int): The maximum number of restarting worker processes.
             openmpi_oversubscribe (bool): adds the `--oversubscribe` command flag (OpenMPI and SLURM) - default False
 
         """
@@ -162,6 +166,9 @@ def __init__(
         resource_dict.update(
             {k: v for k, v in default_resource_dict.items() if k not in resource_dict}
         )
+        check_restart_limit(
+            restart_limit=restart_limit, block_allocation=block_allocation
+        )
         if not disable_dependencies:
             super().__init__(
                 executor=DependencyTaskScheduler(
@@ -175,6 +182,7 @@ def __init__(
                         init_function=init_function,
                         log_obj_size=log_obj_size,
                         wait=wait,
+                        restart_limit=restart_limit,
                     ),
                     max_cores=max_cores,
                     refresh_rate=refresh_rate,
@@ -197,6 +205,7 @@ def __init__(
                     init_function=init_function,
                     log_obj_size=log_obj_size,
                     wait=wait,
+                    restart_limit=restart_limit,
                 )
             )
 
@@ -397,6 +406,7 @@ def create_single_node_executor(
     init_function: Optional[Callable] = None,
     log_obj_size: bool = False,
     wait: bool = True,
+    restart_limit: int = 0,
 ) -> Union[OneProcessTaskScheduler, BlockAllocationTaskScheduler]:
     """
     Create a single node executor
@@ -432,6 +442,7 @@ def create_single_node_executor(
         init_function (None): optional function to preset arguments for functions which are submitted later
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
         wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
+        restart_limit (int): The maximum number of restarting worker processes.
 
     Returns:
         InteractiveStepExecutor/ InteractiveExecutor
@@ -466,6 +477,7 @@ def create_single_node_executor(
             ),
             executor_kwargs=resource_dict,
             spawner=MpiExecSpawner,
+            restart_limit=restart_limit,
         )
     else:
         return OneProcessTaskScheduler(

src/executorlib/executor/slurm.py

@@ -6,6 +6,7 @@
     check_log_obj_size,
     check_plot_dependency_graph,
     check_refresh_rate,
+    check_restart_limit,
     check_wait_on_shutdown,
     validate_number_of_cores,
 )
@@ -288,6 +289,7 @@ class SlurmJobExecutor(BaseExecutor):
         export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
         wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
+        restart_limit (int): The maximum number of restarting worker processes.
         openmpi_oversubscribe (bool): adds the `--oversubscribe` command flag (OpenMPI and SLURM) - default False
 
     Examples:
@@ -328,6 +330,7 @@ def __init__(
         export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
         wait: bool = True,
+        restart_limit: int = 0,
         openmpi_oversubscribe: bool = False,
     ):
         """
@@ -378,6 +381,7 @@ def __init__(
             export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
             wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
+            restart_limit (int): The maximum number of restarting worker processes.
             openmpi_oversubscribe (bool): adds the `--oversubscribe` command flag (OpenMPI and SLURM) - default False
 
         """
@@ -394,6 +398,9 @@ def __init__(
         resource_dict.update(
             {k: v for k, v in default_resource_dict.items() if k not in resource_dict}
         )
+        check_restart_limit(
+            restart_limit=restart_limit, block_allocation=block_allocation
+        )
         if not disable_dependencies:
             super().__init__(
                 executor=DependencyTaskScheduler(
@@ -408,6 +415,7 @@ def __init__(
                         init_function=init_function,
                         log_obj_size=log_obj_size,
                         wait=wait,
+                        restart_limit=restart_limit,
                     ),
                     max_cores=max_cores,
                     refresh_rate=refresh_rate,
@@ -431,6 +439,7 @@ def __init__(
                     init_function=init_function,
                     log_obj_size=log_obj_size,
                     wait=wait,
+                    restart_limit=restart_limit,
                 )
             )
 
@@ -446,6 +455,7 @@ def create_slurm_executor(
     init_function: Optional[Callable] = None,
     log_obj_size: bool = False,
     wait: bool = True,
+    restart_limit: int = 0,
 ) -> Union[OneProcessTaskScheduler, BlockAllocationTaskScheduler]:
     """
     Create a SLURM executor
@@ -487,6 +497,7 @@ def create_slurm_executor(
         init_function (None): optional function to preset arguments for functions which are submitted later
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
         wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
+        restart_limit (int): The maximum number of restarting worker processes.
 
     Returns:
         InteractiveStepExecutor/ InteractiveExecutor
@@ -517,6 +528,7 @@ def create_slurm_executor(
             max_workers=max_workers,
             executor_kwargs=resource_dict,
             spawner=SrunSpawner,
+            restart_limit=restart_limit,
         )
     else:
         return OneProcessTaskScheduler(

src/executorlib/standalone/inputcheck.py

@@ -158,6 +158,13 @@ def check_hostname_localhost(hostname_localhost: Optional[bool]) -> None:
         )
 
 
+def check_restart_limit(restart_limit: int, block_allocation: bool = True) -> None:
+    if not block_allocation and restart_limit != 0:
+        raise ValueError(
+            "The option to specify a restart limit for worker processes is only available with block_allocation=True."
+        )
+
+
 def check_pmi_mode(pmi_mode: Optional[str]) -> None:
     if pmi_mode is not None:
         raise ValueError(

src/executorlib/task_scheduler/interactive/blockallocation.py

@@ -38,6 +38,7 @@ class BlockAllocationTaskScheduler(TaskSchedulerBase):
         max_workers (int): defines the number workers which can execute functions in parallel
         executor_kwargs (dict): keyword arguments for the executor
         spawner (BaseSpawner): interface class to initiate python processes
+        restart_limit (int): The maximum number of restarting worker processes.
 
     Examples:
 
@@ -65,13 +66,15 @@ def __init__(
         max_workers: int = 1,
         executor_kwargs: Optional[dict] = None,
         spawner: type[BaseSpawner] = MpiExecSpawner,
+        restart_limit: int = 0,
     ):
         if executor_kwargs is None:
             executor_kwargs = {}
         super().__init__(max_cores=executor_kwargs.get("max_cores"))
         executor_kwargs["future_queue"] = self._future_queue
         executor_kwargs["spawner"] = spawner
         executor_kwargs["queue_join_on_shutdown"] = False
+        executor_kwargs["restart_limit"] = restart_limit
         self._process_kwargs = executor_kwargs
         self._max_workers = max_workers
         self_id = random.getrandbits(128)

tests/unit/standalone/test_inputcheck.py

@@ -14,6 +14,7 @@
     check_refresh_rate,
     check_resource_dict,
     check_resource_dict_is_empty,
+    check_restart_limit,
     check_pmi_mode,
     check_max_workers_and_cores,
     check_hostname_localhost,
@@ -70,6 +71,10 @@ def test_check_pmi(self):
         with self.assertRaises(ValueError):
             check_pmi(backend="flux_allocation", pmi="test")
 
+    def test_check_restart_limit(self):
+        with self.assertRaises(ValueError):
+            check_restart_limit(restart_limit=1, block_allocation=False)
+
     def test_check_nested_flux_executor(self):
         with self.assertRaises(ValueError):
             check_nested_flux_executor(nested_flux_executor=True)