oh-my-batch 0.0.1.dev0__tar.gz → 0.1.0.dev1__tar.gz

oh_my_batch-0.1.0.dev1/PKG-INFO

@@ -0,0 +1,130 @@
+Metadata-Version: 2.1
+Name: oh-my-batch
+Version: 0.1.0.dev1
+Summary: 
+License: GPL
+Author: weihong.xu
+Author-email: xuweihong.cn@gmail.com
+Requires-Python: >=3.8,<4.0
+Classifier: License :: Other/Proprietary License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: fire (>=0.7.0,<0.8.0)
+Description-Content-Type: text/markdown
+
+# oh-my-batch
+A simple tool to manipulate batch tasks.
+
+The goal of this to is to provide a handy command tool for
+* `omb combo`: generate folders/files from different combinations of parameters
+* `omb batch`: generate batch scripts from multiple working directories
+* `omb job`: track the state of job in job schedular
+
+## Install
+```bash
+pip install oh-my-batch
+```
+
+## Use cases
+
+### Generate files from different combinations of parameters
+
+It's common to generate files with different combinations of parameters in scientific computing.
+For example, you have 3 LAMMPS data files in `tmp` directory: `tmp/1.data`, `tmp/2.data`, `tmp/3.data`.
+And you want to generate a series of input files with different parameters,
+for example, different temperatures 300K, 400K, 500K, against each data file.
+
+In this case, you can use `omb combo` command to generate a series of input files for you.
+
+```bash
+#! /bin/bash
+# prepare fake data files
+mkdir -p tmp/
+touch tmp/1.data tmp/2.data tmp/3.data
+
+# prepare a lammps input file template
+cat > tmp/in.lmp.tmp <<EOF
+read_data $DATA_FILE
+velocity all create $TEMP $RANDOM
+run 1000
+EOF
+
+# prepare a run script template
+cat > tmp/run.sh.tmp <<EOF
+cat in.lmp  # simulate running lammps
+EOF
+
+# generate input files
+omb combo \
+    add_files DATA_FILE tmp/*.data - \
+    add_var TEMP 300 400 500 - \
+    add_randint RANDOM -n 3 -a 1 -b 1000 --broadcast - \
+    make_files tmp/in.lmp.tmp tmp/tasks/{i}-T-{TEMP}/in.lmp - \
+    make_files tmp/run.sh.tmp tmp/tasks/{i}-T-{TEMP}/run.sh --mode 755 - \
+    done
+```
+
+The above script will generate 9 folders in `tmp/tasks` directory
+with names from `0-T-300`, `1-T-400`, `2-T-500`, `3-T-300` to `8-T-500`.
+Each folder will contain a `in.lmp` file and a `run.sh` file.
+
+The 9 folders are the combinations of 3 data files and 3 temperatures,
+and each input file will have a independent random number between 1 and 1000 as `RANDOM`.
+
+You can run the about script by `./examples/omb-combo.sh`,
+and you can also run `omb combo --help` to see the detailed usage of `combo` command.
+
+### Generate batch scripts from multiple working directories
+It's common to submit a lot of jobs to a job scheduler. `omb batch` is designed to help you generate batch scripts from multiple working directories and package them into several batch scripts.
+
+Let's continue the above example, now you have 9 folders in `tmp/tasks` directory.
+You want to package them into 2 batch scripts to submit to a job scheduler.
+
+You can use `omb batch` to generate batch scripts for you like this:
+
+```bash
+#! /bin/bash
+cat > tmp/lammps_header.sh <<EOF
+#!/bin/bash
+#SBATCH -J lmp
+#SBATCH -n 1
+#SBATCH -t 1:00:00
+EOF
+
+omb batch \
+    add_work_dir tmp/tasks/* - \
+    add_header_file tmp/lammps_header.sh - \
+    add_command "checkpoint lmp.done ./run.sh" - \
+    make tmp/lmp-{i}.slurm --concurrency 2
+```
+
+You will find batch scripts `tmp/lmp-0.slurm` and `tmp/lmp-1.slurm` in `tmp` directory.
+
+`omb batch` will provide some useful functions in the batch script.
+For example, `checkpoint` will check if the job is done and skip the job if it's done.
+
+You can run the above script by `./examples/omb-batch.sh`,
+
+### Track the state of job in job schedular
+
+Let's continue the above example, now you have submitted the batch scripts to the job scheduler.
+
+You can use `omb job` to track the state of the jobs.
+
+```bash
+
+omb job slurm \
+    submit tmp/*.slurm --max_tries 3 --wait --recovery lammps-jobs.json 
+```
+
+The above command will submit the batch scripts to the job scheduler,
+and wait for the jobs to finish. If the job fails, it will retry for at most 3 times.
+
+The `--recovery` option will save the job information to `lammps-jobs.json` file,
+if `omb job` is interrupted, you can run the exact same command to recover the job status, 
+so that you don't need to resubmit the jobs that are already submitted.
+

oh_my_batch-0.1.0.dev1/README.md

@@ -0,0 +1,111 @@
+# oh-my-batch
+A simple tool to manipulate batch tasks.
+
+The goal of this to is to provide a handy command tool for
+* `omb combo`: generate folders/files from different combinations of parameters
+* `omb batch`: generate batch scripts from multiple working directories
+* `omb job`: track the state of job in job schedular
+
+## Install
+```bash
+pip install oh-my-batch
+```
+
+## Use cases
+
+### Generate files from different combinations of parameters
+
+It's common to generate files with different combinations of parameters in scientific computing.
+For example, you have 3 LAMMPS data files in `tmp` directory: `tmp/1.data`, `tmp/2.data`, `tmp/3.data`.
+And you want to generate a series of input files with different parameters,
+for example, different temperatures 300K, 400K, 500K, against each data file.
+
+In this case, you can use `omb combo` command to generate a series of input files for you.
+
+```bash
+#! /bin/bash
+# prepare fake data files
+mkdir -p tmp/
+touch tmp/1.data tmp/2.data tmp/3.data
+
+# prepare a lammps input file template
+cat > tmp/in.lmp.tmp <<EOF
+read_data $DATA_FILE
+velocity all create $TEMP $RANDOM
+run 1000
+EOF
+
+# prepare a run script template
+cat > tmp/run.sh.tmp <<EOF
+cat in.lmp  # simulate running lammps
+EOF
+
+# generate input files
+omb combo \
+    add_files DATA_FILE tmp/*.data - \
+    add_var TEMP 300 400 500 - \
+    add_randint RANDOM -n 3 -a 1 -b 1000 --broadcast - \
+    make_files tmp/in.lmp.tmp tmp/tasks/{i}-T-{TEMP}/in.lmp - \
+    make_files tmp/run.sh.tmp tmp/tasks/{i}-T-{TEMP}/run.sh --mode 755 - \
+    done
+```
+
+The above script will generate 9 folders in `tmp/tasks` directory
+with names from `0-T-300`, `1-T-400`, `2-T-500`, `3-T-300` to `8-T-500`.
+Each folder will contain a `in.lmp` file and a `run.sh` file.
+
+The 9 folders are the combinations of 3 data files and 3 temperatures,
+and each input file will have a independent random number between 1 and 1000 as `RANDOM`.
+
+You can run the about script by `./examples/omb-combo.sh`,
+and you can also run `omb combo --help` to see the detailed usage of `combo` command.
+
+### Generate batch scripts from multiple working directories
+It's common to submit a lot of jobs to a job scheduler. `omb batch` is designed to help you generate batch scripts from multiple working directories and package them into several batch scripts.
+
+Let's continue the above example, now you have 9 folders in `tmp/tasks` directory.
+You want to package them into 2 batch scripts to submit to a job scheduler.
+
+You can use `omb batch` to generate batch scripts for you like this:
+
+```bash
+#! /bin/bash
+cat > tmp/lammps_header.sh <<EOF
+#!/bin/bash
+#SBATCH -J lmp
+#SBATCH -n 1
+#SBATCH -t 1:00:00
+EOF
+
+omb batch \
+    add_work_dir tmp/tasks/* - \
+    add_header_file tmp/lammps_header.sh - \
+    add_command "checkpoint lmp.done ./run.sh" - \
+    make tmp/lmp-{i}.slurm --concurrency 2
+```
+
+You will find batch scripts `tmp/lmp-0.slurm` and `tmp/lmp-1.slurm` in `tmp` directory.
+
+`omb batch` will provide some useful functions in the batch script.
+For example, `checkpoint` will check if the job is done and skip the job if it's done.
+
+You can run the above script by `./examples/omb-batch.sh`,
+
+### Track the state of job in job schedular
+
+Let's continue the above example, now you have submitted the batch scripts to the job scheduler.
+
+You can use `omb job` to track the state of the jobs.
+
+```bash
+
+omb job slurm \
+    submit tmp/*.slurm --max_tries 3 --wait --recovery lammps-jobs.json 
+```
+
+The above command will submit the batch scripts to the job scheduler,
+and wait for the jobs to finish. If the job fails, it will retry for at most 3 times.
+
+The `--recovery` option will save the job information to `lammps-jobs.json` file,
+if `omb job` is interrupted, you can run the exact same command to recover the job status, 
+so that you don't need to resubmit the jobs that are already submitted.

