flux-batch 0.0.0__py3-none-any.whl

flux_batch/__init__.py

@@ -0,0 +1,7 @@
+from .jobspec import BatchJobspecV1
+from .models import BatchAttributesV1, BatchJobV1
+from .submit import submit
+
+__all__ = ["BatchJobV1", "BatchAttributesV1", "BatchJobspecV1", "submit"]
+
+from .version import __version__  # noqa

flux_batch/jobspec.py

@@ -0,0 +1,132 @@
+import shlex
+from typing import List
+
+import flux_batch.models as models
+
+
+class BatchJobspecV1:
+    """
+    A BatchJobspecV1 mirrors a JobspecV1. We need to:
+
+    1. Add some number of commands or a script
+    2. Add optional services (start/stop)
+    """
+
+    def __init__(self, attributes: models.BatchAttributesV1 = None):
+        self.attributes = attributes or models.BatchAttributesV1()
+        self.commands: List[str] = []
+        self.prologs: List[str] = []
+        self.epilogs: List[str] = []
+        self.services: List[str] = []
+
+    @classmethod
+    def from_command(cls, command: List[str], **kwargs):
+        inst = cls(models.BatchAttributesV1(**kwargs))
+        inst.commands = [shlex.join(command)]
+        return inst
+
+    @classmethod
+    def from_jobs(cls, batch: models.BatchJobV1, **kwargs):
+        """
+        Generate the batch script from a set of jobs.
+
+        With more than one job, we assume we are waiting.
+        """
+        inst = cls(models.BatchAttributesV1(**kwargs))
+        if len(batch.jobs) > 1:
+            for job_str in batch.jobs:
+                inst.commands.append(f"flux submit --wait {job_str}")
+            # Assume we want to wait for all jobs
+            inst.commands.append("flux job wait --all")
+        else:
+            inst.commands = batch.jobs
+        return inst
+
+    def add_service(self, service: str):
+        self.services.append(service)
+
+    def add_prolog(self, cmd: str):
+        self.prologs.append(cmd)
+
+    def add_epilog(self, cmd: str):
+        self.epilogs.append(cmd)
+
+    def get_cli_flags(self) -> List[str]:
+        """
+        Converts BatchAttributesV1 into a list of strings for subprocess.
+        """
+        flags = []
+        attr = self.attributes
+
+        # Mapping table for simple flags
+        mapping = {
+            "nslots": "-n",
+            "cores_per_slot": "-c",
+            "gpus_per_slot": "-g",
+            "nodes": "-N",
+            "bank": "-B",
+            "queue": "-q",
+            "time_limit": "-t",
+            "urgency": "--urgency",
+            "job_name": "--job-name",
+            "cwd": "--cwd",
+            "dependency": "--dependency",
+            "requires": "--requires",
+            "begin_time": "--begin-time",
+            "signal": "--signal",
+            "broker_opts": "--broker-opts",
+            "dump": "--dump",
+            "flags": "--flags",
+        }
+
+        for field_name, flag in mapping.items():
+            val = getattr(attr, field_name)
+            if val is not None:
+                flags.extend([flag, str(val)])
+
+        # Boolean flags
+        if attr.exclusive:
+            flags.append("-x")
+        if attr.unbuffered:
+            flags.append("-u")
+        if attr.wrap:
+            flags.append("--wrap")
+
+        # Multi-use flags
+        multi_mapping = {
+            "setopt": "-o",
+            "setattr": "-S",
+            "add_file": "--add-file",
+            "env": "--env",
+            "env_remove": "--env-remove",
+            "env_file": "--env-file",
+            "rlimit": "--rlimit",
+            "conf": "--conf",
+        }
+        for field_name, flag in multi_mapping.items():
+            for val in getattr(attr, field_name):
+                flags.extend([flag, str(val)])
+
+        return flags
+
+    def generate_wrapper_script(self) -> str:
+        """
+        Generate the wrapper script.
+
+        1. Start with hashbang!
+        2. Add prologs
+        3. Add services start
+        4. Add jobs/commands
+        5. Stop services
+        6. And epilogs
+        """
+
+        lines = ["#!/bin/bash"]
+        lines.extend(self.prologs)
+        for s in self.services:
+            lines.append(f"systemctl --user start {s}")
+        lines.extend(self.commands)
+        for s in reversed(self.services):
+            lines.append(f"systemctl --user stop {s}")
+        lines.extend(self.epilogs)
+        return "\n".join(lines)

flux_batch/logger/__init__.py

@@ -0,0 +1 @@
+from .logger import LogColors, logger, setup_logger

flux_batch/logger/generate.py

@@ -0,0 +1,208 @@
+from random import choice
+
+
+class JobNamer:
+    _descriptors = [
+        "chunky",
+        "buttery",
+        "delicious",
+        "scruptious",
+        "dinosaur",
+        "boopy",
+        "lovely",
+        "carnivorous",
+        "hanky",
+        "loopy",
+        "doopy",
+        "astute",
+        "gloopy",
+        "outstanding",
+        "stinky",
+        "conspicuous",
+        "fugly",
+        "frigid",
+        "angry",
+        "adorable",
+        "sticky",
+        "moolicious",
+        "cowy",
+        "spicy",
+        "grated",
+        "crusty",
+        "stanky",
+        "blank",
+        "bumfuzzled",
+        "fuzzy",
+        "hairy",
+        "peachy",
+        "tart",
+        "creamy",
+        "arid",
+        "strawberry",
+        "butterscotch",
+        "wobbly",
+        "persnickety",
+        "nerdy",
+        "dirty",
+        "placid",
+        "bloated",
+        "swampy",
+        "pusheena",
+        "hello",
+        "goodbye",
+        "milky",
+        "purple",
+        "rainbow",
+        "bricky",
+        "muffled",
+        "anxious",
+        "misunderstood",
+        "eccentric",
+        "quirky",
+        "lovable",
+        "reclusive",
+        "faux",
+        "evasive",
+        "confused",
+        "crunchy",
+        "expensive",
+        "ornery",
+        "fat",
+        "phat",
+        "joyous",
+        "expressive",
+        "psycho",
+        "chocolate",
+        "salted",
+        "gassy",
+        "red",
+        "blue",
+    ]
+
+    _nouns = [
+        "squidward",
+        "hippo",
+        "butter",
+        "animal",
+        "peas",
+        "lettuce",
+        "carrot",
+        "onion",
+        "peanut",
+        "cupcake",
+        "muffin",
+        "buttface",
+        "leopard",
+        "parrot",
+        "parsnip",
+        "poodle",
+        "itch",
+        "punk",
+        "kerfuffle",
+        "soup",
+        "noodle",
+        "avocado",
+        "peanut-butter",
+        "latke",
+        "milkshake",
+        "banana",
+        "lizard",
+        "lemur",
+        "lentil",
+        "bits",
+        "house",
+        "leader",
+        "toaster",
+        "signal",
+        "pancake",
+        "kitty",
+        "cat",
+        "cattywampus",
+        "poo",
+        "malarkey",
+        "general",
+        "rabbit",
+        "chair",
+        "staircase",
+        "underoos",
+        "snack",
+        "lamp",
+        "eagle",
+        "hobbit",
+        "diablo",
+        "earthworm",
+        "pot",
+        "plant",
+        "leg",
+        "arm",
+        "bike",
+        "citrus",
+        "dog",
+        "puppy",
+        "blackbean",
+        "ricecake",
+        "gato",
+        "nalgas",
+        "lemon",
+        "caramel",
+        "fudge",
+        "cherry",
+        "sundae",
+        "truffle",
+        "cinnamonbun",
+        "pastry",
+        "egg",
+        "omelette",
+        "fork",
+        "knife",
+        "spoon",
+        "salad",
+        "train",
+        "car",
+        "motorcycle",
+        "bicycle",
+        "platanos",
+        "mango",
+        "taco",
+        "pedo",
+        "nunchucks",
+        "destiny",
+        "hope",
+        "despacito",
+        "frito",
+        "chip",
+    ]
+
+    def generate(self, delim="-", length=4, chars="0123456789"):
+        """
+        Generate a robot name. Inspiration from Haikunator, but much more
+                 poorly implemented ;)
+
+        Parameters
+        ==========
+        delim: Delimiter
+        length: TokenLength
+        chars: TokenChars
+        """
+
+        descriptor = self._select(self._descriptors)
+        noun = self._select(self._nouns)
+        numbers = "".join((self._select(chars) for _ in range(length)))
+        return delim.join([descriptor, noun, numbers])
+
+    def _select(self, select_from):
+        """select an element from a list using random.choice
+
+        Parameters
+        ==========
+        should be a list of things to select from
+        """
+        if not select_from:
+            return ""
+
+        return choice(select_from)
+
+
+def generate_name():
+    namer = JobNamer()
+    return namer.generate()

flux_batch/logger/logger.py

@@ -0,0 +1,194 @@
+import inspect
+import logging as _logging
+import os
+import platform
+import sys
+import threading
+
+from fastmcp.utilities.logging import get_logger
+
+
+class LogColors:
+    PURPLE = "\033[95m"
+    OKBLUE = "\033[94m"
+    OKCYAN = "\033[96m"
+    OKGREEN = "\033[92m"
+    WARNING = "\033[93m"
+    RED = "\033[91m"
+    ENDC = "\033[0m"
+    BOLD = "\033[1m"
+    UNDERLINE = "\033[4m"
+
+
+class ColorizingStreamHandler(_logging.StreamHandler):
+    BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
+    RESET_SEQ = "\033[0m"
+    COLOR_SEQ = "\033[%dm"
+    BOLD_SEQ = "\033[1m"
+
+    colors = {
+        "WARNING": YELLOW,
+        "INFO": GREEN,
+        "DEBUG": BLUE,
+        "CRITICAL": RED,
+        "ERROR": RED,
+    }
+
+    def __init__(self, nocolor=False, stream=sys.stderr, use_threads=False):
+        super().__init__(stream=stream)
+        self._output_lock = threading.Lock()
+        self.nocolor = nocolor or not self.can_color_tty()
+
+    def can_color_tty(self):
+        if "TERM" in os.environ and os.environ["TERM"] == "dumb":
+            return False
+        return self.is_tty and not platform.system() == "Windows"
+
+    @property
+    def is_tty(self):
+        isatty = getattr(self.stream, "isatty", None)
+        return isatty and isatty()
+
+    def emit(self, record):
+        with self._output_lock:
+            try:
+                self.format(record)  # add the message to the record
+                self.stream.write(self.decorate(record))
+                self.stream.write(getattr(self, "terminator", "\n"))
+                self.flush()
+            except BrokenPipeError as e:
+                raise e
+            except (KeyboardInterrupt, SystemExit):
+                # ignore any exceptions in these cases as any relevant messages have been printed before
+                pass
+            except Exception:
+                self.handleError(record)
+
+    def decorate(self, record):
+        message = record.message
+        message = [message]
+        if not self.nocolor and record.levelname in self.colors:
+            message.insert(0, self.COLOR_SEQ % (30 + self.colors[record.levelname]))
+            message.append(self.RESET_SEQ)
+        return "".join(message)
+
+
+class Logger:
+    def __init__(self):
+        self.logger = get_logger(__name__)
+        self.log_handler = [self.text_handler]
+        self.stream_handler = None
+        self.printshellcmds = False
+        self.quiet = False
+        self.logfile = None
+        self.last_msg_was_job_info = False
+        self.logfile_handler = None
+
+    def cleanup(self):
+        if self.logfile_handler is not None:
+            self.logger.removeHandler(self.logfile_handler)
+            self.logfile_handler.close()
+        self.log_handler = [self.text_handler]
+
+    def handler(self, msg):
+        for handler in self.log_handler:
+            handler(msg)
+
+    def set_stream_handler(self, stream_handler):
+        if self.stream_handler is not None:
+            self.logger.removeHandler(self.stream_handler)
+        self.stream_handler = stream_handler
+        self.logger.addHandler(stream_handler)
+
+    def set_level(self, level):
+        self.logger.setLevel(level)
+
+    def location(self, msg):
+        callerframerecord = inspect.stack()[1]
+        frame = callerframerecord[0]
+        info = inspect.getframeinfo(frame)
+        self.debug("{}: {info.filename}, {info.function}, {info.lineno}".format(msg, info=info))
+
+    def info(self, msg):
+        self.handler(dict(level="info", msg=msg))
+
+    def warning(self, msg):
+        self.handler(dict(level="warning", msg=msg))
+
+    def debug(self, msg):
+        self.handler(dict(level="debug", msg=msg))
+
+    def error(self, msg):
+        self.handler(dict(level="error", msg=msg))
+
+    def exit(self, msg, return_code=1):
+        self.handler(dict(level="error", msg=msg))
+        sys.exit(return_code)
+
+    def progress(self, done=None, total=None):
+        self.handler(dict(level="progress", done=done, total=total))
+
+    def shellcmd(self, msg):
+        if msg is not None:
+            msg = dict(level="shellcmd", msg=msg)
+            self.handler(msg)
+
+    def success(self, message):
+        """
+        Wrapper to add success to output for LLM.
+        """
+        return "✅ SUCCESS: " + message
+
+    def failure(self, message):
+        """
+        Wrapper to add success to output for LLM.
+        """
+        return "❌ FAILED: " + message
+
+    def text_handler(self, msg):
+        """The default snakemake log handler.
+        Prints the output to the console.
+        Args:
+            msg (dict):     the log message dictionary
+        """
+        level = msg["level"]
+        if level == "info" and not self.quiet:
+            self.logger.info(msg["msg"])
+        if level == "warning":
+            self.logger.warning(msg["msg"])
+        elif level == "error":
+            self.logger.error(msg["msg"])
+        elif level == "debug":
+            self.logger.debug(msg["msg"])
+        elif level == "progress" and not self.quiet:
+            done = msg["done"]
+            total = msg["total"]
+            p = done / total
+            percent_fmt = ("{:.2%}" if p < 0.01 else "{:.0%}").format(p)
+            self.logger.info("{} of {} steps ({}) done".format(done, total, percent_fmt))
+        elif level == "shellcmd":
+            if self.printshellcmds:
+                self.logger.warning(msg["msg"])
+
+
+logger = Logger()
+
+
+def setup_logger(
+    quiet=False,
+    printshellcmds=False,
+    nocolor=False,
+    stdout=False,
+    debug=False,
+    use_threads=False,
+):
+    # console output only if no custom logger was specified
+    stream_handler = ColorizingStreamHandler(
+        nocolor=nocolor,
+        stream=sys.stdout if stdout else sys.stderr,
+        use_threads=use_threads,
+    )
+    logger.set_stream_handler(stream_handler)
+    logger.set_level(_logging.DEBUG if debug else _logging.INFO)
+    logger.quiet = quiet
+    logger.printshellcmds = printshellcmds

