graphviper.graph_tools.process_with_mpi

MPI execution backend for GraphVIPER map/reduce graphs.

This is an optional alternative to the Dask backend (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 graphviper.graph_tools.map.map() and 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 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 map() adapts an explicit-signature node task into a functools.wraps closure (see 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.

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

force_exit_after(seconds[, note])

Arm (or re-arm) a daemon watchdog that hard-exits this process after

processes_with_mpi(viper_graph[, cluster_setup])

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

Module Contents

force_exit_after(seconds, note='')[source]

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 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 threading.Timer (or None if cancelling).

processes_with_mpi(viper_graph, cluster_setup=None)[source]

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.

Parameters:

  • viper_graph (dict) – Graph from graphviper.graph_tools.map.map() and optionally graphviper.graph_tools.reduce.reduce().

  • cluster_setup (dict, optional) –

    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 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.

Returns:

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

Return type:

object