Skip to content

snap_pour_point_d8: no memory guard on H*W working arrays #1362

@brendancol

Description

@brendancol

Summary

snap_pour_point_d8 allocates three full-resolution float64 buffers on the numpy and cupy backends with no memory guard. On a large grid this can exhaust host or device memory before the kernel completes.

Where

xrspatial/hydro/snap_pour_point_d8.py -- public dispatch in snap_pour_point_d8(...).

CPU dispatch (isinstance(fa_data, np.ndarray) branch):

fa = fa_data.astype(np.float64)              # H*W * 8 B
pp = np.asarray(pp_data, dtype=np.float64)   # H*W * 8 B
out = _snap_pour_point_cpu(fa, pp, ...)      # H*W * 8 B

CuPy dispatch (_snap_pour_point_cupy) does the same on device:

fa = flow_accum_data.astype(cp.float64)
pp = pour_points_data.astype(cp.float64)
out = cp.full((H, W), cp.nan, dtype=cp.float64)

Bytes per pixel

backend working buffers total
numpy fa(8) + pp(8) + out(8) 24 B/px
cupy fa(8) + pp(8) + out(8) 24 B/px

The dask and dask+cupy paths process windows around sparse pour points and do not allocate H*W working buffers, so they should skip the guard.

Worked example

A 30000 x 30000 float64 grid:

30000 * 30000 * 24 = 21.6 GB working memory

That fits on a workstation with 32 GB free, fails on most laptops, and can OOM-kill the process on shared hosts.

Proposed fix

Same pattern as #1318 / #1319 (flow_accumulation_d8) and the follow-on guards in #1350 / #1355 / #1357:

  • Add _check_memory(H, W) and _check_gpu_memory(H, W) helpers in snap_pour_point_d8.py, with _BYTES_PER_PIXEL = 24 and _GPU_BYTES_PER_PIXEL = 24.
  • Threshold: raise MemoryError if required > 50% of MemAvailable (psutil fallback) for CPU, or 50% of cupy.cuda.runtime.memGetInfo()[0] for GPU.
  • Wire the checks into the public snap_pour_point_d8 dispatch before the eager numpy / cupy allocations. Dask and dask+cupy paths skip the guard.
  • Error message points the user at the dask backend.

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