pspipe 0.12__tar.gz → 0.14__tar.gz

{pspipe-0.12 → pspipe-0.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pspipe
-Version: 0.12
+Version: 0.14
 Summary: Power spectra pipeline for Cosmic Dawn, Epoch of Reionization radio interferometric experiments
 License-Expression: GPL-3.0-or-later
 License-File: LICENSE
@@ -18,7 +18,7 @@ Provides-Extra: docs
 Provides-Extra: ml-gpr
 Requires-Dist: astropy (>6.0)
 Requires-Dist: click (>=8.0)
-Requires-Dist: libpipe (>=0.8)
+Requires-Dist: libpipe (>=0.9)
 Requires-Dist: myst-parser (>=2) ; extra == "docs"
 Requires-Dist: ps-eor (>=1.0)
 Requires-Dist: ps-eor[ml-gpr] (>=1.0) ; extra == "ml-gpr"

{pspipe-0.12 → pspipe-0.14}/pspipe/settings.py

@@ -5,13 +5,20 @@
 import os
 
 import astropy.units as u
-from libpipe import futils
+from libpipe import futils, worker
 from libpipe.settings import BaseSettings
 
 # Lazy loaded:
 # - ps_eor.datacube in Settings.validate_vis_cube()
 from . import utils
 
+
+def worker_resources(s_worker):
+    '''worker.resources as {category: {'cpus': .., 'mem': ..}}, or None.'''
+    resources = s_worker.get('resources') or {}
+    return {k: dict(v) for k, v in resources.items()} or None
+
+
 TEMPLATE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'templates')
 
 TEMPLATE_NAMES = ['a12', 'nenufar', 'hba', 'alo']
@@ -51,9 +58,19 @@ class Settings(BaseSettings):
     def validate_worker(self, val_fct):
         val_fct(isinstance(self.worker.max_concurrent, int), 'max_concurrent')
         val_fct(not self.worker.env_file or futils.is_file(self.worker.env_file), 'env_file')
-        val_fct(self.worker.get('max_time') is None
+        val_fct(not self.worker.get('max_time')
                 or (isinstance(self.worker.max_time, (int, float)) and self.worker.max_time > 0),
                 'max_time')
+        slurm_nodes = self.worker.get('slurm_nodes') or 0
+        val_fct(isinstance(slurm_nodes, int) and slurm_nodes >= 0, 'slurm_nodes')
+        try:
+            resources = worker.resolve_resources(worker_resources(self.worker))
+        except (ValueError, TypeError):
+            resources = None
+        val_fct(resources is not None, 'resources')
+        categories = self.worker.get('categories') or {}
+        val_fct(resources is None or all(c in resources for c in categories.values()),
+                'categories')
 
     def validate_merge_ms(self, val_fct):
         val_fct(isinstance(self.merge_ms.apply_aoflagger, bool), 'apply_aoflagger')

{pspipe-0.12 → pspipe-0.14}/pspipe/tasks.py

@@ -14,7 +14,7 @@ import astropy.io.fits as pf
 import numpy as np
 from libpipe import futils, worker
 
-from . import database, utils
+from . import database, settings, utils
 
 CWD = os.getcwd()
 
@@ -48,6 +48,11 @@ def get_target_host(in_file, worker_settings):
 
 class AbstractSubTask:
 
+    # Task type (worker.categories key, usage summary) and resource category as a job step of an
+    # allocation (worker.slurm_nodes): small, medium, high_mem, cpu or heavy
+    label = None
+    category = 'heavy'
+
     def __init__(self, settings_obj, exec_name):
         self.settings = settings_obj
         self.exec_name = exec_name
@@ -74,23 +79,24 @@ class AbstractTask:
         self.name = name
         self.settings = settings_obj.copy()
         self.settings.worker.update(worker_kargs)
-        s_worker = self.settings.worker
 
         if not self.settings.validate('worker'):
             raise ValueError('Invalid worker settings.')
 
-        self.worker = worker.get_worker_pool(
+        self.worker = self._worker_pool(name)
+
+    def _worker_pool(self, name):
+        s_worker = self.settings.worker
+        # 0 and '' (the defaults) leave max_time unset and slurm_options from the environment
+        return worker.get_worker_pool(
                 name, nodes=s_worker.nodes, max_concurrent=s_worker.max_concurrent,
-                env_file=s_worker.env_file, max_time=s_worker.get('max_time'),
-                slurm_options=s_worker.get('slurm_options'))
+                env_file=s_worker.env_file, max_time=s_worker.get('max_time') or None,
+                slurm_options=s_worker.get('slurm_options') or None,
+                slurm_nodes=s_worker.get('slurm_nodes') or None,
+                resources=settings.worker_resources(s_worker))
 
     def init_for_new_task(self, post_task_name):
-        s_worker = self.settings.worker
-        self.worker = worker.get_worker_pool(post_task_name, nodes=s_worker.nodes,
-                                             max_concurrent=s_worker.max_concurrent,
-                                             env_file=s_worker.env_file,
-                                             max_time=s_worker.get('max_time'),
-                                             slurm_options=s_worker.get('slurm_options'))
+        self.worker = self._worker_pool(post_task_name)
 
     def add_sub_task(self, sub_task, run_on_host=None):
         log_file = sub_task.get_log_file()
@@ -103,8 +109,10 @@ class AbstractTask:
         cmd = sub_task.build_command()
         if self.settings.debug:
             print(cmd)
+        label = sub_task.label or sub_task.exec_name
+        category = (self.settings.worker.get('categories') or {}).get(label, sub_task.category)
         self.worker.add(cmd, output_file=log_file, done_callback=sub_task.done_callback,
-                        run_on_host=run_on_host)
+                        run_on_host=run_on_host, category=category, label=label)
 
     def run(self, obs_ids):
         raise NotImplementedError()
