dry_run_stats objec it now part of Node even when distributing on the…

… fly, added more explicit tests
BlueBrain · Jul 4, 2024 · 2f8a763 · 2f8a763
1 parent 0c9d78b
commit 2f8a763
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 4 deletions.
diff --git a/neurodamus/node.py b/neurodamus/node.py
@@ -381,12 +381,13 @@ def compute_load_balance(self):
                 alloc = import_allocation_stats(filename, self._cycle_i)
             else:
                 logging.warning("Allocation file not found. Generating on-the-fly.")
-                dry_run_stats = DryRunStats()
-                dry_run_stats.try_import_cell_memory_usage()
+                self._dry_run_stats = DryRunStats()
+                self._dry_run_stats.try_import_cell_memory_usage()
                 cell_distributor = CellDistributor(circuit, self._target_manager, self._run_conf)
                 cell_distributor.load_nodes(None, loader_opts={"load_mode": "load_nodes_metype",
-                                                               "dry_run_stats": dry_run_stats})
-                alloc, _, _ = dry_run_stats.distribute_cells(
+                                                               "dry_run_stats": self._dry_run_stats}
+                                            )
+                alloc, _, _ = self._dry_run_stats.distribute_cells(
                     MPI.size,
                     SimConfig.modelbuilding_steps,
                     DryRunStats._MEMORY_USAGE_PER_METYPE_FILENAME

diff --git a/tests/integration-e2e/test_dry_run_workflow.py b/tests/integration-e2e/test_dry_run_workflow.py
@@ -152,3 +152,21 @@ def test_dynamic_distribute():
                     modelbuilding_steps=2,
                     lb_mode="Memory")
     nd.run()
+
+    rank_allocation, _, _ = nd._dry_run_stats.distribute_cells(1, 2)
+    rank_allocation_standard = convert_to_standard_types(rank_allocation)
+
+    expected_items = {
+        'default': {
+            (0, 0): [
+                62798, 62946, 63257, 63699, 64164, 64862, 65916, 65952, 66069, 66106
+            ],
+            (0, 1): [
+                66141, 66497, 66872, 67667, 68224, 68354, 68533, 68581, 68942, 69531,
+                69840, 63623, 64234, 64666, 64788, 64936, 69878, 65821, 67078, 68856
+            ]
+        }
+    }
+
+    for key, sub_value in rank_allocation_standard['default'].items():
+        assert set(sub_value) == set(expected_items['default'][key])