mpiprof 0.2.2
Profiler for MPI python programs
License: MIT
Requires Python: >=3.8
Keywords: MPI, profiling, cProfile, mpi4py

mpiprof.py: Profiler for MPI python programs
============================================

Enables:

1. using cProfile on each launched MPI process.
2. measuring the communication overhead of your MPI setup.

mpiprof provides two complementary pieces:

1. A drop-in wrapper for cProfile that you can invoke as a module,
   ``mpiexec -n 4 python3 -m mpiprof your_script.py arg1 arg2``.
   This writes a separate .pstats file for each rank.

2. An optional ``MPIProfiler`` class that wraps a ``mpi4py.MPI.Comm`` and
   records basic timing/count statistics of MPI calls and where they came from.
   You can write the results to disk by calling ``write_statistics()``.

Installation
------------

- From pypi::

    pip install mpiprof

- From a local checkout, by cloning this repository::

    git clone https://github.com/JohannesBuchner/mpiprof.py
    cd mpiprof.py/
    pip install .

Requirements
------------
- Python 3.8+
- mpi4py
- An MPI runtime if you plan to launch under mpiexec/srun/etc.

Usage: per-rank cProfile runner
-------------------------------
Run your script under mpiexec (or srun) using the module form:

- ``mpiexec -n 4 python3 -m mpiprof your_script.py arg1 arg2``

This writes one profile file per rank named ``mpiprof.<rank>.pstats``
in the current working directory.

Options:

- ``-o / --outfile`` sets the output path or pattern (default:
  ``mpiprof.{rank}.pstats``). The literal substring ``{rank}`` will be
  replaced with the MPI rank.

Signal handling:

- mpiprof installs SIGINT and SIGTERM handlers to dump the profile
  before exiting, so Ctrl-C or a clean termination should still produce
  output. If the MPI launcher escalates to SIGKILL immediately, no tool
  can save a profile.

Usage: MPIProfiler for mpi4py
-----------------------------
The ``MPIProfiler`` wraps a communicator and measures wall-clock time
and call counts for common blocking MPI calls. It is intentionally
lightweight and safe to leave in production runs (low overhead on the
measured calls). Nonblocking calls are counted but not timed precisely
(the time recorded is call overhead, not the transfer time).

Example:

.. code-block:: python

    from mpi4py import MPI
    from mpiprof import MPIProfiler

    comm = MPIProfiler(MPI.COMM_WORLD)

    # Your MPI code as usual, using comm instead of MPI.COMM_WORLD:
    rank = comm.Get_rank()
    size = comm.Get_size()
    data = rank

    # Simple collective
    total = comm.allreduce(data, op=MPI.SUM)

    # Point-to-point
    if rank == 0:
        comm.send(b"hello", dest=1, tag=0)
    elif rank == 1:
        msg = comm.recv(source=0, tag=0)

    # Write stats at the end (one file per rank)
    comm.write_statistics()  # default name: mpiprof.stats.<rank>.json

Notes:

- The wrapper exposes common methods (e.g., ``send``, ``recv``, ``bcast``,
  ``reduce``, ``allreduce``, ``gather``, ``scatter``, ``barrier``) and
  forwards any other attributes to the underlying communicator. It tries
  to be case-insensitive to match mpi4py idioms (both ``Send`` and
  ``send`` are supported).

- For nonblocking operations (``Isend``, ``Irecv``), the wrapper records
  the call count but cannot attribute data transfer time unless you also
  wrap and time ``Wait``/``Waitall``. A simple ``wait`` wrapper is
  provided to time individual requests returned by the wrapper.

Rank detection
--------------
The runner tries to detect the rank via common environment variables:

- ``OMPI_COMM_WORLD_RANK``, ``PMIX_RANK`` (Open MPI)
- ``PMI_RANK`` (MPICH, Intel MPI)
- ``MV2_COMM_WORLD_RANK`` (MVAPICH2)
- ``SLURM_PROCID`` (scheduler fallback)
- default: 0 if none found

Output files
------------
- Runner: ``mpiprof.<rank>.pstats`` (or the pattern you set with ``-o``).
  You can analyze it with ``pstats`` or tools like ``snakeviz``:

  - ``python3 -m pstats mpiprof.0.pstats``
  - ``snakeviz mpiprof.0.pstats``

- MPIProfiler: ``MPIprofile.<rank>.out`` with operation counts and
  total wall-clock time per operation. Example::

    Function: scatter
    Call stack:
            surveymcmc.py:466
            surveymcmc.py:431 sampler.advance_one_step(False)
            mcaeis.py:107 local_coords = self.comm.scatter(chunks, root=0)
    Number of calls: 2
    Duration During Call: 98.990435s
    Duration Before Call: 0.000014s

    Function: scatter
    Call stack:
            surveymcmc.py:466
            surveymcmc.py:448 sampler.advance_one_step()
            mcaeis.py:75 self._init()
            mcaeis.py:55 self.log_probs = self._evaluate_log_probs(self.coords)
            mcaeis.py:64 local_coords = self.comm.scatter(chunks, root=0)
    Number of calls: 39
    Duration During Call: 2.228181s
    Duration Before Call: 0.000025s
    
    Total MPI Time: 388.922329s
    Total Non-MPI Time: 206.503627s

Results are sorted (descending) by duration during call.

Limitations
-----------
- The runner cannot save profiles if the process is killed by SIGKILL.
- MPIProfiler’s accounting for nonblocking calls is approximate unless
  you consistently call ``wait``/``waitall`` on the requests returned
  by the wrapper’s nonblocking methods.

License
-------
MIT