Rectified differences between threads (OpenMP) and processes (MPI) (#73)

* Rectify difference between threads (OpenMP) and processes (MPI)

* Transparently handle MPI, set OpenMP from config

* Catch mpi4py missing. Added optional MPI dependencies to setup.

* Removed obsolete reference to `nodes`

README.md

@@ -90,3 +90,10 @@ When modifying the config object through scripts and then saving it to file, you
 the original configuration file text, and you won't actually serialize the modified object
 
 We will fix this by version 4.0
+
+## If MPI is installed but mpi4py is not undefined behavior may occur
+
+The amount of NEST virtual processes is determined by using mpi4py to get the amount of
+MPI processes. But if the package is not installed it is assumed no MPI simulations will
+be ran and the amount of virtual processes might be lower than expected when used in
+combination with OpenMP. Be sure to `pip install` using the `[MPI]` requirement tag.

bsb/simulators/nest.py

@@ -12,6 +12,13 @@
 from itertools import chain
 from sklearn.neighbors import KDTree
 
+try:
+    import mpi4py.MPI
+
+    _MPI_processes = mpi4py.MPI.COMM_WORLD.Get_size()
+except ImportError:
+    _MPI_processes = 1
+
 LOCK_ATTRIBUTE = "dbbs_scaffold_lock"
 
 
@@ -276,7 +283,7 @@ class NestAdapter(SimulatorAdapter):
         "entities": NestEntity,
     }
 
-    casts = {"threads": int, "virtual_processes": int, "modules": list}
+    casts = {"threads": int, "modules": list}
 
     defaults = {
         "default_synapse_model": "static_synapse",
@@ -415,7 +422,7 @@ def delete_lock(self):
     def reset_kernel(self):
         self.nest.set_verbosity(self.verbosity)
         self.nest.ResetKernel()
-        self.set_threads(self.threads)
+        self.reset_processes(self.threads)
         self.nest.SetKernelStatus(
             {
                 "resolution": self.resolution,
@@ -438,24 +445,22 @@ def get_master_seed(self):
         # Use a constant reproducible master seed
         return 1989
 
-    def set_threads(self, threads, virtual=None):
+    def reset_processes(self, threads):
         master_seed = self.get_master_seed()
-        # Update the internal reference to the amount of threads
-        if virtual is None:
-            virtual = threads
+        total_num = _MPI_processes * threads
         # Create a range of random seeds and generators.
-        random_generator_seeds = range(master_seed, master_seed + virtual)
+        random_generator_seeds = range(master_seed, master_seed + total_num)
         # Create a different range of random seeds for the kernel.
-        thread_seeds = range(master_seed + virtual + 1, master_seed + 1 + 2 * virtual)
+        thread_seeds = range(master_seed + 1 + total_num, master_seed + 1 + 2 * total_num)
         success = True
         try:
             # Update the kernel with the new RNG and thread state.
             self.nest.SetKernelStatus(
                 {
-                    "grng_seed": master_seed + virtual,
+                    "grng_seed": master_seed + total_num,
                     "rng_seeds": thread_seeds,
                     "local_num_threads": threads,
-                    "total_num_virtual_procs": virtual,
+                    "total_num_virtual_procs": total_num,
                 }
             )
         except Exception as e:
@@ -471,8 +476,8 @@ def set_threads(self, threads, virtual=None):
             else:
                 raise
         if success:
-            self.threads = threads
-            self.virtual_processes = virtual
+            self.threads_per_node = threads
+            self.virtual_processes = total_num
             self.random_generators = [
                 np.random.RandomState(seed) for seed in random_generator_seeds
             ]

setup.py

@@ -45,5 +45,6 @@
     extras_require={
         "dev": ["sphinx", "sphinx_rtd_theme>=0.4.3", "pyarmor", "pre-commit", "black"],
         "NEURON": ["dbbs_models>=0.4.4"],
+        "MPI": ["mpi4py"],
     },
 )