Issue #1696 Optimize modelsplitter

[Manangka]

Fixes #1696

Description

This PR implements some of the key insights @Huite found while optimzing the splitting of Models. Although there are some differences in his implementation and this PR. The optimizatiob exists of 2 parts

The First part has to do with the way models and packages are split. As @Huite pointed out the current implementation leads to a package being loaded and unloaded several times, which leads to a huge performance penalty because the package needs to be loaded from disk each time.
This has been solved by reverting the order of looping. In the original implementation there was a loop over the models, foreach partition a deepcopy was made of the original model. Then afterwards we loop over each model and access each package in the model to determine if it belongs to the same domain as the model. This is expensive.
This PR reverts the loop. Frist a shallow copy is made of all the models. Then we loop over each package of the original package and we determine if belongs to once of the submodels. This way we only need the open a package once.

The second part is the way the SSM and BUY packages are updated after the splitting has happend. There was a pience of code that looped over all the flow and transport models to determine which models share the same domain (are in the same partition). This is an expensive operation because the grid data has to be loaded into memory over and over.
In this PR this has been solved by bringing the update of the SSM and BUY package to the place where that information is already know i.e. the model splitter. This way we can avoid the expensive loading of the grid data.
Also by moving it to the modelsplitter we can avoid forcing the load of the grid data manually

With these changes the splitting of the mode @Huite provided me takes now around 5m:30s

Checklist

• Links to correct issue
• Update changelog, if changes affect users
• PR title starts with Issue #nr, e.g. Issue #737
• Unit tests were added
• If feature added: Added/extended example
• If feature added: Added feature to API documentation
• If pixi.lock was changed: Ran pixi run generate-sbom and committed changes

[Huite]

I'm probably a bit too early FWIW, but I'll comment anyway; I think I recognize the predicate variable approach I suggested in a comment.

The activity count function in my draft PR has this:

    # Use ellipsis to reduce over ALL dimensions except label dims
    dims_to_aggregate = [dim for dim in sample.dims if dim not in label_dims]
    counts = sample.notnull().sum(dim=dims_to_aggregate).groupby(labels).sum()
    return {label: int(n) for label, n in enumerate(counts.data)}

This is more robust in terms of IO-overhead, since if the data is lazy, we require a single load of the data to do the groupby. Whereas doing a clip / isel kind of thing creates a necessity to load for each partition created.

Of course, the IO per partition is inescapable if you start writing the models, hence my spatial chunking (after reordering via space filling curve) prior to split & write.

Even so, I think the groupby on the unpartitioned data is worthwhile since it will lead to less IO regardless of whether the data is spatially chunked.

[Manangka]

Hi @Huite, I tackled the problem in a slightly different way that you do. Or at least I think I'm doing the same but in a different order. As you correctly asset one of the biggest problem with the model splitter is that it is loading and unloading the same package data for each partition.

This is happening because due the way the current modelsplitter is designed. It makes a deepcopy of the original model, including all the packages, and then goes over each partitioned model and assesses if the packages in that partitioned model actually belongs to the domein of that model. (You of course already know this but just stating it for future references)

If i read your code correctly you solved this by creating a boundary polygon for each partition first. Then when looping over the models and assessing if a package belongs to that model you use the polygon test. This works great because you;ve checked before if a package belongs to a partition or not.

I Solved the issue in a different way. Basically your approach and the original code do a double loop with the model as the outer loop and the packages as the inner loop. I turn this around. I first create all the partioned models but without packages. Then I start looping over the packages (outer loop). The package is loaded and stays in memory. Then i loop over the models and check if the package belongs to the model. The latter is done by comparing the a known domain value of the package to the active_domain of the model partition. Doing it this way I ensure that the package data is only loaded once.

With this approach the split takes between 4m30 and 5m30 for the big model that you shared.

What do you think of this approach?
If you have time can you run it on your machine and test if you also see the same gain in performance

[Huite]

Alright! The most important part to me is:

The latter is done by comparing the a known domain value of the package to the active_domain of the model partition.

This sounds like a good approach!

Doing it this way I ensure that the package data is only loaded once.

I'm not entirely sure this is the case, though?

The key lines are:

                active_package_domain = package[
                    self._pkg_id_to_var_mapping[pkg_id]
                ].notnull()

Unless we force a .compute() here, I think the active_package_domain could be a dask backed xarray DataArray. If that's the case, the following part is where the compute is forced:

               has_overlap = (
                    active_package_domain & partition_info.active_domain.astype(bool)
                ).any()  # type: ignore

Since it's the any() reduction which definitely forces a computation.

I think there are two minor but possibly impactful improvements: Reduce the active_package_domain to spatial dimensions only immediately in _get_package_domain, i.e. get rid of time. This single spatial grid should easily fit into memory, so we can call a .compute() without worry and ensure we incur IO only once.

[JoerivanEngelen]

Looks good, I have a few comments. Thanks for going through the effort of creating this clean solution and adding a separate ModelSplitter class for it. Depending on capacity we'll see whether you or I will pick up my comments.

[JoerivanEngelen]

I don't really see this being used?

[JoerivanEngelen]

I think this can be simplified to:

Suggested change

	if isinstance(package, IAgnosticPackage):
	if sliced_package["index"].size == 0:
	sliced_package = None
	if isinstance(package, IAgnosticPackage) and sliced_package.is_empty:
	sliced_package = None

In this way, you are not referring to the "index" variable, which might be or not be present. I think the check for IAgnosticPackage prevents calling the AllNoDataScheme on large, chunked grids, which .is_empty does.

[JoerivanEngelen]

I find this method name a bit generic, how about:

Suggested change

	def update_packages(self) -> None:
	def update_packages_after_split(self) -> None:

Perhaps CoPilot has a briefer suggestion

[JoerivanEngelen]

I resorted to "update_dependent_packages" in the end

[JoerivanEngelen]

I miss the remove_expanded_auxiliary_variables_from_dataset(package) preceding this, that was there before. This is used for multi-species simulations

[JoerivanEngelen]

Forget what I said, I see a split with multispecies is tested in test_split_flow_and_transport_model, so this aux expansion isn't necessary here.

[JoerivanEngelen]

I think there are two minor but possibly impactful improvements: Reduce the active_package_domain to spatial dimensions only immediately in _get_package_domain, i.e. get rid of time. This single spatial grid should easily fit into memory, so we can call a .compute() without worry and ensure we incur IO only once.

FYI: We can call the trim_time_dimension function to get rid of the time dimension.

[JoerivanEngelen]

One more comment

[JoerivanEngelen]

I don't think layer can be dropped here: What if there is no data present in layer 1, but there is in layer 2? What does xarray do when you then drop_vars on layer? I think we either have to aggregate by summing all notnull across the layer dim, or just leave the layer dim in here and let xarray broadcast.

[JoerivanEngelen]

Adding "pass" is a code smell. I refactored this to a separate method _is_package_to_skip, as the same logic is also used in _get_package_domain.

[JoerivanEngelen]

I resorted to "update_dependent_packages" in the end

[sonarqubecloud]

Quality Gate passed

Issues
1 New issue
0 Accepted issues

Measures
0 Security Hotspots
0.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

[JoerivanEngelen]

I've added functionality to ignore time, and it reduced the splitting operation to 1.5 minute.