oh_my_batch-0.1.0.dev1/oh_my_batch/__main__.py

@@ -0,0 +1,5 @@
+from .cli import main
+
+
+if __name__ == '__main__':
+    main()

oh_my_batch-0.1.0.dev1/oh_my_batch/assets/__init__.py

@@ -0,0 +1,7 @@
+import os
+
+DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+def get_asset(rel_path: str):
+    return os.path.join(DIR, rel_path)

oh_my_batch-0.1.0.dev1/oh_my_batch/assets/functions.sh

@@ -0,0 +1,20 @@
+
+checkpoint() {
+    # Usage: checkpoint <flag_file> <command> [arg1] [arg2] ...
+    local flag_file="$1"
+    shift  # Remove the first argument so $@ contains only the command and its arguments
+    if [ -f "$flag_file" ]; then
+        cat "$flag_file"
+    else
+        "$@"  # Execute the command
+        local exit_code=$?
+        if [ $exit_code -eq 0 ]; then
+            local current_time=$(date '+%Y-%m-%d %H:%M:%S')
+            printf 'Command succeeded at %s\n' "$current_time" > "$flag_file"
+            echo "Created flag file '$flag_file' with timestamp: $current_time"
+        else
+            echo "Command `$@` failed with exit code $exit_code"
+            return $exit_code
+        fi
+    fi
+}

oh_my_batch-0.1.0.dev1/oh_my_batch/batch.py

