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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: oh-my-batch
-Version: 0.2.1
+Version: 0.2.3
 Summary: 
 License: GPL
 Author: weihong.xu
@@ -13,6 +13,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: fire (>=0.7.0,<0.8.0)
 Description-Content-Type: text/markdown
 
@@ -28,27 +29,29 @@ A toolkit to manipulate batch tasks with command line. Designed for scientific c
 * `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)
+* `omb misc`: miscellaneous commands
 
 ## Install
 ```bash
 pip install oh-my-batch
 ```
 
+## 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)
+* [Run multiple line shell script with ssh](https://stackoverflow.com/a/32082912/3099733)
+
 ## Use cases
-### Load useful functions to your shell script
+### Load functions in 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
+omb misc export-shell-func > omb-func.sh && source omb-func.sh && rm omb-func.sh
 ```
 
-This will load some useful functions to your shell script, for example, `checkpoint` function.
+This will load some functions to your shell script, for example, `checkpoint`.
 
 ### Generate files from different combinations of parameters
 
@@ -66,8 +69,8 @@ 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
+read_data @DATA_FILE
+velocity all create @TEMP @RANDOM
 run 1000
 EOF
 

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

@@ -10,27 +10,29 @@ A toolkit to manipulate batch tasks with command line. Designed for scientific c
 * `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)
+* `omb misc`: miscellaneous commands
 
 ## Install
 ```bash
 pip install oh-my-batch
 ```
 
+## 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)
+* [Run multiple line shell script with ssh](https://stackoverflow.com/a/32082912/3099733)
+
 ## Use cases
-### Load useful functions to your shell script
+### Load functions in 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
+omb misc export-shell-func > omb-func.sh && source omb-func.sh && rm omb-func.sh
 ```
 
-This will load some useful functions to your shell script, for example, `checkpoint` function.
+This will load some functions to your shell script, for example, `checkpoint`.
 
 ### Generate files from different combinations of parameters
 
@@ -48,8 +50,8 @@ 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
+read_data @DATA_FILE
+velocity all create @TEMP @RANDOM
 run 1000
 EOF
 

{oh_my_batch-0.2.1 → oh_my_batch-0.2.3}/oh_my_batch/combo.py

@@ -1,9 +1,10 @@
 from itertools import product
 from string import Template
 import random
+import json
 import os
 
-from .util import expand_globs, mode_translate
+from .util import expand_globs, mode_translate, ensure_dir
 
 class ComboMaker:
 
@@ -67,7 +68,7 @@ class ComboMaker:
         self.add_var(key, *args, broadcast=broadcast)
         return self
 
