westpa 2022.13__cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

westpa/work_managers/mpi.py

@@ -0,0 +1,318 @@
+"""
+A work manager which uses MPI to distribute tasks and collect results.
+"""
+
+import sys
+import time
+import traceback
+import logging
+import threading
+from collections import deque
+from mpi4py import MPI
+from westpa.work_managers import WorkManager, WMFuture
+
+log = logging.getLogger(__name__)
+
+
+# +------+
+# | Task |
+# +------+
+class Task:
+    """Tasks are tuples of (task_id, function, args, keyword args)"""
+
+    def __init__(self, task_id, fn, args, kwargs):
+        self.task_id = task_id
+        self.fn = fn
+        self.args = args if args is not None else ()
+        self.kwargs = kwargs if kwargs is not None else {}
+
+    def __repr__(self):
+        return '<Task {self.task_id}: {self.fn!r}(*{self.args!r}, **{self.kwargs!r})>'.format(self=self)
+
+
+# +----------------+
+# | MPIWorkManager |
+# +----------------+
+class MPIWorkManager(WorkManager):
+    """MPIWorkManager factory."""
+
+    @classmethod
+    def from_environ(cls, wmenv=None):
+        return cls()
+
+    def __new__(cls):
+        """Creates a Serial WorkManager if size is 1.  Otherwise creates a
+        single Manager and size-1 Worker.
+        """
+        log.debug('MPIWorkManager.__new__()')
+        assert MPI.Is_initialized()
+        assert MPI.Is_thread_main()
+
+        rank = MPI.COMM_WORLD.Get_rank()
+        size = MPI.COMM_WORLD.Get_size()
+
+        if size == 1:
+            return super().__new__(Serial)
+        elif rank == 0:
+            return super().__new__(Manager)
+        else:
+            return super().__new__(Worker)
+
+    def __init__(self):
+        """Initialize info shared by Manager and Worker classes."""
+        log.debug('MPIWorkManager.__init__()')
+
+        super().__init__()
+        comm = MPI.COMM_WORLD
+        self.comm = MPI.COMM_WORLD
+        self.rank = comm.Get_rank()
+        self.size = comm.Get_size()
+        self.name = MPI.Get_processor_name()
+
+        # some tags
+        self.task_tag = 110  # tag for server to client msgs
+        self.result_tag = 120  # tag for client to server msgs
+        self.shutdown_tag = 130  # tag for server to client to stop
+
+        self.managerID = 0
+
+    def submit(self, fn, args=None, kwargs=None):
+        """Adhere to WorkManager interface.  This method should never be
+        called.
+        """
+        assert False
+
+
+# +--------+
+# | Serial |
+# +--------+
+# TODO: no need for the code replication here, just use the original serial wm
+class Serial(MPIWorkManager):
+    """Replication of the serial work manager.  This is a fallback for MPI runs
+    that request only 1 (size=1) processor.
+    """
+
+    def __init__(self):
+        super().__init__()
+        log.debug('Serial.__init__()')
+
+    def submit(self, fn, args=None, kwargs=None):
+        log.debug('Serial.__init__()')
+
+        ft = WMFuture()
+        try:
+            result = fn(*(args if args is not None else ()), **(kwargs if kwargs is not None else {}))
+        except Exception as e:
+            ft._set_exception(e, sys.exc_info()[2])
+        else:
+            ft._set_result(result)
+        return ft
+
+
+# +---------+
+# | Manager |
+# +---------+
+class Manager(MPIWorkManager):
+    """Manager of the MPIWorkManage.  Distributes tasks to Worker as they are
+    received from the sim_manager.  In addition to the main thread, this class
+    spawns two threads, a receiver and a dispatcher.
+    """
+
+    def __init__(self):
+        """Initialize different state variables used by Manager."""
+        super().__init__()
+        log.debug('Manager__init__()')
+
+        # number of workers
+        self.nworkers = self.size - 1
+
+        # list of worker ranks
+        self.workerIDs = range(1, self.size)
+        assert self.nworkers == len(self.workerIDs)
+
+        # deque of idle workers
+        self.dests = deque(self.workerIDs)
+
+        # deque of tesks
+        self.tasks = deque()
+
+        # number of unmatched send/receive pairs
+        self.nPending = 0
+
+        # thread shutdown sentinel
+        self.shutItDown = False
+
+        # task_id, future key value pair
+        self.pending_futures = dict()
+
+        # list of manager threads
+        self.workers = []
+
+        # thread lock
+        self.lock = threading.Lock()
+
+    def startup(self):
+        """Spawns the dispatcher and receiver threads."""
+        log.debug('Manager.startup()')
+
+        if not self.running:
+            self.workers.append(threading.Thread(name='dispatcher', target=self._dispatcher))
+            self.workers.append(threading.Thread(name='receiver', target=self._receiver))
+
+            for t in self.workers:
+                t.start()
+                log.info('Started thread: %s' % t.getName())
+
+            self.running = True
+
+    def _dispatcher(self):
+        """Continuously dispatches tasks to idle destinations until the
+        shutdown sentinel is set.
+        """
+        log.debug('Manager._dispatcher()')
+        assert MPI.Is_thread_main() is False
+        assert threading.currentThread().getName() == "dispatcher"
+
+        while not self.shutItDown:
+            req = []
+            # do we have work and somewhere to send it?
+            while self.tasks and self.dests:
+                with self.lock:
+                    task = self.tasks.popleft()
+                    sendTo = self.dests.popleft()
+                    self.nPending += 1
+
+                req.append(self.comm.isend(task, dest=sendTo, tag=self.task_tag))
+
+            # make sure all sends completed
+            MPI.Request.Waitall(requests=req)
+
+            # force context switching ( 1ms )
+            time.sleep(0.001)
+
+    def _receiver(self):
+        """Continuously receives futures from workers until the shutdown
+        sentinel is set.
+        """
+        log.debug('Manager._receiver()')
+        assert MPI.Is_thread_main() is False
+        assert threading.currentThread().getName() == "receiver"
+
+        while not self.shutItDown:
+            # are we waiting on any results?
+            while self.nPending:
+                stat = MPI.Status()
+                (tid, msg, val) = self.comm.recv(source=MPI.ANY_SOURCE, tag=self.result_tag, status=stat)
+                log.debug('Manager._receiver received task: %s' % tid)
+
+                # update future
+                ft = self.pending_futures.pop(tid)
+                if msg == 'exception':
+                    ft._set_exception(*val)
+                else:
+                    ft._set_result(val)
+
+                with self.lock:
+                    self.dests.append(stat.Get_source())
+                    self.nPending -= 1
+
+            # force context switching ( 1ms )
+            time.sleep(0.001)
+
+    def submit(self, fn, args=None, kwargs=None):
+        """Receive task from simulation manager and add it to pending_futures."""
+        log.debug('Manager.submit()')
+
+        ft = WMFuture()
+        task_id = ft.task_id
+        with self.lock:
+            self.tasks.append(Task(task_id, fn, args, kwargs))
+            self.pending_futures[task_id] = ft
+
+        return ft
+
+    def shutdown(self):
+        """Send shutdown tag to all worker processes, and set the shutdown
+        sentinel to stop the receiver and dispatcher loops.
+        """
+        log.debug('Manager.shutdown()')
+
+        # wait on any unfinished work
+        while self.pending_futures:
+            pass
+
+        # send shutdown msg to all workers
+        req = [MPI.REQUEST_NULL] * self.nworkers
+        for rank in self.workerIDs:
+            req[rank - 1] = self.comm.isend(MPI.BOTTOM, dest=rank, tag=self.shutdown_tag)
+
+        MPI.Request.Waitall(requests=req)
+
+        # stop threads
+        self.shutItDown = True
+
+        for t in self.workers:
+            t.join()
+            log.info('Stopped thread: %s' % t.getName())
+
+        self.running = False
+
+
+# +--------+
+# | Worker |
+# +--------+
+class Worker(MPIWorkManager):
+    """Client class for executing tasks as distributed by the Manager in the
+    MPI Work Manager
+    """
+
+    def __init__(self):
+        super().__init__()
+        log.debug('Worker.__init__() %s' % self.rank)
+
+    def startup(self):
+        """Clock the worker in for work."""
+        log.debug('Worker.startup() %s' % self.rank)
+        if not self.running:
+            self.clockIn()
+
+            self.running = True
+
+    def clockIn(self):
+        """Do each task as it comes in.  The completion of a task is
+        notice to the manager that more work is welcome.
+        """
+        log.info('Worker %s clocking in.' % self.rank)
+
+        comm = self.comm
+
+        while True:
+            stat = MPI.Status()
+            task = comm.recv(source=self.managerID, tag=MPI.ANY_TAG, status=stat)
+
+            tag = stat.Get_tag()
+
+            if tag == self.task_tag:
+                log.debug('Worker %s received task: %s' % (self.rank, task.task_id))
+
+                # do the work
+                try:
+                    rv = task.fn(*task.args, **task.kwargs)
+                except BaseException as e:
+                    ro = (task.task_id, 'exception', (e, traceback.format_exc()))
+                else:
+                    ro = (task.task_id, 'result', rv)
+
+                # send result back to manager
+                comm.send(ro, dest=self.managerID, tag=self.result_tag)
+
+            if tag == self.shutdown_tag:
+                log.info('Worker %s clocking out.' % self.rank)
+                return
+
+    @property
+    def is_master(self):
+        """Worker processes need to be marked as not manager.  This ensures that
+        the proper branching is followed in w_run.py.
+        """
+        return False

