Skip to content

Helper function for MPI-accelerated GST fits#724

Open
rileyjmurray wants to merge 13 commits intodevelopfrom
easy-mpi-run-gst-2
Open

Helper function for MPI-accelerated GST fits#724
rileyjmurray wants to merge 13 commits intodevelopfrom
easy-mpi-run-gst-2

Conversation

@rileyjmurray
Copy link
Contributor

@rileyjmurray rileyjmurray commented Mar 17, 2026

This PR adds Protocol.run_mpi() to run MPI-accelerated protocols launched as a subprocess. This function can be called directly from a Jupyter notebook, an interactive Python session, or a plain script. It removes the need for the user to know about mpi4py (besides its existence) and the need to write their own MPI launcher scripts.

Key changes

Added Protocol.run_mpi(...) to pygsti/protocols/protocol.py. Notable features of this function:

  • Accepts num_ranks, an optional work_dir, forwarded **run_kwargs, and several keyword arguments for controlling launcher selection, thread counts, and dry-run behavior.
  • When num_ranks=1, it short-circuits to a plain self.run() call with no subprocess overhead.
  • Auto-detects the MPI launcher (mpiexec, mpirun, or mpiexec.hydra) when mpiexec='auto'.
  • Auto-computes BLAS thread counts per rank (OMP_NUM_THREADS, MKL_NUM_THREADS, etc.) based on physical core count, with a blas_threads_per_rank override.
  • Supports dry_run=True for generating the launch command and writing data to a permanent work_dir without executing. This is useful for embedding into HPC batch scripts.

This new function is tested in test/integration/test_gst_run_mpi.py and demonstrated in
jupyter_notebooks/Examples/MPI-RunningGST-Automagic.ipynb.

Minor change: tests in test/unit/mpi/test_mpi.py have been made a little more robust. The MPI launcher checks both mpiexec and mpirun. We pass --oversubscribe as appropriate so our use of 4 MPI ranks doesn't break on systems with < 4 physical CPU cores.

Incidental changes

  • pygsti/tools/optools.py: extract real part of trace for Hermitian matrices (makes warning messages clearer)
  • pygsti/report/colormaps.py, plotly_plot_ex.py: robustness against plotly version differences.
  • numpy.core namespace access removed (deprecated in NumPy 2.x).
  • Removed a spurious warning about non-MPI-accelerated gauge optimization.

@rileyjmurray rileyjmurray marked this pull request as ready for review March 17, 2026 18:27
@rileyjmurray rileyjmurray requested review from a team as code owners March 17, 2026 18:27
@rileyjmurray rileyjmurray requested a review from sserita March 17, 2026 18:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant