MPI Message Size Fix (NGWPC-11089)

[idtodd]

MPI broadcasts have a data size limit of ~2 GB due to a 32-bit integer being used for messaging the number of bytes passed. Some of the xarray Datasets we are passing surpass this size limit, leading to an exception being thrown in some areas.

The changes made address this by checking the size of the Dataset before passing, and if it exceeds 1.5 GB, the Dataset will be pickled by MPI root, split into 1.5 GB chunks, then broadcasted to the other ranks to be merged and depickled. Attempts were made in this process to limit memory usage during this process as, at minimum, you're duplicating an object that is at least 1.5 GB in size.

The 1.5 GB size was chosen to attempt a sizable buffer before hitting the 32-bit size limit. MPI seems to have a custom pickling system when they transfer python objects, and I'm unsure how much overhead gets added to an object during the pickling process. I've been able to confirm that 1.5 GB is a very safe number to use to fix the problem, but it might be worth experimenting with larger numbers if we notice a significant slowdown from the chunking overhead.

Additions

Removals

Changes

• Chunk dataset MPI transfers if the dataset exceeds 1.5 GB in size.

Testing

Screenshots

Notes

Todos

• Because of the incoming refactoring, debug logs were not added to prevent adding logging structures that will soon become outdated. In the future, debug logging things like when chunking is being done and the size of those chunks might be helpful.

Checklist

• PR has an informative and human-readable title
• Changes are limited to a single goal (no scope creep)
• Code can be automatically merged (no conflicts)
• Code follows project standards (link if applicable)
• Passes all existing automated tests
• Any change in functionality is tested
• New functions are documented (with a description, list of inputs, and expected output)
• Placeholder code is flagged / future todos are captured in comments
• Visually tested in supported browsers and devices (see checklist below 👇)
• Project documentation has been updated (including the "Unreleased" section of the CHANGELOG)
• Reviewers requested with the Reviewers tool ➡️

Testing checklist

Target Environment support

• Linux

[mxkpp]

Would Bcast (uppercase B) help vs bcast since the former does not use pickling and is more direct if I understand correctly?

[mxkpp]

Would Bcast (uppercase B) help vs bcast since the former does not use pickling and is more direct if I understand correctly?

The PR description confused me since the picking/depickling described made me think bcast was being used, but I see Bcast is being leveraged in the proposed code, with chunking. Is the pickling/depickling actually needed with Bcast? (I understand if chunking is needed due to memory constraints, but I wonder if pickling could be avoided).

Here is another uses of Bcast in the codebase for example. This probably does not experience the magnitude of data that is being experienced in the part of the code affected by the current PR though I realize.

ngen-forcing/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py

Lines 360 to 381 in a0f217f

	def broadcast_parameter(self, value_broadcast, config_options, param_type):
	"""Broadcast a single parameter value to all processors.
	Generic function for sending a parameter value out to the processors.
	:param value_broadcast:
	:param config_options:
	:return:
	"""
	dtype = np.dtype(param_type)
	if self.rank == 0:
	param = np.asarray(value_broadcast, dtype=dtype)
	else:
	param = np.empty(dtype=dtype, shape=())
	try:
	self.comm.Bcast(param, root=0)
	except MPI.Exception:
	config_options.errMsg = "Unable to broadcast single value from rank 0."
	err_handler.log_critical(config_options, self)
	return None
	return param.item(0)

[mxkpp]

Here's another existing Bcast call (without explicit picking/depickling), this one closer to current PR use case since it's operating on a 2D array:

ngen-forcing/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py

Lines 417 to 441 in a0f217f

	# Broadcast the global array to the child processors, then
	if self.rank == 0:
	arrayGlobalTmp = array_broadcast
	else:
	if data_type_flag == 1:
	arrayGlobalTmp = np.empty(
	[geoMeta.ny_global, geoMeta.nx_global], np.float32
	)
	else: # data_type_flag == 2:
	arrayGlobalTmp = np.empty(
	[geoMeta.ny_global, geoMeta.nx_global], np.float64
	)
	try:
	self.comm.Bcast(arrayGlobalTmp, root=0)
	except Exception:
	ConfigOptions.errMsg = (
	"Unable to broadcast a global numpy array from rank 0"
	)
	err_handler.log_critical(ConfigOptions, self)
	return None
	arraySub = arrayGlobalTmp[
	geoMeta.y_lower_bound : geoMeta.y_upper_bound,
	geoMeta.x_lower_bound : geoMeta.x_upper_bound,
	]
	return arraySub

[idtodd]

Bcast can only be used with contiguous data of a single size. bcast is used to let the MPI module handle pickling and unpickling complex data to and from contiguous byte arrays.

In this particular place, ds is an xarray.Dataset, meaning it's a complex object of references to arrays and strings. You cannot directly Bcast this since it isn't represented by contiguous data. At some point, the dataset needs to be pickled so it can be properly messaged through the MPI C functions. The point of the change here is to manually pickle if the size of the dataset is expected to exceed the maximum number of bytes allowed to be transferred in a single MPI broadcast call.

For the 2D numpy array, numpy arrays are always stored in contiguous C arrays. If you have an array of shape (5, 5), the underlying data is stored as a 25 length C array. When you access the data with arr[2, 3] in python, it's converting that to arr[2 * 5 + 3] for finding the location in the C array. That means a single numpy array of any shape can be Bcast, but if you start collecting multiple arrays together in pandas or xarray datasets, you need to pickle it first.

[mxkpp]

It's nice that a xarray.Dataset is picklable so bcast (lowercase b) is easy to use with it. If we end up running into performance bottlenecks, we could Bcast (uppercase B) the individual Bcastable components of the Dataset even though the Dataset object itself cannot be Bcasted.

[mxkpp]

This approach is reasonable, approving, but is somewhat wasteful of memory. I needed to increase my environment's memory to be able to run the forcing pytest for the AnA case with mpirun -n 2.

1. It is creating a pickled representation of the entire Dataset
2. It is creating a list of all the chunks before iterating through them.

If the increased memory usage needs to be mitigated, I think 2) should be revised so that the chunks are processed JIT.

[idtodd]

Update made to let it iterate over the pickled bytes instead of building a list. This is a double-edged sword since we're no longer clearing memory after we use it, but it does decrease overhead of creating copies of the end data as it's chunked.

Primarily, I've increased the cutoff for when chunking happens and decreased the chunk size significantly. This will mean more MPI messages need to be sent, but it prevents needing 1.5 GB of memory to be allocated for each chunk.

[idtodd]

On the topic of trying to use Bcast, first bcast will ultimately pickle the object passed, so our pickling shouldn't increase memory usage relative to the default bcast to any significant degree.

As for attempting to Bcast the xarray.Dataset, it looks to be made up of several private dictionaries at its base level. I'm hesitant to make a custom Bcast wrapper around the innards of a 3rd party library object since those are not meant to be stable properties between versions.