flux_batch/models.py

@@ -0,0 +1,63 @@
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
+
+
+@dataclass
+class BatchJobV1:
+    """
+    Commands to be run within the batch wrapper.
+
+    This should mirror JobspecV1
+    """
+
+    jobs: List[str] = field(default_factory=list)
+
+    def add_job(self, command: List[str]):
+        import shlex
+
+        self.jobs.append(shlex.join(command))
+
+
+@dataclass
+class BatchAttributesV1:
+    """
+    Explicitly defined arguments allowed by flux batch for V1 spec
+    """
+
+    # Resources
+    nslots: Optional[int] = None  # -n
+    cores_per_slot: Optional[int] = None  # -c
+    gpus_per_slot: Optional[int] = None  # -g
+    nodes: Optional[int] = None  # -N
+    exclusive: bool = False  # -x
+
+    # Basic Options
+    bank: Optional[str] = None  # -B
+    queue: Optional[str] = None  # -q
+    time_limit: Optional[str] = None  # -t
+    urgency: Optional[int] = None  # --urgency
+    job_name: Optional[str] = None  # --job-name
+    cwd: Optional[str] = None  # --cwd
+
+    # More complex options
+    setopt: List[str] = field(default_factory=list)  # -o
+    setattr: List[str] = field(default_factory=list)  # -S
+    add_file: List[str] = field(default_factory=list)  # --add-file
+    env: List[str] = field(default_factory=list)  # --env
+    env_remove: List[str] = field(default_factory=list)  # --env-remove
+    env_file: List[str] = field(default_factory=list)  # --env-file
+    rlimit: List[str] = field(default_factory=list)  # --rlimit
+    conf: List[str] = field(default_factory=list)  # --conf
+
+    # Other Attributes
+    dependency: Optional[str] = None  # --dependency
+    requires: Optional[str] = None  # --requires
+    begin_time: Optional[str] = None  # --begin-time
+    signal: Optional[str] = None  # --signal
+    broker_opts: Optional[str] = None  # --broker-opts
+    dump: Optional[str] = None  # --dump
+
+    # Flags
+    unbuffered: bool = False  # -u
+    wrap: bool = False  # --wrap
+    flags: Optional[str] = None  # --flags (debug, waitable, etc)

flux_batch/service/__init__.py

@@ -0,0 +1,31 @@
+import os
+import subprocess
+import sys
+
+from .scribe import SERVICE_TEMPLATE as scribe_template
+
+# Lookup of known services
+services = {"scribe": scribe_template}
+
+
+def ensure_user_service(service_name: str):
+    """
+    Checks for the existence of a systemd service file in the user's home.
+    If it doesn't exist, it creates it and reloads the daemon.
+    """
+    user_systemd_dir = os.path.expanduser("~/.config/systemd/user")
+    os.makedirs(user_systemd_dir, exist_ok=True)
+    service_path = os.path.join(user_systemd_dir, f"{service_name}.service")
+
+    if not os.path.exists(service_path):
+        if service_name in services:
+            template = services[service_name]
+            print(f"[*] Provisioning {service_name} at {service_path}")
+            with open(service_path, "w") as f:
+                f.write(template.format(python_path=sys.executable))
+
+        else:
+            print(f"[*] Service {service_name} is not known, assuming exists.")
+
+        # Reload the user-session manager to recognize the new unit
+        subprocess.run(["systemctl", "--user", "daemon-reload"], check=True)

flux_batch/service/scribe.py

@@ -0,0 +1,12 @@
+# Template for the Scribe Journal Consumer
+SERVICE_TEMPLATE = """[Unit]
+Description=Flux Scribe Journal Consumer
+After=network.target
+
+[Service]
+ExecStart={python_path} -m flux_mcp_server.scribe
+Restart=on-failure
+
+[Install]
+WantedBy=default.target
+"""

flux_batch/submit.py

