Skip to content

diffusion: unbounded allocations and unbounded steps loop #1267

Description

@brendancol

Summary

xrspatial.diffusion.diffuse has no memory guard on its core allocations, and steps has no upper bound. Large rasters or large step counts can OOM the host or pin a CPU.

What goes wrong

  1. diffuse() always materializes a full float64 alpha raster via np.full(agg.shape, alpha_scalar, dtype=np.float64) (diffusion.py line 305), even when the input is a dask DataArray with scalar diffusivity. A 100000x100000 input forces an 80 GB numpy allocation in the public API before any backend dispatch.

  2. _diffuse_step_numpy allocates np.empty((rows, cols)) every step (line 57) with no memory check.

  3. _diffuse_cupy allocates cp.empty_like(u) every step (line 123) with no GPU memory check.

  4. _validate_scalar(steps, ..., min_val=1) only sets a lower bound. steps=10**12 is accepted and would loop forever.

Fix

  • Add a peak-memory guard like the cost_distance / morphology one. Budget ~32 bytes per pixel (u, out, alpha, padded copy as float64) and raise MemoryError before the first allocation in the eager numpy / cupy paths.
  • For the dask + scalar case, don't materialize a full numpy alpha array. Pass the scalar through.
  • Cap steps so a single call can't pin a CPU forever.

Acceptance criteria

  • Oversize raster raises MemoryError before allocating.
  • Scalar diffusivity on a dask raster doesn't materialize a full numpy alpha array.
  • Out-of-range steps raises a clean ValueError.
  • Existing tests still pass.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't workinghigh-priorityinput-validationInput validation and error messagesoomOut-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