Add concurrent limitation

Allow user to define how many workers to run in parallel rather than attempting to apply workers to every job in the array

Added validation for settings of max concurrent (So only positive int or None)

Author: calvinp0

arc/job/pipe/pipe_run.py

@@ -17,6 +17,7 @@
 import stat
 import sys
 import time
+from numbers import Integral
 from typing import Dict, List, Optional
 
 import arc.parser.parser as parser
@@ -51,16 +52,30 @@ class PipeRun:
         run_id (str): Unique identifier for this pipe run.
         tasks (List[TaskSpec]): Task specifications to execute.
         cluster_software (str): Cluster scheduler type.
-        max_workers (int): Maximum number of concurrent array workers.
+        max_workers (int): Maximum total array worker slots (array size).
+        max_concurrent (Optional[int]): Max workers running simultaneously
+            (like PBS ``%N``). ``None`` disables throttling.
         max_attempts (int): Maximum retry attempts per task.
     """
 
+    @staticmethod
+    def _validate_max_concurrent(max_concurrent: Optional[int]) -> None:
+        """Accept ``None`` or a positive integer for throttle settings."""
+        if max_concurrent is None:
+            return
+        if isinstance(max_concurrent, bool) or not isinstance(max_concurrent, Integral):
+            raise ValueError('PipeRun max_concurrent must be None or a positive integer.')
+        if max_concurrent > 0:
+            return
+        raise ValueError('PipeRun max_concurrent must be None or a positive integer.')
+
     def __init__(self,
                  project_directory: str,
                  run_id: str,
                  tasks: List[TaskSpec],
                  cluster_software: str,
                  max_workers: int = 100,
+                 max_concurrent: Optional[int] = None,
                  max_attempts: int = 3,
                  pipe_root: Optional[str] = None,
                  ):
@@ -69,6 +84,8 @@ def __init__(self,
         self.tasks = tasks
         self.cluster_software = cluster_software
         self.max_workers = max_workers
+        self._validate_max_concurrent(max_concurrent)
+        self.max_concurrent = None if max_concurrent is None else int(max_concurrent)
         self.max_attempts = max_attempts
         self.pipe_root = pipe_root if pipe_root is not None \
             else os.path.join(project_directory, 'calcs', 'pipe_' + run_id)
@@ -103,6 +120,7 @@ def _save_run_metadata(self) -> None:
             'status': self.status.value,
             'cluster_software': self.cluster_software,
             'max_workers': self.max_workers,
+            'max_concurrent': self.max_concurrent,
             'max_attempts': self.max_attempts,
             'task_family': task_family,
             'engine': engine,
@@ -146,6 +164,7 @@ def from_dir(cls, pipe_root: str) -> 'PipeRun':
             tasks=tasks,
             cluster_software=data['cluster_software'],
             max_workers=data.get('max_workers', 100),
+            max_concurrent=data.get('max_concurrent'),
             max_attempts=data.get('max_attempts', 3),
             pipe_root=pipe_root,
         )
@@ -192,12 +211,41 @@ def _submission_resources(self):
         Derive resource settings from the homogeneous task list.
 
         Returns:
-            Tuple[int, int, int]: ``(cpus, memory_mb, array_size)``
+            Tuple[int, int, int, Optional[int]]:
+                ``(cpus, memory_mb, array_size, throttle)`` where ``throttle``
+                caps workers running simultaneously (clamped to ``array_size``),
+                or ``None`` if unthrottled.
         """
         cpus = self.tasks[0].required_cores if self.tasks else 1
         memory_mb = self.tasks[0].required_memory_mb if self.tasks else 4096
         array_size = min(self.max_workers, len(self.tasks)) if self.tasks else self.max_workers
