vortex-nwp 2.0.0__py3-none-any.whl

vortex/nwp/tools/partitioning.py

@@ -0,0 +1,258 @@
+"""
+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,71 @@
+"""
+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,212 @@
+"""
+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"))