westpa/work_managers/processes.py

@@ -0,0 +1,201 @@
+import os
+import sys
+import random
+import logging
+import threading
+import traceback
+import multiprocessing
+from multiprocessing.queues import Empty
+
+import westpa.work_managers as work_managers
+from .core import WorkManager, WMFuture
+
+log = logging.getLogger(__name__)
+
+# Tasks are tuples ('task', task_id, fn, args, kwargs).
+# Results are tuples (rtype, task_id, payload) where rtype is 'result' or 'exception' and payload is the return value
+# or exception, respectively.
+
+task_shutdown_sentinel = ('shutdown', None, None, (), {})
+result_shutdown_sentinel = ('shutdown', None, None)
+
+
+class ProcessWorkManager(WorkManager):
+    '''A work manager using the ``multiprocessing`` module.
+
+    Notes
+    -----
+
+    On MacOS, as of Python 3.8 the default start method for multiprocessing launching new processes was changed from fork to spawn.
+    On Linux, as of Python 3.14, the default start method for multiprocessing launching new processes was changed from fork to spawn.
+    In general, spawn is more robust and efficient, however it requires serializability of everything being passed to the child process.
+    In contrast, fork is much less memory efficient, as it makes a full copy of everything in the parent process.
+    However, it does not require picklability.
+
+    So, on MacOS and Linux, the method for launching new processes is explicitly changed to fork from the (UNIX-specific) default of spawn.
+    Unix should default to fork.
+
+    See https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods and
+    https://docs.python.org/3/library/multiprocessing.html#the-spawn-and-forkserver-start-methods for more details.
+    '''
+
+    @classmethod
+    def from_environ(cls, wmenv=None):
+        if wmenv is None:
+            wmenv = work_managers.environment.default_env
+        return cls(wmenv.get_val('n_workers', multiprocessing.cpu_count(), int))
+
+    def __init__(self, n_workers=None, shutdown_timeout=1):
+        super().__init__()
+
+        try:
+            if sys.platform in ['darwin', 'linux']:  # UNIX Platforms
+                multiprocessing.set_start_method('fork')
+                log.debug('setting multiprocessing start method to fork')
+        except RuntimeError:
+            log.debug('failed to set start method to fork')
+
+        self.n_workers = n_workers or multiprocessing.cpu_count()
+        self.workers = None
+        self.task_queue = multiprocessing.Queue()
+        self.result_queue = multiprocessing.Queue()
+        self.receive_thread = None
+        self.pending = None
+
+        self.shutdown_received = threading.Event()
+        self.shutdown_timeout = shutdown_timeout or 1
+
+    def task_loop(self):
+        # Close standard input, so we don't get SIGINT from ^C
+        try:
+            sys.stdin.close()
+        except Exception as e:
+            log.info("can't close stdin: {}".format(e))
+
+        # (re)initialize random number generator in this process
+        random.seed()
+
+        while not self.shutdown_received.is_set():
+            if not self.task_queue.empty():
+                message, task_id, fn, args, kwargs = self.task_queue.get()[:5]
+
+                if message == 'shutdown':
+                    break
+                try:
+                    result = fn(*args, **kwargs)
+                except BaseException as e:
+                    result_tuple = ('exception', task_id, (e, traceback.format_exc()))
+                else:
+                    result_tuple = ('result', task_id, result)
+                self.result_queue.put(result_tuple)
+
+        log.debug('exiting task_loop')
+        return
+
+    def results_loop(self):
+        while not self.shutdown_received.is_set():
+            if not self.result_queue.empty():
+                message, task_id, payload = self.result_queue.get()[:3]
+
+                if message == 'shutdown':
+                    break
+                elif message == 'exception':
+                    future = self.pending.pop(task_id)
+                    future._set_exception(*payload)
+                elif message == 'result':
+                    future = self.pending.pop(task_id)
+                    future._set_result(payload)
+                else:
+                    raise AssertionError('unknown message {!r}'.format((message, task_id, payload)))
+
+        log.debug('exiting results_loop')
+        return
+
+    def submit(self, fn, args=None, kwargs=None):
+        ft = WMFuture()
+        log.debug('dispatching {!r}'.format(fn))
+        self.pending[ft.task_id] = ft
+        self.task_queue.put(('task', ft.task_id, fn, args or (), kwargs or {}))
+        return ft
+
+    def startup(self):
+        from . import environment
+
+        if not self.running:
+            log.debug('starting up work manager {!r}'.format(self))
+            self.running = True
+            self.workers = [
+                multiprocessing.Process(target=self.task_loop, name='worker-{:d}-{:x}'.format(i, id(self)))
+                for i in range(self.n_workers)
+            ]
+
+            pi_name = '{}_PROCESS_INDEX'.format(environment.WMEnvironment.env_prefix)
+            for iworker, worker in enumerate(self.workers):
+                os.environ[pi_name] = str(iworker)
+                worker.start()
+            try:
+                del os.environ[pi_name]
+            except KeyError:
+                pass
+
+            self.pending = dict()
+
+            self.receive_thread = threading.Thread(target=self.results_loop, name='receiver')
+            self.receive_thread.daemon = True
+            self.receive_thread.start()
+
+    def _empty_queues(self):
+        '''Empty self.task_queue and self.result_queue until queue is empty'''
+        try:
+            while True:
+                self.task_queue.get_nowait()
+        except (Empty, ValueError):
+            log.debug('Empty task_queue')
+
+        try:
+            while True:
+                self.result_queue.get_nowait()
+        except (Empty, ValueError):
+            log.debug('Empty result_queue')
+
+    def shutdown(self):
+        while self.running:
+            log.debug('shutting down {!r}'.format(self))
+
+            # Empty queues
+            self._empty_queues()
+            for _i in range(self.n_workers):
+                self.task_queue.put_nowait(task_shutdown_sentinel)
+            self.result_queue.put_nowait(result_shutdown_sentinel)
+
+            # Signal shutdown Event to stop queue loops
+            self.shutdown_received.set()
+            self._empty_queues()
+
+            # Terminating all workers
+            for worker in self.workers:
+                worker.join(self.shutdown_timeout)
+                if worker.is_alive():
+                    log.debug('sending SIGINT to worker process {:d}'.format(worker.pid))
+                    worker.terminate()
+                    worker.join(self.shutdown_timeout)
+                    if worker.is_alive():
+                        log.warning('sending SIGKILL to worker process {:d}'.format(worker.pid))
+                        worker.kill()
+                        worker.join()
+
+                    log.debug('worker process {:d} terminated with code {:d}'.format(worker.pid, worker.exitcode))
+                else:
+                    log.debug('worker process {:d} terminated gracefully with code {:d}'.format(worker.pid, worker.exitcode))
+
+                try:
+                    worker.close()  # Release all resources
+                except ValueError:
+                    try:
+                        if worker.is_alive():
+                            log.debug('worker process {:d} could not be closed'.format(worker.pid))
+                    except ValueError:
+                        pass  # Already closed.
+
+            self.running = False
+            log.debug('Done shutting down the processes work manager')

