feat: Add simulation's controllers

packages/bsb-arbor/bsb_arbor/adapter.py

@@ -408,11 +408,13 @@ def prepare(self, simulation: "ArborSimulation") -> ArborSimulationData:
             report("MPI processes:", context.ranks, level=2)
             report("Threads per process:", context.threads, level=2)
             recipe = self.get_recipe(simulation, simdata)
+            self._duration = simulation.duration
             # Gap junctions are required for domain decomposition
             self.domain = arbor.partition_load_balance(recipe, context)
             self.gids = set(it.chain.from_iterable(g.gids for g in self.domain.groups))
             simdata.arbor_sim = arbor.simulation(recipe, context, self.domain)
-            self.prepare_samples(simulation, simdata)
+            self.implement_components(simulation)
+            self.load_controllers(simulation)
             report("prepared simulation", level=1)
             return simdata
         except Exception:
@@ -422,10 +424,6 @@ def prepare(self, simulation: "ArborSimulation") -> ArborSimulationData:
     def get_gid_manager(self, simulation, simdata):
         return GIDManager(simulation, simdata)
 
-    def prepare_samples(self, simulation, simdata):
-        for device in simulation.devices.values():
-            device.prepare_samples(simdata, comm=self.comm)
-
     def run(self, *simulations):
         if len(simulations) != 1:
             raise RuntimeError(
@@ -446,14 +444,19 @@ def run(self, *simulations):
 
             start = time.time()
             report("running simulation", level=1)
-            arbor_sim.run(simulation.duration * U.ms, dt=simulation.resolution * U.ms)
-            report(f"completed simulation. {time.time() - start:.2f}s", level=1)
+
+            for t, checkpoint_controllers in self.get_next_checkpoint():
+                arbor_sim.run(t * U.ms, dt=simulation.resolution * U.ms)
+                self.execute_checkpoints(checkpoint_controllers)
+            report(f"Completed simulation. {time.time() - start:.2f}s", level=1)
             if simulation.profiling and arbor.config()["profiling"]:
                 report("printing profiler summary", level=2)
                 report(arbor.profiler_summary(), level=1)
             return [simdata.result]
         finally:
+            results = [self.simdata[sim].result for sim in simulations]
             del self.simdata[simulation]
+        return results
 
     def get_recipe(self, simulation, simdata=None):
         if simdata is None:

packages/bsb-arbor/bsb_arbor/device.py

@@ -24,7 +24,7 @@ def __boot__(self):
     def register_probe_id(self, gid, tag):
         self._probe_ids.append((gid, tag))
 
-    def prepare_samples(self, simdata, comm):
+    def implement(self, adapter, simulation, simdata):
         self._handles = [
             self.sample(simdata.arbor_sim, probe_id) for probe_id in self._probe_ids
         ]

packages/bsb-arbor/bsb_arbor/devices/probe.py

@@ -16,15 +16,10 @@ def validate_specifics(self):
                 f"`{self.probe_type}` is not a valid probe type for `{self.name}`"
             )
 
-    def implement(self, target):
-        probe_args = ("where", "mechanism", "ion", "state")
-        kwargs = dict((k, getattr(self, k)) for k in probe_args if hasattr(self, k))
-        return [getattr(arbor, self.get_probe_name())(**kwargs)]
-
-    def prepare_samples(self, sim, comm):
-        super().prepare_samples(sim, comm)
+    def implement(self, adapter, simulation, simdata):
+        super().implement(adapter, simulation, simdata)
         for probe_id, handle in zip(self._probe_ids, self._handles, strict=False):
-            self.adapter.result.add(ProbeRecorder(self, sim, probe_id, handle))
+            self.adapter.result.add(ProbeRecorder(self, simdata, probe_id, handle))
 
 
 class ProbeRecorder:

packages/bsb-arbor/bsb_arbor/devices/spike_recorder.py

@@ -9,9 +9,9 @@ class SpikeRecorder(ArborDevice, classmap_entry="spike_recorder"):
     def boot(self):
         self._gids = set()
 
-    def prepare_samples(self, simdata, comm):
-        super().prepare_samples(simdata, comm)
-        if not comm.get_rank():
+    def implement(self, adapter, simulation, simdata):
+        super().implement(adapter, simulation, simdata)
+        if not adapter.comm.get_rank():
 
             def record_device_spikes(segment):
                 spiketrain = list()

packages/bsb-core/bsb/__init__.py

@@ -133,7 +133,6 @@ def __dir__():
   import bsb.voxels
 
 AdapterError: type["bsb.exceptions.AdapterError"]
-AdapterProgress: type["bsb.simulation.adapter.AdapterProgress"]
 AfterConnectivityHook: type["bsb.postprocessing.AfterConnectivityHook"]
 AfterPlacementHook: type["bsb.postprocessing.AfterPlacementHook"]
 AllToAll: type["bsb.connectivity.general.AllToAll"]
