Skip to content

Mf6Splitter peak memory is dominated by per-cell python dicts #2813

Description

@jdhughes-dev

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions