Source code for hnn_core.parallel_backends

"""Parallel backends"""

# Authors: Blake Caldwell <blake_caldwell@brown.edu>
#          Mainak Jas <mainakjas@gmail.com>

import os
import sys
import re
import multiprocessing
import shlex
import pickle
import base64
from warnings import warn
from subprocess import Popen, PIPE, TimeoutExpired
import binascii
from queue import Queue, Empty
from threading import Thread, Event

from .cell_response import CellResponse
from .dipole import Dipole
from .network_builder import _simulate_single_trial

_BACKEND = None


def _thread_handler(event, out, queue):
    while not event.is_set():
        line = out.readline()
        if line == '':
            break
        queue.put(line)


def _gather_trial_data(sim_data, net, n_trials, postproc):
    """Arrange data by trial

    To be called after simulate(). Returns list of Dipoles, one for each trial,
    and saves spiking info in net (instance of Network).
    """
    dpls = list()

    # Create array of equally sampled time points for simulating currents
    cell_type_names = list(net.cell_types.keys())
    cell_response = CellResponse(times=sim_data[0]['times'],
                                 cell_type_names=cell_type_names)
    net.cell_response = cell_response

    for idx in range(n_trials):

        # cell response
        net.cell_response._spike_times.append(sim_data[idx]['spike_times'])
        net.cell_response._spike_gids.append(sim_data[idx]['spike_gids'])
        net.cell_response.update_types(net.gid_ranges)
        net.cell_response._vsec.append(sim_data[idx]['vsec'])
        net.cell_response._isec.append(sim_data[idx]['isec'])

        # extracellular array
        for arr_name, arr in net.rec_arrays.items():
            # voltages is a n_trials x n_contacts x n_samples array
            arr._data.append(sim_data[idx]['rec_data'][arr_name])
            arr._times = sim_data[idx]['rec_times'][arr_name]

        # dipole
        dpl = Dipole(times=sim_data[idx]['times'],
                     data=sim_data[idx]['dpl_data'])

        N_pyr_x = net._params['N_pyr_x']
        N_pyr_y = net._params['N_pyr_y']
        dpl._baseline_renormalize(N_pyr_x, N_pyr_y)  # XXX cf. #270
        dpl._convert_fAm_to_nAm()  # always applied, cf. #264
        if postproc:
            window_len = net._params['dipole_smooth_win']  # specified in ms
            fctr = net._params['dipole_scalefctr']
            if window_len > 0:  # param files set this to zero for no smoothing
                dpl.smooth(window_len=window_len)
            if fctr > 0:
                dpl.scale(fctr)
        dpls.append(dpl)

    return dpls


def _get_mpi_env():
    """Set some MPI environment variables."""
    my_env = os.environ.copy()
    if 'win' not in sys.platform:
        my_env["OMPI_MCA_btl_base_warn_component_unused"] = '0'

    if 'darwin' in sys.platform:
        my_env["PMIX_MCA_gds"] = "^ds12"  # open-mpi/ompi/issues/7516
        my_env["TMPDIR"] = "/tmp"  # open-mpi/ompi/issues/2956
    return my_env