-        return cpus, memory_mb, array_size
+        throttle = None
+        if self.max_concurrent is not None and self.max_concurrent > 0:
+            throttle = min(self.max_concurrent, array_size)
+        return cpus, memory_mb, array_size, throttle
+
+    def _render_throttle(self, array_size: int, throttle: Optional[int]) -> Dict[str, str]:
+        """
+        Render scheduler-specific array-range and extra-directives strings.
+
+        SLURM/PBS encode the throttle as an inline ``%K`` suffix on the range.
+        SGE uses a separate ``-tc`` directive. HTCondor uses ``max_materialize``
+        and takes a bare count (not a range) for ``queue``.
+        """
+        cs = 'sge' if self.cluster_software == 'oge' else self.cluster_software
+        if cs == 'htcondor':
+            array_range = str(array_size)
+            extra = f'max_materialize = {throttle}' if throttle else ''
+        elif cs == 'sge':
+            array_range = f'1-{array_size}'
+            extra = f'#$ -tc {throttle}' if throttle else ''
+        elif cs in ('slurm', 'pbs'):
+            suffix = f'%{throttle}' if throttle else ''
+            array_range = f'1-{array_size}{suffix}'
+            extra = ''
+        else:
+            raise NotImplementedError(f'No throttle rendering for {self.cluster_software}')
+        return {'array_range': array_range, 'extra_directives': extra}
 
     def write_submit_script(self) -> str:
         """
@@ -215,7 +263,8 @@ def write_submit_script(self) -> str:
             raise NotImplementedError(
                 f'No pipe submit template for cluster software: {self.cluster_software}. '
                 f'Available templates: {list(pipe_submit.keys())}')
-        cpus, memory_mb, array_size = self._submission_resources()
+        cpus, memory_mb, array_size, throttle = self._submission_resources()
+        throttle_fields = self._render_throttle(array_size, throttle)
         server = servers_dict.get('local', {})
         queue, _ = next(iter(server.get('queues', {}).items()), ('', None))
         engine = self.tasks[0].engine if self.tasks else ''
@@ -226,7 +275,8 @@ def write_submit_script(self) -> str:
             env_setup = f'{env_setup}\n{scratch_export}' if env_setup else scratch_export
         content = pipe_submit[template_key].format(
             name=f'pipe_{self.run_id}',
-            max_task_num=array_size,
+            array_range=throttle_fields['array_range'],
+            extra_directives=throttle_fields['extra_directives'],
             pipe_root=self.pipe_root,
             python_exe=sys.executable,
             cpus=cpus,

arc/job/pipe/pipe_run_test.py

@@ -11,6 +11,7 @@
 import tempfile
 import time
 import unittest
+import uuid
 
 from arc.job.adapters.mockter import MockAdapter
 from arc.job.pipe.pipe_state import TaskState, PipeRunState, TaskSpec, read_task_state, update_task_state
@@ -134,12 +135,14 @@ def setUp(self):
     def tearDown(self):
         shutil.rmtree(self.tmpdir, ignore_errors=True)
 
-    def _make_run(self, cluster_software, max_workers=10, n_tasks=None):
+    def _make_run(self, cluster_software, max_workers=10, n_tasks=None,
+                  max_concurrent=None, run_id=None):
         n = n_tasks if n_tasks is not None else max_workers
         tasks = [_make_spec(f't_{i}') for i in range(n)]
-        run = PipeRun(project_directory=self.tmpdir, run_id='sub_test',
+        run = PipeRun(project_directory=self.tmpdir,
+                      run_id=run_id or f'sub_test_{uuid.uuid4().hex[:8]}',
                       tasks=tasks, cluster_software=cluster_software,
-                      max_workers=max_workers)
+                      max_workers=max_workers, max_concurrent=max_concurrent)
         run.stage()
         return run
 
@@ -170,6 +173,104 @@ def test_htcondor_content(self):
             content = f.read()
         self.assertIn('queue 12', content)
 
+    def test_slurm_throttle(self):
+        run = self._make_run('slurm', max_workers=100, n_tasks=100, max_concurrent=8)
+        with open(run.write_submit_script()) as f:
+            content = f.read()
+        self.assertIn('#SBATCH --array=1-100%8', content)
+
+    def test_pbs_throttle(self):
+        run = self._make_run('pbs', max_workers=50, n_tasks=50, max_concurrent=4)
+        with open(run.write_submit_script()) as f:
+            content = f.read()
+        self.assertIn('#PBS -J 1-50%4', content)
+
+    def test_sge_throttle_uses_tc_directive(self):
+        run = self._make_run('sge', max_workers=20, n_tasks=20, max_concurrent=5)
+        with open(run.write_submit_script()) as f:
+            content = f.read()
+        self.assertIn('#$ -t 1-20', content)
+        self.assertIn('#$ -tc 5', content)
+
+    def test_htcondor_throttle_uses_max_materialize(self):
+        run = self._make_run('htcondor', max_workers=12, n_tasks=12, max_concurrent=3)
+        with open(run.write_submit_script()) as f:
+            content = f.read()
+        self.assertIn('queue 12', content)
+        self.assertIn('max_materialize = 3', content)
+
+    def test_throttle_clamped_to_array_size(self):
+        # max_concurrent > array_size should clamp to array_size (no-op throttle).
+        run = self._make_run('slurm', max_workers=6, n_tasks=6, max_concurrent=99)
+        with open(run.write_submit_script()) as f:
+            content = f.read()
+        self.assertIn('#SBATCH --array=1-6%6', content)
+
+    def test_unthrottled_has_no_throttle_markers(self):
+        """Regression guard: unthrottled submit scripts contain no throttle syntax."""
+        array_line_markers = {
+            'slurm':    '#SBATCH --array=',
+            'pbs':      '#PBS -J ',
+            'sge':      '#$ -t ',
+            'htcondor': 'queue ',
+        }
+        for cs, marker in array_line_markers.items():
+            with self.subTest(cluster_software=cs):
+                run = self._make_run(cs, max_workers=10, n_tasks=10, max_concurrent=None)
+                with open(run.write_submit_script()) as f:
+                    content = f.read()
+                array_line = next(line for line in content.splitlines() if marker in line)
+                self.assertNotIn('%', array_line)
+                self.assertNotIn('-tc ', content)
+                self.assertNotIn('max_materialize', content)
+
+    def test_oge_routes_to_sge_throttle(self):
+        """cluster_software='oge' should use the SGE throttle directive."""
+        run = self._make_run('oge', max_workers=15, n_tasks=15, max_concurrent=4)
+        with open(run.write_submit_script()) as f:
+            content = f.read()
+        self.assertIn('#$ -t 1-15', content)
+        self.assertIn('#$ -tc 4', content)
+
+    def test_invalid_max_concurrent_values_raise(self):
+        """Only None and positive integers are accepted."""
+        for max_concurrent in (-1, 0, -2, 1.5, '3', True, False):
+            with self.subTest(max_concurrent=max_concurrent):
+                with self.assertRaisesRegex(ValueError, 'max_concurrent'):
+                    PipeRun(project_directory=self.tmpdir,
+                            run_id=f'invalid_{uuid.uuid4().hex[:8]}',
+                            tasks=[_make_spec('t_0')],
+                            cluster_software='slurm',
+                            max_concurrent=max_concurrent)
+
+    def test_render_throttle_branching_matrix(self):
+        """Unit-test _render_throttle directly across scheduler × throttle combos."""
+        run = self._make_run('slurm', max_workers=1, n_tasks=1)
+        cases = [
+            ('slurm',    100, None, {'array_range': '1-100',   'extra_directives': ''}),
+            ('slurm',    100, 8,    {'array_range': '1-100%8', 'extra_directives': ''}),
+            ('pbs',      50,  4,    {'array_range': '1-50%4',  'extra_directives': ''}),
+            ('sge',      20,  None, {'array_range': '1-20',    'extra_directives': ''}),
+            ('sge',      20,  5,    {'array_range': '1-20',    'extra_directives': '#$ -tc 5'}),
+            ('oge',      20,  5,    {'array_range': '1-20',    'extra_directives': '#$ -tc 5'}),
+            ('htcondor', 12,  None, {'array_range': '12',      'extra_directives': ''}),
+            ('htcondor', 12,  3,    {'array_range': '12',      'extra_directives': 'max_materialize = 3'}),
+        ]
+        for cs, array_size, throttle, expected in cases:
+            with self.subTest(cluster_software=cs, throttle=throttle):
+                run.cluster_software = cs
+                self.assertEqual(run._render_throttle(array_size, throttle), expected)
+
+    def test_from_dir_round_trip_preserves_max_concurrent(self):
+        """Persist max_concurrent through run.json so crash-recovered runs keep their throttle."""
+        run = self._make_run('slurm', max_workers=20, n_tasks=20, max_concurrent=7)
+        run.write_submit_script()
+        reloaded = PipeRun.from_dir(run.pipe_root)
+        self.assertEqual(reloaded.max_concurrent, 7)
+        with open(reloaded.write_submit_script()) as f:
+            content = f.read()
+        self.assertIn('#SBATCH --array=1-20%7', content)
+
     def test_overwrite_is_safe(self):
         run = self._make_run('slurm')
         p1 = run.write_submit_script()