Allow timeout to acquiring the lock
MODIFY
pangeo_forge_recipes/recipes/xarray_zarr.py
pangeo_forge_recipes/recipes/xarray_zarr.py
@@ -84,6 +84,7 @@ class XarrayZarrRecipe(BaseRecipe):
`(ds: xr.Dataset, filename: str) -> ds: xr.Dataset`.
:param process_chunk: Function to call on each concatenated chunk, with signature
`(ds: xr.Dataset) -> ds: xr.Dataset`.
+ :param lock_timeout: The default timeout for acquiring a chunk lock.
"""
file_pattern: FilePattern
@@ -101,6 +102,7 @@ class XarrayZarrRecipe(BaseRecipe):
fsspec_open_kwargs: dict = field(default_factory=dict)
process_input: Optional[Callable[[xr.Dataset, str], xr.Dataset]] = None
process_chunk: Optional[Callable[[xr.Dataset], xr.Dataset]] = None
+ lock_timeout: Optional[int] = None
# internal attributes not meant to be seen or accessed by user
_concat_dim: Optional[str] = None
@@ -322,7 +324,7 @@ class XarrayZarrRecipe(BaseRecipe):
zarr_region = tuple(write_region.get(dim, slice(None)) for dim in var.dims)
lock_keys = [f"{vname}-{c}" for c in conflicts]
logger.debug(f"Acquiring locks {lock_keys}")
- with lock_for_conflicts(lock_keys):
+ with lock_for_conflicts(lock_keys, timeout=self.lock_timeout):
logger.info(
f"Storing variable {vname} chunk {chunk_key} "
f"to Zarr region {zarr_region}"

MODIFY
pangeo_forge_recipes/utils.py
pangeo_forge_recipes/utils.py
@@ -84,7 +84,13 @@ def chunk_bounds_and_conflicts(
@contextmanager
# TODO: use a recipe-specific base_name to handle multiple recipes potentially
# running at the same time
-def lock_for_conflicts(conflicts, base_name="pangeo-forge"):
+def lock_for_conflicts(conflicts, base_name="pangeo-forge", timeout=None):
+ """
+ Parameters
+ ----------
+ timeout : int, optional
+ The time to wait *for each lock*.
+ """
try:
global_client = get_client()
@@ -98,7 +104,7 @@ def lock_for_conflicts(conflicts, base_name="pangeo-forge"):
locks = [Lock(f"{base_name}-{c}", global_client) for c in conflicts]
for lock in locks:
logger.debug(f"Acquiring lock {lock.name}...")
- lock.acquire()
+ lock.acquire(timeout=timeout)
logger.debug(f"Acquired lock {lock.name}")
else:
logger.debug(f"Asked to lock {conflicts} but no Dask client found.")

MODIFY
tests/test_recipes.py
tests/test_recipes.py
from contextlib import nullcontext as does_not_raise
+from unittest.mock import patch
import pytest
import xarray as xr
@@ -6,6 +7,8 @@ import xarray as xr
# need to import this way (rather than use pytest.lazy_fixture) to make it work with dask
from pytest_lazyfixture import lazy_fixture
+from tests.conftest import execute_recipe
+
all_recipes = [
lazy_fixture("netCDFtoZarr_sequential_recipe"),
lazy_fixture("netCDFtoZarr_sequential_multi_variable_recipe"),
@@ -154,3 +157,15 @@ def test_chunks(
ds_actual.load()
print(ds_actual)
xr.testing.assert_identical(ds_actual, ds_expected)
+
+
+def test_lock_timeout(netCDFtoZarr_sequential_recipe, execute_recipe):
+ RecipeClass, file_pattern, kwargs, ds_expected, target = netCDFtoZarr_sequential_recipe
+ recipe = RecipeClass(file_pattern=file_pattern, lock_timeout=1, **kwargs)
+
+ with patch("pangeo_forge_recipes.recipes.xarray_zarr.lock_for_conflicts") as p:
+ execute_recipe(recipe)
+
+ #
+ if p.call_count:
+ assert p.call_args[1]["timeout"] == 1

Allow timeout to acquiring the lock