@@ -0,0 +1,96 @@
+import shlex
+import os
+
+from .util import split_list, ensure_dir, expand_globs, mode_translate
+from .assets import get_asset
+
+
+class BatchMaker:
+
+    def __init__(self):
+        self._work_dirs = []
+        self._script_header = []
+        self._script_bottom = []
+        self._command = []
+
+    def add_work_dir(self, *dir: str):
+        """
+        Add working directories
+
+        :param dir: Directories to work on, can be glob patterns
+        """
+        self._work_dirs.extend(expand_globs(dir))
+        return self
+
+    def add_header_file(self, file: str, encoding='utf-8'):
+        """
+        Add script header from files
+
+        :param file: File path
+        :param encoding: File encoding
+        """
+        with open(file, 'r', encoding=encoding) as f:
+            self._script_header.append(f.read())
+        return self
+
+    def add_bottom_file(self, file: str, encoding='utf-8'):
+        """
+        Add script bottom from files
+
+        :param file: File path
+        :param encoding: File encoding
+        """
+        with open(file, 'r', encoding=encoding) as f:
+            self._script_bottom.append(f.read())
+
+    def add_command_file(self, file: str, encoding='utf-8'):
+        """
+        Add commands from files to run under every working directory
+
+        :param file: File path
+        :param encoding: File encoding
+        """
+        with open(file, 'r', encoding=encoding) as f:
+            self._command.append(f.read())
+        return self
+
+    def add_command(self, *cmd: str):
+        """
+        add commands to run under every working directory
+
+        :param cmd: Commands to run, can be multiple
+        """
+        self._command.extend(cmd)
+        return self
+
+    def make(self, path: str, concurrency=1, encoding='utf-8', mode='755'):
+        """
+        Make batch script files from the previous setup
+
+        :param path: Path to save batch script files, use {i} to represent index
+        :param concurrency: Number of concurrent commands to run
+        """
+        # inject pre-defined functions
+        self.add_header_file(get_asset('functions.sh'))
+
+        header = '\n'.join(self._script_header)
+        bottom = '\n'.join(self._script_bottom)
+        for i, work_dirs in enumerate(split_list(self._work_dirs, concurrency)):
+            body = []
+            work_dirs_arr = "\n".join(shlex.quote(w) for w in work_dirs)
+            body.extend([
+                '[ -n "$PBS_O_WORKDIR" ] && cd $PBS_O_WORKDIR  # fix PBS',
+                f'work_dirs=({work_dirs_arr})',
+                '',
+                'for work_dir in "${work_dirs[@]}"; do',
+                'pushd $work_dir',
+                *self._command,
+                'popd',
+                'done'
+            ])
+            script = '\n'.join([header, *body, bottom])
+            out_path = path.format(i=i)
+            ensure_dir(out_path)
+            with open(out_path, 'w', encoding=encoding) as f:
+                f.write(script)
+            os.chmod(out_path, mode_translate(str(mode)))

oh_my_batch-0.1.0.dev1/oh_my_batch/cli.py

@@ -0,0 +1,25 @@
+import fire
+
+class JobCli:
+
+    def slurm(self):
+        from .job import Slurm
+        return Slurm
+
+
+class OhMyBatch:
+
+    def combo(self):
+        from .combo import ComboMaker
+        return ComboMaker
+
+    def batch(self):
+        from .batch import BatchMaker
+        return BatchMaker
+
+    def job(self):
+        return JobCli()
+
+
+def main():
+    fire.Fire(OhMyBatch)

oh_my_batch-0.1.0.dev1/oh_my_batch/combo.py

@@ -0,0 +1,202 @@
+from itertools import product
+from string import Template
+import random
+import os
+
+from .util import expand_globs, mode_translate
+
+class ComboMaker:
+
+    def __init__(self, seed=None):
+        """
+        ComboMaker constructor
+
+        :param seed: Seed for random number generator
+        """
+        self._product_vars = {}
+        self._broadcast_vars = {}
+        if seed is not None:
+            random.seed(seed)
+        self._combos = []
+
+    def add_seq(self, key: str, start: int, stop: int, step: int=1, broadcast=False):
+        """
+        Add a variable with sequence of integer values
+
+        :param key: Variable name
+        :param start: Start value
+        :param stop: Stop value
+        :param step: Step
+        :param broadcast: If True, values are broadcasted, otherwise they are producted when making combos
+        """
+        args = list(range(start, stop, step))
+        self.add_var(key, *args, broadcast=broadcast)
+        return self
+
+    def add_randint(self, key: str, n: int, a: int, b: int, broadcast=False, seed=None):
+        """
+        Add a variable with random integer values
+
+        :param key: Variable name
+        :param n: Number of values
+        :param a: Lower bound
+        :param b: Upper bound
+        :param broadcast: If True, values are broadcasted, otherwise they are producted when making combos
+        :param seed: Seed for random number generator
+        """
+        if seed is not None:
+            random.seed(seed)
+        args = [random.randint(a, b) for _ in range(n)]
+        self.add_var(key, *args, broadcast=broadcast)
+        return self
+
+    def add_rand(self, key: str, n: int, a: float, b: float, broadcast=False, seed=None):
+        """
+        Add a variable with random float values
+
+        :param key: Variable name
+        :param n: Number of values
+        :param a: Lower bound
+        :param b: Upper bound
+        :param broadcast: If True, values are broadcasted, otherwise they are producted when making combos
+        :param seed: Seed for random number generator
+        """
+        if seed is not None:
+            random.seed(seed)
+        args = [random.uniform(a, b) for _ in range(n)]
+        self.add_var(key, *args, broadcast=broadcast)
+        return self
+
+    def add_files(self, key: str, *path: str, broadcast=False, abs=False):
+        """
+        Add a variable with files by glob pattern
+        For example, suppose there are 3 files named 1.txt, 2.txt, 3.txt in data directory,
+        then calling add_files('DATA_FILE', 'data/*.txt') will add list ["data/1.txt", "data/2.txt", "data/3.txt"]
+        to the variable DATA_FILE.
+
+        :param key: Variable name
+        :param path: Path to files, can include glob pattern
+        :param broadcast: If True, values are broadcasted, otherwise they are producted when making combos
+        :param abs: If True, path will be turned into absolute path
+        """
+        args = expand_globs(path, raise_invalid=True)
+        if not args:
+            raise ValueError(f"No files found for {path}")
+        if abs:
+            args = [os.path.abspath(p) for p in args]
+        self.add_var(key, *args, broadcast=broadcast)
+        return self
+
+    def add_files_as_one(self, key: str, path: str, broadcast=False, sep=' ', abs=False):
+        """
+        Add a variable with files by glob pattern as one string
+        Unlike add_files, this function joins the files with a delimiter.
+        For example, suppose there are 1.txt, 2.txt, 3.txt in data directory,
+        then calling add_files_as_one('DATA_FILE', 'data/*.txt') will add string "data/1.txt data/2.txt data/3.txt"
+        to the variable DATA_FILE.
+
+        :param key: Variable name
+        :param path: Path to files, can include glob pattern
+        :param broadcast: If True, values are broadcasted, otherwise they are producted when making combos
+        :param sep: Separator to join files
+        :param abs: If True, path will be turned into absolute path
+        """
+        args = expand_globs(path, raise_invalid=True)
+        if not args:
+            raise ValueError(f"No files found for {path}")
+        if abs:
+            args = [os.path.abspath(p) for p in args]
+        self.add_var(key, sep.join(args), broadcast=broadcast)
+        return self
+
+    def add_var(self, key: str, *args, broadcast=False):
+        """
+        Add a variable with values
+
+        :param key: Variable name
+        :param args: Values
+        :param broadcast: If True, values are broadcasted, otherwise they are producted when making combos
+        """
+
+        if key == 'i':
+            raise ValueError("Variable name 'i' is reserved")
+
+        if broadcast:
+            if key in self._product_vars:
+                raise ValueError(f"Variable {key} already defined as product variable")
+            self._broadcast_vars.setdefault(key, []).extend(args)
+        else:
+            if key in self._broadcast_vars:
+                raise ValueError(f"Variable {key} already defined as broadcast variable")
+            self._product_vars.setdefault(key, []).extend(args)
+        return self
+
+    def shuffle(self, *keys: str, seed=None):
+        """
+        Shuffle variables
+        :param keys: Variable names to shuffle
+        :param seed: Seed for random number generator
+        """
+        if seed is not None:
+            random.seed(seed)
+
+        for key in keys:
+            if key in self._product_vars:
+                random.shuffle(self._product_vars[key])
+            elif key in self._broadcast_vars:
+                random.shuffle(self._broadcast_vars[key])
+            else:
+                raise ValueError(f"Variable {key} not found")
+        return self
+
+    def make_files(self, template: str, dest: str, delimiter='$', mode=None):
+        """
+        Make files from template
+        The template file can include variables with delimiter.
+        For example, if delimiter is '$', then the template file can include $var1, $var2, ...
+
+        The destination can also include variables in string format style.
+        For example, if dest is 'output/{i}.txt', then files are saved as output/0.txt, output/1.txt, ...
+
+        :param template: Path to template file
+        :param dest: Path pattern to destination file
+        :param delimiter: Delimiter for variables in template, default is '$',
+        can be changed to other character, e.g $$, @, ...
+        """
+        _delimiter = delimiter
+
+        class _Template(Template):
+            delimiter = _delimiter
+
+        combos = self._make_combos()
+        for i, combo in enumerate(combos):
+            with open(template, 'r') as f:
+                template_text = f.read()
+            text = _Template(template_text).safe_substitute(combo)
+            _dest = dest.format(i=i, **combo)
+            os.makedirs(os.path.dirname(_dest), exist_ok=True)
+            with open(_dest, 'w') as f:
+                f.write(text)
+            if mode is not None:
+                os.chmod(_dest, mode_translate(str(mode)))
+        return self
+
+    def done(self):
+        """
+        End of command chain
+        """
+        pass
+
+    def _make_combos(self):
+        if not self._product_vars and not self._broadcast_vars:
+            return self._combos
+        keys = self._product_vars.keys()
+        values_list = product(*self._product_vars.values())
+        combos = [ dict(zip(keys, values)) for values in values_list ]
+        for i, combo in enumerate(combos):
+            for k, v in self._broadcast_vars.items():
+                combo[k] = v[i % len(v)]
+        self._combos.extend(combos)
+        self._product_vars = {}
+        self._broadcast_vars = {}
+        return self._combos

