oh-my-batch 0.2.2__py3-none-any.whl → 0.2.3__py3-none-any.whl

oh_my_batch/combo.py

@@ -4,7 +4,7 @@ import random
 import json
 import os
 
-from .util import expand_globs, mode_translate
+from .util import expand_globs, mode_translate, ensure_dir
 
 class ComboMaker:
 
@@ -161,9 +161,9 @@ 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, ...
 
@@ -175,6 +175,8 @@ class ComboMaker:
         :param dest: Path pattern to destination file
         :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
 
@@ -187,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):
         """
@@ -201,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/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.2.dist-info → oh_my_batch-0.2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: oh-my-batch
-Version: 0.2.2
+Version: 0.2.3
 Summary: 
 License: GPL
 Author: weihong.xu

{oh_my_batch-0.2.2.dist-info → oh_my_batch-0.2.3.dist-info}/RECORD

@@ -4,12 +4,12 @@ oh_my_batch/assets/__init__.py,sha256=Exub46UbQaz2V2eXpQeiVfnThQpXaNeuyjlGY6gBSZ
 oh_my_batch/assets/functions.sh,sha256=dvPGpOKz4CyUSlE5IAewQc8HUrQJuaZ0j-WZIuxB3Tg,1002
 oh_my_batch/batch.py,sha256=6qnaXEVyA493heGzzbCrdZXCcnYk8zgl7WP0rmo7KlU,3690
 oh_my_batch/cli.py,sha256=jMSM_orAamNoEP7jPsbp0UYPX-UDY1-kHSHO31zELHs,548
-oh_my_batch/combo.py,sha256=UQYyZAGFi-siUN13XVcpQwuqSZ0ZYPJl3hdVCORECng,8407
-oh_my_batch/job.py,sha256=8kZnWtvpr1rAl4tc9I_Vlhi-T0o3rh4RQZZgMNBCGho,5800
+oh_my_batch/combo.py,sha256=-qgqT7XSvEeyefMPD3-HuX-aozefjSKhbZbPm1pm_18,9324
+oh_my_batch/job.py,sha256=q-1WxZQxQNqDV4xZXz5V9RQQ787rhdX18tQWTNBKBjQ,6268
 oh_my_batch/misc.py,sha256=G_iOovRCrShBJJCc82QLN0CvMqW4adOefEoY1GedEiw,452
 oh_my_batch/util.py,sha256=okg_kY8dJouyJ2BYCXRl7bxDUAtNH6GLh2UjXRnkoW0,2385
-oh_my_batch-0.2.2.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-oh_my_batch-0.2.2.dist-info/METADATA,sha256=7Y9HPRolb4gljmDpHLaqkfdh1bfPYR4WdzFuh8DeKVw,5496
-oh_my_batch-0.2.2.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-oh_my_batch-0.2.2.dist-info/entry_points.txt,sha256=ZY2GutSoNjjSyJ4qO2pTeseKUFgoTYdvmgkuZZkwi68,77
-oh_my_batch-0.2.2.dist-info/RECORD,,
+oh_my_batch-0.2.3.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+oh_my_batch-0.2.3.dist-info/METADATA,sha256=sRJJinfObbY4AVBpO1i86Ds2ognYtQ7upeJ4Q_q3WG8,5496
+oh_my_batch-0.2.3.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+oh_my_batch-0.2.3.dist-info/entry_points.txt,sha256=ZY2GutSoNjjSyJ4qO2pTeseKUFgoTYdvmgkuZZkwi68,77
+oh_my_batch-0.2.3.dist-info/RECORD,,