diff --git a/xee/ext.py b/xee/ext.py index 6b8b203..f5d4a97 100644 --- a/xee/ext.py +++ b/xee/ext.py @@ -147,6 +147,8 @@ def open( ee_init_kwargs: Optional[Dict[str, Any]] = None, ee_init_if_necessary: bool = False, executor_kwargs: Optional[Dict[str, Any]] = None, + compute_pixels_max_retries: Optional[int] = None, + compute_pixels_initial_delay: Optional[int] = None, ) -> 'EarthEngineStore': if mode != 'r': raise ValueError( @@ -168,6 +170,8 @@ def open( ee_init_kwargs=ee_init_kwargs, ee_init_if_necessary=ee_init_if_necessary, executor_kwargs=executor_kwargs, + compute_pixels_max_retries=compute_pixels_max_retries, + compute_pixels_initial_delay=compute_pixels_initial_delay, ) def __init__( @@ -186,6 +190,8 @@ def __init__( ee_init_kwargs: Optional[Dict[str, Any]] = None, ee_init_if_necessary: bool = False, executor_kwargs: Optional[Dict[str, Any]] = None, + compute_pixels_max_retries: Optional[int] = None, + compute_pixels_initial_delay: Optional[int] = None, ): self.ee_init_kwargs = ee_init_kwargs self.ee_init_if_necessary = ee_init_if_necessary @@ -195,6 +201,17 @@ def __init__( executor_kwargs = {} self.executor_kwargs = executor_kwargs + # Here 6 & 500 is default value. + # (https://github.com/pydata/xarray/blob/main/xarray/backends/common.py#L181). + self.compute_pixels_max_retries = ( + 6 if compute_pixels_max_retries is None else compute_pixels_max_retries + ) + self.compute_pixels_initial_delay = ( + 500 + if compute_pixels_initial_delay is None + else compute_pixels_initial_delay + ) + self.image_collection = image_collection if n_images != -1: self.image_collection = image_collection.limit(n_images) @@ -482,7 +499,11 @@ def image_to_array( **kwargs, } raw = common.robust_getitem( - pixels_getter, params, catch=ee.ee_exception.EEException + pixels_getter, + params, + catch=ee.ee_exception.EEException, + max_retries=self.compute_pixels_max_retries, + initial_delay=self.compute_pixels_initial_delay, ) # Extract out the shape information from EE response. @@ -965,6 +986,8 @@ def open_dataset( ee_init_if_necessary: bool = False, ee_init_kwargs: Optional[Dict[str, Any]] = None, executor_kwargs: Optional[Dict[str, Any]] = None, + compute_pixels_max_retries: Optional[int] = None, + compute_pixels_initial_delay: Optional[int] = None, ) -> xarray.Dataset: # type: ignore """Open an Earth Engine ImageCollection as an Xarray Dataset. @@ -1037,6 +1060,10 @@ def open_dataset( executor_kwargs (optional): A dictionary of keyword arguments to pass to the ThreadPoolExecutor that handles the parallel computation of pixels i.e. {'max_workers': 2}. + compute_pixels_max_retries (optional): The maximum number of retry + attempts for calling ee.data.computePixels(). + compute_pixels_initial_delay (optional): The initial delay in milliseconds + before the first retry of calling ee.data.computePixels(). Returns: An xarray.Dataset that streams in remote data from Earth Engine. @@ -1067,6 +1094,8 @@ def open_dataset( ee_init_kwargs=ee_init_kwargs, ee_init_if_necessary=ee_init_if_necessary, executor_kwargs=executor_kwargs, + compute_pixels_max_retries=compute_pixels_max_retries, + compute_pixels_initial_delay=compute_pixels_initial_delay, ) store_entrypoint = backends_store.StoreBackendEntrypoint()