[Feature] Implement run time limits (#930)

* [Feature] Implement run time limits

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

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

* extend tests

* another fix

* Test flux job is reaching the run time limit

* Fix error import

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

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

* Add to __all__

* fix broken import

* fix run time pysqa

---------

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

Author: jan-janssen

src/executorlib/api.py

@@ -8,6 +8,7 @@
 from executorlib.executor.single import TestClusterExecutor
 from executorlib.standalone.command import get_command_path
 from executorlib.standalone.interactive.communication import (
+    ExecutorlibSocketError,
     SocketInterface,
     interface_bootup,
     interface_connect,
@@ -32,4 +33,5 @@
     "MpiExecSpawner",
     "SocketInterface",
     "SubprocessSpawner",
+    "ExecutorlibSocketError",
 ]

src/executorlib/executor/flux.py

@@ -45,6 +45,7 @@ class FluxJobExecutor(BaseExecutor):
                               - 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
+                              - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
         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.
@@ -136,6 +137,7 @@ def __init__(
                                                       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.
+                                  - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
             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.
@@ -246,6 +248,7 @@ class FluxClusterExecutor(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.
+                              - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
         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
@@ -333,6 +336,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.
+                                  - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
             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
@@ -364,6 +368,7 @@ def __init__(
             "cwd": None,
             "openmpi_oversubscribe": False,
             "slurm_cmd_args": [],
+            "run_time_limit": None,
         }
         if resource_dict is None:
             resource_dict = {}
@@ -478,6 +483,7 @@ def create_flux_executor(
                                                   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.
+                              - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
         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.

src/executorlib/executor/single.py

@@ -42,6 +42,7 @@ class SingleNodeExecutor(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.
+                              - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
         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
@@ -126,6 +127,7 @@ def __init__(
                                   - 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
+                                  - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
             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
@@ -219,6 +221,7 @@ class TestClusterExecutor(BaseExecutor):
                               - cwd (str/None): current working directory where the parallel python task is executed
                               - error_log_file (str): Name of the error log file to use for storing exceptions raised
                                                       by the Python functions submitted to the Executor.
+                              - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
         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
@@ -296,6 +299,7 @@ def __init__(
                                   - cwd (str/None): current working directory where the parallel python task is executed
                                   - error_log_file (str): Name of the error log file to use for storing exceptions
                                                           raised by the Python functions submitted to the Executor.
+                                  - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
             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

src/executorlib/executor/slurm.py

@@ -45,6 +45,7 @@ class SlurmClusterExecutor(BaseExecutor):
                               - 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
+                              - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
         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
@@ -132,6 +133,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.
+                                  - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
             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
@@ -267,6 +269,7 @@ class SlurmJobExecutor(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.
+                              - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
         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
                                       context of an HPC cluster this essential to be able to communicate to an
@@ -356,6 +359,7 @@ def __init__(
                                                       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.
+                                  - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
             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
                                       context of an HPC cluster this essential to be able to communicate to an
@@ -469,6 +473,7 @@ def create_slurm_executor(
                                                   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.
+                              - run_time_limit (int): The maximum runtime in seconds for each task. Default: None
         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
                                   context of an HPC cluster this essential to be able to communicate to an

src/executorlib/standalone/command.py

@@ -126,6 +126,7 @@ def generate_slurm_command(
     openmpi_oversubscribe: bool = False,
     slurm_cmd_args: Optional[list[str]] = None,
     pmi_mode: Optional[str] = None,
+    run_time_limit: Optional[int] = None,
 ) -> list[str]:
     """
     Generate the command list for the SLURM interface.
@@ -140,6 +141,7 @@ def generate_slurm_command(
         openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
         slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].
         pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
+        run_time_limit (int): The maximum runtime in seconds for each task. Default: None
 
     Returns:
         list[str]: The generated command list.
@@ -159,6 +161,8 @@ def generate_slurm_command(
         command_prepend_lst += ["--exact"]
     if openmpi_oversubscribe:
         command_prepend_lst += ["--oversubscribe"]
+    if run_time_limit is not None:
+        command_prepend_lst += ["--time=" + str(run_time_limit // 60 + 1)]
     if slurm_cmd_args is not None and len(slurm_cmd_args) > 0:
         command_prepend_lst += slurm_cmd_args
     return command_prepend_lst

src/executorlib/task_scheduler/file/spawner_pysqa.py

@@ -69,6 +69,7 @@ def execute_with_pysqa(
             "command": " ".join(command),
             "dependency_list": [str(qid) for qid in task_dependent_lst],
             "working_directory": os.path.abspath(cwd),
+            "run_time_max": resource_dict.get("run_time_limit"),
         }
         if "cwd" in resource_dict:
             del resource_dict["cwd"]

src/executorlib/task_scheduler/interactive/spawner_flux.py

@@ -44,6 +44,7 @@ class FluxPythonSpawner(BaseSpawner):
         flux_executor (flux.job.FluxExecutor, optional): The FluxExecutor instance. Defaults to None.
         flux_executor_nesting (bool, optional): Whether to use nested FluxExecutor. Defaults to False.
         flux_log_files (bool, optional): Write flux stdout and stderr files. Defaults to False.
+        run_time_limit (int): The maximum runtime in seconds for each task. Default: None
     """
 
     def __init__(
@@ -61,6 +62,7 @@ def __init__(
         flux_executor: Optional[flux.job.FluxExecutor] = None,
         flux_executor_nesting: bool = False,
         flux_log_files: bool = False,
+        run_time_limit: Optional[int] = None,
     ):
         super().__init__(
             cwd=cwd,
@@ -78,6 +80,7 @@ def __init__(
         self._flux_log_files = flux_log_files
         self._priority = priority
         self._future = None
+        self._run_time_limit = run_time_limit
 
     def bootup(
         self,
@@ -128,6 +131,8 @@ def bootup(
         if self._cwd is not None:
             jobspec.cwd = self._cwd
             os.makedirs(self._cwd, exist_ok=True)
+        if self._run_time_limit is not None:
+            jobspec.duration = self._run_time_limit
         file_prefix = "flux_" + str(self._worker_id)
         if self._flux_log_files and self._cwd is not None:
             jobspec.stderr = os.path.join(self._cwd, file_prefix + ".err")

src/executorlib/task_scheduler/interactive/spawner_pysqa.py

@@ -31,6 +31,7 @@ def __init__(
         pmi_mode: Optional[str] = None,
         config_directory: Optional[str] = None,
         backend: Optional[str] = None,
+        run_time_limit: Optional[int] = None,
         **kwargs,
     ):
         """
@@ -50,6 +51,7 @@ def __init__(
             pmi_mode (str, optional): PMI interface to use (OpenMPI v5 requires pmix) default is None
             config_directory (str, optional): path to the pysqa config directory (only for pysqa based backend).
             backend (str): name of the backend used to spawn tasks.
+            run_time_limit (int): The maximum runtime in seconds for each task. Default: None
         """
         super().__init__(
             cwd=cwd,
@@ -68,6 +70,7 @@ def __init__(
         self._pysqa_submission_kwargs = kwargs
         self._process: Optional[int] = None
         self._queue_adapter: Optional[QueueAdapter] = None
+        self._run_time_limit = run_time_limit
 
     def bootup(
         self,
@@ -191,6 +194,7 @@ def _start_process_helper(
             command=" ".join(self.generate_command(command_lst=command_lst)),
             working_directory=working_directory,
             cores=int(self._cores * self._threads_per_core),
+            run_time_max=self._run_time_limit,
             **self._pysqa_submission_kwargs,
         )
 

src/executorlib/task_scheduler/interactive/spawner_slurm.py

@@ -32,6 +32,7 @@ def __init__(
         openmpi_oversubscribe: bool = False,
         slurm_cmd_args: Optional[list[str]] = None,
         pmi_mode: Optional[str] = None,
+        run_time_limit: Optional[int] = None,
     ):
         """
         Srun interface implementation.
@@ -47,6 +48,7 @@ def __init__(
             openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
             slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].
             pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
+            run_time_limit (int): The maximum runtime in seconds for each task. Default: None
         """
         super().__init__(
             cwd=cwd,
@@ -60,6 +62,7 @@ def __init__(
         self._num_nodes = num_nodes
         self._exclusive = exclusive
         self._pmi_mode = pmi_mode
+        self._run_time_limit = run_time_limit
 
     def generate_command(self, command_lst: list[str]) -> list[str]:
         """
@@ -81,6 +84,7 @@ def generate_command(self, command_lst: list[str]) -> list[str]:
             openmpi_oversubscribe=self._openmpi_oversubscribe,
             slurm_cmd_args=self._slurm_cmd_args,
             pmi_mode=self._pmi_mode,
+            run_time_limit=self._run_time_limit,
         )
         return super().generate_command(
             command_lst=command_prepend_lst + command_lst,

tests/unit/executor/test_flux_job.py

@@ -1,9 +1,11 @@
 import os
 import unittest
+from time import sleep
 
 import numpy as np
 
 from executorlib import FluxJobExecutor
+from executorlib.api import ExecutorlibSocketError
 
 
 try:
@@ -20,6 +22,11 @@ def calc(i):
     return i
 
 
+def delayed_calc(i):
+    sleep(2)
+    return i
+
+
 def mpi_funct(i):
     from mpi4py import MPI
 
@@ -110,6 +117,24 @@ def test_single_task(self):
             [[(1, 2, 0), (1, 2, 1)], [(2, 2, 0), (2, 2, 1)], [(3, 2, 0), (3, 2, 1)]],
         )
 
+    def test_run_time_limit(self):
+        with FluxJobExecutor(
+            max_cores=1,
+            resource_dict={"cores": 1},
+            flux_executor=self.executor,
+            block_allocation=False,
+            pmi_mode=pmi,
+        ) as p:
+            f1 = p.submit(delayed_calc, 1, resource_dict={"run_time_limit": 1})
+            f2 = p.submit(delayed_calc, 2, resource_dict={"run_time_limit": 5})
+            self.assertFalse(f1.done())
+            self.assertFalse(f2.done())
+            self.assertEqual(f2.result(), 2)
+            self.assertTrue(f1.done())
+            self.assertTrue(f2.done())
+            with self.assertRaises(ExecutorlibSocketError):
+                f1.result()
+
     def test_output_files_cwd(self):
         dirname = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
         os.makedirs(dirname, exist_ok=True)