def run_subprocess(command, obj, timeout, proc_queue=None, *args, **kwargs):
    """Run process and communicate with it.
    Parameters
    ----------
    command : list of str | str
        Command to run as subprocess (see subprocess.Popen documentation).
    obj : object
        The object to write to stdin after starting child process
        with MPI command.
    timeout : float
        The number of seconds to wait for a process without output.
    *args, **kwargs : arguments
        Additional arguments to pass to subprocess.Popen.
    Returns
    -------
    child_data : object
        The data returned by the child process.
    """
    proc_data_bytes = b''
    # each loop while waiting will involve two Queue.get() timeouts, each
    # 0.01s. This caclulation will error on the side of a longer timeout
    # than is specified because more is done each loop that just Queue.get()
    timeout_cycles = timeout / 0.02

    pickled_obj = base64.b64encode(pickle.dumps(obj))

    # non-blocking adapted from https://stackoverflow.com/questions/375427/non-blocking-read-on-a-subprocess-pipe-in-python#4896288  # noqa: E501
    out_q = Queue()
    err_q = Queue()

    threads_started = False

    try:
        proc = Popen(command, stdin=PIPE, stdout=PIPE, stderr=PIPE, *args,
                     **kwargs)

        # now that the process has started, add it to the queue
        # used by MPIBackend.terminate()
        if proc_queue is not None:
            proc_queue.put(proc)

        # set up polling first so all of child's stdout/stderr
        # gets captured
        event = Event()
        out_t = Thread(target=_thread_handler,
                       args=(event, proc.stdout, out_q))
        err_t = Thread(target=_thread_handler,
                       args=(event, proc.stderr, err_q))
        out_t.start()
        err_t.start()
        threads_started = True
        data_received = False
        sent_network = False
        count_since_last_output = 0

        # loop while the process is running the simulation
        while True:
            child_terminated = proc.poll() is not None

            if not data_received:
                if _echo_child_output(out_q):
                    count_since_last_output = 0
                else:
                    count_since_last_output += 1
                # look for data in stderr and print child stdout
                data_len, proc_data_bytes = _get_data_from_child_err(err_q)
                if data_len > 0:
                    data_received = True
                    _write_child_exit_signal(proc.stdin)
                elif child_terminated:
                    # child terminated early, and we already
                    # captured output left in queues
                    warn("Child process failed unexpectedly")
                    kill_proc_name('nrniv')
                    break

            if not sent_network:
                # Send network object to child so it can start
                try:
                    _write_net(proc.stdin, pickled_obj)
                except BrokenPipeError:
                    # child failed during _write_net(). get the
                    # output and break out of loop on the next
                    # iteration
                    warn("Received BrokenPipeError exception. "
                         "Child process failed unexpectedly")
                    continue
                else:
                    sent_network = True
                    # This is not the same as "network received", but we
                    # assume it was successful and move on to waiting for
                    # data in the next loop iteration.

            if child_terminated and data_received:
                # both exit conditions have been met (also we know that
                # the network has been sent)
                break

            if not child_terminated and \
                    count_since_last_output > timeout_cycles:
                warn("Timeout exceeded while waiting for child process output"
                     ". Terminating...")
                kill_proc_name('nrniv')
                break
    except KeyboardInterrupt:
        warn("Received KeyboardInterrupt. Stopping simulation process...")

    if threads_started:
        # stop the threads
        event.set()  # close signal
        out_t.join()
        err_t.join()

    # wait for the process to terminate. we need use proc.communicate to
    # read any output at its end of life.
    try:
        outs, errs = proc.communicate(timeout=1)
    except TimeoutExpired:
        proc.kill()
        # wait for output again after kill signal
        outs, errs = proc.communicate(timeout=1)

    sys.stdout.write(outs)
    sys.stdout.write(errs)

    if proc.returncode is None:
        # It's theoretically possible that we have received data
        # and exited the loop above, but the child process has not
        # yet terminated. This is unexpected unless KeyboarInterrupt
        # is caught
        proc.terminate()
        try:
            proc.wait(1)  # wait maximum of 1s
        except TimeoutExpired:
            warn("Could not kill python subprocess: PID %d" % proc.pid)

    if not proc.returncode == 0:
        # simulation failed with a numeric return code
        raise RuntimeError("MPI simulation failed. Return code: %d" %
                           proc.returncode)

    child_data = _process_child_data(proc_data_bytes, data_len)

    # clean up the queue
    try:
        proc_queue.get_nowait()
    except Empty:
        pass

    return proc, child_data