-    def add_files(self, key: str, *path: str, broadcast=False, abs=False):
+    def add_files(self, key: str, *path: str, broadcast=False, abs=False, raise_invalid=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,
@@ -78,8 +79,9 @@ class ComboMaker:
         :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
+        :param raise_invalid: If True, will raise error if no file found for a glob pattern
         """
-        args = expand_globs(path, raise_invalid=True)
+        args = expand_globs(path, raise_invalid=raise_invalid)
         if not args:
             raise ValueError(f"No files found for {path}")
         if abs:
@@ -87,7 +89,8 @@ class ComboMaker:
         self.add_var(key, *args, broadcast=broadcast)
         return self
 
-    def add_files_as_one(self, key: str, path: str, broadcast=False, sep=' ', abs=False):
+    def add_files_as_one(self, key: str, *path: str, broadcast=False, format=None,
+                         sep=' ', abs=False, raise_invalid=False):
         """
         Add a variable with files by glob pattern as one string
         Unlike add_files, this function joins the files with a delimiter.
@@ -98,15 +101,25 @@ class ComboMaker:
         :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 format: the way to format the files, can be None, 'json-list','json-item'
         :param sep: Separator to join files
         :param abs: If True, path will be turned into absolute path
+        :param raise_invalid: If True, will raise error if no file found for a glob pattern
         """
-        args = expand_globs(path, raise_invalid=True)
+        args = expand_globs(path, raise_invalid=raise_invalid)
         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)
+        if format is None:
+            value = sep.join(args)
+        elif format == 'json-list':
+            value = json.dumps(args)
+        elif format == 'json-item':
+            value = json.dumps(args).strip('[]')
+        else:
+            raise ValueError(f"Invalid format: {format}")
+        self.add_var(key, value, broadcast=broadcast)
         return self
 
     def add_var(self, key: str, *args, broadcast=False):
@@ -148,19 +161,22 @@ class ComboMaker:
                 raise ValueError(f"Variable {key} not found")
         return self
 
-    def make_files(self, template: str, dest: str, delimiter='$', mode=None):
+    def make_files(self, template: str, dest: str, delimiter='@', mode=None, encoding='utf-8'):
         """
-        Make files from template
+        Make files from template against each combo
         The template file can include variables with delimiter.
-        For example, if delimiter is '$', then the template file can include $var1, $var2, ...
+        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, ...
+        For example, if dest is 'output/{i}-{TEMP}.txt', 
+        then files are saved as output/0-300K.txt, output/1-400K.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 $$, @, ...
+        :param delimiter: Delimiter for variables in template, default is '@', as '$' is popular in shell scripts
+        can be changed to other character, e.g $, $$, ...
+        :param mode: File mode, e.g. 755, 644, ...
+        :param encoding: File encoding
         """
         _delimiter = delimiter
 
@@ -173,12 +189,34 @@ class ComboMaker:
                 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:
+            ensure_dir(_dest)
+            with open(_dest, 'w', encoding=encoding) as f:
                 f.write(text)
             if mode is not None:
                 os.chmod(_dest, mode_translate(str(mode)))
         return self
+    
+    def print(self, *line: str, file: str = '', mode=None, encoding='utf-8'):
+        """
+        Print lines to a file against each combo
+
+        :param line: Lines to print, can include format style variables, e.g. {i}, {i:03d}, {TEMP}
+        :param file: File to save the output, if not provided, print to stdout
+        """
+        combos = self._make_combos()
+        out_lines = []
+        for i, combo in enumerate(combos):
+            out_lines.extend(l.format(i=i, **combo) for l in line)
+        out = '\n'.join(out_lines)
+        if file:
+            ensure_dir(file)
+            with open(file, 'w', encoding=encoding) as f:
+                f.write(out)
+            if mode is not None:
+                os.chmod(file, mode_translate(str(mode)))
+        else:
+            print(out)
+        return self
 
     def done(self):
         """
@@ -187,7 +225,7 @@ class ComboMaker:
         pass
 
     def _make_combos(self):
-        if not self._product_vars and not self._broadcast_vars:
+        if not self._product_vars:
             return self._combos
         keys = self._product_vars.keys()
         values_list = product(*self._product_vars.values())

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

@@ -68,7 +68,7 @@ class BaseJobManager:
 
         current = time.time()
         while True:
-            self._update_jobs(jobs, max_tries, opts)
+            jobs = self._update_jobs(jobs, max_tries, opts)
             if recovery:
                 ensure_dir(recovery)
                 with open(recovery, 'w', encoding='utf-8') as f:
@@ -77,7 +77,7 @@ class BaseJobManager:
             if not wait:
                 break
 
-            # stop if all jobs are terminal and not job to be submitted
+            # stop if all jobs are terminal and no job to be submitted
             if (all(JobState.is_terminal(j['state']) for j in jobs) and
                     not any(should_submit(j, max_tries) for j in jobs)):
                 break
@@ -88,7 +88,30 @@ class BaseJobManager:
 
             time.sleep(interval)
 
+
     def _update_jobs(self, jobs: List[dict], max_tries: int, submit_opts: str):
+        jobs = self._update_state(jobs)
+
+        # 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
+                job_id = self._submit_job(job, submit_opts)
+                if job_id:
+                    job['state'] = JobState.PENDING
+                    job['id'] = job_id
+                    logger.info('Job %s submitted', job['id'])
+                else:
+                    job['state'] = JobState.FAILED
+                    logger.error('Failed to submit job %s', job['script'])
+        return jobs
+
+    def _update_state(self, jobs: List[dict]):
+        raise NotImplementedError
+    
+    def _submit_job(self, job: dict, submit_opts: str):
         raise NotImplementedError
 
 
@@ -97,8 +120,7 @@ class Slurm(BaseJobManager):
         self._sbatch_bin = sbatch
         self._sacct_bin = sacct
 
-    def _update_jobs(self, jobs: List[dict], max_tries: int, submit_opts: str):
-        # query job status
+    def _update_state(self, jobs: List[dict]):
         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)}'
@@ -106,8 +128,9 @@ class Slurm(BaseJobManager):
             if cp.returncode != 0:
                 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'))
+            out = cp.stdout.decode('utf-8')
+            logger.info('Job status:\n%s', out)
+            new_state = parse_csv(out)
         else:
             new_state = []
 
@@ -122,23 +145,19 @@ class Slurm(BaseJobManager):
                     break
             else:
                 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', log_cp(cp))
-                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'])
+        return jobs 
+
+    def _submit_job(self, job:dict, submit_opts:str):
+        submit_cmd = f'{self._sbatch_bin} {submit_opts} {job["script"]}'
+        cp = shell_run(submit_cmd)
+        if cp.returncode != 0:
+            logger.error('Failed to submit job: %s', log_cp(cp))
+            return ''
+        out = cp.stdout.decode('utf-8')
+        job_id  = self._parse_job_id(out)
+        if not job_id:
+            raise ValueError(f'Unexpected sbatch output: {out}')
+        return job_id
 
     def _map_state(self, state: str):
         if state.startswith('CANCELLED'):

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

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