diff --git a/neurodamus/cell_distributor.py b/neurodamus/cell_distributor.py index 9579e2b1..9c5458f6 100644 --- a/neurodamus/cell_distributor.py +++ b/neurodamus/cell_distributor.py @@ -202,7 +202,8 @@ def load_nodes(self, load_balancer=None, *, _loader=None, loader_opts=None): logging.info("Reading Nodes (METype) info from '%s'", conf.CellLibraryFile) if load_balancer and \ - getattr(load_balancer, 'population', None) != self._target_spec.population: + hasattr(load_balancer, 'population') and \ + load_balancer.population != self._target_spec.population: log_verbose("Load balance object doesn't apply to '%s'", self._target_spec.population) load_balancer = None if not load_balancer or SimConfig.dry_run: @@ -252,6 +253,7 @@ def _load_nodes_balance_mem(self, loader_f, load_balancer): population = targetspec.population all_gids = load_balancer[population][MPI.rank] + logging.debug("Loading %d cells in rank %d", len(all_gids), MPI.rank) total_cells = len(all_gids) gidvec, me_infos, full_size = loader_f(self._circuit_conf, all_gids) return gidvec, me_infos, total_cells, full_size diff --git a/neurodamus/node.py b/neurodamus/node.py index e6a7b03b..a904fd47 100644 --- a/neurodamus/node.py +++ b/neurodamus/node.py @@ -379,6 +379,9 @@ def compute_load_balance(self): elif lb_mode == LoadBalanceMode.Memory: logging.info("Load Balancing ENABLED. Mode: Memory") alloc = import_allocation_stats("allocation.pkl.gz") + for pop, ranks in alloc.items(): + for rank, gids in ranks.items(): + logging.debug(f"Population: {pop}, Rank: {rank}, Number of GIDs: {len(gids)}") if MPI.rank == 0: unique_ranks = set(rank for pop in alloc.values() for rank in pop.keys()) logging.debug("Unique ranks in allocation file: %s", len(unique_ranks))