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
23 changes: 21 additions & 2 deletions xee/ext.py
Original file line number Diff line number Diff line change
@@ -102,6 +102,11 @@ class EarthEngineStore(common.AbstractDataStore):
'height': 256,
}

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

SCALE_UNITS: Dict[str, int] = {
'degree': 1,
'metre': 10_000,
@@ -147,6 +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,
getitem_kwargs: Optional[Dict[str, int]] = None,
) -> 'EarthEngineStore':
if mode != 'r':
raise ValueError(
@@ -168,6 +174,7 @@ def open(
ee_init_kwargs=ee_init_kwargs,
ee_init_if_necessary=ee_init_if_necessary,
executor_kwargs=executor_kwargs,
getitem_kwargs=getitem_kwargs,
)

def __init__(
@@ -186,6 +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,
getitem_kwargs: Optional[Dict[str, int]] = None,
):
self.ee_init_kwargs = ee_init_kwargs
self.ee_init_if_necessary = ee_init_if_necessary
@@ -195,6 +203,8 @@ def __init__(
executor_kwargs = {}
self.executor_kwargs = executor_kwargs

self.getitem_kwargs = {**self.GETITEM_KWARGS, **(getitem_kwargs or {})}

self.image_collection = image_collection
if n_images != -1:
self.image_collection = image_collection.limit(n_images)
@@ -482,7 +492,10 @@ def image_to_array(
**kwargs,
}
raw = common.robust_getitem(
dabhicusp marked this conversation as resolved.
Show resolved Hide resolved
pixels_getter, params, catch=ee.ee_exception.EEException
pixels_getter,
params,
catch=ee.ee_exception.EEException,
**self.getitem_kwargs,
)

# Extract out the shape information from EE response.
@@ -965,6 +978,7 @@ def open_dataset(
ee_init_if_necessary: bool = False,
ee_init_kwargs: Optional[Dict[str, Any]] = None,
executor_kwargs: Optional[Dict[str, Any]] = None,
getitem_kwargs: Optional[Dict[str, int]] = None,
) -> xarray.Dataset: # type: ignore
"""Open an Earth Engine ImageCollection as an Xarray Dataset.

@@ -1037,7 +1051,11 @@ 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}.

getitem_kwargs (optional): Exponential backoff kwargs passed into the
xarray function to index the array (`robust_getitem`).
- 'max_retries', the maximum number of retry attempts. Defaults to 6.
- 'initial_delay', the initial delay in milliseconds before the first
retry. Defaults to 500.
Returns:
An xarray.Dataset that streams in remote data from Earth Engine.
"""
@@ -1067,6 +1085,7 @@ def open_dataset(
ee_init_kwargs=ee_init_kwargs,
ee_init_if_necessary=ee_init_if_necessary,
executor_kwargs=executor_kwargs,
getitem_kwargs=getitem_kwargs,
)

store_entrypoint = backends_store.StoreBackendEntrypoint()
15 changes: 15 additions & 0 deletions xee/ext_integration_test.py
Original file line number Diff line number Diff line change
@@ -69,6 +69,7 @@ def setUp(self):
'2017-01-01', '2017-01-03'
),
n_images=64,
getitem_kwargs={'max_retries': 10, 'initial_delay': 1500},
)
self.lnglat_store = xee.EarthEngineStore(
ee.ImageCollection.fromImages([ee.Image.pixelLonLat()]),
@@ -80,6 +81,7 @@ def setUp(self):
'2020-03-30', '2020-04-01'
),
n_images=64,
getitem_kwargs={'max_retries': 9},
)
self.all_img_store = xee.EarthEngineStore(
ee.ImageCollection('LANDSAT/LC08/C01/T1').filterDate(
@@ -255,6 +257,19 @@ def __getitem__(self, params):

self.assertEqual(getter.count, 3)

def test_getitem_kwargs(self):
arr = xee.EarthEngineBackendArray('B4', self.store)
self.assertEqual(arr.store.getitem_kwargs['initial_delay'], 1500)
self.assertEqual(arr.store.getitem_kwargs['max_retries'], 10)

arr1 = xee.EarthEngineBackendArray('longitude', self.lnglat_store)
self.assertEqual(arr1.store.getitem_kwargs['initial_delay'], 500)
self.assertEqual(arr1.store.getitem_kwargs['max_retries'], 6)

arr2 = xee.EarthEngineBackendArray('spi2y', self.conus_store)
self.assertEqual(arr2.store.getitem_kwargs['initial_delay'], 500)
self.assertEqual(arr2.store.getitem_kwargs['max_retries'], 9)


class EEBackendEntrypointTest(absltest.TestCase):

Loading