Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 12 commits into from
Mar 15, 2024
Prev Previous commit
Next Next commit
code updated after suggestions.
  • Loading branch information
dabhicusp committed Mar 14, 2024
commit 41b12e2a0de12d45df3440bd2485469b9d3348d6
27 changes: 14 additions & 13 deletions xee/ext.py
Original file line number Diff line number Diff line change
@@ -102,6 +102,11 @@ class EarthEngineStore(common.AbstractDataStore):
'height': 256,
}

TILE_FETCH_KWARGS: Dict[str, int] = {
'max_retries': 6,
'initial_delay': 500,
}

SCALE_UNITS: Dict[str, int] = {
'degree': 1,
'metre': 10_000,
@@ -147,10 +152,7 @@ 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,
},
tile_fetch_kwargs: Dict[str, int] = TILE_FETCH_KWARGS,
) -> 'EarthEngineStore':
if mode != 'r':
raise ValueError(
@@ -191,10 +193,7 @@ 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,
},
tile_fetch_kwargs: Dict[str, int] = TILE_FETCH_KWARGS,
):
self.ee_init_kwargs = ee_init_kwargs
self.ee_init_if_necessary = ee_init_if_necessary
@@ -204,9 +203,7 @@ 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.tile_fetch_kwargs = tile_fetch_kwargs

self.image_collection = image_collection
if n_images != -1:
@@ -498,8 +495,12 @@ def image_to_array(
pixels_getter,
params,
catch=ee.ee_exception.EEException,
max_retries=self.tile_fetch_max_retries,
initial_delay=self.tile_fetch_initial_delay,
max_retries=self.tile_fetch_kwargs.get(
dabhicusp marked this conversation as resolved.
Show resolved Hide resolved
'max_retries', self.TILE_FETCH_KWARGS.get('max_retries')
),
initial_delay=self.tile_fetch_kwargs.get(
'initial_delay', self.TILE_FETCH_KWARGS.get('initial_delay')
),
)

# Extract out the shape information from EE response.
12 changes: 8 additions & 4 deletions xee/ext_integration_test.py
Original file line number Diff line number Diff line change
@@ -81,6 +81,7 @@ def setUp(self):
'2020-03-30', '2020-04-01'
),
n_images=64,
tile_fetch_kwargs={'max_retries': 9},
)
self.all_img_store = xee.EarthEngineStore(
ee.ImageCollection('LANDSAT/LC08/C01/T1').filterDate(
@@ -258,13 +259,16 @@ def __getitem__(self, params):

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)
self.assertEqual(arr.store.tile_fetch_kwargs['initial_delay'], 1500)
self.assertEqual(arr.store.tile_fetch_kwargs['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)
self.assertEqual(arr1.store.tile_fetch_kwargs['initial_delay'], 500)
self.assertEqual(arr1.store.tile_fetch_kwargs['max_retries'], 6)

arr1 = xee.EarthEngineBackendArray('spi2y', self.conus_store)
self.assertNotIn('initial_delay', arr1.store.tile_fetch_kwargs)
self.assertEqual(arr1.store.tile_fetch_kwargs['max_retries'], 9)

class EEBackendEntrypointTest(absltest.TestCase):

Loading