Skip to content

stream_link_mfd: no memory guard on H*W working arrays #1337

@brendancol

Description

@brendancol

Summary

stream_link_mfd allocates several H*W working arrays plus an (8, H, W) fractions copy on the numpy and cupy backends with no memory check before allocation. A large input DataArray triggers an unbounded host or device allocation that crashes the process or the calling notebook.

Same pattern fixed in #1318 (flow_accumulation_d8), #1319 (flow_accumulation_mfd), #1335 (stream_link_d8). MFD adds the (8, H, W) fractions input copy on top of the d8 working set.

Affected backends

  • numpy (_stream_link_mfd_cpu): peak working set ~97 B/px

    • input copy frac (8, H, W) float64 -> 64
    • stream_mask int8 -> 1
    • link_id float64 -> 8
    • in_degree int32 -> 4
    • orig_indeg int32 -> 4
    • queue_r int64 -> 8
    • queue_c int64 -> 8
  • cupy (_stream_link_mfd_cupy): peak device working set ~93 B/px

    • fractions_f64 (8, H, W) float64 -> 64
    • stream_mask_i8 int8 -> 1
    • in_degree int32 -> 4
    • orig_indeg int32 -> 4
    • state int32 -> 4
    • link_id float64 -> 8
    • plus fa_cp input copy float64 -> 8

The dask and dask+cupy paths process per-tile via _stream_link_mfd_tile_kernel, so per-call allocation is bounded by chunk shape and the guard does not need to run on those paths.

Worked example

import numpy as np
import xarray as xr
from xrspatial.hydro import stream_link_mfd

H = W = 35_000  # ~118 GB on the numpy backend at 97 B/px
frac = xr.DataArray(np.zeros((8, H, W), dtype=np.float64), dims=['neighbor', 'y', 'x'])
fa = xr.DataArray(np.full((H, W), 200.0, dtype=np.float64))
stream_link_mfd(frac, fa, threshold=100)

The kernel OOMs before returning.

Proposed fix

Mirror the helper layout from #1335:

  • per-module _BYTES_PER_PIXEL = 97, _GPU_BYTES_PER_PIXEL = 100
  • _available_memory_bytes, _available_gpu_memory_bytes, _check_memory, _check_gpu_memory (50% threshold, MemoryError with dimensions and a pointer to the dask path)
  • wire the checks into stream_link_mfd before any H*W allocation on the numpy and cupy branches; leave dask and dask+cupy untouched

Acceptance

  • numpy and cupy paths raise MemoryError mentioning grid dimensions and the dask alternative when projected usage exceeds 50% of available memory
  • dask and dask+cupy paths unaffected
  • existing test suite passes

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't workinghigh-priorityoomOut-of-memory risk with large datasets

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions