oh-my-batch 0.1.0.dev3__py3-none-any.whl → 0.2.0__py3-none-any.whl

oh_my_batch/__init__.py

@@ -0,0 +1,4 @@
+if __name__ == '__main__':
+    import fire
+    from .cli import OhMyBatch
+    fire.Fire(OhMyBatch)

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/combo.py

@@ -117,7 +117,6 @@ class ComboMaker:
         :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")
 

oh_my_batch/job.py

@@ -1,5 +1,4 @@
 from typing import List
-from enum import Enum
 
 import logging
 import json
@@ -7,11 +6,12 @@ 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__)
 
+
 class JobState:
     NULL = 0
     PENDING = 1
@@ -59,7 +59,7 @@ class BaseJobManager:
         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))
+        scripts = set(norm_path(s) for s in expand_globs(script, raise_invalid=True))
         logger.info('Scripts to submit: %s', scripts)
 
         for script_file in scripts:
@@ -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)
 
@@ -101,20 +102,18 @@ class Slurm(BaseJobManager):
         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'))
+                logger.error('Failed to query job status: %s', log_cp(cp))
                 return jobs
-            logger.info('Job status: %s', cp.stdout.decode('utf-8'))
+            logger.info('Job status:\n%s', cp.stdout.decode('utf-8'))
             new_state = parse_csv(cp.stdout.decode('utf-8'))
         else:
             new_state = []
 
         for job in jobs:
+            if not job['id']:
+                continue
             for row in new_state:
                 if job['id'] == row['JobID']:
                     job['state'] = self._map_state(row['State'])
@@ -122,8 +121,7 @@ class Slurm(BaseJobManager):
                         logger.warning('Unknown job %s state: %s',row['JobID'], row['State'])
                     break
             else:
-                if job['id']:
-                    logger.error('Job %s not found in sacct output', job['id'])
+                logger.error('Job %s not found in sacct output', job['id'])
 
         # check if there are jobs to be (re)submitted
         for job in jobs:
@@ -135,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'
@@ -169,3 +167,7 @@ def should_submit(job: dict, max_tries: int):
     if job['tries'] >= max_tries:
         return False
     return state != JobState.COMPLETED
+
+
+def norm_path(path: str):
+    return os.path.normpath(os.path.abspath(path))

oh_my_batch/util.py