def _process_child_data(data_bytes, data_len):
    """Process the data returned by child process.

    Parameters
    ----------
    data_bytes : str
        The data bytes

    Returns
    -------
    data_unpickled : object
        The unpickled data.
    """
    if not data_len == len(data_bytes):
        # This is indicative of a failure. For debugging purposes.
        warn("Length of received data unexpected. Expecting %d bytes, "
             "got %d" % (data_len, len(data_bytes)))

    if len(data_bytes) == 0:
        raise RuntimeError("MPI simulation didn't return any data")

    # decode base64 byte string
    try:
        data_pickled = base64.b64decode(data_bytes, validate=True)
    except binascii.Error:
        # This is here for future debugging purposes. Unit tests can't
        # reproduce an incorrectly padded string, but this has been an
        # issue before
        raise ValueError("Incorrect padding for data length %d bytes" %
                         len(data_len) + " (mod 4 = %d)" %
                         (len(data_len) % 4))

    # unpickle the data
    return pickle.loads(data_pickled)


def _echo_child_output(out_q):
    out = ''
    while True:
        try:
            out += out_q.get(timeout=0.01)
        except Empty:
            break

    if len(out) > 0:
        sys.stdout.write(out)
        return True
    return False


def _get_data_from_child_err(err_q):
    err = ''
    data_length = 0
    data_bytes = b''

    while True:
        try:
            err += err_q.get(timeout=0.01)
        except Empty:
            break

    # check for data signal
    extracted_data = _extract_data(err, 'data')
    if len(extracted_data) > 0:
        # _extract_data only returns data when signals on
        # both sides were seen

        err = err.replace('@start_of_data@', '')
        err = err.replace(extracted_data, '')
        data_length = _extract_data_length(err, 'data')
        err = err.replace('@end_of_data:%d@\n' % data_length, '')
        data_bytes = extracted_data.encode()

    # print the rest of the child's stderr to our stdout
    sys.stdout.write(err)

    return data_length, data_bytes


def _has_mpi4py():
    """Determine if mpi4py is present."""
    try:
        import mpi4py  # noqa
    except ImportError:
        return False
    else:
        return True


def _has_psutil():
    """Determine if psutil is present."""
    try:
        import psutil  # noqa
    except ImportError:
        return False
    else:
        return True


def requires_mpi4py(function):
    """Decorator for testing functions that require MPI."""
    import pytest

    try:
        import mpi4py
        assert hasattr(mpi4py, '__version__')
        skip = False
    except (ImportError, ModuleNotFoundError) as err:
        if "TRAVIS_OS_NAME" not in os.environ:
            skip = True
        else:
            raise ImportError(err)
    reason = 'mpi4py not available'
    return pytest.mark.skipif(skip, reason=reason)(function)


def requires_psutil(function):
    """Decorator for testing functions that require psutil."""
    import pytest

    try:
        import psutil
        assert hasattr(psutil, '__version__')
        skip = False
    except (ImportError, ModuleNotFoundError) as err:
        if "TRAVIS_OS_NAME" not in os.environ:
            skip = True
        else:
            raise ImportError(err)
    reason = 'psutil not available'
    return pytest.mark.skipif(skip, reason=reason)(function)


def _extract_data_length(data_str, object_name):
    data_len_match = re.search('@end_of_%s:' % object_name + r'(\d+)@',
                               data_str)
    if data_len_match is not None:
        return int(data_len_match.group(1))
    else:
        raise ValueError("Couldn't find data length in string")


def _extract_data(data_str, object_name):
    start_idx = 0
    end_idx = 0
    start_match = re.search('@start_of_%s@' % object_name, data_str)
    if start_match is not None:
        start_idx = start_match.end()
    else:
        # need start signal
        return ''

    end_match = re.search('@end_of_%s:' % object_name + r'\d+@', data_str)
    if end_match is not None:
        end_idx = end_match.start()

    return data_str[start_idx:end_idx]


# Next 3 functions are from HNN. Will move here. They require psutil
def _kill_procs(procs):
    """Tries to terminate processes in a list before sending kill signal"""
    from psutil import wait_procs, NoSuchProcess

    # try terminate first
    for p in procs:
        try:
            p.terminate()
        except NoSuchProcess:
            pass
    _, alive = wait_procs(procs, timeout=3)

    # now try kill
    for p in alive:
        p.kill()
    _, alive = wait_procs(procs, timeout=3)

    return alive


