Refactored ArraySpec

Author: Ladme

src/qq_lib/batch/interface/interface.py

@@ -12,7 +12,6 @@
 from typing import TYPE_CHECKING, Any, cast
 
 from qq_lib.batch.interface._arfe import _AtomicRemoteFileEditor
-from qq_lib.batch.interface.array_spec import ArraySpec
 from qq_lib.core.common import convert_absolute_to_relative
 from qq_lib.core.config import CFG
 from qq_lib.core.error import QQError
@@ -27,6 +26,7 @@
 if TYPE_CHECKING:
     from collections.abc import Callable
 
+    from qq_lib.core.array_spec import ArraySpec
     from qq_lib.properties.depend import Depend
     from qq_lib.properties.resources import Resources
 
@@ -35,7 +35,7 @@
 """
 Type alias for a batch system class.
 """
-type AnyBatchClass = type[BatchInterface[Any, Any, Any, Any]]
+type AnyBatchClass = type[BatchInterface[Any, Any, Any]]
 
 
 class _BatchMeta(ABCMeta):
@@ -137,7 +137,6 @@ class BatchInterface[
     TBatchJob: BatchJobInterface = BatchJobInterface,
     TBatchQueue: BatchQueueInterface = BatchQueueInterface,
     TBatchNode: BatchNodeInterface = BatchNodeInterface,
-    TArraySpec: ArraySpec = ArraySpec,
 ](ABC, metaclass=_BatchMeta):
     """
     Abstract base class for batch system integrations.
@@ -222,7 +221,7 @@ def job_submit(
         depend: list[Depend],
         env_vars: dict[str, str],
         account: str | None = None,
-        array: TArraySpec | None = None,
+        array: ArraySpec | None = None,
         server: str | None = None,
         remote_host: str | None = None,
     ) -> str:
@@ -239,7 +238,7 @@ def job_submit(
             depend (list[Depend]): List of job dependencies.
             env_vars (dict[str, str]): Dictionary of environment variables to propagate to the job.
             account (str | None): Optional account name to use for the job.
-            array (TArraySpec | None): Optional array job specification.
+            array (ArraySpec | None): Optional array job specification.
             server (str | None): Optional name of the server to submit the job to.
             remote_host (str | None): Optional name of the machine to submit the job from.
 

src/qq_lib/batch/pbs/array_spec.py

@@ -12,10 +12,10 @@
 from pathlib import Path
 
 from qq_lib.batch.interface import BatchInterface
-from qq_lib.batch.pbs.array_spec import PBSArraySpec
 from qq_lib.batch.pbs.common import parse_multi_pbs_dump_to_dictionaries
 from qq_lib.batch.pbs.node import PBSNode
 from qq_lib.batch.pbs.queue import PBSQueue
+from qq_lib.core.array_spec import ArraySpec
 from qq_lib.core.common import equals_normalized
 from qq_lib.core.config import CFG
 from qq_lib.core.error import QQError
@@ -29,7 +29,7 @@
 logger = get_logger(__name__)
 
 
-class PBS(BatchInterface[PBSJob, PBSQueue, PBSNode, PBSArraySpec]):
+class PBS(BatchInterface[PBSJob, PBSQueue, PBSNode]):
     """
     Implementation of BatchInterface for PBS Pro batch system.
     """
@@ -76,7 +76,7 @@ def job_submit(
         depend: list[Depend],
         env_vars: dict[str, str],
         account: str | None = None,
-        array: PBSArraySpec | None = None,
+        array: ArraySpec | None = None,
         server: str | None = None,
         remote_host: str | None = None,
     ) -> str:
@@ -589,7 +589,7 @@ def _translate_submit(
         job_name: str,
         depend: list[Depend],
         env_vars: dict[str, str],
-        array: PBSArraySpec | None,
+        array: ArraySpec | None,
     ) -> str:
         """
         Generate the PBS submission command for a job.
