Skip to content

feat: add additional transformers #7

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion fractale/cli/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,10 @@ def get_parser():
default="kubernetes",
)
transform.add_argument(
"-f", "--from", dest="from_transformer", help="transform from this jobspec", default="flux"
"-f",
"--from",
dest="from_transformer",
help="transform from this jobspec",
)
transform.add_argument(
"--pretty",
Expand Down
6 changes: 5 additions & 1 deletion fractale/cli/transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from rich import print
from rich.pretty import pprint

from fractale.transformer import get_transformer
from fractale.transformer import detect_transformer, get_transformer


def main(args, extra, **kwargs):
Expand All @@ -18,6 +18,10 @@ def main(args, extra, **kwargs):
if not os.path.exists(args.jobspec):
sys.exit(f"JobSpec {args.jobspec} does not exist.")

# If no from transformer defined, try to detect
if args.from_transformer is None:
args.from_transformer = detect_transformer(args.jobspec)

# No selector or solver, just manual transform
from_transformer = get_transformer(args.from_transformer)
to_transformer = get_transformer(args.to_transformer)
Expand Down
32 changes: 32 additions & 0 deletions fractale/transformer/__init__.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,47 @@
import fractale.utils as utils

from .cobalt import Transformer as CobaltTransformer
from .flux import Transformer as FluxTransformer
from .kubernetes import Transformer as KubernetesTransformer
from .lsf import Transformer as LSFTransformer
from .oar import Transformer as OARTransformer
from .pbs import Transformer as PBSTransformer
from .slurm import Transformer as SlurmTransformer

plugins = {
"kubernetes": KubernetesTransformer,
"flux": FluxTransformer,
"slurm": SlurmTransformer,
"pbs": PBSTransformer,
"lsf": LSFTransformer,
"oar": OARTransformer,
"cobalt": CobaltTransformer,
}


def get_transformer(name, selector="random", solver=None):
if name not in plugins:
raise ValueError(f"{name} is not a valid transformer.")
return plugins[name](selector, solver)


def detect_transformer(jobspec):
"""
Quick and dirty detection.
"""
content = utils.read_file(jobspec)
if "#FLUX" in content and "FLUX_CAPACITOR" not in content:
return "flux"
if "#SBATCH " in content:
return "slurm"
if "kind:" in content and "Job" in content:
return "kubernetes"
if "#PBS " in content:
return "pbs"
if "#BSUB" in content:
return "lsf"
if "#OAR" in content:
return "oar"
if "#COBALT" in content:
return "cobalt"
raise ValueError("Unkown transformer.")
53 changes: 51 additions & 2 deletions fractale/transformer/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ class TransformerBase:
This can be very manual, or use an LLM.
"""

def __init__(self, selector, solver):
def __init__(self, selector="random", solver=None):
"""
Create a new transformer backend, accepting any options type.

Expand All @@ -21,12 +21,18 @@ def __init__(self, selector, solver):
self.selector = get_selector(selector)
self.solver = solver

def parse(self, *args, **kwargs):
def _parse(self, *args, **kwargs):
"""
Parse converts the native jobspec to the standard JobSpec
"""
raise NotImplementedError

def parse(self, filename):
return self._parse(filename)

def unhandled(self, filename):
return self._parse(filename, return_unhandled=True)

def convert(self, *args, **kwargs):
"""
Convert a normalized jobspec to the format here.
Expand All @@ -42,3 +48,46 @@ def render(self, matches, jobspec):
"""
js = utils.load_jobspec(jobspec)
return self.run(matches, js)


class Script:
"""
A helper class to build a batch script line by line.
"""

def __init__(self, directive=""):
self.script_lines = ["#!/bin/bash"]
self.directive = directive

def newline(self):
self.script_lines.append("")

def add_line(self, line: str):
"""
Add a custom line to the script.
"""
self.script_lines.append(line)

def add(self, name: str, value=None):
"""
Add a Flux directive, e.g., #FLUX: --job-name=my-job or #FLUX: -N 4.
Handles both short and long options.
"""
if value is None:
return

# Determine if it's a short (-n) or long (--tasks) option
prefix = "-" if len(name) == 1 else "--"
self.script_lines.append(f"{self.directive}: {prefix}{name}={value}")

def add_flag(self, name: str):
"""
Add a boolean flag (e.g., --exclusive).
"""
self.script_lines.append(f"{self.directive}: --{name}")

def render(self) -> str:
"""
Return the complete script as a single string.
"""
return "\n".join(self.script_lines)
3 changes: 3 additions & 0 deletions fractale/transformer/cobalt/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
from .transform import CobaltTransformer as Transformer

assert Transformer
Loading