@@ -128,6 +136,9 @@ class AbstractTask:
 
 class ImagingSubTask(AbstractSubTask):
 
+    label = 'image'
+    category = 'heavy'
+
     def __init__(self, settings_obj, obs_id, inp_ms_files, even=False, odd=False, index=0,
                  even_odd=False):
         AbstractSubTask.__init__(self, settings_obj, settings_obj.image.imager)
@@ -180,10 +191,8 @@ class ImagingSubTask(AbstractSubTask):
                       odd=self.odd if odd is None else odd, img_name=img_name)
         meta.save()
 
-    def done_callback(self):
-        for file in glob.glob(f'{self.get_fits_prefix()}*-MFS-*.fits'):
-            os.remove(file)
-
+    def fits_keywords(self):
+        '''Integration time, total time and channel width of the images, for their headers.'''
         img = self.settings.image
 
         if img.time_start_index != img.time_end_index:
@@ -193,6 +202,15 @@ class ImagingSubTask(AbstractSubTask):
         else:
             total_time = self.obs_meta.total_time
 
+        return {'PEINTTIM': int(self.obs_meta.int_time), 'PETOTTIM': int(total_time),
+                'PECHWIDT': int(self.obs_meta.chan_width)}
+
+    def done_callback(self):
+        for file in glob.glob(f'{self.get_fits_prefix()}*-MFS-*.fits'):
+            os.remove(file)
+
+        img = self.settings.image
+
         if self.even_odd:
             parts = [('-even', True, False), ('-odd', False, True)]
         else:
@@ -209,10 +227,12 @@ class ImagingSubTask(AbstractSubTask):
                 files = glob.glob(f'{prefix}*.fits')
                 img_name = 'image'
 
-            for file in files:
-                pf.setval(file, 'PEINTTIM', value=int(self.obs_meta.int_time))
-                pf.setval(file, 'PETOTTIM', value=int(total_time))
-                pf.setval(file, 'PECHWIDT', value=int(self.obs_meta.chan_width))
+            # psimager writes the keywords with the files (-fits-keyword)
+            if img.imager != 'psimager':
+                keywords = self.fits_keywords()
+                for file in files:
+                    with pf.open(file, mode='update') as hdul:
+                        hdul[0].header.update(keywords)
 
             if not img.lst_bins:
                 self.register_files(files, img_name=img_name, even=even, odd=odd)
