kind of working
MODIFY
pangeo_forge/recipe.py
pangeo_forge/recipe.py
@@ -16,7 +16,12 @@ from rechunker.types import MultiStagePipeline, ParallelPipelines, Stage
from .patterns import ExplicitURLSequence, VariableSequencePattern
from .storage import AbstractTarget, UninitializedTarget
-from .utils import chunked_iterable, fix_scalar_attr_encoding
+from .utils import (
+ calc_chunk_conflicts,
+ chunked_iterable,
+ fix_scalar_attr_encoding,
+ lock_for_conflicts,
+)
logger = logging.getLogger(__name__)
@@ -181,24 +186,19 @@ class NetCDFtoZarrRecipe(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
- target_chunks: Optional[Dict[str, int]] = None
+ target_chunks: Dict[str, int] = field(default_factory=dict)
def __post_init__(self):
self._cache_metadata = self.nitems_per_input is None
- if self.target_chunks:
- sequence_dim_chunks = self.target_chunks.get(self.sequence_dim)
- if sequence_dim_chunks:
- if sequence_dim_chunks > self.nitems_per_input:
- raise ValueError(
- f"`sequence_dim_chunks` ({sequence_dim_chunks}) must be ",
- f"<= `nitems_per_input` ({self.nitems_per_input})",
- )
- if self.nitems_per_input % sequence_dim_chunks > 0:
- raise ValueError(
- "`sequence_dim_chunks` ({sequence_dim_chunks}) must "
- f"evenly divide `nitems_per_input` ({self.nitems_per_input})"
- )
+ target_sequence_dim_chunks = self.target_chunks.get(self.sequence_dim)
+ if (self.nitems_per_input is None) and (target_sequence_dim_chunks is None):
+ raise ValueError("Unable to determine target chunks.")
+ elif target_sequence_dim_chunks:
+ self.sequence_dim_chunks = target_sequence_dim_chunks
+ else:
+ self.sequence_dim_chunks = self.nitems_per_input * self.inputs_per_chunk
+
# TODO: more input validation
@property
@@ -285,7 +285,8 @@ class NetCDFtoZarrRecipe(BaseRecipe):
ds_chunk = drop_vars(ds_chunk)
target_mapper = self.target.get_mapper()
- write_region = self.region_for_chunk(chunk_key)
+ write_region, conflicts = self.region_and_conflicts_for_chunk(chunk_key)
+ with lock_for_conflicts(conflicts):
logger.info(f"Storing chunk '{chunk_key}' to Zarr region {write_region}")
ds_chunk.to_zarr(target_mapper, region=write_region)
@@ -397,13 +398,14 @@ class NetCDFtoZarrRecipe(BaseRecipe):
for input in self.inputs_for_chunk(chunk_key):
yield input
- def region_for_chunk(self, chunk_key):
+ def region_and_conflicts_for_chunk(self, chunk_key):
# return a dict suitable to pass to xr.to_zarr(region=...)
# specifies where in the overall array to put this chunk's data
if self.nitems_per_input:
stride = self.nitems_per_input * self.inputs_per_chunk
start = self.chunk_position(chunk_key) * stride
stop = start + stride
+ this_chunk_conflicts = []
else:
input_keys = self.inputs_for_chunk(chunk_key)
input_sequence_lens = json.loads(
@@ -412,9 +414,13 @@ class NetCDFtoZarrRecipe(BaseRecipe):
start = sum(input_sequence_lens[: self.input_position(input_keys[0])])
chunk_len = sum([input_sequence_lens[self.input_position(k)] for k in input_keys])
stop = start + chunk_len
+ all_chunk_conflicts = calc_chunk_conflicts(
+ input_sequence_lens, self.sequence_dim_chunks
+ )
+ this_chunk_conflicts = [all_chunk_conflicts[self.input_position(k)] for k in input_keys]
region_slice = slice(start, stop)
- return {self.sequence_dim: region_slice}
+ return {self.sequence_dim: region_slice}, this_chunk_conflicts
def iter_chunks(self):
for k in self._chunks_inputs:

MODIFY
pangeo_forge/utils.py
pangeo_forge/utils.py
import itertools
+from contextlib import contextmanager
from typing import Iterable, List, Tuple
import numpy as np
+from dask.distributed import Lock
# https://alexwlchan.net/2018/12/iterating-in-fixed-size-chunks/
@@ -57,3 +59,15 @@ def calc_chunk_conflicts(chunks: Iterable[int], zchunks: int) -> List[Tuple[int]
chunk_conflicts.append(tuple(conflicts))
return chunk_conflicts
+
+
+@contextmanager
+def lock_for_conflicts(conflicts, base_name="pangeo-forge"):
+ locks = [Lock(f"{base_name}-{c}") for c in conflicts]
+ for lock in locks:
+ lock.acquire()
+ try:
+ yield
+ finally:
+ for lock in locks:
+ lock.release()

MODIFY
setup.cfg
setup.cfg
@@ -3,7 +3,7 @@ max-line-length = 100
[isort]
known_first_party=pangeo_forge
-known_third_party=aiohttp,fsspec,numpy,pandas,pkg_resources,prefect,pytest,rechunker,setuptools,sphinx_book_theme,xarray,zarr
+known_third_party=aiohttp,dask,fsspec,numpy,pandas,pkg_resources,prefect,pytest,rechunker,setuptools,sphinx_book_theme,xarray,zarr
multi_line_output=3
include_trailing_comma=True
force_grid_wrap=0

MODIFY
tests/test_utils.py
tests/test_utils.py
+import dask
import pytest
+from dask import delayed
+from dask.distributed import Client, LocalCluster
from pangeo_forge import utils
@@ -23,3 +26,19 @@ def test_chunk_conflicts():
assert utils.calc_chunk_conflicts([9, 10], zchunks) == [(0,), (0,)]
assert utils.calc_chunk_conflicts([10, 9, 11, 10], zchunks) == [(), (1,), (1,), ()]
assert utils.calc_chunk_conflicts([9, 12, 5], zchunks) == [(0,), (0, 2), (2,)]
+
+
+@pytest.mark.parametrize("conflicts", [{}, {0}, {0, 1}])
+def test_locks(conflicts):
+ # TOOD: move this into a fixture
+ with LocalCluster(n_workers=1, processes=False, threads_per_worker=1,) as cluster, Client(
+ cluster
+ ):
+
+ @delayed
+ def do_stuff():
+ with utils.lock_for_conflicts(conflicts):
+ # todo; how to actually test for concurrency! hard!
+ pass
+
+ dask.compute([do_stuff() for n in range(3)])

kind of working