@@ -603,7 +603,7 @@ def _translate_submit(
             job_name (str): Name of the job.
             depend (list[Depend]): List of dependencies of the job.
             env_vars (dict[str, str]): Dictionary of environment variables to set.
-            array (PBSArraySpec | None): Optional array job specification.
+            array (ArraySpec | None): Optional array job specification.
 
         Returns:
             str: The fully constructed qsub command string.
@@ -612,7 +612,7 @@ def _translate_submit(
 
         # translate array specification
         if array:
-            command += f"-J {array.translate()} "
+            command += f"-J {cls._translate_array(array)} "
 
         # translate environment variables
         if env_vars:
@@ -865,6 +865,28 @@ def _translate_dependencies(cls, depend: list[Depend]) -> str | None:
 
         return ",".join(Depend.to_str(x).replace("=", ":") for x in depend)
 
+    @classmethod
+    def _translate_array(cls, array: ArraySpec) -> str:
+        """
+        Translate an `ArraySpec` object into a PBS-compatible array job specification string.
+
+        Args:
+            array (ArraySpec): The array job specification to translate.
+
+        Returns:
+            str: The translated PBS array job specification string.
+        """
+        parts: list[str] = []
+        for elem in array.elements:
+            match elem:
+                case int(index):
+                    parts.append(str(index))
+                case (int(start), int(stop)):
+                    parts.append(f"{start}-{stop}")
+                case (int(start), int(stop), int(step)):
+                    parts.append(f"{start}-{stop}:{step}")
+        return ",".join(parts)
+
     @classmethod
     def _collect_ams_env_vars(cls) -> dict[str, str]:
         """

src/qq_lib/batch/pbs/pbs.py

@@ -9,7 +9,7 @@
 
 from qq_lib.batch.interface import BatchInterface
 from qq_lib.batch.pbs.pbs import PBS
-from qq_lib.batch.slurm.array_spec import SlurmArraySpec
+from qq_lib.core.array_spec import ArraySpec
 from qq_lib.core.config import CFG
 from qq_lib.core.error import QQError
 from qq_lib.core.logger import get_logger
@@ -28,7 +28,7 @@
 logger = get_logger(__name__)
 
 
-class Slurm(BatchInterface[SlurmJob, SlurmQueue, SlurmNode, SlurmArraySpec]):
+class Slurm(BatchInterface[SlurmJob, SlurmQueue, SlurmNode]):
     """
     Implementation of BatchInterface for Slurm batch system.
     """
@@ -59,7 +59,7 @@ def job_submit(
         depend: list[Depend],
         env_vars: dict[str, str],
         account: str | None = None,
-        array: SlurmArraySpec | None = None,
+        array: ArraySpec | None = None,
         server: str | None = None,
         remote_host: str | None = None,
     ) -> str:
@@ -411,7 +411,7 @@ def _translate_submit(
         depend: list[Depend],
         env_vars: dict[str, str],
         account: str | None,
-        array: SlurmArraySpec | None,
+        array: ArraySpec | None,
     ) -> str:
         """
         Generate the Slurm submission command for a job.
@@ -425,7 +425,7 @@ def _translate_submit(
             depend (list[Depend]): List of dependencies of the job.
             env_vars (dict[str, str]): Dictionary of environment variables and their values to propagate to the job's environment.
             account (str | None): Optional name of the account to use for the job.
-            array (SlurmArraySpec | None): Optional array job specification.
+            array (ArraySpec | None): Optional array job specification.
 
         Returns:
             str: The fully constructed sbatch command string.
@@ -434,7 +434,7 @@ def _translate_submit(
         command = f"sbatch -J {job_name} -p {queue} -e {qq_output} -o {qq_output} "
 
         if array:
-            command += f"--array={array.translate()} "
+            command += f"--array={cls._translate_array(array)} "
 
         if account:
             command += f"--account {account} "
@@ -582,6 +582,20 @@ def _translate_dependencies(cls, depend: list[Depend]) -> str | None:
 
         return ",".join(Depend.to_str(x).replace("=", ":") for x in depend)
 
+    @classmethod
+    def _translate_array(cls, array: ArraySpec) -> str:
+        """
+        Translate an `ArraySpec` object into a Slurm-compatible array job specification string.
+
+        Args:
+            array (ArraySpec): The array job specification to translate.
+
+        Returns:
+            str: The translated Slurm array job specification string.
+        """
+        # Slurm array syntax is the same as PBS syntax
+        return PBS._translate_array(array)
+
     @classmethod
     def _get_default_server_resources(cls) -> Resources:
         """

src/qq_lib/batch/slurm/array_spec.py

@@ -8,9 +8,9 @@
 from typing import cast
 
 from qq_lib.batch.slurm import Slurm
-from qq_lib.batch.slurm.array_spec import SlurmArraySpec
 from qq_lib.batch.slurmit4i import SlurmIT4I
 from qq_lib.batch.slurmlumi.node import SlurmLumiNode
+from qq_lib.core.array_spec import ArraySpec
 from qq_lib.core.config import CFG
 from qq_lib.core.error import QQError
 from qq_lib.core.logger import get_logger
@@ -46,7 +46,7 @@ def job_submit(
         depend: list[Depend],
         env_vars: dict[str, str],
         account: str | None = None,
-        array: SlurmArraySpec | None = None,
+        array: ArraySpec | None = None,
         server: str | None = None,
         remote_host: str | None = None,
     ) -> str:

src/qq_lib/batch/slurm/slurm.py

@@ -1,19 +1,14 @@
 # Released under MIT License.
 # Copyright (c) 2025-2026 Ladislav Bartos and Robert Vacha Lab
 
-from abc import ABC, abstractmethod
-
 from qq_lib.core.error import QQError
 
 type ArrayElement = int | tuple[int, int] | tuple[int, int, int]
 
 
-class ArraySpec(ABC):
+class ArraySpec:
     """
-    Abstract specification for job-array task indices.
-
-    Each scheduler backend provides a concrete subclass that knows how
-    to translate index specifications into the syntax that particular scheduler expects.
+    Specification for job-array task indices.
 
     Args:
         elements (list[ArrayElement]): Non-empty list of indices and ranges.
@@ -24,17 +19,7 @@ class ArraySpec(ABC):
 
     def __init__(self, elements: list[ArrayElement]):
         _validate_elements(elements)
-        self._elements = _merge_elements(elements)
-
-    @abstractmethod
-    def translate(self) -> str:
-        """
-        Translate this specification into the scheduler's native syntax.
-
-        Returns:
-            str: A string suitable for passing to the scheduler's array flag
-            (e.g. `--array=` for Slurm, `-J` for PBS).
-        """
+        self.elements = _merge_elements(elements)
 
 
 def _validate_elements(elements: list[ArrayElement]) -> None:

src/qq_lib/batch/slurmlumi/slurm.py

@@ -149,14 +149,15 @@ def submit(self, remote: str | None = None) -> str:
             QQError: If job submission fails.
         """
         job_id = self._batch_system.job_submit(
-            self._resources,
-            self._queue,
-            self._script,
-            self._job_name,
-            self._depend,
-            self._create_env_vars_dict(),
-            self._account,
-            self._server,
+            res=self._resources,
+            queue=self._queue,
+            script=self._script,
+            job_name=self._job_name,
+            depend=self._depend,
+            env_vars=self._create_env_vars_dict(),
+            account=self._account,
+            array=None,
+            server=self._server,
             remote_host=remote,
         )