@@ -0,0 +1,53 @@
+import os
+import stat
+import subprocess
+import tempfile
+
+import flux
+import flux.job
+
+import flux_batch.models as models
+import flux_batch.utils as utils
+from flux_batch.service import ensure_user_service
+
+
+def submit(handle: flux.Flux, spec: models.BatchJobV1, dry_run=False) -> int:
+    """
+    Orchestrates the submission process:
+    1. Provisions any required user-space services.
+    2. Generates the wrapper shell script.
+    3. Uses 'flux batch --dryrun' to compile the Jobspec JSON.
+    4. Submits the Jobspec to the Flux instance.
+    """
+
+    # Provision services (like flux-scribe) if requested
+    for service in spec.services:
+        ensure_user_service(service)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Write the wrapper script (handling prologs, services, and jobs)
+        wrapper_path = os.path.join(tmpdir, "wrapper.sh")
+
+        # dry run here just displays it
+        script = spec.generate_wrapper_script()
+        if dry_run:
+            return script
+
+        utils.write_file(script, wrapper_path)
+
+        # Make the script executable so 'flux batch' can analyze it
+        os.chmod(wrapper_path, os.stat(wrapper_path).st_mode | stat.S_IEXEC)
+
+        # Generate the RFC 25 Jobspec JSON via the Flux CLI
+        # This handles all resource mapping (-N, -n, etc.)
+        cmd = ["flux", "batch"] + spec.get_cli_flags() + ["--dry-run", wrapper_path]
+
+        try:
+            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        except subprocess.CalledProcessError as e:
+            print(f"Error during flux batch dryrun: {e.stderr}")
+            raise
+
+        # Submit the JSON string to the Flux instance
+        # The result.stdout contains the raw JSON Jobspec
+        return flux.job.submit(handle, result.stdout)

flux_batch/utils/__init__.py

@@ -0,0 +1,2 @@
+from .fileio import *
+from .text import *

flux_batch/utils/fileio.py

@@ -0,0 +1,203 @@
+import json
+import os
+import platform
+import re
+import stat
+import subprocess
+import tempfile
+from contextlib import contextmanager
+
+import yaml
+
+
+def run_sync(coroutine):
+    """
+    Runs an async coroutine synchronously.
+    Patches the loop if running inside IPython/Jupyter.
+
+    Note that I'm not currently using this - keeping here if need.
+    """
+    import asyncio
+
+    import nest_asyncio
+
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+
+    if loop and loop.is_running():
+        # We aren't in async -> patch it
+        nest_asyncio.apply(loop)
+        return loop.run_until_complete(coroutine)
+    else:
+        # We are in standard Python script -> Standard run
+        return asyncio.run(coroutine)
+
+
+def get_local_cluster():
+    """
+    Guess the local cluster based on the hostname
+    """
+    return platform.node().split("-")[0]
+
+
+def read_json(filename):
+    """
+    Read json from file
+    """
+    return json.loads(read_file(filename))
+
+
+def write_json(obj, filename):
+    with open(filename, "w") as fd:
+        fd.write(json.dumps(obj, indent=4))
+
+
+def load_jobspec(filename):
+    """
+    Load a jobspec. First try yaml and fall back to json
+    """
+    # It is already loaded!
+    if isinstance(filename, dict):
+        return filename
+    if isinstance(filename, str) and not os.path.exists(filename):
+        return yaml.safe_load(filename)
+    try:
+        return read_yaml(filename)
+    except:
+        return read_json(filename)
+
+
+def read_file(filename):
+    """
+    Read in a file content
+    """
+    with open(filename, "r") as fd:
+        content = fd.read()
+    return content
+
+
+def make_executable(path):
+    """
+    Adds execute permission to a file.
+    """
+    current_mode = os.stat(path).st_mode
+
+    # Add execute permission for owner, group, and others
+    new_mode = current_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
+
+    # Set the new permissions
+    os.chmod(path, new_mode)
+
+
+def recursive_find(base, pattern="[.]py"):
+    """recursive find will yield python files in all directory levels
+    below a base path.
+
+    Arguments:
+      - base (str) : the base directory to search
+      - pattern: a pattern to match, defaults to *.py
+    """
+    for root, _, filenames in os.walk(base):
+        for filename in filenames:
+            filepath = os.path.join(root, filename)
+            if not re.search(pattern, filepath):
+                continue
+            yield filepath
+
+
+def get_tmpfile(tmpdir=None, prefix="", suffix=None):
+    """
+    Get a temporary file with an optional prefix.
+    """
+    # First priority for the base goes to the user requested.
+    tmpdir = get_tmpdir(tmpdir)
+
+    # If tmpdir is set, add to prefix
+    if tmpdir:
+        prefix = os.path.join(tmpdir, os.path.basename(prefix))
+
+    fd, tmp_file = tempfile.mkstemp(prefix=prefix, suffix=suffix)
+    os.close(fd)
+
+    return tmp_file
+
+
+def get_tmpdir(tmpdir=None, prefix="", create=True):
+    """
+    Get a temporary directory for an operation.
+    """
+    tmpdir = tmpdir or tempfile.gettempdir()
+    prefix = prefix or "jobspec"
+    prefix = "%s.%s" % (prefix, next(tempfile._get_candidate_names()))
+    tmpdir = os.path.join(tmpdir, prefix)
+
+    if not os.path.exists(tmpdir) and create is True:
+        os.mkdir(tmpdir)
+
+    return tmpdir
+
+
+def read_yaml(filename):
+    """
+    Read yaml from file
+    """
+    with open(filename, "r") as fd:
+        content = yaml.safe_load(fd)
+    return content
+
+
+def write_file(content, filename):
+    """
+    Write content to file
+    """
+    with open(filename, "w") as fd:
+        fd.write(content)
+
+
+def write_yaml(obj, filename):
+    """
+    Read yaml to file
+    """
+    with open(filename, "w") as fd:
+        yaml.dump(obj, fd)
+
+
+@contextmanager
+def workdir(dirname):
+    """
+    Provide context for a working directory, e.g.,
+
+    with workdir(name):
+       # do stuff
+    """
+    here = os.getcwd()
+    os.chdir(dirname)
+    try:
+        yield
+    finally:
+        os.chdir(here)
+
+
+def run_command(cmd, stream=False, check_output=False, return_code=0):
+    """
+    use subprocess to send a command to the terminal.
+
+    If check_output is True, check against an expected return code.
+    """
+    stdout = subprocess.PIPE if not stream else None
+    output = subprocess.Popen(cmd, stderr=subprocess.STDOUT, stdout=stdout, env=os.environ.copy())
+    t = output.communicate()[0], output.returncode
+    output = {"message": t[0], "return_code": t[1]}
+
+    if isinstance(output["message"], bytes):
+        output["message"] = output["message"].decode("utf-8")
+
+    # Check the output and raise an error if not success
+    if check_output and t[1] != return_code:
+        if output["message"]:
+            raise ValueError(output["message"].strip())
+        else:
+            raise ValueError(f"Failed execution, return code {t[1]}")
+    return output

