graphviper.graph_tools.process_with_mpi
=======================================

.. py:module:: graphviper.graph_tools.process_with_mpi

.. autoapi-nested-parse::

   MPI execution backend for GraphVIPER map/reduce graphs.

   This is an *optional* alternative to the Dask backend
   (:func:`graphviper.graph_tools.generate_dask_workflow.generate_dask_workflow`
   followed by ``dask.compute``).  The GraphVIPER architecture keeps the graph a
   plain, backend-agnostic description (``map`` / optional ``reduce`` /optional
   ``load`` stages produced by :func:`graphviper.graph_tools.map.map` and
   :func:`graphviper.graph_tools.reduce.reduce`); ``dask.delayed`` only ever appears
   inside ``generate_dask_workflow``.  Because of that modularity the *same* graph
   can be executed by a different engine.  This module provides an MPI engine so a
   caller can swap::

       viper_graph = map(...)
       viper_graph = reduce(...)
       dask_graph  = generate_dask_workflow(viper_graph)   # Dask backend
       return_dict = dask.compute(dask_graph)[0]

   for::

       viper_graph = map(...)
       viper_graph = reduce(...)
       return_dict = processes_with_mpi(viper_graph, cluster_setup)   # MPI backend

   Execution model (manager / worker pool)
   ---------------------------------------
   The backend uses :class:`mpi4py.futures.MPIPoolExecutor` in the **static**
   manager-worker model.  The program must therefore be launched as::

       mpiexec -n <N> python -m mpi4py.futures <script>.py ...
       # or, on TACC:
       ibrun python -m mpi4py.futures <script>.py ...

   With ``-m mpi4py.futures`` exactly one rank (rank 0) runs the user script
   (``__main__``) -- it is the *manager*: it builds the graph and calls
   ``processes_with_mpi``.  The remaining ranks never run the script body; they sit
   in a worker server loop and execute map (and, optionally, reduce) node tasks
   dispatched by the manager.  This deliberately mirrors the "dedicated driver
   node" pattern: the manager does the orchestration (and the small final reduce)
   while every other rank does the heavy per-chunk compute -- but unlike a Dask
   in-process scheduler there is no large task graph to hold, and unlike a
   dedicated *Dask* scheduler node the manager rank shares its node with worker
   ranks (only one rank, not a whole node, is spent on orchestration).

   Pickling
   --------
   GraphVIPER's :func:`map` adapts an explicit-signature node task into a
   ``functools.wraps`` *closure* (see
   :func:`graphviper.graph_tools.map.make_graph_node_task`), and plain ``pickle``
   cannot serialise closures.  mpi4py is therefore reconfigured to pickle with
   ``cloudpickle`` on the manager side; cloudpickle's output is loadable by the
   workers' stdlib ``pickle.loads`` (cloudpickle only customises *dumps*), so the
   manager->worker direction (the closure node task + its ``input_params``) works
   with only the manager reconfigured.

   The worker->manager *result* direction is NOT symmetric: workers stay on stdlib
   ``pickle`` (they are already in the server loop and cannot be reconfigured from
   here), so node-task **return values must be stdlib-pickle-serialisable**.  This
   differs from the Dask backend (which serialises worker results with cloudpickle):
   a result that is only cloudpickle-serialisable -- a closure/lambda, or an instance
   of a class defined in ``__main__`` or a local scope -- works under
   ``dask.compute`` but raises ``PicklingError`` here.  This is fine for the imaging
   use case (results are small stdlib-picklable metadata).  To make results fully
   cloudpickle-capable, the launching program must call
   ``MPI.pickle.__init__(cloudpickle.dumps, cloudpickle.loads)`` at import time on
   *every* rank (before ``mpi4py.futures`` starts the worker loop), not just inside
   this function.

   .. rubric:: Notes

   * Map node tasks write their result chunk straight to the shared (Lustre) output
     store and return only small metadata; the manager creates/initialises that
     store *before* calling this function, so all workers see it.  Results are
     returned in input order.
   * A ``load`` stage (disk-chunk coalescing) is **not** specially handled here:
     each map task self-loads its data (``input_data is None``), exactly as in the
     no-load Dask path.  Results are identical; only the cross-task I/O sharing
     optimisation is skipped.  A warning is logged if a ``load`` stage is present.