@@ -229,6 +228,7 @@ def __dir__():
 FixedIndegree: type["bsb.connectivity.general.FixedIndegree"]
 FixedOutdegree: type["bsb.connectivity.general.FixedOutdegree"]
 FixedPositions: type["bsb.placement.strategy.FixedPositions"]
+FixedStepProgressController: type["bsb.simulation.adapter.FixedStepProgressController"]
 FractionFilter: type["bsb.simulation.targetting.FractionFilter"]
 GatewayError: type["bsb.exceptions.GatewayError"]
 GeneratedMorphology: type["bsb.storage.interfaces.GeneratedMorphology"]

packages/bsb-core/bsb/_options.py

@@ -92,6 +92,25 @@ def get_default(self):
         return "network_configuration.json"
 
 
+class SimulationProgress(
+    BsbOption,
+    name="sim_console_progress",
+    cli=("scp", "sim_console_progress"),
+    project=("sim_console_progress",),
+    script=("sim_console_progress",),
+    env=("BSB_SIM_PROGRESS",),
+):
+    """
+    Activate reports during simulations, set the time steps of the report.
+    """
+
+    def setter(self, value):
+        return float(value)
+
+    def getter(self, value):
+        return float(value)
+
+
 class ProfilingOption(
     BsbOption,
     name="profiling",
@@ -157,6 +176,10 @@ def config():
     return ConfigOption
 
 
+def sim_console_progress():
+    return SimulationProgress
+
+
 def profiling():
     return ProfilingOption
 

packages/bsb-core/bsb/options.py

@@ -32,7 +32,7 @@ def get_default(self):
 import contextlib
 import functools
 
-from ._options import ProfilingOption, VerbosityOption
+from ._options import ProfilingOption, SimulationProgress, VerbosityOption
 
 # Store the module magic for unpolluted namespace copy
 _module_magic = globals().copy()
@@ -388,6 +388,7 @@ def __delattr__(self, attr):
 
 register_option("verbosity", VerbosityOption())
 register_option("profiling", ProfilingOption())
+register_option("sim_console_progress", SimulationProgress())
 
 # Static public API
 __all__ = [

packages/bsb-core/bsb/simulation/adapter.py

@@ -1,49 +1,68 @@
 import abc
-import itertools
-import types
+import os
+import sys
 import typing
 from contextlib import ExitStack
 from time import time
 
-import numpy as np
+from tqdm import tqdm
+
+from bsb import AttributeMissingError, SimulationResult, options
 
 from ..services.mpi import MPIService
-from .results import SimulationResult
 
 if typing.TYPE_CHECKING:
     from ..storage import PlacementSet
     from .cell import CellModel
     from .simulation import Simulation
 
 
-class AdapterProgress:
-    def __init__(self, duration):
-        self._duration = duration
+class FixedStepProgressController:
+    def __init__(self, simulations, adapter, step=1):
+        self._status = 0
+        self._adapter = adapter
         self._start = self._last_tick = time()
-        self._ticks = 0
-
-    def tick(self, step):
-        """
-        Report simulation progress.
-        """
+        self._step = step
+        self._sim_name = [sim._name for sim in simulations]
+        self._use_tty = os.isatty(sys.stdout.fileno()) and sum(os.get_terminal_size())
+
+        def silent():
+            self._status = self._adapter.current_checkpoint
+
+        if self._use_tty:
+            if not self._adapter.comm.get_rank():
+                self.run_checkpoint = self.use_bar
+            else:
+                self.run_checkpoint = silent
+        else:
+            if self._adapter.comm.get_rank():
+                self.run_checkpoint = silent
+
+    def get_next_checkpoint(self):
+        return self._status + self._step
+
+    def run_checkpoint(self):
         now = time()
+        self._status = self._adapter.current_checkpoint
         tic = now - self._last_tick
-        self._ticks += 1
-        el = now - self._start
-        progress = types.SimpleNamespace(
-            progression=step, duration=self._duration, time=time(), tick=tic, elapsed=el
-        )
+        el_time = now - self._start
+        duration = self._adapter._duration
+        msg = f"Simulation {self._sim_name} | progress: {self._status:.2f} "
+        msg += f"elapsed: {el_time:.2f}s - last step time: {tic:.2f}s - "
+        msg += f"exectuted: {(self._status / duration) * 100:.2f}%"
+        print(msg)
         self._last_tick = now
-        return progress
-
-    def steps(self, step=1):
-        steps = itertools.chain(np.arange(0, self._duration, step), (self._duration,))
-        a, b = itertools.tee(steps)
-        next(b, None)
-        yield from zip(a, b, strict=False)
 
-    def complete(self):
-        return
+    def use_bar(self):
+        if not self._status:
+            self._progress_bar = tqdm(total=self._adapter._duration)
+            self._progress_bar.update(self._adapter.current_checkpoint - self._status)
+        elif self._adapter.current_checkpoint == self._adapter._duration:
+            self._progress_bar.update(self._adapter.current_checkpoint - self._status)
+            self._progress_bar.close()
+        else:
+            self._progress_bar.update(self._adapter.current_checkpoint - self._status)
+        self._status = self._adapter.current_checkpoint
 
 
 class SimulationData:
@@ -66,9 +85,12 @@ def __init__(self, comm=None):
         :param comm: The mpi4py MPI communicator to use. Only nodes in the communicator
           will participate in the simulation. The first node will idle as the main node.
         """
-        self._progress_listeners = []
+
         self.simdata: dict[Simulation, SimulationData] = dict()
         self.comm = MPIService(comm)
+        self._controllers = []
+        self._duration = None
+        self.current_checkpoint = 0
 
     def simulate(self, *simulations, post_prepare=None):
         """
@@ -84,6 +106,12 @@ def simulate(self, *simulations, post_prepare=None):
             for simulation in simulations:
                 context.enter_context(simulation.scaffold.storage.read_only())
             alldata = []
+            if options.sim_console_progress:
+                listener = FixedStepProgressController(
+                    simulations, self, options.sim_console_progress
+                )
+                self._controllers.append(listener)
+
             for simulation in simulations:
                 data = self.prepare(simulation)
                 alldata.append(data)
@@ -118,6 +146,30 @@ def run(self, *simulations):
         """
         pass
 
+    def get_next_checkpoint(self):
+        while self.current_checkpoint < self._duration:
+            checkpoints = [
+                controller.get_next_checkpoint() for controller in self._controllers
+            ]
+            # Filter out invalid "regressive" checkpoints,
+            # and default to the end of the simulation
+            chkp_noregressive = [
+                checkpoint
+                for checkpoint in checkpoints
+                if checkpoint > self.current_checkpoint
+            ]
+            self.current_checkpoint = min(chkp_noregressive, default=self._duration)
+            participants = [
+                self._controllers[i]
+                for i, checkpoint in enumerate(checkpoints)
+                if checkpoint == self.current_checkpoint
+            ]
+            yield (self.current_checkpoint, participants)
+
+    def execute_checkpoints(self, controllers):
+        for controller in controllers:
+            controller.run_checkpoint()
+
     def collect(self, results):
         """
         Collect the output the simulations that completed.
@@ -129,8 +181,25 @@ def collect(self, results):
             result.flush()
         return results
 
-    def add_progress_listener(self, listener):
-        self._progress_listeners.append(listener)
-
-
-__all__ = ["AdapterProgress", "SimulationData", "SimulatorAdapter"]
+    def implement_components(self, simulation):
+        simdata = self.simdata[simulation]
+        for component in simulation.get_components():
+            component.implement(self, simulation, simdata)
+
+    def load_controllers(self, simulation):
+        for component in simulation.get_components():
+            if hasattr(component, "get_next_checkpoint"):
+                if hasattr(component, "run_checkpoint"):
+                    self._controllers.append(component)
+                else:
+                    raise AttributeMissingError(
+                        f"Device {component.name} is configured to be a controller "
+                        f"but run_checkpoint is not defined"
+                    )
+
+
+__all__ = [
+    "FixedStepProgressController",
+    "SimulationData",
+    "SimulatorAdapter",
+]

packages/bsb-core/bsb/simulation/component.py

@@ -14,5 +14,10 @@ def simulation(self):
     def __str__(self):
         return f"'{self.name}'"
 
+    def implement(self, adapter, simulation, simdata):
+        """Method called when simulation is being set up. Can be used for components to
+        set themselves up and store the context they need to operate."""
+        pass
+
 
 __all__ = ["SimulationComponent"]

packages/bsb-core/bsb/simulation/simulation.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import itertools
 import typing
 
 from .. import config
@@ -84,5 +85,12 @@ def get_connectivity_sets(
             for model in sorted(self.connection_models.values())
         }
 
+    def get_components(self):
+        return itertools.chain(
+            self.cell_models.values(),
+            self.connection_models.values(),
+            self.devices.values(),
+        )
+
 
 __all__ = ["ProgressEvent", "Simulation"]

packages/bsb-core/pyproject.toml

@@ -58,6 +58,7 @@ verbosity = "bsb._options:verbosity"
 sudo = "bsb._options:sudo"
 version = "bsb._options:version"
 config = "bsb._options:config"
+sim_console_progress= "bsb._options:sim_console_progress"
 profiling = "bsb._options:profiling"
 debug_pool = "bsb._options:debug_pool"