quimb.linalg.mpi_launcher ========================= .. py:module:: quimb.linalg.mpi_launcher .. autoapi-nested-parse:: Manages the spawning of mpi processes to send to the various solvers. Attributes ---------- .. autoapisummary:: quimb.linalg.mpi_launcher.QUIMB_MPI_LAUNCHED quimb.linalg.mpi_launcher.ALLOW_SPAWN quimb.linalg.mpi_launcher.NUM_MPI_WORKERS quimb.linalg.mpi_launcher.eigs_slepc_mpi quimb.linalg.mpi_launcher.eigs_slepc_spawn quimb.linalg.mpi_launcher.svds_slepc_mpi quimb.linalg.mpi_launcher.svds_slepc_spawn quimb.linalg.mpi_launcher.mfn_multiply_slepc_mpi quimb.linalg.mpi_launcher.mfn_multiply_slepc_spawn quimb.linalg.mpi_launcher.ssolve_slepc_mpi quimb.linalg.mpi_launcher.ssolve_slepc_spawn Classes ------- .. autoapisummary:: quimb.linalg.mpi_launcher.SyncroFuture quimb.linalg.mpi_launcher.SynchroMPIPool quimb.linalg.mpi_launcher.CachedPoolWithShutdown quimb.linalg.mpi_launcher.GetMPIBeforeCall quimb.linalg.mpi_launcher.SpawnMPIProcessesFunc Functions --------- .. autoapisummary:: quimb.linalg.mpi_launcher.can_use_mpi_pool quimb.linalg.mpi_launcher.bcast quimb.linalg.mpi_launcher.get_mpi_pool Module Contents --------------- .. py:data:: QUIMB_MPI_LAUNCHED .. py:data:: ALLOW_SPAWN .. py:data:: NUM_MPI_WORKERS .. py:function:: can_use_mpi_pool() Function to determine whether we are allowed to call `get_mpi_pool`. .. py:function:: bcast(result, comm, result_rank) Broadcast a result to all workers, dispatching to proper MPI (rather than pickled) communication if the result is a numpy array. .. py:class:: SyncroFuture(result, result_rank, comm) .. py:attribute:: _result .. py:attribute:: result_rank .. py:attribute:: comm .. py:method:: result() .. py:method:: cancel() :staticmethod: .. py:class:: SynchroMPIPool An object that looks like a ``concurrent.futures`` executor but actually distributes tasks in a round-robin fashion based to MPI workers, before broadcasting the results to each other. .. py:attribute:: comm .. py:attribute:: size .. py:attribute:: rank .. py:attribute:: counter .. py:attribute:: _max_workers .. py:method:: submit(fn, *args, **kwargs) .. py:method:: shutdown() .. py:class:: CachedPoolWithShutdown(pool_fn) Decorator for caching the mpi pool when called with the equivalent args, and shutting down previous ones when not needed. .. py:attribute:: _settings :value: '__UNINITIALIZED__' .. py:attribute:: _pool_fn .. py:method:: __call__(num_workers=None, num_threads=1) .. py:function:: get_mpi_pool(num_workers=None, num_threads=1) Get the MPI executor pool, with specified number of processes and threads per process. .. py:class:: GetMPIBeforeCall(fn) Bases: :py:obj:`object` Wrap a function to automatically get the correct communicator before its called, and to set the `comm_self` kwarg to allow forced self mode. This is called by every mpi process before the function evaluation. .. py:attribute:: fn .. py:method:: __call__(*args, comm_self=False, wait_for_workers=None, **kwargs) :param args: Supplied to self.fn :param comm_self: Whether to force use of MPI.COMM_SELF :type comm_self: bool, optional :param wait_for_workers: If set, wait for the communicator to have this many workers, this can help to catch some errors regarding expected worker numbers. :type wait_for_workers: int, optional :param kwargs: Supplied to self.fn .. py:class:: SpawnMPIProcessesFunc(fn) Bases: :py:obj:`object` Automatically wrap a function to be executed in parallel by a pool of mpi workers. This is only called by the master mpi process in manual mode, only by the (non-mpi) spawning process in automatic mode, or by all processes in syncro mode. .. py:attribute:: fn .. py:method:: __call__(*args, num_workers=None, num_threads=1, mpi_pool=None, spawn_all=USE_SYNCRO or not ALREADY_RUNNING_AS_MPI, **kwargs) :param args: Supplied to `self.fn`. :param num_workers: How many total process should run function in parallel. :type num_workers: int, optional :param num_threads: How many (OMP) threads each process should use :type num_threads: int, optional :param mpi_pool: If not None (default), submit function to this pool. :type mpi_pool: pool-like, optional :param spawn_all: Whether all the parallel processes should be spawned (True), or num_workers - 1, so that the current process can also do work. :type spawn_all: bool, optional :param kwargs: Supplied to `self.fn`. :rtype: `fn` output from the master process. .. py:data:: eigs_slepc_mpi .. py:data:: eigs_slepc_spawn .. py:data:: svds_slepc_mpi .. py:data:: svds_slepc_spawn .. py:data:: mfn_multiply_slepc_mpi .. py:data:: mfn_multiply_slepc_spawn .. py:data:: ssolve_slepc_mpi .. py:data:: ssolve_slepc_spawn