Skip to content

watershed_dinf: no memory guard on H*W working arrays #1345

@brendancol

Description

@brendancol

Summary

watershed_dinf in xrspatial/hydro/watershed_dinf.py allocates a handful of H*W float64/int8/int64 arrays eagerly in its numpy and cupy paths with no upfront check that the projected working set fits in available memory. Pass a large enough input and the OS kills the process instead of letting Python raise a clean error.

Same pattern already fixed for the d8 and mfd siblings:

Bytes per pixel

CPU peak working set per pixel for the numpy dispatch plus _watershed_dinf_cpu:

  • fd (float64 cast) -> 8
  • labels (float64) -> 8
  • state (int8) -> 1
  • path_r (int64) -> 8
  • path_c (int64) -> 8

Total: 33 bytes/pixel.

D-inf encodes the downstream direction as a single real-valued angle, so the per-pixel CPU footprint matches d8 (one float64 channel) rather than mfd's eight-channel fractions buffer.

GPU peak working set per pixel for _watershed_dinf_cupy, which copies the device input to host, runs the CPU kernel, then ships the result back:

  • caller's flow_dir_data device float64 -> 8
  • caller's pour_points_data device float64 -> 8
  • cp.asarray(out) device float64 -> 8

Total: 24 bytes/pixel on the device.

Worked example

A 30000 x 30000 D-inf flow direction grid pulls roughly:

  • CPU: 30000 * 30000 * 33 = 29.7 GB working memory
  • GPU: 30000 * 30000 * 24 = 21.6 GB device memory

On a 32 GB workstation the CPU path will quietly eat nearly all RAM and trip the OOM killer. Most consumer GPUs can't fit 21 GB at all.

Fix

Same per-module helpers as #1330 / #1342:

  • _BYTES_PER_PIXEL = 33
  • _GPU_BYTES_PER_PIXEL = 24
  • _available_memory_bytes, _available_gpu_memory_bytes
  • _check_memory, _check_gpu_memory raising MemoryError when the projected allocation exceeds 50% of available RAM
  • Guards run in the public watershed_dinf dispatch ahead of the numpy and cupy allocations
  • Dask paths skip the guard; per-tile allocations are bounded by chunk size

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions