vortex-nwp 2.0.0b1__py3-none-any.whl

vortex/nwp/tools/partitioning.py

@@ -0,0 +1,234 @@
+"""
+Compute simple domain partitionings.
+
+The partitioning classes can be used on their own. Alternatively, the
+:meth:`setup_partitioning_in_namelist` method can be used to interact with
+namelist's Content objects.
+"""
+
+import functools
+import math
+import re
+
+from bronx.syntax import iterators as b_iter
+from bronx.fancies import loggers
+
+
+logger = loggers.getLogger(__name__)
+
+__all__ = ['PartitioningError',
+           'Rectangular2DPartitioner',
+           'setup_partitioning_in_namelist']
+
+
+class PartitioningError(ValueError):
+    """Any error raised during domain partitionings."""
+    pass
+
+
+class AbstratctPartitioner:
+    """The base class for any concrete partitioning class.
+
+    Provides the mechanism to filter the partitioning methods and
+    cache the results.
+    """
+
+    _REGISTERED_METHODS = ()
+
+    def __init__(self, p_method_specification):
+        """
+        :param p_method_specification: The partitioning method definition
+        """
+        # Process the partitioning method specification string
+        p_method_parts = p_method_specification.lower().split('_')
+        self.p_method_name = None
+        self.p_method_args = ()
+        for a_method, a_method_n_args in self._REGISTERED_METHODS:
+            if p_method_parts[0] == a_method:
+                self.p_method_name = p_method_parts[0]
+                if len(p_method_parts) - 1 != a_method_n_args:
+                    raise ValueError('Erroneous number of interger args ' +
+                                     'for the {:s} p_method ({:d} required)'
+                                     .format(a_method, a_method_n_args))
+                self.p_method_args = tuple([int(s) for s in p_method_parts[1:]])
+        # Unknown method -> crash
+        if self.p_method_name is None:
+            raise PartitioningError("Unknown partitioning method ({:s})."
+                                    .format(p_method_specification))
+        # The actual class' method that will be used to compute a and b
+        self.p_method = functools.partial(getattr(self, '_' + self.p_method_name),
+                                          * self.p_method_args)
+        # Implement a caching mechanism
+        self._p_cache = dict()
+
+    def __call__(self, ntasks):
+        """Return the appropriate partitioning given **ntasks**."""
+        if ntasks not in self._p_cache:
+            self._p_cache[ntasks] = self.p_method(ntasks)
+        return self._p_cache[ntasks]
+
+
+class Rectangular2DPartitioner(AbstratctPartitioner):
+    """Find an adequat 2D simple partitioning given the number of tasks.
+
+    Here is a description of the problem :
+
+    * Let ``D`` be a 2D array/field;
+    * Let ``N`` be the number of partitions you want to create;
+    * What are ``x`` and ``x`` that satisfies ``x * y = N`` (so that the
+      ``D`` array can be partitionned in ``x`` (reps. ``y``) chunks
+      in its first (resp. second) dimension ?
+
+    For exemple, with N=16, an array can be split in 4 chunks in each
+    dimension. It can also be partitioned in 2 chunks in the first
+    dimension and 8 in the second...  There is no unique solution.
+    Consequently, the user needs to provide a partitioning method.
+
+    Example::
+
+        # Look for a partitioning around a given fixed value
+        # e.g with xcloseto_16, the x value will be close to 16
+        >>> Rectangular2DPartitioner('xcloseto_16')(128)
+        (16, 8)
+        >>> Rectangular2DPartitioner('xcloseto_16')(990)
+        (15, 66)
+        >>> Rectangular2DPartitioner('xcloseto_16')(500)
+        (20, 25)
+
+        # e.g with ycloseto_16, the y value will be close to 16
+        >>> Rectangular2DPartitioner('ycloseto_16')(128)
+        (8, 16)
+        >>> Rectangular2DPartitioner('ycloseto_16')(990)
+        (66, 15)
+        >>> Rectangular2DPartitioner('ycloseto_16')(500)
+        (25, 20)
+
+        # Squar-est partition of the domain: y and y as close
+        # as possible
+        >>> Rectangular2DPartitioner('square')(16)
+        (4, 4)
+        >>> Rectangular2DPartitioner('square')(12)
+        (3, 4)
+        >>> Rectangular2DPartitioner('square')(7)
+        (1, 7)
+
+        # Try to find x and y so that a given aspect ratio is preserved
+        # e.g with aspect_16_9, x / y should roughly be equal to 16 / 9
+        >>> Rectangular2DPartitioner('aspect_2_1')(32)
+        (8, 4)
+        >>> Rectangular2DPartitioner('aspect_2_1')(27)
+        (9, 3)
+        >>> Rectangular2DPartitioner('aspect_16_9')(28)  # roughly 16/9e like a TV...
+        (7, 4)
+
+    """
+
+    _REGISTERED_METHODS = (
+        ('xcloseto', 1),
+        ('ycloseto', 1),
+        ('square', 0),
+        ('aspect', 2)
+    )
+
+    @staticmethod
+    def _test_and_return(ntasks, guesses):
+        found = 1
+        for i_guess in guesses:
+            if ntasks % i_guess == 0:
+                found = i_guess
+                break
+        return found, ntasks // found
+
+    def _xcloseto(self, close_to_what, ntasks):
+        """Find ``x`` as the closest possible value to **close_to_what**."""
+        guesses = b_iter.interleave(
+            range(close_to_what, 0, -1),
+            range(close_to_what + 1, min(close_to_what * 2,
+                                         ntasks))
+        )
+        return self._test_and_return(ntasks, guesses)
+
+    def _ycloseto(self, close_to_what, ntasks):
+        """Find ``y`` as the closest possible value to **close_to_what**."""
+        y_value, x_value = self._xcloseto(close_to_what, ntasks)
+        return x_value, y_value
+
+    def _square(self, ntasks):
+        """Find ``x`` and ``y`` so that they are close to the square root of ``N``.
+
+        With this method, ``x`` is always the smalest value.
+        """
+        guesses = range(int(math.sqrt(ntasks)), 0, -1)
+        return self._test_and_return(ntasks, guesses)
+
+    def _aspect(self, x_spec, y_spec, ntasks):
+        """Find ``x`` and ``y`` so that ``x / y =~ x_spec / y_spec``."""
+        aspect_ratio = x_spec / y_spec
+        return self._xcloseto(int(math.sqrt(ntasks * aspect_ratio)),
+                              ntasks)
+
+
+_PARTITIONERS_CACHE = dict()
+
+
+def setup_partitioning_in_namelist(namcontents,
+                                   effective_tasks,
+                                   effective_threads,
+                                   namlocal=None):
+    """Look in a namelist Content object and replace the macros related to partitioning.
+
+    :param nwp.data.namelists.NamelistContent namcontents: The namelist's Content
+                                                              object to work with
+    :param int effective_tasks: The number of tasks that will be used when computing
+                                the partitioning
+    :param int effective_threads: The number of threads that will be used when computing
+                                  the partitioning
+    :param str namlocal: The namelist's file name
+    :return: ``True`` if the namelist's Contents object has been modified
+    :rtype: bool
+
+    This function will detect namelist macros like ``PART_TASKS2D_X_SQUARE`` where:
+
+    * ``TASKS`` tells that **effective_tasks** will be used to compute the
+      decomposition (alternatively, ``THREADS`` can be used.
+    * ``2D`` tells that the :class:`Rectangular2DPartitioner` class will be used
+      to compute the partitioning. For now, ``2D`` is the only available option.
+    * ``X`` tells that the user wants to get the X value of the computed partioning.
+      Alternatively, ``Y`` can be used.
+    * ``SQUARE`` refers to the partitioning method that will be used by the
+      partitioning class. Any value that is accepted by the partitioning class is
+      fine.
+    """
+    macrovalid = re.compile('PART_' +
+                            '(?P<what>TASKS|THREADS)(?P<cls>2D)_' +
+                            '(?P<dim>[XY])_(?P<def>.*)$')
+    partitioning_classes = {'2D': Rectangular2DPartitioner}
+    namw = False
+    # Find the list of existing macros
+    all_macros = set()
+    for nam in namcontents.values():
+        all_macros.update(nam.macros())
+    # Consider only relevant macros
+    for macroname in all_macros:
+        macroname_re = macrovalid.match(macroname)
+        if macroname_re:
+            cache_key = (macroname_re.group('cls'), macroname_re.group('def'))
+            if cache_key not in _PARTITIONERS_CACHE:
+                partitioning_class = partitioning_classes[macroname_re.group('cls')]
+                _PARTITIONERS_CACHE[cache_key] = partitioning_class(macroname_re.group('def'))
+            effective_n = dict(TASKS=effective_tasks,
+                               THREADS=effective_threads)[macroname_re.group('what')]
+            part_x, part_y = _PARTITIONERS_CACHE[cache_key](effective_n)
+            final_result = part_x if macroname_re.group('dim') == 'X' else part_y
+            if namlocal:
+                logger.info('Setup macro %s=%s in %s', macroname, final_result, namlocal)
+            else:
+                logger.info('Setup macro %s=%s', macroname, final_result)
+            namcontents.setmacro(macroname, final_result)
+            namw = True
+    return namw
+
+
+if __name__ == '__main__':
+    import doctest
+    doctest.testmod()

vortex/nwp/tools/satrad.py

@@ -0,0 +1,56 @@
+"""
+Common interest classes to help setup the RTTOV/IFS environment.
+"""
+
+import re
+
+from bronx.fancies import loggers
+
+from vortex.algo.components import AlgoComponentDecoMixin, AlgoComponentError
+from vortex.algo.components import algo_component_deco_mixin_autodoc
+
+#: No automatic export
+__all__ = []
+
+logger = loggers.getLogger(__name__)
+
+
+@algo_component_deco_mixin_autodoc
+class SatRadDecoMixin(AlgoComponentDecoMixin):
+    """RTTOV settings + Satellites related stuffs.
+
+    This mixin class is intended to be used with AlgoComponent classes. It will
+    automatically set up the path to RTTOV coefficient files
+    (:meth:`_satrad_coeffdir_setup`).
+
+    In addition it provides the :meth:`setchannels` utility method (that have to
+    be called manually if needed).
+    """
+
+    def _satrad_coeffdir_setup(self, rh, opts):  # @UnusedVariable
+        """Look for RTTOV coefficient files and act on it."""
+        rtcoefs = self.context.sequence.effective_inputs(role='RtCoef', kind='rtcoef')
+        if rtcoefs:
+            sh = self.system
+            rtpaths = {sh.path.dirname(sh.path.realpath(rtcoef.rh.container.localpath()))
+                       for rtcoef in rtcoefs}
+            if len(rtpaths) != 1:
+                raise AlgoComponentError('The Radiative Transfer Coefficients are scattered in' +
+                                         'several directories: {!s}'.format(rtpaths))
+            rtpath = rtpaths.pop()
+            logger.info('Setting %s = %s', 'RTTOV_COEFDIR', rtpath)
+            self.env['RTTOV_COEFDIR'] = rtpath
+
+    _MIXIN_PREPARE_HOOKS = (_satrad_coeffdir_setup, )
+
+    def setchannels(self):
+        """Look up for channels namelists in effective inputs."""
+        namchan = [
+            x.rh for x in self.context.sequence.effective_inputs(kind='namelist')
+            if 'channel' in x.rh.options
+        ]
+        for thisnam in namchan:
+            thisloc = re.sub(r'\d+$', '', thisnam.options['channel']) + 'channels'
+            if thisloc != thisnam.container.localpath():
+                logger.info('Linking < %s > to < %s >', thisnam.container.localpath(), thisloc)
+                self.system.softlink(thisnam.container.localpath(), thisloc)