westpa/work_managers/serial.py

@@ -0,0 +1,28 @@
+import logging
+import sys
+
+from .core import WorkManager, WMFuture
+
+
+log = logging.getLogger(__name__)
+
+
+class SerialWorkManager(WorkManager):
+    @classmethod
+    def from_environ(cls, wmenv=None):
+        return cls()
+
+    def __init__(self):
+        log.debug('initializing serial work manager')
+        super().__init__()
+        self.n_workers = 1
+
+    def submit(self, fn, args=None, kwargs=None):
+        ft = WMFuture()
+        try:
+            result = fn(*(args if args is not None else ()), **(kwargs if kwargs is not None else {}))
+        except Exception as e:
+            ft._set_exception(e, sys.exc_info()[2])
+        else:
+            ft._set_result(result)
+        return ft

westpa/work_managers/threads.py

@@ -0,0 +1,79 @@
+import logging
+import multiprocessing
+import queue
+import sys
+import threading
+
+from .core import WorkManager, WMFuture
+import westpa.work_managers as work_managers
+
+log = logging.getLogger(__name__)
+
+
+class Task:
+    def __init__(self, fn, args, kwargs, future):
+        self.fn = fn
+        self.args = args
+        self.kwargs = kwargs
+        self.future = future
+
+    def run(self):
+        try:
+            result = self.fn(*self.args, **self.kwargs)
+        except Exception as e:
+            self.future._set_exception(e, sys.exc_info()[2])
+        else:
+            self.future._set_result(result)
+
+
+ShutdownSentinel = object()
+
+
+class ThreadsWorkManager(WorkManager):
+    '''A work manager using threads.'''
+
+    @classmethod
+    def from_environ(cls, wmenv=None):
+        if wmenv is None:
+            wmenv = work_managers.environment.default_env
+        return cls(wmenv.get_val('n_workers', multiprocessing.cpu_count(), int))
+
+    def __init__(self, n_workers=None):
+        super().__init__()
+        self.n_workers = n_workers or multiprocessing.cpu_count()
+        self.workers = []
+        self.task_queue = queue.Queue()
+
+    def runtask(self, task_queue):
+        while True:
+            task = task_queue.get()
+            if task is ShutdownSentinel:
+                return
+            else:
+                task.run()
+
+    def submit(self, fn, args=None, kwargs=None):
+        ft = WMFuture()
+        task = Task(fn, args if args is not None else (), kwargs if kwargs is not None else {}, ft)
+        self.task_queue.put(task)
+        return ft
+
+    def startup(self):
+        if not self.running:
+            self.running = True
+            self.workers = [
+                threading.Thread(target=self.runtask, args=[self.task_queue], name='worker-{:d}'.format(i))
+                for i in range(0, self.n_workers)
+            ]
+            for thread in self.workers:
+                log.debug('starting thread {!r}'.format(thread))
+                thread.start()
+
+    def shutdown(self):
+        if self.running:
+            # Put one sentinel on the queue per worker, then wait for threads to terminate
+            for i in range(0, self.n_workers):
+                self.task_queue.put(ShutdownSentinel)
+            for thread in self.workers:
+                thread.join()
+            self.running = False

westpa/work_managers/zeromq/__init__.py

@@ -0,0 +1,20 @@
+from .core import ZMQWMError, ZMQWMTimeout, ZMQWMEnvironmentError, ZMQWorkerMissing, ZMQCore
+from .node import ZMQNode
+from .worker import ZMQWorker
+from .work_manager import ZMQWorkManager
+
+import atexit
+
+__all__ = [
+    'ZMQWMError',
+    'ZMQWMTimeout',
+    'ZMQWMEnvironmentError',
+    'ZMQWorkerMissing',
+    'ZMQCore',
+    'ZMQNode',
+    'ZMQWorker',
+    'ZMQWorkManager',
+]
+
+atexit.register(ZMQCore.remove_ipc_endpoints)
+del atexit