Merge pull request #119 from quantifyearth/mwd-cache-limits

Enforce cache limits on GDAL and MLX

Author: mdales

CHANGES.md

@@ -1,8 +1,9 @@
-## v1.12.2 (26/1/2026)
+## v1.12.2 (27/1/2026)
 
 ### Added
 
 * Added `yg.sum`, `yg.all`, `yg.any` to build layers from lists of layers.
+* Cache limits applied to GDAL and MLX which tend to assume that they are the only thing using memory. You can override these by adjusting `yg.constants.GDAL_CACHE_LIMIT` and `yg.constants.MLX_CACHE_LIMIT`.
 
 ## v1.12.1 (21/1/2026)
 

yirgacheffe/_backends/mlx.py

@@ -2,12 +2,14 @@
 
 from typing import Callable
 
+from osgeo import gdal
 import numpy as np
 import mlx.core as mx # type: ignore
 import mlx.nn
 
 from .enumeration import operators as op
 from .enumeration import dtype
+from ..  import constants
 
 array_t = mx.array
 float_t = mx.float32
@@ -53,6 +55,10 @@
 round_op = mx.round
 ceil_op = mx.ceil
 
+def init() -> None:
+    gdal.SetCacheMax(constants.GDAL_CACHE_LIMIT)
+    mx.set_cache_limit(constants.MLX_CACHE_LIMIT)
+
 def sum_op(a):
     # By default the type promotion rules for sum in MLX are not the same as with Numpy. E.g.,
     #

yirgacheffe/_backends/numpy.py

@@ -2,11 +2,13 @@
 
 from typing import Callable
 
+from osgeo import gdal
 import numpy as np
 import torch
 
 from .enumeration import operators as op
 from .enumeration import dtype
+from .. import constants
 
 array_t = np.ndarray
 float_t = np.float64
@@ -57,6 +59,9 @@
 round_op = np.round
 ceil_op = np.ceil
 
+def init() -> None:
+    gdal.SetCacheMax(constants.GDAL_CACHE_LIMIT)
+
 def conv2d_op(data, weights):
     # torch wants to process dimensions of channels of width of height
     # Which is why both the data and weights get nested into two arrays here,

yirgacheffe/_operators/__init__.py

@@ -811,6 +811,8 @@ def sum(self) -> float:
 
         cse_cache = CSECacheTable(self, computation_window)
 
+        backend.init()
+
         for yoffset in range(0, computation_window.ysize, self.ystep):
             cse_cache.reset_cache()
             step=self.ystep
@@ -837,6 +839,8 @@ def min(self) -> float:
 
         cse_cache = CSECacheTable(self, computation_window)
 
+        backend.init()
+
         for yoffset in range(0, computation_window.ysize, self.ystep):
             cse_cache.reset_cache()
             step=self.ystep
@@ -863,6 +867,8 @@ def max(self) -> float:
 
         cse_cache = CSECacheTable(self, computation_window)
 
+        backend.init()
+
         for yoffset in range(0, computation_window.ysize, self.ystep):
             cse_cache.reset_cache()
             step=self.ystep
@@ -898,6 +904,8 @@ def unique(self, return_counts:bool=False) -> np.ndarray | tuple[np.ndarray,np.n
 
         cse_cache = CSECacheTable(self, computation_window)
 
+        backend.init()
+
         for yoffset in range(0, computation_window.ysize, self.ystep):
             cse_cache.reset_cache()
             step=self.ystep
@@ -984,6 +992,8 @@ def save(self, destination_layer, and_sum=False, callback=None, band=1) -> float
 
         cse_cache = CSECacheTable(self, computation_window)
 
+        backend.init()
+
         for yoffset in range(0, computation_window.ysize, self.ystep):
 
             cse_cache.reset_cache()
@@ -1024,6 +1034,8 @@ def _parallel_worker(
         # the cache build once per worker
         cse_cache = CSECacheTable(self, computation_window)
 
+        backend.init()
+
         arr = np.ndarray((self.ystep, width), dtype=np_dtype, buffer=shared_mem.buf) # type: ignore[var-annotated]
         projection = self.map_projection
         # TODO: the `save` method does more sanity checking that parallel save!

yirgacheffe/constants.py

@@ -3,6 +3,13 @@
 YSTEP = 512
 MINIMUM_CHUNKS_PER_THREAD = 1
 
+# Both GDAL and MLX assume that there is one instance running and it has the right to use
+# all the memory it can. In general Yirgacheffe's chunking and it's own caching is what we shou
+# be relying on, so we set some limits here. These are applied before each calculation, and so
+# in theory they can be tweaked if necessary on demand.
+GDAL_CACHE_LIMIT = 1 * 1024 * 1024 * 1024
+MLX_CACHE_LIMIT = 1 * 1024 * 1024 * 1024
+
 # I don't really want this here, but it's just too useful having it exposed
 # This used to be a fixed string, but now it is at least programmatically generated
 WGS_84_PROJECTION = pyproj.CRS.from_epsg(4326).to_wkt(version='WKT1_GDAL')