vortex/nwp/util/__init__.py

@@ -0,0 +1,6 @@
+"""
+Utilities mostly for NWP.
+"""
+
+#: No automatic export
+__all__ = []

vortex/nwp/util/async.py

@@ -0,0 +1,184 @@
+"""
+Callback functions for Jeeves.
+If needed, VORTEX must be loaded via a VortexWorker in this context.
+"""
+
+from vortex.tools import compression, systems
+from vortex.util.worker import VortexWorker
+
+#: No automatic export
+__all__ = []
+
+
+def _double_ssh(sh, loginnode, transfernode):
+    """Applies a double ssh to retrieve the effective name of a machine.
+
+    This trick enables the load balancing and node crash recovery
+    capabilities handled by the network teams through DNS names.
+
+    May return None when network problems occur.
+    """
+    cmd = ['ssh', '-x', loginnode, 'ssh', '-x', transfernode, 'hostname', '-s']
+    rc = sh.spawn(cmd, shell=False, output=True, fatal=False)
+    if not rc:
+        return None
+    return rc[0]
+
+
+def system_ftput(pnum, ask, config, logger, **opts):
+    """Ftp transfer to some archive host.
+
+    Removes the source on success.
+    In phase mode, raw ftp is not allowed, and the hostname is dynamically
+    obtained by a double ssh.
+    """
+
+    logger.info('System', todo=ask.todo, pnum=pnum, opts=opts)
+    value = dict(rpool='retry')
+
+    phasemode = opts.get('phasemode', False)
+    nbtries = opts.get('attempts', 1)
+    if phasemode:
+        rawftput = False
+    else:
+        rawftput = opts.get('rawftput', False)
+    trynum = 0
+
+    profile = config['driver'].get('profile', None)
+    with VortexWorker(logger=logger, profile=profile) as vwork:
+        sh = vwork.session.sh
+        sh.trace = 'log'
+        sh.ftpflavour = systems.FTP_FLAVOUR.STD  # Because errors are handled directly by jeeves
+
+        data = vwork.get_dataset(ask)
+        logger.info('ftput', source=data.source, destination=data.destination)
+        if not sh.path.exists(data.source):
+            logger.error('The source file is missing - sorry')
+            return pnum, False, dict(rpool='error')
+
+        if phasemode:
+            data.hostname = _double_ssh(sh, data.phase_loginnode, data.phase_transfernode)
+            if data.hostname is None:
+                return pnum, False, dict(rpool='retry')
+
+        cpipeline = (None if not hasattr(data, 'cpipeline') or not data.cpipeline
+                     else compression.CompressionPipeline(sh, data.cpipeline))
+
+        logger.info('FTPut host', hostname=data.hostname, logname=data.logname)
+        logger.info('FTPut data', source=data.source, destination=data.destination)
+        while trynum < nbtries:
+            trynum += 1
+            if nbtries > 1:
+                logger.info('FTPut loop', attempt=trynum)
+            try:
+                if rawftput:
+                    putrc = sh.rawftput(data.source, data.destination, hostname=data.hostname,
+                                        logname=data.logname, cpipeline=cpipeline,
+                                        fmt=data.fmt)
+                else:
+                    putrc = sh.ftput(data.source, data.destination, hostname=data.hostname,
+                                     logname=data.logname, cpipeline=cpipeline,
+                                     fmt=data.fmt)
+            except Exception as e:
+                logger.warning('FTPut failed', attempt=trynum, error=e)
+                putrc = False
+            if putrc:
+                value = dict(clear=sh.rm(data.source, fmt=data.fmt))
+                break
+
+    return pnum, putrc and vwork.rc, value
+
+
+def system_cp(pnum, ask, config, logger, **opts):
+    """Local transfers (between filesystems) on a given host."""
+
+    logger.info('System', todo=ask.todo, pnum=pnum, opts=opts)
+    value = dict(rpool='retry')
+
+    profile = config['driver'].get('profile', None)
+    with VortexWorker(logger=logger, profile=profile) as vwork:
+        sh = vwork.session.sh
+        sh.trace = 'log'
+        data = vwork.get_dataset(ask)
+        logger.info('cp', source=data.source, destination=data.destination)
+        if not sh.path.exists(data.source):
+            logger.error('The source file is missing - sorry')
+            return pnum, False, dict(rpool='error')
+
+        try:
+            rc = sh.cp(data.source, data.destination, fmt=data.fmt)
+        except Exception as e:
+            logger.warning('cp failed', error=e)
+            rc = False
+        if rc:
+            value = dict(clear=sh.rm(data.source, fmt=data.fmt))
+
+    return pnum, rc and vwork.rc, value
+
+
+def system_scp(pnum, ask, config, logger, **opts):
+    """Scp transfer to some archive host.
+
+    Removes the source on success.
+    In phase mode, raw ftp is not allowed, and the hostname is dynamically
+    obtained by a double ssh.
+    """
+    logger.info('System', todo=ask.todo, pnum=pnum, opts=opts)
+    value = dict(rpool='retry')
+
+    phasemode = opts.get('phasemode', False)
+
+    profile = config['driver'].get('profile', None)
+    with VortexWorker(logger=logger, profile=profile) as vwork:
+        sh = vwork.session.sh
+        sh.trace = 'log'
+
+        data = vwork.get_dataset(ask)
+        logger.info('scp', source=data.source, destination=data.destination)
+        if not sh.path.exists(data.source):
+            logger.error('The source file is missing - sorry')
+            return pnum, False, dict(rpool='error')
+
+        if phasemode:
+            data.hostname = _double_ssh(sh, data.phase_loginnode, data.phase_transfernode)
+            if data.hostname is None:
+                return pnum, False, value
+        logger.info('scp host', hostname=data.hostname, logname=data.logname)
+        logger.info('scp data', source=data.source, destination=data.destination)
+        try:
+            putrc = sh.scpput(data.source, data.destination, hostname=data.hostname,
+                              logname=data.logname, fmt=data.fmt)
+        except Exception as e:
+            logger.warning('scp failed', error=e)
+            putrc = False
+        if putrc:
+            value = dict(clear=sh.rm(data.source, fmt=data.fmt))
+
+    return pnum, putrc and vwork.rc, value
+
+
+def system_noop(pnum, ask, config, logger, **opts):
+    """A callback able to do nothing, but cleanly.
+
+    Used to desactivate jeeves when mirroring the operational suite.
+    """
+    logger.info('Noop', todo=ask.todo, pnum=pnum, opts=opts)
+
+    profile = config['driver'].get('profile', None)
+    with VortexWorker(logger=logger, profile=profile) as vwork:
+        sh = vwork.session.sh
+        sh.trace = 'log'
+        data = vwork.get_dataset(ask)
+        value = dict(clear=sh.rm(data.source, fmt=data.fmt))
+
+    return pnum, vwork.rc, value
+
+
+if __name__ == '__main__':
+    import vortex
+
+    t = vortex.ticket()
+    main_sh = t.sh
+    main_sh.trace = True
+    main_sh.verbose = True
+    print(_double_ssh(main_sh, 'beaufixoper', 'beaufixtransfert-agt'))

