Feature Update: Retry logic for the ee.data.computePixels update.

xee/ext.py

@@ -147,6 +147,10 @@ def open(
       ee_init_kwargs: Optional[Dict[str, Any]] = None,
       ee_init_if_necessary: bool = False,
       executor_kwargs: Optional[Dict[str, Any]] = None,
+      tile_fetch_kwargs: Dict[str, int] = {
+          'max_retries': 6,
+          'initial_delay': 500,
+      },
   ) -> 'EarthEngineStore':
     if mode != 'r':
       raise ValueError(
@@ -168,6 +172,7 @@ def open(
         ee_init_kwargs=ee_init_kwargs,
         ee_init_if_necessary=ee_init_if_necessary,
         executor_kwargs=executor_kwargs,
+        tile_fetch_kwargs=tile_fetch_kwargs,
     )
 
   def __init__(
@@ -186,6 +191,10 @@ def __init__(
       ee_init_kwargs: Optional[Dict[str, Any]] = None,
       ee_init_if_necessary: bool = False,
       executor_kwargs: Optional[Dict[str, Any]] = None,
+      tile_fetch_kwargs: Dict[str, int] = {
+          'max_retries': 6,
+          'initial_delay': 500,
+      },
   ):
     self.ee_init_kwargs = ee_init_kwargs
     self.ee_init_if_necessary = ee_init_if_necessary
@@ -195,6 +204,10 @@ def __init__(
       executor_kwargs = {}
     self.executor_kwargs = executor_kwargs
 
+    # Default value: (https://github.com/pydata/xarray/blob/main/xarray/backends/common.py#L181).
+    self.tile_fetch_max_retries = tile_fetch_kwargs.pop('max_retries', 6)
+    self.tile_fetch_initial_delay = tile_fetch_kwargs.pop('initial_delay', 500)
+
     self.image_collection = image_collection
     if n_images != -1:
       self.image_collection = image_collection.limit(n_images)
@@ -482,7 +495,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.tile_fetch_max_retries,
+        initial_delay=self.tile_fetch_initial_delay,
     )
 
     # Extract out the shape information from EE response.
@@ -965,6 +982,10 @@ def open_dataset(
       ee_init_if_necessary: bool = False,
       ee_init_kwargs: Optional[Dict[str, Any]] = None,
       executor_kwargs: Optional[Dict[str, Any]] = None,
+      tile_fetch_kwargs: Dict[str, int] = {
+          'max_retries': 6,
+          'initial_delay': 500,
+      },
   ) -> xarray.Dataset:  # type: ignore
     """Open an Earth Engine ImageCollection as an Xarray Dataset.
 
@@ -1037,7 +1058,14 @@ 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}.
-
+      tile_fetch_kwargs (Dict): The necessary kwargs like `max_retries`,
+        `initial_delay` which helps while fetching data through calling
+        ee.data.computePixels(). i.e. {'max_retries' : 6, 'initial_delay': 500}.
+        - max_retries is maximum number of retry attempts for calling
+        ee.data.computePixels().By default, it is 6.
+        - initial_delay is the initial delay in milliseconds before the first
+        retry of calling ee.data.computePixels(). By default, it is 500.
+        (https://github.com/pydata/xarray/blob/main/xarray/backends/common.py#L181).
     Returns:
       An xarray.Dataset that streams in remote data from Earth Engine.
     """
@@ -1067,6 +1095,7 @@ def open_dataset(
         ee_init_kwargs=ee_init_kwargs,
         ee_init_if_necessary=ee_init_if_necessary,
         executor_kwargs=executor_kwargs,
+        tile_fetch_kwargs=tile_fetch_kwargs,
     )
 
     store_entrypoint = backends_store.StoreBackendEntrypoint()

xee/ext_integration_test.py

@@ -69,6 +69,7 @@ def setUp(self):
             '2017-01-01', '2017-01-03'
         ),
         n_images=64,
+        tile_fetch_kwargs={'max_retries': 10, 'initial_delay': 1500},
     )
     self.lnglat_store = xee.EarthEngineStore(
         ee.ImageCollection.fromImages([ee.Image.pixelLonLat()]),
@@ -255,6 +256,15 @@ def __getitem__(self, params):
 
     self.assertEqual(getter.count, 3)
 
+  def test_tile_fetch_kwargs(self):
+    arr = xee.EarthEngineBackendArray('B4', self.store)
+    self.assertEqual(arr.store.tile_fetch_initial_delay, 1500)
+    self.assertEqual(arr.store.tile_fetch_max_retries, 10)
+
+    arr1 = xee.EarthEngineBackendArray('longitude', self.lnglat_store)
+    self.assertEqual(arr1.store.tile_fetch_initial_delay, 500)
+    self.assertEqual(arr1.store.tile_fetch_max_retries, 6)
+
 
 class EEBackendEntrypointTest(absltest.TestCase):