def _get_procs_running(proc_name):
    """Return a list of processes currently running"""
    from psutil import process_iter

    process_list = []
    for p in process_iter(attrs=["name", "exe", "cmdline"]):
        if proc_name == p.info['name'] or \
                (p.info['exe'] is not None and
                 os.path.basename(p.info['exe']) == proc_name) or \
                (p.info['cmdline'] and
                 p.info['cmdline'][0] == proc_name):
            process_list.append(p)
    return process_list


def kill_proc_name(proc_name):
    """Make best effort to kill processes

    Parameters
    ----------
    proc_name : str
        A string to match process names against and kill all matches

    Returns
    -------
    killed_procs : bool
        True if any processes were killed
    """

    killed_procs = False
    procs = _get_procs_running(proc_name)
    if len(procs) > 0:
        running = _kill_procs(procs)
        if len(running) > 0:
            if len(running) < len(procs):
                killed_procs = True
            pids = [str(proc.pid) for proc in running]
            warn("Failed to kill nrniv process(es) %s" %
                 ','.join(pids))
        else:
            killed_procs = True

    return killed_procs


def _write_net(stream, pickled_net):
    stream.flush()
    stream.write('@start_of_net@')
    stream.write(pickled_net.decode())
    stream.write('@end_of_net:%d@\n' % len(pickled_net))
    stream.flush()


def _write_child_exit_signal(stream):
    stream.flush()
    stream.write('@data_received@\n')
    stream.flush()