flux_batch/utils/text.py

@@ -0,0 +1,22 @@
+import re
+
+
+def format_rules(rules):
+    return "\n".join([f"- {r}" for r in rules])
+
+
+def get_code_block(content, code_type):
+    """
+    Parse a code block from the response
+    """
+    pattern = f"```(?:{code_type})?\n(.*?)```"
+    match = re.search(pattern, content, re.DOTALL)
+    if match:
+        return match.group(1).strip()
+    if content.startswith(f"```{code_type}"):
+        content = content[len(f"```{code_type}") :]
+    if content.startswith("```"):
+        content = content[len("```") :]
+    if content.endswith("```"):
+        content = content[: -len("```")]
+    return content.strip()

flux_batch/utils/timer.py

@@ -0,0 +1,24 @@
+import time
+
+
+class Timer:
+    """
+    A context timer! You can use in context OR explicitly start/stop.
+    t = Timer()
+    with t:
+       do_stuff()
+    """
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop()
+
+    def stop(self):
+        self.end_time = time.perf_counter()
+        self.elapsed_time = self.end_time - self.start_time
+
+    def start(self):
+        self.start_time = time.perf_counter()

flux_batch/version.py

@@ -0,0 +1,16 @@
+__version__ = "0.0.0"
+AUTHOR = "Vanessa Sochat"
+AUTHOR_EMAIL = "vsoch@users.noreply.github.com"
+NAME = "flux-batch"
+PACKAGE_URL = "https://github.com/converged-computing/flux-batch"
+KEYWORDS = "flux, flux framework, hpc, batch, workloads"
+DESCRIPTION = "Python SDK for flux batch jobs and services"
+LICENSE = "LICENSE"
+
+INSTALL_REQUIRES = (
+    ("pyyaml", {"min_version": None}),
+    ("ply", {"min_version": None}),
+)
+
+TESTS_REQUIRES = (("pytest", {"min_version": "4.6.2"}),)
+INSTALL_REQUIRES_ALL = INSTALL_REQUIRES + TESTS_REQUIRES

flux_batch-0.0.0.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022-2023 LLNS, LLC and other HPCIC DevTools Developers.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

flux_batch-0.0.0.dist-info/METADATA

