vortex-nwp 2.0.0b1__py3-none-any.whl

vortex/nwp/tools/odb.py

@@ -0,0 +1,548 @@
+"""
+Common interest classes to help setup the ODB software environment.
+"""
+
+import re
+
+from bronx.fancies import loggers
+from bronx.stdtypes import date as bdate
+import footprints
+
+from vortex.algo.components import AlgoComponentDecoMixin, AlgoComponentError, algo_component_deco_mixin_autodoc
+from vortex.layout.dataflow import intent
+
+from ..syntax.stdattrs import ArpIfsSimplifiedCycle
+
+#: No automatic export
+__all__ = []
+
+logger = loggers.getLogger(__name__)
+
+
+class TimeSlots:
+    """Handling of assimilation time slots."""
+
+    def __init__(self, nslot=7, start='-PT3H', window='PT6H', chunk=None, center=True):
+        if isinstance(nslot, str):
+            info = [x.strip() for x in nslot.split('/')]
+            nslot = info[0]
+            if len(info) > 1:
+                start = info[1]
+            if len(info) > 2:
+                window = info[2]
+            if len(info) > 3:
+                if re.match('^regular', info[3]):
+                    center = False
+                else:
+                    chunk = info[3]
+        self.nslot = int(nslot)
+        self.center = center if self.nslot > 1 else False
+        self.start = bdate.Period(start)
+        self.window = bdate.Period(window)
+        if chunk is None:
+            cslot = self.nslot - 1 if self.center else self.nslot
+            chunk = 'PT' + str((self.window.length // max(1, cslot)) // 60) + 'M'
+        self.chunk = self.window if self.nslot < 2 else bdate.Period(chunk)
+
+    def __eq__(self, other):
+        if isinstance(other, str):
+            try:
+                other = TimeSlots(other)
+            except ValueError:
+                pass
+        return (isinstance(other, TimeSlots) and
+                self.nslot == other.nslot and self.center == other.center and
+                self.start == other.start and self.window == other.window and
+                self.chunk == other.chunk)
+
+    def __str__(self):
+        chunky = self.chunk.isoformat() if self.center else 'regular'
+        return '{0.nslot:d}/{1:s}/{2:s}/{3:s}'.format(self,
+                                                      self.start.isoformat(),
+                                                      self.window.isoformat(),
+                                                      chunky)
+
+    def __repr__(self, *args, **kwargs):
+        return super().__repr__()[:-1] + ' | {!s}>'.format(self)
+
+    def as_slots(self):
+        """Return a list of slots in seconds."""
+        if self.center:
+            slots = [self.chunk.length, ] * self.nslot
+            nb = self.window.length // self.chunk.length
+            if nb != self.nslot:
+                slots[0] = slots[-1] = self.chunk.length // 2
+        else:
+            islot = self.window.length // self.nslot
+            slots = [islot, ] * self.nslot
+        return slots
+
+    def as_centers_fromstart(self):
+        """Return time slots centers as a list of Period objects."""
+        slots = self.as_slots()
+        fromstart = []
+        acc = 0
+        for i in range(len(slots)):
+            fromstart.append(acc + slots[i] / 2)
+            acc += slots[i]
+        if self.center and (self.window.length // self.chunk.length != self.nslot):
+            fromstart[0] = 0
+            fromstart[-1] = self.window.length
+        return [bdate.Period(seconds=t) for t in fromstart]
+
+    def as_bounds(self, date):
+        """Return time slots as a list of compact date values."""
+        date = bdate.Date(date)
+        boundlist = [date + self.start, ]
+        for x in self.as_slots():
+            boundlist.append(boundlist[-1] + x)
+        boundlist = [x.compact() for x in boundlist]
+        return boundlist
+
+    @property
+    def leftmargin(self):
+        """Return length in minutes from left margin of the window."""
+        return int(self.start.total_seconds()) // 60
+
+    @property
+    def rightmargin(self):
+        """Return length in minutes from rigth margin of the window."""
+        return int((self.start + self.window).total_seconds()) // 60
+
+    def as_environment(self):
+        """Return a dictionary of ready-to-export variables that describe the timeslots."""
+        thelen = self.chunk.length // 60 if self.center and self.nslot > 1 else 0
+        return dict(BATOR_WINDOW_LEN=self.window.length // 60,
+                    BATOR_WINDOW_SHIFT=int(self.start.total_seconds()) // 60,
+                    BATOR_SLOT_LEN=thelen, BATOR_CENTER_LEN=thelen)
+
+    def as_file(self, date, filename):
+        """Fill the specified ``filename`` wih the current list of time slots at this ``date``."""
+        with open(filename, 'w') as fd:
+            for x in self.as_bounds(date):
+                fd.write(str(x) + '\n')
+            nbx = fd.tell()
+        return nbx
+
+
+class OdbDriver:
+    """A dedicated class for handling some ODB settings."""
+
+    def __init__(self, cycle, sh=None, env=None):
+        """
+        A quite challenging initialisation since cycle, sh, env and target
+        should be provided...
+        """
+        self.cycle = cycle
+        self.sh = sh
+        if self.sh is None:
+            logger.critical('%s created with a proper shell access [%s]', self.__class__, self)
+        self.env = env
+        if self.env is None:
+            logger.critical('%s created with a proper environment access [%s]', self.__class__, self)
+
+    def setup(self, date, npool=1, nslot=1, iomethod=1, layout='ecma'):
+        """Setup given environment with default ODB env variables."""
+
+        logger.info('ODB: generic setup called.'),
+        self.env.update(
+            ODB_CMA=layout.upper(),
+            ODB_IO_METHOD=iomethod,
+        )
+
+        self.env.default(
+            ODB_DEBUG=0,
+            ODB_CTX_DEBUG=0,
+            ODB_REPRODUCIBLE_SEQNO=4,
+            ODB_STATIC_LINKING=1,
+            ODB_ANALYSIS_DATE=date.ymd,
+            ODB_ANALYSIS_TIME=date.hm + '00',
+            TO_ODB_ECMWF=0,
+            TO_ODB_SWAPOUT=0,
+        )
+
+        if iomethod == 4:
+            self.env.default(
+                ODB_IO_GRPSIZE=npool,
+                ODB_IO_FILESIZE=128,
+            )
+
+        if self.sh.path.exists('IOASSIGN'):
+            self.env.default(
+                IOASSIGN=self.sh.path.abspath('IOASSIGN'),
+            )
+
+    def force_overwrite_method(self):
+        """Force ODB_OVERWRITE_METHOD if necessary."""
+        if not int(self.env.get('ODB_OVERWRITE_METHOD', 0)):
+            logger.info('ODB: Some input ODB databases are read-only. Setting ODB_OVERWRITE_METHOD to 1.')
+            self.env.ODB_OVERWRITE_METHOD = 1
+
+    def _process_layout_dbpath(self, layout, dbpath=None):
+        """Normalise **layout** and **dbpath**."""
+        layout = layout.upper()
+        thispwd = self.sh.path.abspath(self.sh.getcwd())
+        if dbpath is None:
+            dbpath = self.sh.path.join(thispwd, layout)
+        return layout, dbpath, thispwd
+
+    def fix_db_path(self, layout, dbpath=None, env=None):
+        """Setup the path to the **layout** database."""
+        if env is None:
+            env = self.env
+        layout, dbpath, _ = self._process_layout_dbpath(layout, dbpath)
+        logger.info('ODB: Fix %s path: %s', layout, dbpath)
+        env['ODB_SRCPATH_{:s}'.format(layout)] = dbpath
+        env['ODB_DATAPATH_{:s}'.format(layout)] = dbpath
+
+    @property
+    def _default_iocreate_path(self):
+        """The location to the default create_ioassign utility."""
+        return self.sh.path.join(
+            from_config(section="nwp-tools", key="odb"),
+            get_from_config_w_default(
+                section="nwp-tools",
+                key="iocreate_cmd",
+                default="create_ioassign"
+            ),
+        )
+
+    @property
+    def _default_iomerge_path(self):
+        """The location to the default merge_ioassign utility."""
+        return self.sh.path.join(
+            from_config(section="nwp-tools", key="odb"),
+            get_from_config_w_default(
+                section="nwp-tools",
+                key="iomerge_cmd",
+                default="merge_ioassign"
+            ),
+        )
+
+    def ioassign_create(self, ioassign='ioassign.x', npool=1, layout='ecma',
+                        dbpath=None, iocreate_path=None):
+        """Build IO-Assign table."""
+        layout, dbpath, _ = self._process_layout_dbpath(layout, dbpath)
+        if iocreate_path is None:
+            iocreate_path = self._default_iocreate_path
+        ioassign = self.sh.path.abspath(ioassign)
+        self.sh.xperm(ioassign, force=True)
+        self.sh.mkdir(dbpath)
+        with self.env.clone() as lenv:
+            lenv['ODB_IOASSIGN_BINARY'] = ioassign
+            self.fix_db_path(layout, dbpath, env=lenv)
+            self.sh.spawn([iocreate_path,
+                           '-d' + dbpath,
+                           '-l' + layout,
+                           '-n' + str(npool)],
+                          output=False)
+        return dbpath
+
+    def ioassign_merge(self, ioassign='ioassign.x', layout='ecma', odbnames=None,
+                       dbpath=None, iomerge_path=None, iocreate_path=None):
+        """Build IO-Assign table."""
+        layout, dbpath, thispwd = self._process_layout_dbpath(layout, dbpath)
+        if iomerge_path is None:
+            iomerge_path = self._default_iomerge_path
+        if iocreate_path is None:
+            iocreate_path = self._default_iocreate_path
+        iocmd = [iomerge_path]
+        ioassign = self.sh.path.abspath(ioassign)
+        self.sh.xperm(ioassign, force=True)
+        with self.sh.cdcontext(dbpath, create=True):
+            iocmd.extend(['-d', thispwd])
+            for dbname in odbnames:
+                iocmd.extend(['-t', dbname])
+            with self.env.clone() as lenv:
+                lenv['ODB_IOASSIGN_BINARY'] = ioassign
+                if 'ODB_IOCREATE_COMMAND' not in lenv:
+                    lenv['ODB_IOCREATE_COMMAND'] = iocreate_path
+                self.fix_db_path(layout, dbpath, env=lenv)
+                self.sh.spawn(iocmd, output=False)
+        return dbpath
+
+    def _ioassign_process(self, dbpaths, wmode):
+        with open('IOASSIGN', wmode) as fhgather:
+            for dbpath in dbpaths:
+                with open(self.sh.path.join(dbpath, 'IOASSIGN')) as fhlay:
+                    for line in fhlay:
+                        fhgather.write(line)
+
+    def ioassign_gather(self, *dbpaths):
+        """Gather IOASSIGN data from **dbpaths** databases and create a global IOASSIGN file."""
+        logger.info('ODB: creating a global IOASSIGN file from: %s',
+                    ','.join([self.sh.path.basename(db) for db in dbpaths]))
+        self._ioassign_process(dbpaths, 'w')
+
+    def ioassign_append(self, *dbpaths):
+        """Append IOASSIGN data from **dbpaths** databases into the global IOASSIGN file."""
+        logger.info('ODB: extending the IOASSIGN file with: %s',
+                    ','.join([self.sh.path.basename(db) for db in dbpaths]))
+        self._ioassign_process(dbpaths, 'a')
+
+    def shuffle_setup(self, slots, mergedirect=False, ccmadirect=False):
+        """Setup environment variables to control ODB shuffle behaviour.
+
+        :param bool mergedirect: Run the shuffle procedure on the input database.
+        :param bool ccmadirect: Create a CCMA database at the end of the run.
+        """
+        logger.info('ODB: shuffle_setup: mergedirect=%s, ccmadirect=%s',
+                    str(mergedirect), str(ccmadirect))
+        if mergedirect or ccmadirect:
+            self.env.update(
+                ODB_CCMA_TSLOTS=slots.nslot,
+            )
+            self.env.default(
+                ODB_CCMA_LEFT_MARGIN=slots.leftmargin,
+                ODB_CCMA_RIGHT_MARGIN=slots.rightmargin,
+            )
+            if mergedirect:
+                self.env.default(ODB_MERGEODB_DIRECT=1)
+            if ccmadirect:
+                self.env.update(ODB_CCMA_CREATE_DIRECT=1)
+
+    def create_poolmask(self, layout, dbpath=None):
+        """Request the poolmask file creation."""
+        layout, dbpath, _ = self._process_layout_dbpath(layout, dbpath)
+        logger.info('ODB: requesting poolmask file for: %s (layout=%s).', dbpath, layout)
+        self.env.update(
+            ODB_CCMA_CREATE_POOLMASK=1,
+            ODB_CCMA_POOLMASK_FILE=self.sh.path.join(dbpath, layout + '.poolmask'),
+        )
+
+    def change_layout(self, layout, layout_new, dbpath=None):
+        """Make the appropriate renaming of files in ECMA to CCMA."""
+        layout, dbpath, _ = self._process_layout_dbpath(layout, dbpath)
+        layout_new = layout_new.upper()
+        logger.info('ODB: changing layout (%s -> %s) for %s.', layout, layout_new, dbpath)
+        to_cleanup = set()
+        for f in self.sh.ls(dbpath):
+            if self.sh.path.islink(self.sh.path.join(dbpath, f)):
+                fullpath = self.sh.path.join(dbpath, f)
+                target = self.sh.readlink(fullpath)
+                self.sh.unlink(fullpath)
+                self.sh.symlink(target.replace(layout, layout_new),
+                                fullpath.replace(layout, layout_new))
+                continue
+            if f in [n.format(layout) for n in ('{:s}.dd', '{:s}.flags')]:
+                self.sh.mv(self.sh.path.join(dbpath, f),
+                           self.sh.path.join(dbpath, f.replace(layout, layout_new)))
+            if f in [n.format(layout) for n in ('{:s}.iomap', '{:s}.sch',
+                                                '{:s}.IOASSIGN', 'IOASSIGN.{:s}', 'IOASSIGN')]:
+                tmp_target = self.sh.path.join(dbpath, f + '.tmp_new')
+                with open(self.sh.path.join(dbpath, f)) as inodb:
+                    with open(tmp_target, 'w') as outodb:
+                        for line in inodb:
+                            outodb.write(line.replace(layout, layout_new))
+                self.sh.mv(tmp_target,
+                           self.sh.path.join(dbpath, f.replace(layout, layout_new)))
+                if layout in f:
+                    to_cleanup.add(self.sh.path.join(dbpath, f))
+        for f_name in to_cleanup:
+            self.sh.rm(f_name)
+
+
+#: Footprint's attributes needed to ODB to setup properly
+odbmix_attributes = footprints.Footprint(
+    info="Abstract ODB footprints' attributes.",
+    attr=dict(
+        npool=dict(
+            info='The number of pool(s) in the ODB database.',
+            type=int,
+            optional=True,
+            default=1,
+        ),
+        iomethod=dict(
+            info='The io_method of the ODB database.',
+            type=int,
+            optional=True,
+            default=1,
+            doc_zorder=-50,
+        ),
+        slots=dict(
+            info='The timeslots of the assimilation window.',
+            type=TimeSlots,
+            optional=True,
+            default=TimeSlots(7, chunk='PT1H'),
+        ),
+        virtualdb=dict(
+            info='The type of the virtual ODB database.',
+            optional=True,
+            default='ecma',
+            access='rwx',
+            doc_visibility=footprints.doc.visibility.ADVANCED,
+        ),
+        date=dict(
+            info='The current run date.',
+            optional=True,
+            access='rwx',
+            type=bdate.Date,
+            doc_zorder=-50,
+        ),
+        ioassign=dict(
+            info='The path to the ioassign binary (needed for merge/create actions',
+            optional=True,
+            default='ioassign.x',
+        )
+    ))
+
+
+@algo_component_deco_mixin_autodoc
+class OdbComponentDecoMixin(AlgoComponentDecoMixin):
+    """Handle ODB settings in AlgoComponents.
+
+    This mixin class is intended to be used with AlgoComponent classes. It will
+    automatically add footprints' arguments related to ODB
+    (see :data:`odbmix_attributes`), set up generic ODB environment variables
+    (:meth:`_odbobj_setup`) and provides a :attr:`odb` property that gives
+    access to a properly initialised :class:`OdbDriver` object that can be used
+    directly in AlgoComponents.
+
+    In addition it provides directly some utility methods that can be called
+    manually if needed.
+    """
+
+    _MIXIN_EXTRA_FOOTPRINTS = [odbmix_attributes, ]
+
+    def _odbobj_init(self, rh, opts):  # @UnusedVariable
+        """Setup the OdbDriver object."""
+        cycle = ArpIfsSimplifiedCycle('cy01')
+        if rh and hasattr(rh.resource, 'cycle'):
+            cycle = rh.resource.cycle
+        self._odb = OdbDriver(
+            cycle=cycle,
+            sh=self.system,
+            env=self.env,
+        )
+        if self.system.path.exists(self.ioassign):
+            self._x_ioassign = self.system.path.abspath(self.ioassign)
+        else:
+            # Legacy...
+            self._x_ioassign = self.ioassign
+
+    def _odbobj_setup(self, rh, opts):  # @UnusedVariable
+        """Setup the ODB object."""
+        self.odb.setup(
+            layout=self.virtualdb,
+            date=self.date,
+            npool=self.npool,
+            nslot=self.slots.nslot,
+            iomethod=self.iomethod,
+        )
+
+    _MIXIN_PREPARE_PREHOOKS = (_odbobj_init,)
+    _MIXIN_PREPARE_HOOKS = (_odbobj_setup,)
+
+    @property
+    def odb(self):
+        """Access to a properly initialised :class:`OdbDriver` object."""
+        if not hasattr(self, '_odb'):
+            raise RuntimeError('Uninitialised *odb* object.')
+        return self._odb
+
+    def lookupodb(self, fatal=True):
+        """Return a list of effective input resources which are odb observations."""
+        allodb = [
+            x for x in self.context.sequence.effective_inputs(kind='observations')
+            if x.rh.container.actualfmt == 'odb'
+        ]
+        allodb.sort(key=lambda s: s.rh.resource.part)
+        if not allodb and fatal:
+            logger.critical('Missing ODB input data for %s', self.fullname())
+            raise ValueError('Missing ODB input data')
+        return allodb
+
+    def odb_date_and_layout_from_sections(self, odbsections):
+        """
+        Look into the **odsections** section list in order to find the current
+        run date and ODB database layout.
+        """
+        alllayouts = {s.rh.resource.layout for s in odbsections}
+        alldates = {s.rh.resource.date for s in odbsections}
+        if len(alllayouts) != 1:
+            raise AlgoComponentError("Inconsistent ODB layouts")
+        if len(alldates) != 1:
+            raise AlgoComponentError("Inconsistent ODB dates")
+        self.virtualdb = alllayouts.pop()
+        self.date = alldates.pop()
+        logger.info('ODB: Detected from ODB database(s). self.date=%s, self.virtualdb=%s.',
+                    self.date.stdvortex, self.virtualdb)
+
+    def _odb_find_ioassign_script(self, purpose):
+        """Look for ioassign script of *purpose" attribute, return path."""
+        scripts = [x.rh.container.abspath
+                   for x in self.context.sequence.effective_inputs(kind='ioassign_script')
+                   if x.rh.resource.purpose == purpose]
+        if len(scripts) > 1:
+            raise AlgoComponentError("More than one purpose={} ioassign_script found in resources.")
+        elif len(scripts) == 1:
+            self.system.xperm(scripts[0], force=True)
+            return scripts[0]
+        else:
+            return None
+
+    def odb_merge_if_needed(self, odbsections, subdir='.'):
+        """
+        If multiple ODB databases are listed in the **odsections** section list,
+        start an ODB merge.
+
+        :return: The path to the ODB database (the virtual database if a merge is
+                 performed).
+        """
+        if len(odbsections) > 1 or self.virtualdb.lower() == 'ecma':
+            logger.info('ODB: merge for: %s.', self.virtualdb)
+            iomerge_p = self._odb_find_ioassign_script('merge')
+            iocreate_p = self._odb_find_ioassign_script('create')
+            with self.system.cdcontext(subdir):
+                virtualdb_path = self.odb.ioassign_merge(
+                    layout=self.virtualdb,
+                    ioassign=self._x_ioassign,
+                    odbnames=[x.rh.resource.part for x in odbsections],
+                    iomerge_path=iomerge_p,
+                    iocreate_path=iocreate_p,
+                )
+        else:
+            virtualdb_path = self.system.path.abspath(odbsections[0].rh.container.localpath())
+        return virtualdb_path
+
+    def odb_create_db(self, layout, dbpath=None):
+        """Create a new empty ODB database.
+
+        :param str layout: The new database layout
+        :param str dbpath: The path to the new database (current_dir/layout if omitted)
+        :return: The path to the new ODB database.
+        """
+        dbout = self.odb.ioassign_create(
+            layout=layout,
+            npool=self.npool,
+            ioassign=self._x_ioassign,
+            dbpath=dbpath,
+            iocreate_path=self._odb_find_ioassign_script('create')
+        )
+        logger.info('ODB: database created: %s (layout=%s).', dbout, layout)
+        return dbout
+
+    def odb_handle_raw_dbs(self):
+        """Look for extras ODB raw databases and fix the environment accordingly."""
+        odbraw = [
+            x.rh for x in self.context.sequence.effective_inputs(kind='odbraw')
+            if x.rh.container.actualfmt == 'odb'
+        ]
+        if not odbraw:
+            logger.error('No ODB raw databases found.')
+        else:
+            for rraw in odbraw:
+                rawpath = rraw.container.localpath()
+                self.odb.fix_db_path(rraw.resource.layout, rawpath)
+                for badlink in [bl
+                                for bl in self.system.glob(self.system.path.join(rawpath, '*.h'))
+                                if self.system.path.islink(bl) and not self.system.path.exists(bl)]:
+                    self.system.unlink(badlink)
+            self.odb.ioassign_append(* [rraw.container.localpath() for rraw in odbraw])
+        return odbraw
+
+    def odb_rw_or_overwrite_method(self, *dbsections):
+        """Are the input databases fetch with intent=inout ?"""
+        needs_work = [s for s in dbsections if s.intent == intent.IN]
+        if needs_work:
+            self.odb.force_overwrite_method()