Refactor rawdem

[rosepearson]

@jennan refactor to try address Dask failures

Specifically:
Create mask arrays instead of repeatably clipping against geometries.
Revisit the save/load
Add a no_values_mask property

• Make code change
• Update tests
  • Update / create new tests
  • Ensure these tests have the expected behavour
  • Test locally and ensure tests are passing
• Ensure tests passing on GH Actions
• Update documentation
  • Doc strings
  • Wiki
• Update package (version.py & pyproject.toml) version (if needed)

[jennan]

Hi, here are my notes about what I have tried so far, in order to improve the behavior of the dask part on the example tile (i.e. getting good cpu utilisation, little transfer overhead reported in the dashboard and no worker timeout errors).

Trying without work stealing (https://distributed.dask.org/en/latest/work-stealing.html)

DASK_DISTRIBUTED__SCHEDULER__WORK_STEALING="False" python -m geofabrics --instruction instruction_lidar.json

• poor utilisation (~65%)
• terrible compute vs. transfer (more transfer than compute)
• so... let's forget it

Trying to change priority (https://distributed.dask.org/en/stable/priority.html)

with dask.annotate(priority=10):
    dem["z"] = dem.z.where(mask , 0)

• 70% to 90% cpu utilisation
• transfer not neglectable but lower than compute

Trying to pre-scatter the mask to clip dem in add_lidar (for foreshore)

from distributed.client import default_client
mask = client.scatter(mask)

• 70% to 90% utilisation
• seems quite slow
• repeated messages about large graph while no transfer (only compute)
• but then timeout messages and transfer increase

Trying to get more reliable comms

DASK_DISTRIBUTED__COMM__RETRY__COUNT=3 python -m geofabrics --instruction instruction_lidar.json

• doesn't seem to avoid timeout errors

Trying to remove all clipping/masking before self.raw_dem.save_dem(cached_file, ...)

• cpu utilisation between 30 and 90%
• no transfer overhead
• no timeout error

Testing with different masking in add_lidar (just testing up to self.raw_dem.save_dem(cached_file, ...) and still without most rio.clip)

• only masking dem.z: timeout errors (run a bit)
• only masking dem.z / only clip_mask: timeout error (run a bit)
• only masking dem.z / only dem.z < 0:
  • no timeout error and no transfer oberhead (run up to 80%)
  • then timeout errors but was playing with Dask profilers...
• masking all dem components / only dem.z < 0: no timeout error and no transfer oberhead
• masking all dem components / xarray.ones_like(dem.z, dtype=bool) | (dem.z <= 0)
  • mostly no timeout error and no transfer oberhead
  • but some timeout errors after 90% completion
• masking all dem components / persist clip_mask: some timeout errors and transfer overhead
• masking all dem components / (dem.z <= 0) | mask and mask not lazy
  • not so good cpu utilisation, transfer overhead
  • no timeout messages
• mask all dem components / (dem.z <= 0) | mask and broadcasted mask
  • see Passing a distributed.Future to the kwargs of apply_ufunc should resolve the future pydata/xarray#6803 (comment)
  • failed with error

ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 1 dimensions. The detected shape was (2,) + inhomogeneous part.

• mask all dem components / (dem.z <= 0) | mask and dask.block

mask = ~clip_mask(dem.z, buffered_foreshore.geometry, None)

arrs = []
for i in range(0, mask.shape[0], self.chunk_size):
    sub_arrs = []
    for j in range(0, mask.shape[1], self.chunk_size):
        chunk = mask[i : i + self.chunk_size, j : j + self.chunk_size].copy()
        sub_arrs.append(chunk)
    arrs.append(sub_arrs)
mask = dask.array.block(arrs)

• transfer overhead, timeout messages
• masking all dem components / clip_mask without compute
  • some timeouts and transfer overhead, not too bad

But adding clip_mask back everywhere, transfer is bad (significant amount vs. compute) and timeout errors for dask workers :(.

TODO export DASK__DISTRIBUTED__COMM__TIMEOUTS__CONNECT="60s"
TODO check scheduling https://distributed.dask.org/en/latest/scheduling-policies.html

[jennan]

I approve on behalf of Rose :-).