@@ -267,6 +287,9 @@ class ImagingSubTask(AbstractSubTask):
 
             args.append(f'-temp-dir {self.temp_dir}')
 
+        if img.imager == 'psimager':
+            args.append(' '.join(f'-fits-keyword {k}={v}' for k, v in self.fits_keywords().items()))
+
         if self.even_odd:
             args.append('-even-odd')
         elif self.even:
@@ -368,6 +391,9 @@ class ImagingTask(AbstractTask):
 
 class GenVisCubeSubTask(AbstractSubTask):
 
+    label = 'gen_vis_cube'
+    category = 'small'
+
     def __init__(self, settings_obj, obs_id, stokes, even=False, odd=False, lst_idx=None):
         AbstractSubTask.__init__(self, settings_obj, 'pstool gen_vis_cube')
         self.obs_id = obs_id
@@ -423,6 +449,9 @@ class GenVisCubeSubTask(AbstractSubTask):
 
 class EvenOddToSumDiffSubTask(AbstractSubTask):
 
+    label = 'even_odd_to_sum_diff'
+    category = 'small'
+
     def __init__(self, settings_obj, obs_id, stokes, lst_idx=None):
         AbstractSubTask.__init__(self, settings_obj, 'pstool even_odd_to_sum_diff')
         self.obs_id = obs_id
@@ -523,6 +552,9 @@ class GenVisCubeTask(AbstractTask):
 
 
 class ConcatMSSubTask(AbstractSubTask):
