Add SLURM_MAX_ARRAY_SIZE env var for to control maximum running slurm jobs for array jobs in cluster_tools (#788)

* add env for slurm max running array jobs

* renamed max_running to max_running_size and SLURM_MAX_RUNNING to SLURM_MAX_RUNNING_SIZE

* add SLURM_MAX_ARRAY_SIZE to documentation

* ran ./format.sh and ./lint.sh

* ran ./lint.sh again with black==21.12b0

Co-authored-by: Philipp Otto <[email protected]>

Author: elhuhdron

cluster_tools/README.md

@@ -26,7 +26,7 @@ if __name__ == '__main__':
 The `cluster_tools` automatically determine the slurm limit for maximum array job size and split up larger job batches into multiple smaller batches.
 Also, the slurm limit for the maximum number of jobs which are allowed to be submitted by a user at the same time is honored by looking up the number of currently submitted jobs and only submitting new batches if they fit within the limit.
 
-If you would like to configure these limits independently, you can do so by setting the `SLURM_MAX_ARRAY_SIZE` and `SLURM_MAX_SUBMIT_JOBS` environment variables.
+If you would like to configure these limits independently, you can do so by setting the `SLURM_MAX_ARRAY_SIZE` and `SLURM_MAX_SUBMIT_JOBS` environment variables. You can also limit the maximum number of simultaneously running tasks within the slurm array job(s) by using the `SLURM_MAX_RUNNING_SIZE` environment variable.
 
 ### Kubernetes
 

cluster_tools/cluster_tools/schedulers/slurm.py

@@ -118,6 +118,18 @@ def get_max_array_size():
             )
         return max_array_size
 
+    @staticmethod
+    @cache_in_production
+    def get_max_running_size():
+        max_running_size_env = os.environ.get("SLURM_MAX_RUNNING_SIZE", None)
+        if max_running_size_env is not None:
+            logging.debug(
+                f"SLURM_MAX_RUNNING_SIZE env variable specified which is {max_running_size_env}."
+            )
+            return int(max_running_size_env)
+
+        return 0
+
     @staticmethod
     @cache_in_production
     def get_max_submit_jobs():
@@ -219,6 +231,10 @@ def inner_submit(
 
         max_array_size = self.get_max_array_size()
         max_submit_jobs = self.get_max_submit_jobs()
+        max_running_size = self.get_max_running_size()
+        slurm_max_running_size_str = (
+            "%{}".format(max_running_size) if max_running_size > 0 else ""
+        )
         # Only ever submit at most max_submit_jobs and max_array_size jobs at once (but at least one).
         batch_size = max(min(max_array_size, max_submit_jobs), 1)
 
@@ -233,7 +249,9 @@ def inner_submit(
 
             job_array_line = ""
             if job_count is not None:
-                job_array_line = "#SBATCH --array=0-{}".format(array_index_end)
+                job_array_line = "#SBATCH --array=0-{}{}".format(
+                    array_index_end, slurm_max_running_size_str
+                )
             script_lines = (
                 [
                     "#!/bin/sh",