"""Running one point's segments: serially, or across a pool of lanes (#563).
Issue #563's implementation proposal ends "Segment simulations can run in parallel", and
they are the one embarrassingly parallel thing in the method: one augmented-model
evaluation is ``m`` spans of one trajectory integrated from ``m`` states the transcription
already knows, with no data flowing between them. This module is that scheduler.
Threads, because the integration releases the GIL
-------------------------------------------------
The choice between threads and processes is not a matter of taste here, and it was
measured rather than assumed. On the motivating model (``Borghans_BiophysChem1997``: 3
species, 21 sensitivity parameters, both axes requested), four warm engine+simulator
replicas driven from four threads integrated 160 segments in 0.057 s against 0.153 s
serially -- a **2.7x** speedup on 4 workers -- and every trajectory column and every entry
of ``d(y)/d(theta)`` came back *bit-identical* to the serial run. So bngsim drops the GIL
inside CVODE, and the arithmetic above this seam does not have to change to accommodate a
scheduler.
Processes were never a real option: a segment costs one integration, and the #563
prototype measured that constructing a sensitivity-bearing simulator costs ~230 ms cold
against ~50 ms for an integration. A process pool would pay that per worker per point (or
pay to pickle an engine model, which does not pickle), which is why
:class:`~pybnf.shooting.backend.SegmentBackend` keeps warm state per parameter point in
the first place.
A lane is a warm engine + simulator, and it is not free
-------------------------------------------------------
Two segments cannot share one ``Simulator``: the backend restarts it from the knot's state
with ``save_concentrations()`` + ``reset()``, so a second segment on the same object would
be integrating from the first's start state. A *lane* is therefore an independent
engine+simulator pair, and running ``L`` segments at once needs ``L`` of them **at that
parameter point** -- so the point pays ``L`` preparations instead of one.
That is the whole cost model, and it decides the default:
parallel wins when (m - 1) * t_integrate > (L - 1) * t_prepare
On Borghans, measured, ``t_prepare`` is ~4.1 ms and ``t_integrate`` is ~1-2.3 ms, so at
``m = L = 4`` parallel segments would *lose* -- the model is small enough that preparing
the extra lanes costs more than the integrations they save. On a model where a segment is
the expensive thing, it wins nearly linearly, and both cost terms move with the model
rather than with the fit: the initial-condition sensitivity system is ``n_species`` wide,
so a segment's integration grows with the state while a lane's preparation does not grow
as fast.
So ``ms_parallel_segments`` defaults to **1** (serial), which is the measured right answer
for the model this feature was built for, and a run that sets it reports what it measured
(:meth:`SegmentPool.describe`) rather than leaving a user to guess whether it helped.
What parallel gives up
----------------------
The serial pass stops at the first segment that fails to integrate -- the rest of that
point's segments are work whose answer is already decided, which is what keeps a search
that has wandered into a non-integrable corner from paying ``m`` simulations per rejected
trial. A submitted future cannot be un-run, so the parallel pass pays all ``m``. The
answer is identical; the *simulation count* a run reports is not, and on a multi-start
sweep over an uninformed box (where most points do not integrate) that difference is the
dominant cost. Stated here rather than discovered from a benchmark table.
"""
import queue
from concurrent.futures import ThreadPoolExecutor
from .backend import SegmentSimulationFailed, trace_from_data
[docs]
class SegmentTask:
"""One span to integrate: everything :meth:`SegmentPool.run` needs and nothing else."""
__slots__ = ('backend', 'times', 'initial_state', 'state_names')
def __init__(self, backend, times, initial_state, state_names):
self.backend = backend
self.times = times
self.initial_state = initial_state
self.state_names = tuple(state_names)
[docs]
class SegmentPool:
"""The segment pass, at one degree of parallelism.
:param n_lanes: How many segments to integrate at once. ``1`` (the default) runs them
on the calling thread, in order, short-circuiting at the first failure -- byte-for-
byte the behaviour that shipped with ADR-0110.
One pool is built per fit and shared by every rung of the ladder and every start, so a
thread pool is created at most once per run.
"""
def __init__(self, n_lanes=1):
self.n_lanes = max(1, int(n_lanes))
self._executor = None
@property
def parallel(self):
return self.n_lanes > 1
def describe(self):
if not self.parallel:
return 'segments simulated serially on the master'
return ('segments simulated %i at a time, on %i warm engine+simulator lane(s) per '
'experiment' % (self.n_lanes, self.n_lanes))
[docs]
def run(self, pset, tasks):
"""Integrate every task at ``pset``.
Returns ``(traces, ok)``: the per-task :class:`~pybnf.shooting.backend.SegmentTrace`
list and whether *every* one integrated to a finite trajectory. A pass with
``ok = False`` returns ``None`` for the traces, because the caller turns any failure
into a non-finite local model and never reads the partial result -- and returning a
half-filled list would invite someone to.
"""
tasks = list(tasks)
if not self.parallel or len(tasks) < 2:
return self._run_serial(pset, tasks)
return self._run_parallel(pset, tasks)
# -- serial -----------------------------------------------------------------
def _run_serial(self, pset, tasks):
traces = []
for task in tasks:
trace = _integrate(task, pset, 0)
if trace is None or not trace.is_finite():
# One unusable segment makes the whole local model unusable, so the rest of
# this point's segments are work whose answer is already decided.
return None, False
traces.append(trace)
return traces, True
# -- parallel ---------------------------------------------------------------
def _run_parallel(self, pset, tasks):
lanes = self._open_lanes(pset, tasks)
executor = self._pool()
futures = []
for task in tasks:
lane_queue = lanes[id(task.backend)]
futures.append(executor.submit(_integrate_in_lane, task, pset, lane_queue))
traces = [future.result() for future in futures]
if any(trace is None or not trace.is_finite() for trace in traces):
return None, False
return traces, True
def _open_lanes(self, pset, tasks):
"""Prepare each backend's lanes **on the calling thread**, and hand back one queue of
lane indices per backend.
Preparation is where a backend touches shared model state (it assigns the parameter
set on the model it owns), so it happens here, once, before any worker starts --
rather than being raced for inside :meth:`SegmentBackend.simulate`. A queue rather
than a modulo of the task index because a backend with more segments than lanes must
make a worker *wait* for a lane instead of two workers sharing one simulator, which
would have each integrating from the other's start state.
"""
lanes = {}
for task in tasks:
key = id(task.backend)
if key in lanes:
continue
available = task.backend.open_lanes(pset, self.n_lanes)
lane_queue = queue.Queue()
for lane in range(max(1, int(available))):
lane_queue.put(lane)
lanes[key] = lane_queue
return lanes
def _pool(self):
if self._executor is None:
self._executor = ThreadPoolExecutor(
max_workers=self.n_lanes, thread_name_prefix='ms-segment')
return self._executor
[docs]
def close(self):
"""Shut the thread pool down. Idempotent; a serial pool never opened one."""
if self._executor is not None:
self._executor.shutdown(wait=True)
self._executor = None
def __repr__(self):
return 'SegmentPool(lanes=%i)' % self.n_lanes
def _integrate(task, pset, lane):
"""One span, in one lane. ``None`` when it did not integrate."""
try:
data = task.backend.simulate(pset, task.times, task.initial_state, lane=lane)
except SegmentSimulationFailed:
return None
return trace_from_data(data, task.state_names)
def _integrate_in_lane(task, pset, lane_queue):
"""Claim a lane for the whole integration, then give it back.
The claim spans the *simulation*, not just the dispatch: a lane is a stateful
engine+simulator that is reset to this segment's start knot, so releasing it early would
let a second segment restart it underneath the first.
"""
lane = lane_queue.get()
try:
return _integrate(task, pset, lane)
finally:
lane_queue.put(lane)
#: The pool a problem built without one uses -- serial, so nothing changes for a caller
#: that does not ask for lanes.
SERIAL = SegmentPool(1)