CoreNeuron Rebalancer

[ferdonline]

Context

CoreNeuron runs, in particular those using multi-cycle, may end up distributing cells unoptimally.

We want to mitigate that by introducing a post processing step which distributes the CoreNeuron input files evenly across ranks, based on their size.

Scope

Added rebalance-corenrn-data.py to neurodamus/tools.

CLI:

neurodamus/tools $ ./rebalance-corenrn-data.py -h
usage: rebalance-corenrn-data.py [-h] [--max-entries MAX_ENTRIES] [--output-file OUTPUT_FILE] [-v] [--histogram] input_file n_ranks

Redistribute CoreNeuron dat files, optimizing for a given number of ranks

positional arguments:
  input_file            Path to the CoreNeuron input file, typically files.dat
  n_ranks               Number of target ranks

options:
  -h, --help            show this help message and exit
  --max-entries MAX_ENTRIES
                        Consider only the first N entries of the input file
  --output-file OUTPUT_FILE
                        The rebalanced output file path
  -v, --verbose         Enable verbose output for debugging.
  --histogram           Additionally display the histogram of the ranks accumulated sizes

Testing

Using /gpfs/bbp.cscs.ch/data/scratch/proj134/home/king/BBPP134-917/o_MultiCycle_Support/output/2557873/coreneuron_input/files.dat for testing

Example

leite@bbpv2 ~/dev/neurodamus/neurodamus-py/tools (leite/corenrn-rebalancer %)$ python rebalance-corenrn-data.py /gpfs/bbp.cscs.ch/data/scratch/proj134/home/king/BBPP134-917/o_MultiCycle_Support/output/2557873/coreneuron_input/files.dat 10 --max-entries=1000
INFO :: Reading from input file: /gpfs/bbp.cscs.ch/data/scratch/proj134/home/king/BBPP134-917/o_MultiCycle_Support/output/2557873/coreneuron_input/files.dat'
WARNING :: files.dat (line 2): reduced number of entries: 1000'
INFO :: Distributing files into 10 buckets...'
INFO :: Processing 1000 entries'
         0 [  0%]
        20 [  2%]
...
       980 [ 98%]
INFO :: Writing out data from 10 buckets to file: rebalanced-files.dat'
INFO :: DONE'

Review

• PR description is complete
• Coding style (imports, function length, New functions, classes or files) are good
• Unit/Scientific test added
• Updated Readme, in-code, developer documentation

[bbpbuildbot]

Logfiles from GitLab pipeline #237055 (:white_check_mark:) have been uploaded here!

Status and direct links:

• ✅ docs
• ✅ test-unit
• ✅ test-integration
• ✅ test-integration-e2e
• ✅ test-scientific
• ✅ test-scientific-ngv
• ✅ set_alt_branches
• ✅ build-package
• ⚙️ publish-package
• ⚙️ publish-docs

[jorblancoa]

Maybe it makes sense to have some constants (for example):

PROGRESS_UPDATE_FREQUENCY = 50
DEFAULT_OUTPUT_FILE = "rebalanced-files.dat"
FILL_VALUE = "-1\n"

Or the names you think makes more sense

[jorblancoa]

If this is a tool to be used in the future and not only a script that will be discarded soon maybe makes sense to have a class?
something like

class RebalancerConfig:
...
class CoreNeuronRebalancer:
    def __init__(self, config: RebalancerConfig):
        self.config = config
        self.logger = self._setup_logger()
        self.buckets = []
        ...

def main():
        # Parse arguments and create config
        args = parse_args()
        config = RebalancerConfig.from_args(args)

        # Create and run rebalancer
        rebalancer = CoreNeuronRebalancer(config)
        rebalancer.rebalance()