openbraininstitute / neurodamus / 24997006652

91%

push

github

New CLI option to allow dry run and memeory load balancing based on heap memory (#507)

## Context
Fix #363, improve memory load balancing to match dry-run memory
estimates
1. In the `--dry-run` mode, cells and synapses are included into the
estimated memory, and neurodamus produces the allocation file
`allocation_r<No_Ranks>_c1.pkl.gz` as well as `cell_memory_usage.json`.
In the `--lb-mode=Memory` mode, if the allocation file is missing, run
the cell distribution on-the-fly just as in dry run mode. This PR fixes
an issue in the `--lb-mode=Memory` mode, that is , when neither
allocation file or `cell_memory_usage.json` exists, distributing cells
on the fly should include the memory estimates for synapses, just as in
dry run.
2. Enable memory measurement by RSS or Heap size via CLI option
`--memory-tracker`.
3. When distributing cells across ranks, distribute cells from the most
memory consuming to the least.


## Scope
1. In `memory.py:_memory_mode_load_balancing` function: create synapse
manager for all the connections when estimate memory for each metype
combination.
2. In `export_cell_memory_usage` function: to reuse
`cell_memory_usage.json` to generate the allocation file without
instantiating cells and synapses, dump the data dict
`metype_cell_syn_average` and `pop_metype_gids` into
`cell_memory_usage.json`, together with the existing data
`metype_memory`.
4. in `memory.py: distribute_cells`, sort `metype_memory_usage` by its
values `total_mem_per_cell` in descending order, and distribute cells
from most memory consuming to the least.
5. CLI option `--memory-tracker` to either measure rss or use memray to
measure the heap size per cell metype. Measuring rss is faster, our
default option. Measuring heap requires additional dependency (memray),
but it would bring indentical measurement from run to run, used for
testing.

## Testing
New test `test_lb_mode_memory_from_scratch` in
`tests/unit-mpi/test_dry_run.py`

## Review
* [x] PR description is com... (continued)

104 of 105 new or added lines in 5 files covered. (99.05%)

7849 of 8599 relevant lines covered (91.28%)

2.84 hits per line