Functions
---------

.. autoapisummary::

   graphviper.graph_tools.process_with_mpi.force_exit_after
   graphviper.graph_tools.process_with_mpi.processes_with_mpi


Module Contents
---------------

.. py:function:: force_exit_after(seconds, note='')

   Arm (or re-arm) a daemon watchdog that hard-exits this process after
   ``seconds`` -- a guard against the mpi4py.futures / MPI shutdown hanging
   after all work is done.

   Rationale.  In the ``python -m mpi4py.futures`` manager-worker model the
   global stop handshake and ``MPI_Finalize`` run at *interpreter exit*.  A
   single worker rank wedged during its own shutdown (typically a non-daemon
   thread stuck in an uninterruptible filesystem syscall) blocks the
   effectively-collective finalize on every rank, idling the whole allocation
   until walltime.  An ``atexit`` hook CANNOT guard against this: CPython
   joins non-daemon threads (including ``concurrent.futures`` pools) *before*
   running atexit callbacks, which is exactly where the wedge occurs.  The
   only robust guard is a daemon thread armed while the interpreter is still
   healthy -- this function.

   Call it once the application's results are safely persisted (or use the
   ``teardown_force_exit_seconds`` option of :func:`processes_with_mpi`).
   Calling again re-arms the countdown; the previous timer is cancelled.
   ``seconds`` of ``None``/``0`` just cancels any armed watchdog.

   The forced exit uses ``os._exit(0)`` (no cleanup, no ``MPI_Finalize``); the
   MPI launcher then tears down the remaining ranks.  The launcher may log the
   exit as unclean -- harmless next to an allocation idling for hours.

   Returns the armed :class:`threading.Timer` (or ``None`` if cancelling).


.. py:function:: processes_with_mpi(viper_graph, cluster_setup=None)

   Execute a GraphVIPER map/reduce graph with an MPI manager-worker pool.

   Drop-in replacement for ``generate_dask_workflow`` + ``dask.compute``: returns
   exactly what ``dask.compute(generate_dask_workflow(viper_graph))[0]`` would --
   the reduced result when a ``reduce`` stage is present, otherwise the list of
   per-map-task results (in input order).

   Must be launched in the static mpi4py.futures manager-worker model
   (``... python -m mpi4py.futures <script>``); see the module docstring.  Only
   the manager rank (rank 0) ever reaches this function.

   :param viper_graph: Graph from :func:`graphviper.graph_tools.map.map` and optionally
                       :func:`graphviper.graph_tools.reduce.reduce`.
   :type viper_graph: dict
   :param cluster_setup: MPI execution options:

                         * ``max_workers`` (int or None) -- cap the pool size; ``None`` (default)
                           uses every available worker rank.
                         * ``chunksize`` (int) -- ``MPIPoolExecutor.map`` chunk size; ``1``
                           (default) gives the best dynamic load balancing for long, uneven node
                           tasks (per-task dispatch overhead is negligible next to ~100 s tasks).
                         * ``reduce_in_pool`` (bool) -- if ``True`` run the reduce tree on the
                           worker pool; default ``False`` reduces on the manager (the right choice
                           when map results are small metadata, as in imaging).
                         * ``use_cloudpickle`` (bool) -- reconfigure mpi4py to pickle with
                           cloudpickle so closure node tasks serialise; default ``True``.
                         * ``progress_every`` (int or None) -- if set, log a progress line every
                           this many completed map tasks.
                         * ``teardown_force_exit_seconds`` (float or None) -- if set, arm
                           :func:`force_exit_after` with this grace period when the compute
                           returns, guarding against the mpi4py.futures / MPI_Finalize shutdown
                           hang at interpreter exit (one wedged worker rank can idle the whole
                           allocation until walltime).  Default ``None`` (off): only enable in
                           run-one-graph-then-exit programs (batch jobs); a long-lived
                           application that keeps working after this call would be killed
                           mid-flight unless it re-arms or cancels via
                           ``force_exit_after(None)``.  A subsequent ``processes_with_mpi``
                           call re-arms the countdown.
   :type cluster_setup: dict, optional

   :returns: The reduced result, or (no reduce stage) the list of map-task results.
   :rtype: object