+
+    label = 'concat_ms'
+    category = 'medium'
     """
     Subtask that calls the CLI 'mstool concat <MS...> <OUT.MS>'
     to concatenate MS files in time.
@@ -604,6 +636,9 @@ class ConcatMSTask(AbstractTask):
 
 class MergeMSSubTask(AbstractSubTask):
 
+    label = 'merge_ms'
+    category = 'heavy'
+
     def __init__(self, settings_obj, ms_files, obs_id_out):
         AbstractSubTask.__init__(self, settings_obj, 'DP3')
         self.db = database.Database(self.settings)
@@ -718,6 +753,9 @@ class MergeMSTask(AbstractTask):
 
 class AOFlaggerSubTask(AbstractSubTask):
 
+    label = 'ao_flag'
+    category = 'medium'
+
     def __init__(self, settings_obj, ms_file):
         AbstractSubTask.__init__(self, settings_obj, 'DP3')
         self.ms_file = ms_file
@@ -771,6 +809,9 @@ class AOFlaggerTask(AbstractTask):
 
 class DP3SubTask(AbstractSubTask):
 
+    label = 'dp3'
+    category = 'heavy'
+
     def __init__(self, settings_obj, ms_file):
         AbstractSubTask.__init__(self, settings_obj, 'DP3')
         self.ms_file = ms_file
@@ -818,6 +859,9 @@ class DP3Task(AbstractTask):
 
 class MakePSSubTask(AbstractSubTask):
 
+    label = 'make_ps'
+    category = 'small'
+
     def __init__(self, settings_obj, obs_id, lst_idx=None):
         AbstractSubTask.__init__(self, settings_obj, 'pstool make_ps')
         self.obs_id = obs_id
@@ -881,6 +925,9 @@ class MakePSTask(AbstractTask):
 
 class RunGPRSubTask(AbstractSubTask):
 
+    label = 'run_gpr'
+    category = 'cpu'
+
     def __init__(self, settings_obj, obs_id, lst_idx=None):
         AbstractSubTask.__init__(self, settings_obj, 'pstool run_gpr')
         self.obs_id = obs_id
@@ -951,6 +998,9 @@ class RunGPRTask(AbstractTask):
 
 class RunMLGPRSubTask(AbstractSubTask):
 
+    label = 'run_ml_gpr'
+    category = 'cpu'
+
     def __init__(self, settings_obj, obs_id, lst_idx=None):
         AbstractSubTask.__init__(self, settings_obj, 'pstool run_ml_gpr')
         self.obs_id = obs_id
@@ -1024,6 +1074,9 @@ class RunMLGPRTask(AbstractTask):
 
 class RunMLGPRInjSubTask(AbstractSubTask):
 
+    label = 'run_ml_gpr_inj'
+    category = 'cpu'
+
     def __init__(self, settings_obj, obs_id, var_21_over_noise, x1, x2, lst_idx=None):
         AbstractSubTask.__init__(self, settings_obj, 'pstool run_ml_gpr_inj')
         self.obs_id = obs_id
@@ -1230,6 +1283,9 @@ class CombineTask(AbstractTask):
 
 class RunFlaggerSubTask(AbstractSubTask):
 
+    label = 'run_flagger'
+    category = 'small'
+
     def __init__(self, settings_obj, obs_id):
         AbstractSubTask.__init__(self, settings_obj, 'pstool run_flagger')
         self.vis_rev = database.VisRevision(self.settings)
@@ -1244,6 +1300,9 @@ class RunFlaggerSubTask(AbstractSubTask):
 
 class VisToSphSubTask(AbstractSubTask):
 
+    label = 'vis_to_sph'
+    category = 'small'
+
     def __init__(self, settings_obj, obs_id_out, stokes, h5_input, h5_out, h5_flag):
         AbstractSubTask.__init__(self, settings_obj, 'pstool vis_to_sph')
         self.vis_rev = database.VisRevision(self.settings)
@@ -1318,6 +1377,9 @@ class VisToSphTask(AbstractTask):
 
 class ApplyFlagSubTask(AbstractSubTask):
 
+    label = 'apply_flag'
+    category = 'small'
+
     def __init__(self, settings_obj, ms_file, h5_file):
         AbstractSubTask.__init__(self, settings_obj, 'flagtool restore')
         self.ms_file = ms_file
@@ -1405,6 +1467,9 @@ class RestoreFlagTask(AbstractTask):
 
 class DelayFlaggerSubTask(AbstractSubTask):
 
+    label = 'delay_flagger'
+    category = 'medium'
+
     def __init__(self, settings_obj, ms_file):
         AbstractSubTask.__init__(self, settings_obj, 'flagtool delay_flagger')
         self.ms_file = ms_file
@@ -1462,6 +1527,9 @@ class DelayFlagger(AbstractTask):
 
 class VisFlaggerSubTask(AbstractSubTask):
 
+    label = 'vis_flagger'
+    category = 'medium'
+
     def __init__(self, settings_obj, ms_file):
         AbstractSubTask.__init__(self, settings_obj, 'flagtool vis_flagger')
         self.ms_file = ms_file
@@ -1513,6 +1581,9 @@ class VisFlagger(AbstractTask):
 
 class SSINSSubTask(AbstractSubTask):
 
+    label = 'ssins'
+    category = 'medium'
+
     def __init__(self, settings_obj, ms_file):
         AbstractSubTask.__init__(self, settings_obj, 'flagtool ssins')
         self.ms_file = ms_file

{pspipe-0.12 → pspipe-0.14}/pspipe/templates/default_settings.toml

@@ -10,6 +10,13 @@ max_concurrent = 4
 env_file = ''
 run_on_file_host = false
 run_on_file_host_pattern = ''
+slurm_options = ''
+max_time = 0
+slurm_nodes = 0
+
+[worker.resources]
+
+[worker.categories]
 
 [merge_ms]
 apply_aoflagger = true

{pspipe-0.12 → pspipe-0.14}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "pspipe"
-version = "0.12"
+version = "0.14"
 description = "Power spectra pipeline for Cosmic Dawn, Epoch of Reionization radio interferometric experiments"
 readme = "README.md"
 license = "GPL-3.0-or-later"
@@ -20,7 +20,7 @@ dependencies = [
     "tables>=3.2",
     "tabulate>=0.8",
     "ps-eor>=1.0",
-    "libpipe>=0.8",
+    "libpipe>=0.9",
     "python-casacore>=3.0",
 ]