oh_my_batch-0.1.0.dev1/oh_my_batch/job.py

@@ -0,0 +1,178 @@
+from typing import List
+from enum import Enum
+
+import logging
+import json
+import time
+import os
+import re
+
+from .util import expand_globs, shell_run, parse_csv
+
+
+logger = logging.getLogger(__name__)
+
+
+class JobState(bytes, Enum):
+    """
+    Job state enumeration
+    """
+    def __new__(cls, value: int, terminal: bool, status_name: str) -> "JobState":
+        obj = bytes.__new__(cls, [value])
+        obj._value_ = value
+        obj.terminal = terminal
+        obj.status_name = status_name
+        return obj
+
+    value: int  # type: ignore
+    terminal: bool
+    status_name: str
+
+    NULL = (0, True, "NULL")
+    PENDING = (1, False, "PENDING")
+    RUNNING = (2, False, "RUNNING")
+    CANCELLED = (3, True, "CANCELLED")
+    COMPLETED = (4, True, "COMPLETED")
+    FAILED = (5, True, "FAILED")
+    UNKNOWN = (6, False, "UNKNOWN")
+
+
+def new_job(script: str):
+    return {
+        'id': '', 'script': script, 'state': JobState.NULL, 'tries': 0,
+    }
+
+
+class BaseJobManager:
+
+    def submit(self, *script: str, recovery: str = '', wait=False,
+               timeout=None, opts='', max_tries=1, interval=10):
+        """
+        Submit scripts
+
+        :param script: Script files to submit, can be glob patterns.
+        :param recovery: Recovery file to store the state of the submitted scripts
+        :param wait: If True, wait for the job to finish
+        :param timeout: Timeout in seconds for waiting
+        :param opts: Additional options for submit command
+        :param max_tries: Maximum number of tries for each job
+        :param interval: Interval in seconds for checking job status
+        """
+        jobs = []
+        if recovery and os.path.exists(recovery):
+            with open(recovery, 'r', encoding='utf-8') as f:
+                jobs = json.load(f)
+
+        recover_scripts = set(j['script'] for j in jobs)
+        logger.info('Scripts in recovery files: %s', recover_scripts)
+
+        scripts = set(os.path.normpath(s) for s in expand_globs(script))
+        logger.info('Scripts to submit: %s', scripts)
+
+        for script_file in scripts:
+            if script_file not in recover_scripts:
+                jobs.append(new_job(script_file))
+
+        current = time.time()
+        while True:
+            self._update_jobs(jobs, max_tries, opts)
+            if recovery:
+                with open(recovery, 'w', encoding='utf-8') as f:
+                    json.dump(jobs, f, indent=2)
+
+            if not wait:
+                break
+
+            # stop if all jobs are terminal and not job to be submitted
+            if (all(j['state'].terminal for j in jobs) and
+                    not any(should_submit(j, max_tries) for j in jobs)):
+                break
+
+            if timeout and time.time() - current > timeout:
+                logger.error('Timeout, current state: %s', jobs)
+                break
+
+            time.sleep(interval)
+
+    def _update_jobs(self, jobs: List[dict], max_tries: int, submit_opts: str):
+        raise NotImplementedError
+
+
+class Slurm(BaseJobManager):
+    def __init__(self, sbatch='sbatch',  sacct='sacct'):
+        self._sbatch_bin = sbatch
+        self._sacct_bin = sacct
+
+    def _update_jobs(self, jobs: List[dict], max_tries: int, submit_opts: str):
+        # query job status
+        job_ids = [j['id'] for j in jobs if j['id']]
+        if job_ids:
+            query_cmd = f'{self._sacct_bin} -X -P --format=JobID,JobName,State -j {",".join(job_ids)}'
+            user = os.environ.get('USER')
+            if user:
+                query_cmd += f' -u {user}'
+
+            cp = shell_run(query_cmd)
+            if cp.returncode != 0:
+                logger.error('Failed to query job status: %s', cp.stderr.decode('utf-8'))
+                return jobs
+            logger.info('Job status: %s', cp.stdout.decode('utf-8'))
+            new_state = parse_csv(cp.stdout.decode('utf-8'))
+        else:
+            new_state = []
+
+        for job in jobs:
+            for row in new_state:
+                if job['id'] == row['JobID']:
+                    job['state'] = self._map_state(row['State'])
+                    if job['state'] == JobState.UNKNOWN:
+                        logger.warning('Unknown job %s state: %s',row['JobID'], row['State'])
+                    break
+            else:
+                job['state'] = JobState.FAILED
+                logger.error('Job %s not found in sacct output', job['id'])
+
+        # check if there are jobs to be (re)submitted
+        for job in jobs:
+            if should_submit(job, max_tries):
+                job['tries'] += 1
+                job['id'] = ''
+                job['state'] = JobState.NULL
+                submit_cmd = f'{self._sbatch_bin} {submit_opts} {job["script"]}'
+                cp = shell_run(submit_cmd)
+                if cp.returncode != 0:
+                    job['state'] = JobState.FAILED
+                    logger.error('Failed to submit job: %s', cp.stderr.decode('utf-8'))
+                else:
+                    job['id'] = self._parse_job_id(cp.stdout.decode('utf-8'))
+                    assert job['id'], 'Failed to parse job id'
+                    job['state'] = JobState.PENDING
+                    logger.info('Job %s submitted', job['id'])
+
+    def _map_state(self, state: str):
+        if state.startswith('CANCELLED'):
+            return JobState.CANCELLED
+        return {
+            'PENDING': JobState.PENDING,
+            'RUNNING': JobState.RUNNING,
+            'COMPLETED': JobState.COMPLETED,
+            'FAILED': JobState.FAILED,
+            'OUT_OF_MEMORY': JobState.FAILED,
+            'TIMEOUT': JobState.FAILED,
+        }.get(state, JobState.UNKNOWN)
+
+    def _parse_job_id(self, output: str):
+        """
+        Parse job id from sbatch output
+        """
+        m = re.search(r'\d+', output)
+        return m.group(0) if m else ''
+
+
+def should_submit(job: dict, max_tries: int):
+    state: JobState = job['state']
+    if not state.terminal:
+        return False
+    if job['tries'] >= max_tries:
+        return False
+    return state != JobState.COMPLETED

oh_my_batch-0.1.0.dev1/oh_my_batch/util.py

@@ -0,0 +1,86 @@
+from typing import List, Iterable
+import subprocess as sp
+import logging
+import glob
+import csv
+import os
+
+
+logger = logging.getLogger(__name__)
+
+
+def expand_globs(patterns: Iterable[str], raise_invalid=False) -> List[str]:
+    """
+    Expand glob patterns in paths
+
+    :param patterns: list of paths or glob patterns
+    :param raise_invalid: if True, will raise error if no file found for a glob pattern
+    :return: list of expanded paths
+    """
+    paths = []
+    for pattern in patterns:
+        result = glob.glob(pattern, recursive=True) if '*' in pattern else [pattern]
+        if raise_invalid and len(result) == 0:
+            raise FileNotFoundError(f'No file found for {pattern}')
+        for p in result:
+            if p not in paths:
+                paths.append(p)
+            else:
+                logger.warning('path %s already exists in the list', p)
+    return paths
+
+
+def split_list(l, n):
+    """
+    Splits a list into n sub-lists.
+
+    :param l: The list to be split.
+    :param n: The number of sub-lists to create.
+    :return: A list of sub-lists.
+    """
+    if n <= 0:
+        raise ValueError("Number of sub-lists must be a positive integer")
+
+    # Calculate the size of each sublist
+    k, m = divmod(len(l), n)
+
+    for i in range(n):
+        start = i * k + min(i, m)
+        end = (i + 1) * k + min(i + 1, m)
+        if start == end:
+            break
+        yield l[start:end]
+
+
+def ensure_dir(path: str):
+    """
+    Ensure the directory exists
+
+    :param path: Path to directory or file.
+    """
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+
+def mode_translate(mode: str):
+    """
+    Translate mode in decimal to octal
+    For example, convert 777 -> 0o777, 755 -> 0o755
+    """
+    return int(mode, 8)
+
+
+def shell_run(cmd: str):
+    """
+    Run a shell command
+
+    :param cmd: Command to run
+    """
+    return sp.run(cmd, shell=True, stdout=sp.PIPE, stderr=sp.PIPE)
+
+
+def parse_csv(text: str, delimiter="|"):
+    """
+    Parse CSV text to list of dictionaries
+    """
+    reader = csv.DictReader(text.splitlines(), delimiter=delimiter)
+    return list(reader)

