Skip to content

Add memory guard to cost_distance numpy path #1252

Description

@brendancol

Describe the bug

_cost_distance_kernel in xrspatial/cost_distance.py allocates several arrays sized from the raster dimensions with no memory check. Peak working memory is about 37 bytes per pixel once you add up dist (float64), h_keys (float64), h_rows (int64), h_cols (int64), visited (int8), and the float32 output. A 20000x20000 raster (400M pixels) tries to grab ~15 GB and either OOM-kills or thrashes swap.

The dask iterative path already has a guard. See _cost_distance_dask_iterative around line 919: it checks estimated against _available_memory_bytes() and raises MemoryError pointing users at max_cost=. The numpy path has nothing. Same pattern viewshed had in #1229.

To Reproduce

import numpy as np
import xarray as xr
from xrspatial import cost_distance

# Tries to allocate a multi-GB heap before any check
arr = np.zeros((20000, 20000), dtype=np.float64)
arr[0, 0] = 1.0
friction = np.ones_like(arr)
xs = np.arange(20000, dtype=float)
ys = np.arange(20000, dtype=float)
r = xr.DataArray(arr, coords=dict(x=xs, y=ys), dims=["y", "x"], attrs={"res": (1.0, 1.0)})
f = xr.DataArray(friction, coords=dict(x=xs, y=ys), dims=["y", "x"], attrs={"res": (1.0, 1.0)})
cost_distance(r, f)

Expected behavior

MemoryError before allocation, with a message pointing users at max_cost= or a dask-backed array.

Proposed fix

Guard at the top of _cost_distance_numpy: estimate peak working memory (~37 * H * W bytes) via _available_memory_bytes() and raise if it would blow past a safe fraction of RAM. Mirrors the viewshed #1229 guard and pathfinding's _check_memory.

Allocations that add up

  • dist: 8 bytes/pixel
  • h_keys (float64): 8 bytes/pixel
  • h_rows (int64): 8 bytes/pixel
  • h_cols (int64): 8 bytes/pixel
  • visited (int8): 1 byte/pixel
  • out (float32): 4 bytes/pixel
  • Total ~37 bytes/pixel

Sister module to pathfinding and viewshed. Both already have memory guards for this exact class of allocation.

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 workinginput-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