From f11ca45fcef8027e3fc56f0758fbbaed494b4b9b Mon Sep 17 00:00:00 2001 From: drodarie Date: Mon, 7 Apr 2025 15:14:07 +0200 Subject: [PATCH 1/2] fix: forward MPI communicator to simulation adapter. --- bsb_neuron/adapter.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/bsb_neuron/adapter.py b/bsb_neuron/adapter.py index dc9c929..f715b65 100644 --- a/bsb_neuron/adapter.py +++ b/bsb_neuron/adapter.py @@ -4,7 +4,6 @@ import numpy as np from bsb import ( - MPI, AdapterError, AdapterProgress, Chunk, @@ -82,7 +81,7 @@ def prepare(self, simulation, comm=None): self.engine.celsius = simulation.temperature self.engine.tstop = simulation.duration report("Load balancing", level=2) - self.load_balance(simulation) + self.load_balance(simulation, comm) report("Creating neurons", level=2) self.create_neurons(simulation) report("Creating transmitters", level=2) @@ -94,17 +93,17 @@ def prepare(self, simulation, comm=None): del self.simdata[simulation] raise - def load_balance(self, simulation): + def load_balance(self, simulation, comm): simdata = self.simdata[simulation] chunk_stats = simulation.scaffold.storage.get_chunk_stats() - size = MPI.get_size() + size = comm.get_size() all_chunks = [Chunk.from_id(int(chunk), None) for chunk in chunk_stats.keys()] simdata.node_chunk_alloc = [all_chunks[rank::size] for rank in range(0, size)] simdata.chunk_node_map = {} for node, chunks in enumerate(simdata.node_chunk_alloc): for chunk in chunks: simdata.chunk_node_map[chunk] = node - simdata.chunks = simdata.node_chunk_alloc[MPI.get_rank()] + simdata.chunks = simdata.node_chunk_alloc[comm.get_rank()] simdata.placement = { model: model.get_placement_set(chunks=simdata.chunks) for model in simulation.cell_models.values() From 9415bd00b21d8a066e03ae55d9441192ae51c1fd Mon Sep 17 00:00:00 2001 From: drodarie Date: Tue, 8 Apr 2025 14:22:35 +0200 Subject: [PATCH 2/2] fix: move MPI communicator to Simulation adapter --- bsb_neuron/adapter.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/bsb_neuron/adapter.py b/bsb_neuron/adapter.py index f715b65..8262b3b 100644 --- a/bsb_neuron/adapter.py +++ b/bsb_neuron/adapter.py @@ -60,8 +60,8 @@ def fill_parameter_data(parameters, data): class NeuronAdapter(SimulatorAdapter): initial = -65 - def __init__(self): - super().__init__() + def __init__(self, comm=None): + super().__init__(comm=comm) self.network = None self.next_gid = 0 @@ -71,7 +71,7 @@ def engine(self): return engine - def prepare(self, simulation, comm=None): + def prepare(self, simulation): self.simdata[simulation] = NeuronSimulationData( simulation, result=NeuronResult(simulation) ) @@ -81,7 +81,7 @@ def prepare(self, simulation, comm=None): self.engine.celsius = simulation.temperature self.engine.tstop = simulation.duration report("Load balancing", level=2) - self.load_balance(simulation, comm) + self.load_balance(simulation) report("Creating neurons", level=2) self.create_neurons(simulation) report("Creating transmitters", level=2) @@ -93,17 +93,17 @@ def prepare(self, simulation, comm=None): del self.simdata[simulation] raise - def load_balance(self, simulation, comm): + def load_balance(self, simulation): simdata = self.simdata[simulation] chunk_stats = simulation.scaffold.storage.get_chunk_stats() - size = comm.get_size() + size = self.comm.get_size() all_chunks = [Chunk.from_id(int(chunk), None) for chunk in chunk_stats.keys()] simdata.node_chunk_alloc = [all_chunks[rank::size] for rank in range(0, size)] simdata.chunk_node_map = {} for node, chunks in enumerate(simdata.node_chunk_alloc): for chunk in chunks: simdata.chunk_node_map[chunk] = node - simdata.chunks = simdata.node_chunk_alloc[comm.get_rank()] + simdata.chunks = simdata.node_chunk_alloc[self.comm.get_rank()] simdata.placement = { model: model.get_placement_set(chunks=simdata.chunks) for model in simulation.cell_models.values()