{oh_my_batch-0.0.1.dev0 → oh_my_batch-0.1.0.dev1}/pyproject.toml

@@ -1,16 +1,21 @@
 [tool.poetry]
 name = "oh-my-batch"
-version = "0.0.1.dev0"
+version = "0.1.0.dev1"
 description = ""
 authors = ["weihong.xu <xuweihong.cn@gmail.com>"]
 license = "GPL"
 readme = "README.md"
-packages = [{include = "on_my_batch"}]
+packages = [{include = "oh_my_batch"}]
 
 [tool.poetry.dependencies]
 python = "^3.8"
+fire = "^0.7.0"
 
 
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.scripts]
+oh-my-batch = "oh_my_batch.cli:main"
+omb = "oh_my_batch.cli:main"

oh_my_batch-0.0.1.dev0/PKG-INFO

@@ -1,20 +0,0 @@
-Metadata-Version: 2.1
-Name: oh-my-batch
-Version: 0.0.1.dev0
-Summary: 
-License: GPL
-Author: weihong.xu
-Author-email: xuweihong.cn@gmail.com
-Requires-Python: >=3.8,<4.0
-Classifier: License :: Other/Proprietary License
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Description-Content-Type: text/markdown
-
-# oh-my-batch
-A simple tool to manipulate batch tasks.
-

oh_my_batch-0.0.1.dev0/README.md

@@ -1,2 +0,0 @@
-# oh-my-batch
-A simple tool to manipulate batch tasks.