From 724b802093271b33fa0f4a3365e27221110aaa87 Mon Sep 17 00:00:00 2001 From: Justin Braaten Date: Tue, 30 Jun 2026 14:35:28 -0700 Subject: [PATCH] Optimize micro_benchmarks execution by limiting dataset metadata scan PiperOrigin-RevId: 940664266 --- xee/micro_benchmarks.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/xee/micro_benchmarks.py b/xee/micro_benchmarks.py index 3f1479f..c9c78e5 100644 --- a/xee/micro_benchmarks.py +++ b/xee/micro_benchmarks.py @@ -31,7 +31,7 @@ import ee -REPEAT = 10 +REPEAT = 5 LOOPS = 1 PROFILE = False @@ -47,12 +47,17 @@ def init_ee_for_tests(): 'shape_2d': (1440, 720), } +_OPEN_KWARGS = { + 'engine': xee.EarthEngineBackendEntrypoint, + 'n_images': 64, # Limit metadata scan to 64 images to prevent RPC timeouts. + **_GRID_PARAMS, +} + def open_dataset() -> None: _ = xarray.open_dataset( 'NASA/GPM_L3/IMERG_V06', - engine=xee.EarthEngineBackendEntrypoint, - **_GRID_PARAMS, + **_OPEN_KWARGS, ) @@ -60,8 +65,7 @@ def open_and_chunk() -> None: ds = xarray.open_dataset( 'NASA/GPM_L3/IMERG_V06', chunks={'index': 24, 'width': 512, 'height': 512}, - engine=xee.EarthEngineBackendEntrypoint, - **_GRID_PARAMS, + **_OPEN_KWARGS, ) ds.chunk() @@ -71,8 +75,7 @@ def open_and_write() -> None: ds = xarray.open_dataset( 'NASA/GPM_L3/IMERG_V06', chunks={'time': 24, 'lon': 1440, 'lat': 720}, - engine=xee.EarthEngineBackendEntrypoint, - **_GRID_PARAMS, + **_OPEN_KWARGS, ) ds = ds.isel(time=slice(0, 24)) ds.to_zarr(os.path.join(tmpdir, 'imerg.zarr'))