ENH: Add a new environment variable EOREADER_NOF_BANDS_IN_CHUNKS to…

… control the number of the bands in chunks when using `dask`. #180

CHANGES.md

@@ -3,6 +3,7 @@
 ## 0.21.8 (2024-mm-dd)
 
 - ENH: Add a new type (`BandsType`) for list of BandType
+- ENH: Add a new environment variable `EOREADER_NOF_BANDS_IN_CHUNKS` to control the number of the bands in chunks when using `dask`.
 - FIX: Fix stack `save_as_int` to use updated int values - by @TabeaW
 - FIX: Fixed PAZ Product Regex to properly indentify PAZ ST products as `PAZProduct` - by @guillemc23
 - FIX: Fixed PNEO Product Regex to properly indentify PNEO products as `PneoProduct` - by @guillemc23

eoreader/env_vars.py

@@ -58,5 +58,14 @@
 
 TILE_SIZE = "EOREADER_TILE_SIZE"
 """
-If set, overrides the default tile size used in chunking (2048 by default, i.e. default chunk is [1, 2048, 2048]). Only used if :code:`EOREADER_USE_DASK` is set to 1.
+If set, overrides the default tile size used in chunking (1024 by default, i.e. default chunk is {"band": 1, "x": 1024, "y": 1024}).
+Only used if :code:`EOREADER_USE_DASK` is set to 1.
+If 'auto' is set, the value passed as chunks will be 'auto'.
+"""
+
+NOF_BANDS_IN_CHUNKS = "EOREADER_NOF_BANDS_IN_CHUNKS"
+"""
+If set, overrides the default number of bands to be considered used in chunking (1 by default, i.e. default chunk is {"band": 1, "x": 1024, "y": 1024}).
+Only used if :code:`EOREADER_USE_DASK` is set to 1.
+Not used in case of :code:`EOREADER_USE_DASK` set as 'auto'.
 """

eoreader/products/optical/vhr_product.py

@@ -582,6 +582,11 @@ def _warp_band(
                 resolution=pixel_size,
             )
 
+            try:
+                tile_size = int(os.getenv(TILE_SIZE, DEFAULT_TILE_SIZE))
+            except ValueError:
+                tile_size = int(DEFAULT_TILE_SIZE)
+
             vrt_options = {
                 "crs": self.crs(),
                 "transform": utm_tr,
@@ -591,9 +596,7 @@ def _warp_band(
                 "resampling": Resampling.bilinear,
                 "nodata": self._raw_nodata,
                 # Float32 is the max possible
-                "warp_mem_limit": 32
-                * int(os.getenv(TILE_SIZE, DEFAULT_TILE_SIZE)) ** 2
-                / 1e6,
+                "warp_mem_limit": 32 * tile_size**2 / 1e6,
                 "dtype": src.meta["dtype"],
                 "num_threads": utils.get_max_cores(),
             }

eoreader/utils.py

@@ -36,12 +36,13 @@
 
 from eoreader import EOREADER_NAME
 from eoreader.bands import is_index, is_sat_band
-from eoreader.env_vars import TILE_SIZE, USE_DASK
+from eoreader.env_vars import NOF_BANDS_IN_CHUNKS, TILE_SIZE, USE_DASK
 from eoreader.exceptions import InvalidProductError
 from eoreader.keywords import _prune_keywords
 
 LOGGER = logging.getLogger(EOREADER_NAME)
 DEFAULT_TILE_SIZE = 1024
+DEFAULT_NOF_BANDS_IN_CHUNKS = 1
 UINT16_NODATA = rasters.UINT16_NODATA
 
 
@@ -161,7 +162,8 @@ def read(
     window = kwargs.get("window")
 
     # Always use chunks
-    tile_size = int(os.getenv(TILE_SIZE, DEFAULT_TILE_SIZE))
+    tile_size = os.getenv(TILE_SIZE, DEFAULT_TILE_SIZE)
+    nof_bands_in_chunks = os.getenv(NOF_BANDS_IN_CHUNKS, DEFAULT_NOF_BANDS_IN_CHUNKS)
 
     if use_dask():
         chunks = kwargs.get("chunks", {"band": 1, "x": tile_size, "y": tile_size})