Add --throttle option to babs init for SLURM array job throttling

Author: tien-tong

babs/bootstrap.py

@@ -37,6 +37,7 @@ def babs_bootstrap(
         container_name,
         container_config,
         initial_inclusion_df=None,
+        throttle=None,
     ):
         """
         Bootstrap a babs project: initialize datalad-tracked RIAs, generate scripts to be used, etc
@@ -57,6 +58,11 @@ def babs_bootstrap(
             of how to run the BIDS App container
         initial_inclusion_df: pd.DataFrame
             initial inclusion dataframe of subjects/sessions to analyze
+        throttle: int or None, optional
+            Optional throttle value for SLURM array jobs. This limits the number of
+            simultaneously running array tasks. The value will be added to the array
+            specification as `%<throttle>`. Example: `10` will result in
+            `--array=1-${max_array}%10`.
         """
         if op.exists(self.project_root):
             raise FileExistsError(
@@ -79,6 +85,9 @@ def babs_bootstrap(
 
         os.makedirs(self.project_root)
 
+        # Store throttle value for job submission template
+        self.throttle = throttle
+
         # validate `processing_level`:
         self.processing_level = validate_processing_level(processing_level)
 

babs/cli.py

@@ -118,6 +118,14 @@ def _parse_init():
         " Please refer to section below 'What if ``babs init`` fails?' for details.",
         #      ^^ in `babs init.rst`, pointed to below section for more
     )
+    parser.add_argument(
+        '--throttle',
+        type=int,
+        help='Optional throttle value for SLURM array jobs. '
+        'This limits the number of simultaneously running array tasks. '
+        'The value will be added to the array specification as ``%%<throttle>``. '
+        'Example: ``--throttle 10`` will result in ``--array=1-${max_array}%%10``.',
+    )
 
     return parser
 
@@ -148,6 +156,7 @@ def babs_init_main(
     processing_level: str,
     queue: str,
     keep_if_failed: bool,
+    throttle: int | None = None,
 ):
     """This is the core function of babs init.
 
@@ -175,6 +184,11 @@ def babs_init_main(
         sge or slurm
     keep_if_failed: bool
         If `babs init` failed with error, whether to keep the created BABS project.
+    throttle: int or None, optional
+        Optional throttle value for SLURM array jobs. This limits the number of
+        simultaneously running array tasks. The value will be added to the array
+        specification as `%<throttle>`. Example: `10` will result in
+        `--array=1-${max_array}%10`.
     """
 
     from babs import BABSBootstrap
@@ -188,6 +202,7 @@ def babs_init_main(
             container_name,
             container_config,
             list_sub_file,
+            throttle=throttle,
         )
     except Exception as exc:
         print('\n`babs init` failed! Below is the error message:')

babs/container.py

@@ -274,6 +274,9 @@ def generate_job_submit_template(self, yaml_path, babs, system, test=False):
                 array_args = '--array=1'
             else:  # need max_array for for `submit_job_template.yaml`
                 array_args = '--array=1-${max_array}'
+                # Add throttle if specified
+                if hasattr(babs, 'throttle') and babs.throttle is not None:
+                    array_args += f'%{babs.throttle}'
 
         # Render the template
         env = Environment(

docs/babs-init.rst

@@ -104,6 +104,14 @@ a SLURM cluster:
         --queue slurm \
         /path/to/a/folder/holding/BABS/project/my_BABS_project
 
+.. note::
+    **Throttling SLURM array jobs**: If you want to limit the number of simultaneously running
+    array tasks, you can add the ``--throttle`` option. For example, ``--throttle 10`` will
+    limit SLURM to run at most 10 array tasks at the same time. This is useful when you have
+    many jobs but want to avoid overwhelming the cluster or hitting resource limits.
+    The throttle value will be added to the array specification as ``%<throttle>``,
+    e.g., ``--array=1-${max_array}%10``.
+
 
 *********
 Debugging

docs/walkthrough.rst

@@ -354,6 +354,12 @@ and results and provenance are saved. An example command of ``babs init`` is as
         --queue slurm \
         "${HOME}/babs_demo/my_BABS_project"
 
+.. note::
+    **Optional: Throttling array jobs**: If you have many jobs and want to limit how many
+    run simultaneously, you can add ``--throttle <number>`` to the command above.
+    For example, ``--throttle 10`` will limit SLURM to run at most 10 array tasks at once.
+    This helps avoid overwhelming the cluster or hitting resource limits when submitting
+    large numbers of array tasks.
 
 Here you will create a BABS project called ``my_BABS_project`` in directory ``~/babs_demo``.
 The input dataset is specified in the yaml file and no longer specified in the command line.

tests/test_base.py

@@ -1,7 +1,9 @@
 """Test the check_setup functionality."""
 
+import os
 import os.path as op
 import random
+import re
 from pathlib import Path
 from unittest.mock import MagicMock
 
@@ -202,3 +204,54 @@ def test_key_info_full(babs_project_sessionlevel):
     babs_proj.wtf_key_info(flag_output_ria_only=False)
     assert babs_proj.output_ria_data_dir is not None
     assert babs_proj.analysis_dataset_id is not None
+
+
+@pytest.mark.parametrize(
+    ('throttle_value', 'expected_in_template'),
+    [(10, True), (None, False)],
+)
+def test_throttle_in_job_template(
+    tmp_path_factory,
+    templateflow_home,
+    simbids_container_ds,
+    bids_data_singlesession,
+    throttle_value,
+    expected_in_template,
+):
+    """Test that throttle value is correctly included in job submission template."""
+    from conftest import get_config_simbids_path, update_yaml_for_run
+
+    from babs.bootstrap import BABSBootstrap
+
+    os.environ['TEMPLATEFLOW_HOME'] = str(templateflow_home)
+
+    project_base = tmp_path_factory.mktemp('project')
+    project_root = project_base / f'my_babs_project_{throttle_value or "none"}'
+    container_config = update_yaml_for_run(
+        project_base,
+        get_config_simbids_path().name,
+        {'BIDS': bids_data_singlesession},
+    )
+
+    babs_bootstrap = BABSBootstrap(project_root=project_root)
+    babs_bootstrap.babs_bootstrap(
+        processing_level='subject',
+        queue='slurm',
+        container_ds=simbids_container_ds,
+        container_name='simbids-0-0-3',
+        container_config=container_config,
+        initial_inclusion_df=None,
+        throttle=throttle_value,
+    )
+
+    assert babs_bootstrap.throttle == throttle_value
+
+    template_path = op.join(babs_bootstrap.analysis_path, 'code', 'submit_job_template.yaml')
+    with open(template_path) as f:
+        cmd_template = yaml.safe_load(f)['cmd_template']
+
+    assert '--array=1-${max_array}' in cmd_template
+    if expected_in_template:
+        assert f'%{throttle_value}' in cmd_template
+    else:
+        assert not re.search(r'%\d+', cmd_template)