Reverted change on load_nodes check since it didn't work as expected,…

… added some debug prints

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

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))