Mf6Splitter holds two python dicts with one entry per grid cell, so peak
memory is set by python object overhead rather than by the data. Splitting a
210,771,902 cell structured model into 256 parts peaks at 320 GB and is killed
on a 229 GB node; the same information as numpy arrays is about 7 GB.
The two structures are self._node_map, a dict of {node: (model, newnode)}
built at model_splitter.py:1067, and self._connection, the dict returned by
modelgrid.neighbors() at model_splitter.py:1011. Both are keyed by every
cell in the user grid, active or not.
Measurements
CONUS groundwater model, structured, nlay=1, nrow=11419, ncol=18458,
210,771,902 cells, 124,884,583 active, split into 256 parts.
| stage |
peak RSS |
outcome |
optimize_splitting_mask(active_only=True) |
194 GB |
1 h 23 m, completes |
split_model on a 229 GB node |
> 229 GB |
OOM killed after 33 min |
split_model on a 377 GB node |
320 GB |
4 h 15 m, completes |
Memory is flat while the model loads (73 GB) and climbs through _remap_nodes
and split_model.
Where it goes
| structure |
contents |
estimated |
as numpy |
_node_map |
210.8 M dict entries, tuple values |
~44 GB |
2 int32 arrays, 1.7 GB |
_connection |
210.8 M dict entries, list values |
~57 GB |
CSR indptr and indices, ~5 GB |
A dict entry with an int key costs roughly 100 bytes of table, the tuple another
56, and the ints inside it 28 each once they exceed the small integer cache.
Two changes that would recover most of it
1. _node_map as arrays rather than a dict. Two int32 arrays of length
nnodes, one for the model and one for the new node number, replace the dict and
its tuples: ~44 GB to 1.7 GB. There are about 25 references across 11 methods,
all internal, and the lookups in _remap_nodes are already loops over arrays
that would become fancy indexing. save_node_mapping and load_node_mapping
serialize the dict to HDF5, so this needs a format bump or a shim.
2. _connection in CSR form. An indptr and an indices array replace the
dict of lists: ~57 GB to ~5 GB. Concentrated in _remap_nodes, but it needs
modelgrid.neighbors() to grow an array returning mode, so it touches
flopy/discretization/grid.py as well.
Either alone would bring this model onto a 229 GB node.
Both structures are also traversed per cell in _remap_nodes, about 1 G dict
operations for this model, and contribute several hundred million objects for
the cyclic collector to walk, so the change should reduce wall time as well as
memory. split_model took 88 minutes here, against 140 minutes for
write_simulation which this would not affect.
Smaller and independent: model_splitter.py:1022 builds
cellids = list(zip([0] * len(cells[0]), cells[0], cells[1])) and passes it to
get_node(), materializing a python tuple per cell of every part, where the
node number is row * ncol + col in numpy.
Versions: flopy 3.10.0 (develop at 7732f39 is unchanged in these paths),
python 3.12.3, numpy 1.26.4.
Mf6Splitterholds two python dicts with one entry per grid cell, so peakmemory is set by python object overhead rather than by the data. Splitting a
210,771,902 cell structured model into 256 parts peaks at 320 GB and is killed
on a 229 GB node; the same information as numpy arrays is about 7 GB.
The two structures are
self._node_map, a dict of{node: (model, newnode)}built at
model_splitter.py:1067, andself._connection, the dict returned bymodelgrid.neighbors()atmodel_splitter.py:1011. Both are keyed by everycell in the user grid, active or not.
Measurements
CONUS groundwater model, structured,
nlay=1,nrow=11419,ncol=18458,210,771,902 cells, 124,884,583 active, split into 256 parts.
optimize_splitting_mask(active_only=True)split_modelon a 229 GB nodesplit_modelon a 377 GB nodeMemory is flat while the model loads (73 GB) and climbs through
_remap_nodesand
split_model.Where it goes
_node_map_connectionA dict entry with an int key costs roughly 100 bytes of table, the tuple another
56, and the ints inside it 28 each once they exceed the small integer cache.
Two changes that would recover most of it
1.
_node_mapas arrays rather than a dict. Twoint32arrays of lengthnnodes, one for the model and one for the new node number, replace the dict and
its tuples: ~44 GB to 1.7 GB. There are about 25 references across 11 methods,
all internal, and the lookups in
_remap_nodesare already loops over arraysthat would become fancy indexing.
save_node_mappingandload_node_mappingserialize the dict to HDF5, so this needs a format bump or a shim.
2.
_connectionin CSR form. Anindptrand anindicesarray replace thedict of lists: ~57 GB to ~5 GB. Concentrated in
_remap_nodes, but it needsmodelgrid.neighbors()to grow an array returning mode, so it touchesflopy/discretization/grid.pyas well.Either alone would bring this model onto a 229 GB node.
Both structures are also traversed per cell in
_remap_nodes, about 1 G dictoperations for this model, and contribute several hundred million objects for
the cyclic collector to walk, so the change should reduce wall time as well as
memory.
split_modeltook 88 minutes here, against 140 minutes forwrite_simulationwhich this would not affect.Smaller and independent:
model_splitter.py:1022buildscellids = list(zip([0] * len(cells[0]), cells[0], cells[1]))and passes it toget_node(), materializing a python tuple per cell of every part, where thenode number is
row * ncol + colin numpy.Versions: flopy 3.10.0 (
developat 7732f39 is unchanged in these paths),python 3.12.3, numpy 1.26.4.