[docs]class JoblibBackend(object): """The JoblibBackend class. Parameters ---------- n_jobs : int | None The number of jobs to start in parallel. If None, then 1 trial will be started without parallelism Attributes ---------- n_jobs : int The number of jobs to start in parallel """ def __init__(self, n_jobs=1): self.n_jobs = n_jobs def _parallel_func(self, func): if self.n_jobs != 1: try: from joblib import Parallel, delayed except ImportError: warn('joblib not installed. Cannot run in parallel.') self.n_jobs = 1 if self.n_jobs == 1: my_func = func parallel = list else: parallel = Parallel(self.n_jobs) my_func = delayed(func) return parallel, my_func def __enter__(self): global _BACKEND self._old_backend = _BACKEND _BACKEND = self return self def __exit__(self, type, value, traceback): global _BACKEND _BACKEND = self._old_backend
[docs] def simulate(self, net, tstop, dt, n_trials, postproc=False): """Simulate the HNN model Parameters ---------- net : Network object The Network object specifying how cells are connected. n_trials : int Number of trials to simulate. tstop : float The simulation stop time (ms). dt : float The integration time step of h.CVode (ms) postproc : bool If False, no postprocessing applied to the dipole Returns ------- dpl: list of Dipole The Dipole results from each simulation trial """ print(f"Joblib will run {n_trials} trial(s) in parallel by " f"distributing trials over {self.n_jobs} jobs.") parallel, myfunc = self._parallel_func(_simulate_single_trial) sim_data = parallel(myfunc(net, tstop, dt, trial_idx) for trial_idx in range(n_trials)) dpls = _gather_trial_data(sim_data, net=net, n_trials=n_trials, postproc=postproc) return dpls
[docs]class MPIBackend(object): """The MPIBackend class. Parameters ---------- n_procs : int | None The number of MPI processes requested by the user. If None, then will attempt to detect number of cores (including hyperthreads) and start parallel simulation over all of them. mpi_cmd : str The name of the mpi launcher executable. Will use 'mpiexec' (openmpi) by default. Attributes ---------- n_procs : int The number of processes MPI will actually use (spread over cores). If 1 is specified or mpi4py could not be loaded, the simulation will be run with the JoblibBackend mpi_cmd : list of str The mpi command with number of procs and options to be passed to Popen expected_data_length : int Used to check consistency between data that was sent and what MPIBackend received. proc_queue : threading.Queue A Queue object to hold process handles from Popen in a thread-safe way. There will be a valid process handle present the queue when a MPI åsimulation is running. """ def __init__(self, n_procs=None, mpi_cmd='mpiexec'): self.expected_data_length = 0 self.proc = None self.proc_queue = Queue() n_logical_cores = multiprocessing.cpu_count() if n_procs is None: self.n_procs = n_logical_cores else: self.n_procs = n_procs # did user try to force running on more cores than available? oversubscribe = False if self.n_procs > n_logical_cores: oversubscribe = True hyperthreading = False if _has_mpi4py() and _has_psutil(): import psutil n_physical_cores = psutil.cpu_count(logical=False) # detect if we need to use hwthread-cpus with mpiexec if self.n_procs > n_physical_cores: hyperthreading = True else: packages = list() if not _has_mpi4py(): packages += ['mpi4py'] if not _has_psutil(): packages += ['psutil'] packages = ' and '.join(packages) warn(f'{packages} not installed. Will run on single processor') self.n_procs = 1 self.mpi_cmd = mpi_cmd if hyperthreading: self.mpi_cmd += ' --use-hwthread-cpus' if oversubscribe: self.mpi_cmd += ' --oversubscribe' self.mpi_cmd += ' -np ' + str(self.n_procs) self.mpi_cmd += ' nrniv -python -mpi -nobanner ' + \ sys.executable + ' ' + \ os.path.join(os.path.dirname(sys.modules[__name__].__file__), 'mpi_child.py') # Split the command into shell arguments for passing to Popen if 'win' in sys.platform: use_posix = True else: use_posix = False self.mpi_cmd = shlex.split(self.mpi_cmd, posix=use_posix) def __enter__(self): global _BACKEND self._old_backend = _BACKEND _BACKEND = self return self def __exit__(self, type, value, traceback): global _BACKEND _BACKEND = self._old_backend # always kill nrniv processes for good measure if self.n_procs > 1: kill_proc_name('nrniv')
[docs] def simulate(self, net, tstop, dt, n_trials, postproc=False): """Simulate the HNN model in parallel on all cores Parameters ---------- net : Network object The Network object specifying how cells are connected. tstop : float The simulation stop time (ms). dt : float The integration time step of h.CVode (ms) n_trials : int Number of trials to simulate. postproc : bool If False, no postprocessing applied to the dipole Returns ------- dpl : list of Dipole The Dipole results from each simulation trial """ # just use the joblib backend for a single core if self.n_procs == 1: print("MPIBackend is set to use 1 core: tranferring the " "simulation to JoblibBackend....") return JoblibBackend(n_jobs=1).simulate(net, tstop=tstop, dt=dt, n_trials=n_trials, postproc=postproc) if self.n_procs > net._n_cells: raise ValueError(f'More MPI processes were assigned than there ' f'are cells in the network. Please decrease ' f'the number of parallel processes (got n_procs=' f'{self.n_procs}) over which you will ' f'distribute the {net._n_cells} network neurons.') print(f"MPI will run {n_trials} trial(s) sequentially by " f"distributing network neurons over {self.n_procs} processes.") env = _get_mpi_env() self.proc, sim_data = run_subprocess( command=self.mpi_cmd, obj=[net, tstop, dt, n_trials], timeout=30, proc_queue=self.proc_queue, env=env, cwd=os.getcwd(), universal_newlines=True) dpls = _gather_trial_data(sim_data, net, n_trials, postproc) return dpls
[docs] def terminate(self): """Terminate running simulation on this MPIBackend Safe to call from another thread from the one `simulate_dipole` was called from. """ proc = None try: proc = self.proc_queue.get(timeout=1) except Empty: warn("No currently running process to terminate") if proc is not None: proc.terminate() try: proc.wait(5) # wait maximum of 5s except TimeoutExpired: warn("Could not kill python subprocess: PID %d" % proc.pid)