AllenInstitute / bmtk

Brain Modeling Toolkit
https://alleninstitute.github.io/bmtk/
BSD 3-Clause "New" or "Revised" License
275 stars 88 forks source link

Ability to assign cells to different processes using their gids #384

Open tjbanks opened 2 months ago

tjbanks commented 2 months ago

Hey all - I have a unique situation where I would like to to assign cells to specific processes due to (potential) network lag. An example situation may be where I have two regions that communicate often on one host but occasionally communicate over the network. I think I see where cells are assigned here: https://github.com/AllenInstitute/bmtk/blob/develop/bmtk/simulator/bionet/bionetwork.py#L192

But want to be sure that my thinking isn't off and ensure that I can't already assign process via gid.

Thank you again, Tyler

Cc: @V-Marco

kaeldai commented 2 months ago

If I'm understanding you correctly you want to explicitly assign certain cells to run on certain cores (eg. MPI ranks)?

This is possible by modifying the loop on line 188 you pointed to. Current cells are assigned to a rank in a round robin method, so if you are running with 5 cores then rank 0 will contain cell gids 0, 5, 10, ...; rank 1 will have gids 1, 6, 11; rank 2 have gids 2, 7, 12, ...; etc. To change this assigment scheme you'd just need to change the for loop.

For example, if you wanted to make it so that cells 20 & 21 is always on rank/processor 0 the code would have to be modified the somewhat along the lines of the following:

for i, node in enumerate(node_pop):
    if node_pop['node_id'] in [20, 21]:
        # If we come across either cells 20 or 21, then only create them when running on rank 0
        if MPI_rank == 0:
            cell = self._build_cell(bionode=node, population_name=node_pop.name)
            node_ids_map[node.node_id] = cell
            self._rank_node_gids[cell.gid] = cell
    else:
        # for all other cells assign them to ranks as before.
        if i % MPI_SIZE != 0:
            continue

        cell = self._build_cell(bionode=node, population_name=node_pop.name)
        node_ids_map[node.node_id] = cell
        self._rank_node_gids[cell.gid] = cell

There are definitely more elegant ways to implement this, maybe either assign ranks based on a pre-existing table or a customized user function. If you think doing that would be useful please let me know! I think it wouldn't be too hard to implement such a solution.

tjbanks commented 2 months ago

Thanks @kaeldai!

I will run a few tests and see about a merge request if things look good.

Appreciate your time, hope you've been well! Tyler