vortex/nwp/util/beacon.py

@@ -0,0 +1,40 @@
+"""
+Functions to create and write a few information in a file using Vortex
+(FunctionStore).
+"""
+
+import io
+import json
+
+#: No automatic export
+__all__ = []
+
+
+def beaconfunction(options):
+    """Function to create a file and write information in it.
+
+    - model
+    - date
+    - cutoff
+    - vapp
+    - vconf
+    - member (optional)
+    """
+    rst = dict()
+
+    # Find out if a resource handler is present and load the elements to be written
+    rhdict = options.get('rhandler', None)
+    if rhdict:
+        rst['vapp'] = rhdict.get('provider', {}).get('vapp', '')
+        rst['vconf'] = rhdict.get('provider', {}).get('vconf', '')
+        rst['model'] = rhdict.get('resource', {}).get('model', '')
+        rst['date'] = rhdict.get('resource', {}).get('date', '')
+        rst['cutoff'] = rhdict.get('resource', {}).get('cutoff', '')
+        member = rhdict.get('provider', {}).get('member', None)
+        if member is not None:
+            rst['member'] = member
+    else:
+        rst['error'] = 'No resource handler here'
+    outstr = json.dumps(rst)
+    # Return the string, which has to be converted to a file like object
+    return io.BytesIO(outstr.encode(encoding='utf_8'))