oh-my-batch 0.1.1__tar.gz → 0.2.1__tar.gz

{oh_my_batch-0.1.1 → oh_my_batch-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: oh-my-batch
-Version: 0.1.1
+Version: 0.2.1
 Summary: 
 License: GPL
 Author: weihong.xu
@@ -22,19 +22,33 @@ Description-Content-Type: text/markdown
 [![PyPI - Downloads](https://img.shields.io/pypi/dm/oh-my-batch)](https://pypi.org/project/oh-my-batch/)
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/oh-my-batch)](https://pypi.org/project/oh-my-batch/)
 
-A simple tool to manipulate batch tasks designed for scientific computing community.
+A toolkit to manipulate batch tasks with command line. Designed for scientific computing community.
 
 ## Features
 * `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
 
+## Shell Tips
+`oh-my-batch` is intended to help you implement computational workflows with shell scripts.
+To make the best use of `oh-my-batch`, you need to know some shell tips.
+
+* [Retry commands until success in shell script](https://stackoverflow.com/a/79191004/3099733)
+
 ## Install
 ```bash
 pip install oh-my-batch
 ```
 
 ## Use cases
+### Load useful functions to your shell script
+You can load useful functions from `oh-my-batch` this way:
+
+```bash
+omb misc export-shell-func > omb-func.sh && source omb-func.sh
+```
+
+This will load some useful functions to your shell script, for example, `checkpoint` function.
 
 ### Generate files from different combinations of parameters
 
@@ -99,9 +113,9 @@ cat > tmp/lammps_header.sh <<EOF
 EOF
 
 omb batch \
-    add_work_dir tmp/tasks/* - \
-    add_header_file tmp/lammps_header.sh - \
-    add_command "checkpoint lmp.done ./run.sh" - \
+    add_work_dirs tmp/tasks/* - \
+    add_header_files tmp/lammps_header.sh - \
+    add_cmds "checkpoint lmp.done ./run.sh" - \
     make tmp/lmp-{i}.slurm --concurrency 2
 ```
 

{oh_my_batch-0.1.1 → oh_my_batch-0.2.1}/README.md

@@ -4,19 +4,33 @@
 [![PyPI - Downloads](https://img.shields.io/pypi/dm/oh-my-batch)](https://pypi.org/project/oh-my-batch/)
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/oh-my-batch)](https://pypi.org/project/oh-my-batch/)
 
-A simple tool to manipulate batch tasks designed for scientific computing community.
+A toolkit to manipulate batch tasks with command line. Designed for scientific computing community.
 
 ## Features
 * `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
 
+## Shell Tips
+`oh-my-batch` is intended to help you implement computational workflows with shell scripts.
+To make the best use of `oh-my-batch`, you need to know some shell tips.
+
+* [Retry commands until success in shell script](https://stackoverflow.com/a/79191004/3099733)
+
 ## Install
 ```bash
 pip install oh-my-batch
 ```
 
 ## Use cases
+### Load useful functions to your shell script
+You can load useful functions from `oh-my-batch` this way:
+
+```bash
+omb misc export-shell-func > omb-func.sh && source omb-func.sh
+```
+
+This will load some useful functions to your shell script, for example, `checkpoint` function.
 
 ### Generate files from different combinations of parameters
 
@@ -81,9 +95,9 @@ cat > tmp/lammps_header.sh <<EOF
 EOF
 
 omb batch \
-    add_work_dir tmp/tasks/* - \
-    add_header_file tmp/lammps_header.sh - \
-    add_command "checkpoint lmp.done ./run.sh" - \
+    add_work_dirs tmp/tasks/* - \
+    add_header_files tmp/lammps_header.sh - \
+    add_cmds "checkpoint lmp.done ./run.sh" - \
     make tmp/lmp-{i}.slurm --concurrency 2
 ```
 

{oh_my_batch-0.1.1 → oh_my_batch-0.2.1}/oh_my_batch/assets/functions.sh

@@ -1,6 +1,7 @@
+GREEN='\033[0;32m'
+NC='\033[0m'
 
 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
@@ -10,6 +11,7 @@ checkpoint() {
         local exit_code=$?
         if [ $exit_code -eq 0 ]; then
             local current_time=$(date '+%Y-%m-%d %H:%M:%S')
+            mkdir -p "$(dirname "$flag_file")"
             printf 'Command succeeded at %s\n' "$current_time" > "$flag_file"
             echo "Created flag file '$flag_file' with timestamp: $current_time"
         else
@@ -18,3 +20,16 @@ checkpoint() {
         fi
     fi
 }
+
+echo -e "${GREEN}Functiion: checkpoint${NC}"
+cat <<EOF
+Usage:
+    checkpoint <flag_file> <command> [arg1] [arg2] ...
+
+Set a flag file to indicate the completion of a command,
+so that it will not be executed again next time.
+
+Example:
+    checkpoint lmp.done lmp -in in.lmp
+
+EOF

{oh_my_batch-0.1.1 → oh_my_batch-0.2.1}/oh_my_batch/batch.py

@@ -13,7 +13,7 @@ class BatchMaker:
         self._script_bottom = []
         self._command = []
 
-    def add_work_dir(self, *dir: str):
+    def add_work_dirs(self, *dir: str):
         """
         Add working directories
 
@@ -22,39 +22,55 @@ class BatchMaker:
         self._work_dirs.extend(expand_globs(dir))
         return self
 
-    def add_header_file(self, file: str, encoding='utf-8'):
+    def add_header_files(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())
+        self._script_header.extend(load_files(*file, encoding=encoding))
         return self
 
-    def add_bottom_file(self, file: str, encoding='utf-8'):
+    def add_headers(self, *header: str):
+        """
+        Add script header
+
+        :param header: Header lines
+        """
+        self._script_header.extend(header)
+        return self
+
+    def add_bottom_files(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())
+        self._script_bottom.extend(load_files(*file, encoding=encoding))
+        return self
+
+    def add_bottoms(self, *bottom: str):
+        """
+        Add script bottom
 
-    def add_command_file(self, file: str, encoding='utf-8'):
+        :param bottom: Bottom lines
+        """
+        self._script_bottom.extend(bottom)
+        return self
+
+    def add_cmd_files(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())
+        self._command.extend(load_files(*file, encoding=encoding))
         return self
 
-    def add_command(self, *cmd: str):
+    def add_cmds(self, *cmd: str):
         """
         add commands to run under every working directory
 
@@ -68,10 +84,10 @@ class BatchMaker:
         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
+        :param concurrency: Number of scripts to to make
         """
         # inject pre-defined functions
-        self.add_header_file(get_asset('functions.sh'))
+        self.add_header_files(get_asset('functions.sh'))
 
         header = '\n'.join(self._script_header)
         bottom = '\n'.join(self._script_bottom)
@@ -80,10 +96,10 @@ class BatchMaker:
             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})',
+                f'WORK_DIRS=({work_dirs_arr})',
                 '',
-                'for work_dir in "${work_dirs[@]}"; do',
-                'pushd $work_dir',
+                'for WORK_DIR in "${WORK_DIRS[@]}"; do',
+                'pushd $WORK_DIR',
                 *self._command,
                 'popd',
                 'done'
@@ -94,3 +110,17 @@ class BatchMaker:
             with open(out_path, 'w', encoding=encoding) as f:
                 f.write(script)
             os.chmod(out_path, mode_translate(str(mode)))
+
+
+def load_files(*file, encoding='utf-8', raise_invalid=False):
+    """
+    Load files from paths
+
+    :param files: List of file paths
+    :return: List of file contents
+    """
+    result = []
+    for file in expand_globs(file, raise_invalid=raise_invalid):
+        with open(file, 'r', encoding=encoding) as f:
+            result.append(f.read())
+    return result

{oh_my_batch-0.1.1 → oh_my_batch-0.2.1}/oh_my_batch/cli.py

@@ -23,6 +23,9 @@ class OhMyBatch:
     def job(self):
         return JobCli()
 
+    def misc(self):
+        from .misc import Misc
+        return Misc()
 
 def main():
     fire.Fire(OhMyBatch)

{oh_my_batch-0.1.1 → oh_my_batch-0.2.1}/oh_my_batch/job.py

@@ -6,7 +6,7 @@ import time
 import os
 import re
 
-from .util import expand_globs, shell_run, parse_csv
+from .util import expand_globs, shell_run, parse_csv, ensure_dir, log_cp
 
 
 logger = logging.getLogger(__name__)
@@ -70,6 +70,7 @@ class BaseJobManager:
         while True:
             self._update_jobs(jobs, max_tries, opts)
             if recovery:
+                ensure_dir(recovery)
                 with open(recovery, 'w', encoding='utf-8') as f:
                     json.dump(jobs, f, indent=2)
 
@@ -103,7 +104,7 @@ class Slurm(BaseJobManager):
             query_cmd = f'{self._sacct_bin} -X -P --format=JobID,JobName,State -j {",".join(job_ids)}'
             cp = shell_run(query_cmd)
             if cp.returncode != 0:
-                logger.error('Failed to query job status: %s', cp.stderr.decode('utf-8'))
+                logger.error('Failed to query job status: %s', log_cp(cp))
                 return jobs
             logger.info('Job status:\n%s', cp.stdout.decode('utf-8'))
             new_state = parse_csv(cp.stdout.decode('utf-8'))
@@ -132,7 +133,7 @@ class Slurm(BaseJobManager):
                 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'))
+                    logger.error('Failed to submit job: %s', log_cp(cp))
                 else:
                     job['id'] = self._parse_job_id(cp.stdout.decode('utf-8'))
                     assert job['id'], 'Failed to parse job id'

oh_my_batch-0.2.1/oh_my_batch/misc.py

@@ -0,0 +1,14 @@
+from .assets import get_asset
+
+class Misc:
+
+    def export_shell_func(self):
+        """
+        Export shell functions for batch scripts
+
+        For example, you can load them to you shell environment with the following command:
+            omb misc export-shell-func > omb-func.sh && source omb-func.sh
+        """
+        shell_func = get_asset('functions.sh')
+        with open(shell_func, 'r', encoding='utf-8') as f:
+            print(f.read())

{oh_my_batch-0.1.1 → oh_my_batch-0.2.1}/oh_my_batch/util.py

@@ -83,4 +83,18 @@ def parse_csv(text: str, delimiter="|"):
     Parse CSV text to list of dictionaries
     """
     reader = csv.DictReader(text.splitlines(), delimiter=delimiter)
-    return list(reader)
+    return list(reader)
+
+
+def log_cp(cp):
+    """
+    Log child process
+    """
+    log = f'Command: {cp.args}\nReturn code: {cp.returncode}'
+
+    out = cp.stdout.decode('utf-8').strip()
+    if out:
+        log += f'\nSTDOUT:\n{out}'
+    err = cp.stderr.decode('utf-8').strip()
+    if err:
+        log += f'\nSTDERR:\n{err}'

{oh_my_batch-0.1.1 → oh_my_batch-0.2.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "oh-my-batch"
-version = "0.1.1"
+version = "0.2.1"
 description = ""
 authors = ["weihong.xu <xuweihong.cn@gmail.com>"]
 license = "GPL"
@@ -11,7 +11,6 @@ packages = [{include = "oh_my_batch"}]
 python = "^3.8"
 fire = "^0.7.0"
 
-
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"