feat: add moab

Signed-off-by: vsoch <[email protected]>

Author: vsoch

fractale/transformer/__init__.py

@@ -7,6 +7,7 @@
 from .oar import Transformer as OARTransformer
 from .pbs import Transformer as PBSTransformer
 from .slurm import Transformer as SlurmTransformer
+from .moab import Transformer as MoabTransformer
 
 plugins = {
     "kubernetes": KubernetesTransformer,
@@ -16,6 +17,7 @@
     "lsf": LSFTransformer,
     "oar": OARTransformer,
     "cobalt": CobaltTransformer,
+    "moab": MoabTransformer,
 }
 
 
@@ -32,6 +34,8 @@ def detect_transformer(jobspec):
     content = utils.read_file(jobspec)
     if "#FLUX" in content and "FLUX_CAPACITOR" not in content:
         return "flux"
+    if "#MSUB " in content:
+        return "moab"
     if "#SBATCH " in content:
         return "slurm"
     if "kind:" in content and "Job" in content:

fractale/transformer/moab/__init__.py

@@ -0,0 +1,3 @@
+from .transform import MoabTransformer as Transformer
+
+assert Transformer

fractale/transformer/moab/transform.py

@@ -0,0 +1,303 @@
+import re
+import shlex
+from datetime import timedelta
+
+import fractale.utils as utils
+from fractale.logger.generate import JobNamer
+from fractale.transformer.base import Script, TransformerBase
+from fractale.transformer.common import JobSpec
+
+
+class MoabScript(Script):
+    """
+    A helper class to build a Moab (#MSUB) batch script line by line.
+    """
+    def __init__(self):
+        self.script_lines = ["#!/bin/bash"]
+        self.directive = "#MSUB"
+
+    def add(self, flag, value=None):
+        """
+        Add a Moab directive, e.g., #MSUB -N my-job.
+        Format is #MSUB -<flag> <value>
+        """
+        if value is None:
+            return
+        self.script_lines.append(f"{self.directive} -{flag} {str(value)}")
+
+
+def seconds_to_moab_walltime(seconds):
+    """
+    Converts an integer number of seconds into Moab's HH:MM:SS walltime format.
+    """
+    # This shouldn't happen, but we return 0 so we use the default.
+    if not seconds or seconds <= 0:
+        return None
+    
+    # Moab walltime does not typically include days.
+    hours, seconds_rem = divmod(seconds, 3600)
+    minutes, seconds = divmod(seconds_rem, 60)
+
+    return f"{int(hours):02d}:{int(minutes):02d}:{int(seconds):02d}"
+
+
+def epoch_to_moab_begin_time(epoch_seconds: int) -> str:
+    """
+    Converts a Unix epoch timestamp into Moab's required epoch integer string
+    for the '-S' (starttime) flag.
+    """
+    if not isinstance(epoch_seconds, int) or epoch_seconds < 0:
+        raise ValueError("begin_time must be a positive integer (Unix epoch seconds).")
+
+    return str(epoch_seconds)
+
+
+def moab_walltime_to_seconds(time_str):
+    if not time_str:
+        return None
+    
+    h = 0
+    m = 0
+    s = 0
+    parts = time_str.split(":")
+    if len(parts) == 3:
+        h, m, s = map(int, parts)
+    elif len(parts) == 2:
+        m, s = map(int, parts)
+    elif len(parts) == 1:
+        m = int(parts[0])
+    return int(timedelta(hours=h, minutes=m, seconds=s).total_seconds())
+
+
+def parse_moab_command(command_lines, spec):
+    """
+    Parses a moab command into parts.
+    """
+    # We use the last command line as the primary execution logic
+    main_command = command_lines[-1]
+    parts = shlex.split(main_command)
+
+    # Unwrap common launchers like mpiexec
+    if parts and parts[0] in ('mpiexec', 'mpirun'):
+        parts = parts[1:]
+
+    if parts and parts[0] in ("singularity", "apptainer") and parts[1] == "exec":
+        spec.container_image = parts[2]
+        # The rest is the command inside the container
+        parts = parts[3:]
+
+    # Handle input redirection
+    if "<" in parts:
+        try:
+            idx = parts.index("<")
+            spec.input_file = parts[idx + 1]
+            # Remove '<' and the filename from the arguments
+            parts.pop(idx)
+            parts.pop(idx)
+        except (ValueError, IndexError):
+            pass
+    return parts
+
+
+def moab_begin_time_to_epoch(time_str):
+    """
+    Converts a Moab begin time string (epoch) to an integer.
+    """
+    if not time_str:
+        return None
+    try:
+        return int(time_str)
+    except ValueError:
+        return None
+
+
+def priority_to_moab_priority(priority):
+    """
+    Maps a semantic priority string ("high") to a Moab priority value (-1024 to 1023).
+    """
+    # Higher value means HIGHER priority in Moab.
+    mapping = {
+        "low": -500,
+        "normal": 0,
+        "high": 500,
+        "urgent": 1000,
+    }
+    # Default to 'normal' (0) if the string is None or not in the map
+    return mapping.get(priority, 0)
+
+
+def moab_priority_to_priority(moab_priority):
+    """
+    Maps a Moab priority value back to a semantic string ("high").
+    """
+    if moab_priority is None or moab_priority == 0:
+        return "normal"
+    if moab_priority < 0:
+        return "low"
+    if 0 < moab_priority < 1000:
+        return "high"
+    return "urgent" # for priority >= 1000
+
+
+class MoabTransformer(TransformerBase):
+    """
+    A Moab Transformer for converting a generic JobSpec into a Moab batch script.
+
+    This transformer maps the fields of the JobSpec to their corresponding #MSUB
+    directives and constructs a runnable script.
+    """
+
+    def convert(self, spec) -> str:
+        """
+        Convert a normalized jobspec to a Moab batch script.
+        """
+        script = MoabScript()
+
+        # If we don't have a job name, generate one
+        job_name = spec.job_name or JobNamer().generate()
+        script.add("N", job_name)
+
+        # Job Identity & Accounting
+        script.add("A", spec.account)
+
+        # I/O
+        script.add("o", spec.output_file)
+        script.add("e", spec.error_file)
+        # Moab does not have direct mail-type equivalents like Slurm, -m is simpler
+        if spec.mail_user:
+            script.add("m", spec.mail_user)
+
+        # Resource Requests
+        resource_parts = []
+        # Moab's `-l` is key: nodes=X:ppn=Y where ppn is procs per node
+        if spec.num_nodes and spec.cpus_per_task:
+            resource_parts.append(f"nodes={spec.num_nodes}:ppn={spec.cpus_per_task}")
+        
+        if spec.gpus_per_task and spec.gpus_per_task > 0:
+            resource_parts.append(f"gpus={spec.gpus_per_task}")
+        
+        if spec.mem_per_task:
+            resource_parts.append(f"mem={spec.mem_per_task}")
+
+        if spec.exclusive_access:
+            resource_parts.append("naccesspolicy=singlejob")
+        
+        # Scheduling and Constraints
+        wt = seconds_to_moab_walltime(spec.wall_time)
+        if wt:
+            script.add("l", f"walltime={wt}")
+        
+        if resource_parts:
+            script.add("l", " ".join(resource_parts))
+
+        script.add("q", spec.queue)
+
+        moab_prio = priority_to_moab_priority(spec.priority)
+        if moab_prio != 0:
+            script.add("p", moab_prio)
+
+        if spec.begin_time:
+            script.add("S", epoch_to_moab_begin_time(spec.begin_time))
+        
+        script.add("d", spec.working_directory)
+        
+        # Dependencies: `-l depend=...`
+        if spec.depends_on:
+            if isinstance(spec.depends_on, list):
+                dependency_str = ":".join(spec.depends_on)
+                script.add("l", f"depend=afterok:{dependency_str}")
+            else:
+                script.add("l", f"depend={spec.depends_on}")
+
+        # I am just adding this for readability
+        script.newline()
+
+        # Environment Variables
+        if spec.environment:
+            for key, value in spec.environment.items():
+                script.add_line(f"export {key}='{value}'")
+            script.newline()
+
+        # Execution logic
+        container_exec = []
+        if spec.container_image:
+            container_exec = ["singularity", "exec", spec.container_image]
+
+        command_parts = container_exec + spec.script
+
+        if spec.input_file:
+            command_parts.append(f"< {spec.input_file}")
+        for line in command_parts:
+            script.add_line(line)
+        script.newline()
+        return script.render()
+
+    def _parse(self, filename, return_unhandled=False):
+        """
+        Parses the content of a Moab batch script into a JobSpec object.
+        """
+        spec = JobSpec()
+        command_lines = []
+        not_handled = set()
+        
+        # Regex to capture #MSUB directives (simple -f <val> format)
+        msub_re = re.compile(r"#MSUB\s+-(\w+)(?:\s+(.+))?")
+        script_content = utils.read_file(filename)
+
+        for line in script_content.splitlines():
+            if not line.strip():
+                continue
+
+            match = msub_re.match(line)
+            if match:
+                key, value = match.groups()
+                value = value.strip() if value else ""
+
+                if key == 'N': spec.job_name = value
+                elif key == 'A': spec.account = value
+                elif key == 'o': spec.output_file = value
+                elif key == 'e': spec.error_file = value
+                elif key == 'm': spec.mail_user = value
+                elif key == 'q': spec.queue = value
+                elif key == 'd': spec.working_directory = value
+                elif key == 'S': spec.begin_time = moab_begin_time_to_epoch(value)
+                elif key == 'p': spec.priority = moab_priority_to_priority(int(value))
+                elif key == 'l':
+                    # The -l line can contain walltime and other resources
+                    # e.g., -l walltime=01:00:00 -l nodes=4:ppn=8
+                    if 'walltime' in value:
+                        spec.wall_time = moab_walltime_to_seconds(value.split('=', 1)[1])
+                    elif 'nodes' in value:
+                        for part in value.split(':'):
+                            k, v = part.split('=', 1)
+                            if k == 'nodes': spec.num_nodes = int(v)
+                            elif k == 'ppn': spec.cpus_per_task = int(v)
+                    elif 'mem' in value:
+                        spec.mem_per_task = value.split('=', 1)[1]
+                    elif 'gpus' in value:
+                        spec.gpus_per_task = int(value.split('=', 1)[1])
+                    elif 'depend' in value:
+                        spec.depends_on = value.split('=', 1)[1]
+                    elif 'naccesspolicy' in value and 'singlejob' in value:
+                        spec.exclusive_access = True
+                else:
+                    not_handled.add(key)
+                continue
+
+            if line.lower().startswith("export "):
+                env_match = re.match(r"export\s+([^=]+)=(.*)", line)
+                if env_match:
+                    env_key, env_val = env_match.groups()
+                    spec.environment[env_key] = env_val.strip("'\"")
+
+            if line.startswith("#"):
+                continue
+            
+            command_lines.append(line)
+
+        if command_lines:
+            spec.script = command_lines
+
+        if return_unhandled:
+            return not_handled
+        return spec