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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: oh-my-batch
-Version: 0.1.1
+Version: 0.2.0
 Summary: 
 License: GPL
 Author: weihong.xu
@@ -99,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
 ```
 

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

@@ -81,9 +81,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.0}/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.0}/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.1.1 → oh_my_batch-0.2.0}/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.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "oh-my-batch"
-version = "0.1.1"
+version = "0.2.0"
 description = ""
 authors = ["weihong.xu <xuweihong.cn@gmail.com>"]
 license = "GPL"