diff --git a/CHANGES.md b/CHANGES.md index c306e099..6fbe4cc9 100644 --- a/CHANGES.md +++ b/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 diff --git a/eoreader/env_vars.py b/eoreader/env_vars.py index bbccf4c2..a7d614a4 100644 --- a/eoreader/env_vars.py +++ b/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'. """ diff --git a/eoreader/products/optical/vhr_product.py b/eoreader/products/optical/vhr_product.py index 04beca7a..8ced5d38 100644 --- a/eoreader/products/optical/vhr_product.py +++ b/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(), } diff --git a/eoreader/utils.py b/eoreader/utils.py index 3f0699c7..a3c2917a 100644 --- a/eoreader/utils.py +++ b/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})