@@ -19,7 +19,7 @@ def expand_globs(patterns: Iterable[str], raise_invalid=False) -> List[str]:
     """
     paths = []
     for pattern in patterns:
-        result = glob.glob(pattern, recursive=True) if '*' in pattern else [pattern]
+        result = glob.glob(pattern, recursive=True)
         if raise_invalid and len(result) == 0:
             raise FileNotFoundError(f'No file found for {pattern}')
         for p in result:
@@ -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.0.dev3.dist-info → oh_my_batch-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: oh-my-batch
-Version: 0.1.0.dev3
+Version: 0.2.0
 Summary: 
 License: GPL
 Author: weihong.xu
@@ -17,6 +17,11 @@ Requires-Dist: fire (>=0.7.0,<0.8.0)
 Description-Content-Type: text/markdown
 
 # oh-my-batch
+
+[![PyPI version](https://badge.fury.io/py/oh-my-batch.svg)](https://badge.fury.io/py/oh-my-batch)
+[![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.
 
 ## Features
@@ -41,7 +46,6 @@ 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
@@ -87,7 +91,6 @@ 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
@@ -96,9 +99,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
 ```
 
@@ -112,19 +115,16 @@ 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.
+In this case, 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 
+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.
+The `--recovery` option will save the job information to `lammps-jobs.json` file.
+If `omb job` is interrupted, you can rerun the exact same command to recover the job status,
+so that you don't need to resubmit the jobs that are still running or completed.
 

oh_my_batch-0.2.0.dist-info/RECORD

@@ -0,0 +1,14 @@
+oh_my_batch/__init__.py,sha256=BsRNxZbqDWfaIZJGxzIDqCubRWztMGFDceW08TECuFs,98
+oh_my_batch/__main__.py,sha256=sWyFZMwWNvhkanwZSJRGfBBDoIevhC028dTSB67i6yI,61
+oh_my_batch/assets/__init__.py,sha256=Exub46UbQaz2V2eXpQeiVfnThQpXaNeuyjlGY6gBSZc,130
+oh_my_batch/assets/functions.sh,sha256=eORxFefV-XrWbG-2I6u-c8uf1XxOQ31LaeVHBumwzJ4,708
+oh_my_batch/batch.py,sha256=6qnaXEVyA493heGzzbCrdZXCcnYk8zgl7WP0rmo7KlU,3690
+oh_my_batch/cli.py,sha256=uelW9ms1N30DipJOcsiuG5K-5VN8O6yu1RNEqex00GY,475
+oh_my_batch/combo.py,sha256=R_WTO4v-LWHIQ0O46bIgeRlL_RGrFcf8305S9auqeQk,7679
+oh_my_batch/job.py,sha256=8kZnWtvpr1rAl4tc9I_Vlhi-T0o3rh4RQZZgMNBCGho,5800
+oh_my_batch/util.py,sha256=okg_kY8dJouyJ2BYCXRl7bxDUAtNH6GLh2UjXRnkoW0,2385
+oh_my_batch-0.2.0.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+oh_my_batch-0.2.0.dist-info/METADATA,sha256=_qeLyk6LEg2--NU9NqgMoV2WU96125A0eiIOdJFY760,4759
+oh_my_batch-0.2.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+oh_my_batch-0.2.0.dist-info/entry_points.txt,sha256=ZY2GutSoNjjSyJ4qO2pTeseKUFgoTYdvmgkuZZkwi68,77
+oh_my_batch-0.2.0.dist-info/RECORD,,

oh_my_batch-0.1.0.dev3.dist-info/RECORD

@@ -1,14 +0,0 @@
-oh_my_batch/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-oh_my_batch/__main__.py,sha256=sWyFZMwWNvhkanwZSJRGfBBDoIevhC028dTSB67i6yI,61
-oh_my_batch/assets/__init__.py,sha256=Exub46UbQaz2V2eXpQeiVfnThQpXaNeuyjlGY6gBSZc,130
-oh_my_batch/assets/functions.sh,sha256=eORxFefV-XrWbG-2I6u-c8uf1XxOQ31LaeVHBumwzJ4,708
-oh_my_batch/batch.py,sha256=e73N-xwxMvgxnWwFMp33PQD1Dy-T-ATjANlwtPRHPQM,3016
-oh_my_batch/cli.py,sha256=uelW9ms1N30DipJOcsiuG5K-5VN8O6yu1RNEqex00GY,475
-oh_my_batch/combo.py,sha256=AHFD5CLoczqtjcfl2Rb4A2ucoQU40-cWtDOYjtP-yY4,7680
-oh_my_batch/job.py,sha256=_fETBYpuSd_hNHKnXSwYcSU3OXtU7PO-P2QMfhE-Wfs,5788
-oh_my_batch/util.py,sha256=H8B4zVNH5xRp-NG_uypgvtmz2YSpXy_6LK5ROv6SYrc,2116
-oh_my_batch-0.1.0.dev3.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-oh_my_batch-0.1.0.dev3.dist-info/METADATA,sha256=sFfJrqAmymhuKZIv4-ydeDX6Y3oby3NWVfiPjQPUAUM,4453
-oh_my_batch-0.1.0.dev3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-oh_my_batch-0.1.0.dev3.dist-info/entry_points.txt,sha256=ZY2GutSoNjjSyJ4qO2pTeseKUFgoTYdvmgkuZZkwi68,77
-oh_my_batch-0.1.0.dev3.dist-info/RECORD,,