@@ -0,0 +1,153 @@
+Metadata-Version: 2.1
+Name: flux-batch
+Version: 0.0.0
+Summary: Python SDK for flux batch jobs and services
+Home-page: https://github.com/converged-computing/flux-batch
+Author: Vanessa Sochat
+Author-email: vsoch@users.noreply.github.com
+Maintainer: Vanessa Sochat
+License: LICENSE
+Keywords: flux,flux framework,hpc,batch,workloads
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: C
+Classifier: Programming Language :: Python
+Classifier: Topic :: Software Development
+Classifier: Topic :: Scientific/Engineering
+Classifier: Operating System :: Unix
+Classifier: Programming Language :: Python :: 3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+License-File: NOTICE
+Requires-Dist: pyyaml
+Requires-Dist: ply
+Provides-Extra: all
+Requires-Dist: pyyaml ; extra == 'all'
+Requires-Dist: ply ; extra == 'all'
+Requires-Dist: pytest >=4.6.2 ; extra == 'all'
+
+# flux-batch
+
+> Python SDK to generate Flux batch jobs and services
+
+[![PyPI version](https://badge.fury.io/py/flux-batch.svg)](https://badge.fury.io/py/flux-batch)
+
+![https://github.com/converged-computing/flux-batch/raw/main/img/flux-batch-small.png](https://github.com/converged-computing/flux-batch/raw/main/img/flux-batch-small.png)
+
+## Related Projects
+
+- [flux-mcp](https://github.com/converged-computing/flux-mcp): MCP functions for Flux.
+- [flux-mcp-server](https://github.com/converged-computing/flux-mcp-server): MCP server.
+- [fractale-mcp](https://github.com/compspec/fractale-mcp): (fractale) MCP orchestration (agents, databases, ui interfaces).
+- [hpc-mcp](https://github.com/converged-computing/hpc-mcp): HPC tools for a larger set of HPC and converged computing use cases.
+
+## Services
+
+ - **flux-scribe**: Write job events to a local sqlite database via the JournalConsumer (not added yet, written and needs testing)
+
+## Usage
+
+This is a small Flux utility that makes it easy to create Flux batch jobs and services.
+The use case is to submit work (one or more jobs) under an instance, and run a custom service, or epilog and prolog commands. We will provision several services here, and you can also provide your own name to start / stop.
+
+### Setup
+
+Install the library and start (or be in) a flux instance.
+
+```bash
+flux start
+pip install -e . --break-system-packages
+```
+
+### Example
+
+Run the controlled example to see a batch job with prolog and epilog run and complete:
+
+```bash
+python3 ./tests/test_flux_batch.py
+```
+```console
+Flux Batch Module Test
+[OK] Connected to Flux.
+[*] Creating batch jobs...
+[*] Mapping attributes to BatchJobspecV1...
+[*] Previewing submission (Dryrun -> Wrapper)...
+#!/bin/bash
+echo 'Batch Wrapper Starting'
+flux submit --wait /bin/echo 'Job 1 starting'
+flux submit --wait /bin/sleep 5
+flux submit --wait /bin/echo 'Job 2 finished'
+flux job wait --all
+echo 'Batch Wrapper Finished'
+[*] Performing submission (Dryrun -> Wrapper -> Submit)...
+[SUCCESS] Batch submitted! Flux Job ID: ƒMX29AwFu
+```
+```bash
+$ flux jobs -a
+       JOBID USER     NAME       ST NTASKS NNODES     TIME INFO
+   ƒMX29AwFu vscode   test-batch  R      1      1   4.213s 68e8c4399c15
+```
+```bash
+$ flux jobs -a
+       JOBID USER     NAME       ST NTASKS NNODES     TIME INFO
+   ƒMX29AwFu vscode   test-batch CD      1      1   6.354s 68e8c4399c15
+```
+
+Here is an explicit (manual) example to do the same:
+
+```python
+import flux
+import flux_batch
+
+# for pretty printing
+# from rich import print
+
+handle = flux.Flux()
+
+# Create your batch job with some number of commands
+batch = flux_batch.BatchJobV1()
+batch.add_job(["echo", "Job 1 starting"])
+batch.add_job(["sleep", "5"])
+batch.add_job(["echo", "Job 2 finished"])
+
+# Wrap it up into a jobspec
+jobspec = flux_batch.BatchJobspecV1.from_jobs(
+    batch,
+    nodes=1,
+    nslots=1,
+    time_limit="10m",
+    job_name="test-batch"
+)
+
+# Add a prolog and epilog
+jobspec.add_prolog("echo 'Batch Wrapper Starting'")
+jobspec.add_epilog("echo 'Batch Wrapper Finished'")
+
+# Add a service (this assumes user level that exists)
+# jobspec.add_service("my-service'")
+
+# Preview it
+print(flux_batch.submit(handle, jobspec, dry_run=True))
+
+# Submit that bad boi.
+jobid = flux_batch.submit(handle, jobspec)
+```
+
+## TODO
+
+- Option for controlled output (that we can easily get after)
+- Create flux-scribe service, add example (need to test on cluster with systemctl)
+
+## License
+
+HPCIC DevTools is distributed under the terms of the MIT license.
+All new contributions must be made under this license.
+
+See [LICENSE](https://github.com/converged-computing/cloud-select/blob/main/LICENSE),
+[COPYRIGHT](https://github.com/converged-computing/cloud-select/blob/main/COPYRIGHT), and
+[NOTICE](https://github.com/converged-computing/cloud-select/blob/main/NOTICE) for details.
+
+SPDX-License-Identifier: (MIT)
+
+LLNL-CODE- 842614

flux_batch-0.0.0.dist-info/NOTICE

@@ -0,0 +1,21 @@
+This work was produced under the auspices of the U.S. Department of
+Energy by Lawrence Livermore National Laboratory under Contract
+DE-AC52-07NA27344.
+
+This work was prepared as an account of work sponsored by an agency of
+the United States Government. Neither the United States Government nor
+Lawrence Livermore National Security, LLC, nor any of their employees
+makes any warranty, expressed or implied, or assumes any legal liability
+or responsibility for the accuracy, completeness, or usefulness of any
+information, apparatus, product, or process disclosed, or represents that
+its use would not infringe privately owned rights.
+
+Reference herein to any specific commercial product, process, or service
+by trade name, trademark, manufacturer, or otherwise does not necessarily
+constitute or imply its endorsement, recommendation, or favoring by the
+United States Government or Lawrence Livermore National Security, LLC.
+
+The views and opinions of authors expressed herein do not necessarily
+state or reflect those of the United States Government or Lawrence
+Livermore National Security, LLC, and shall not be used for advertising
+or product endorsement purposes.

flux_batch-0.0.0.dist-info/RECORD

@@ -0,0 +1,21 @@
+flux_batch/__init__.py,sha256=ZAZF-E0SbVVY2U1_WkRGZFB4rge5BGqQZJ2MdjloZhY,230
+flux_batch/jobspec.py,sha256=f0C3ba8kR9lS-ej-AH-q3fXd8F_WarB7qCSyIDKhhgk,3891
+flux_batch/models.py,sha256=WACPnAvjEteLOq7Jx2zb8M2gKz5YVhz6V5pKjzRpPXI,1999
+flux_batch/submit.py,sha256=TSDg1Dwa5HKhg7Rj9Due8hTDz5__ihyoUdGGKhtVqWo,1795
+flux_batch/version.py,sha256=DX2m2hMpubLBjIWS8CfP-WZViMaSR2YD6eT8I_Akw8c,533
+flux_batch/logger/__init__.py,sha256=eDdpw_uppR5mPLHE39qT_haqMxu-2wniLlJZDigRC2k,52
+flux_batch/logger/generate.py,sha256=L9JyMY2oapp0ss7f7LGuihbLomzVJsMq7sByy9NhbZI,4017
+flux_batch/logger/logger.py,sha256=HKymVBNcoPdX87QWy69er5wUzHVeriiKp9p0bIYboUo,5927
+flux_batch/service/__init__.py,sha256=8IDdhIZY2B20RuJUcWsuTpwB6fm5BYxtTGot6pmn4Ag,1111
+flux_batch/service/scribe.py,sha256=dY6geiLvXYIRcIzuP_naZscKgzX4Y5dPzxoWf9Wywg0,253
+flux_batch/utils/__init__.py,sha256=CqMhw_mBfR0HBcHwv7LtFITq0J7LBV413VQE9xrz8ks,42
+flux_batch/utils/fileio.py,sha256=Elz8WkNkJ9B6x7WmCwiIBW0GgsRSSFCcbuJh7aqu2z4,4879
+flux_batch/utils/text.py,sha256=Ci1BqHs2IbOSn2o60zhLkT4kIA7CSNuGj8mdiGaDIGk,606
+flux_batch/utils/timer.py,sha256=_Weec7Wd5hWQ1r4ZHjownG4YdoIowpVqilXhvYFmIgA,491
+flux_batch-0.0.0.dist-info/LICENSE,sha256=AlyLB1m_z0CENCx1ob0PedLTTohtH2VLZhs2kfygrfc,1108
+flux_batch-0.0.0.dist-info/METADATA,sha256=JGWbOpKTZIpesJokTG29z9p-x1L-tTajMlQJ7TcjQUQ,4877
+flux_batch-0.0.0.dist-info/NOTICE,sha256=9CR93geVKl_4ZrJORbXN0fzkEM2y4DglWhY1hn9ZwQw,1167
+flux_batch-0.0.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+flux_batch-0.0.0.dist-info/entry_points.txt,sha256=ynoKpD82xn2V2sD-aZIQoq7NnfOu9VEKqW55Y1AoPGI,67
+flux_batch-0.0.0.dist-info/top_level.txt,sha256=jj8zAsZzMmbjiBISJL7lRtA37MSEAQYfObGLUncn9Lw,11
+flux_batch-0.0.0.dist-info/RECORD,,

flux_batch-0.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.43.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flux_batch-0.0.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+flux-mcp = flux_mcp.server.fastmcp.__main__:main

flux_batch-0.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+flux_batch