wip; getting closer
MODIFY
pangeo_forge/recipe.py
pangeo_forge/recipe.py
@@ -190,7 +190,6 @@ class NetCDFtoZarrRecipe(BaseRecipe):
def __post_init__(self):
self._cache_metadata = self.nitems_per_input is None
-
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.")
@@ -401,22 +400,23 @@ class NetCDFtoZarrRecipe(BaseRecipe):
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
+ # also return the conflicts with other chunks
+
+ input_keys = self.inputs_for_chunk(chunk_key)
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 = []
+ input_sequence_lens = (self.nitems_per_input,) * self._n_inputs_along_sequence
else:
- input_keys = self.inputs_for_chunk(chunk_key)
input_sequence_lens = json.loads(
self.metadata_cache.get_mapper()[_GLOBAL_METADATA_KEY]
)["input_sequence_lens"]
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
- )
+
+ 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)
@@ -456,6 +456,8 @@ class NetCDFtoZarrSequentialRecipe(NetCDFtoZarrRecipe):
self._chunks_inputs = {
(k,): v for k, v in enumerate(chunked_iterable(self._inputs, self.inputs_per_chunk))
}
+ self._n_inputs_along_sequence = len(self._inputs)
+
# just the first chunk is needed to initialize the recipe
self._init_chunks = [next(iter(self._chunks_inputs))]
@@ -482,6 +484,7 @@ class NetCDFtoZarrMultiVarSequentialRecipe(NetCDFtoZarrRecipe):
def __post_init__(self):
super().__post_init__()
self._variables = self.input_pattern.keys["variable"]
+ self._other_key = [k for k in self.input_pattern.keys if k != "foo"][0]
self._inputs = {k: v for k, v in self.input_pattern}
# input keys are tuples like
# ("temp", 0)
@@ -502,6 +505,7 @@ class NetCDFtoZarrMultiVarSequentialRecipe(NetCDFtoZarrRecipe):
(vname, input_sequence_key) for input_sequence_key in chunk_inputs
]
self._chunks_inputs = chunks_inputs
+ self._n_inputs_along_sequence = len(self.input_pattern.keys[self._other_key])
# should be the first chunk from each variable
self._init_chunks = [chunk_key for chunk_key in self._chunks_inputs if chunk_key[1] == 0]

MODIFY
pangeo_forge/utils.py
pangeo_forge/utils.py
@@ -3,7 +3,7 @@ from contextlib import contextmanager
from typing import Iterable, List, Tuple
import numpy as np
-from dask.distributed import Lock
+from dask.distributed import Lock, client
# https://alexwlchan.net/2018/12/iterating-in-fixed-size-chunks/
@@ -63,11 +63,19 @@ def calc_chunk_conflicts(chunks: Iterable[int], zchunks: int) -> List[Tuple[int]
@contextmanager
def lock_for_conflicts(conflicts, base_name="pangeo-forge"):
+
+ # https://stackoverflow.com/questions/59070260/dask-client-detect-local-default-cluster-already-running
+ is_distributed = client._get_global_client() is not None
+ # Don't bother with locks if we are not in a distributed context
+ # NOTE! This means we HAVE to use dask.distributed as our parallel execution enviroment
+ # That is compatible with Prefect.
+ if is_distributed:
locks = [Lock(f"{base_name}-{c}") for c in conflicts]
for lock in locks:
lock.acquire()
try:
yield
finally:
+ if is_distributed:
for lock in locks:
lock.release()

MODIFY
tests/test_recipe.py
tests/test_recipe.py
@@ -63,15 +63,26 @@ def test_NetCDFtoZarrSequentialRecipeHttpAuth(
)
@pytest.mark.parametrize("inputs_per_chunk", [1, 2])
@pytest.mark.parametrize(
- "target_chunks,chunk_expectation",
+ "target_chunks,specify_nitems_per_input,chunk_expectation",
[
- (None, does_not_raise()),
- ({"lon": 12}, does_not_raise()),
- ({"lon": 12, "time": 1}, does_not_raise()),
- ({"lon": 12, "time": 3}, pytest.raises(ValueError)),
+ ({}, True, does_not_raise()),
+ ({"lon": 12}, True, does_not_raise()),
+ ({"lon": 12, "time": 1}, True, does_not_raise()),
+ (
+ {"lon": 12, "time": 3},
+ True,
+ does_not_raise(),
+ ), # currently failing because we don't handle last chunk properly
+ ({}, False, pytest.raises(ValueError)), # can't determine target chunks
+ ({"lon": 12}, False, pytest.raises(ValueError)), # can't determine target_chunks
+ ({"lon": 12, "time": 1}, False, does_not_raise()),
+ (
+ {"lon": 12, "time": 3},
+ False,
+ does_not_raise(),
+ ), # currently failing because we don't handle last chunk properly
],
)
-@pytest.mark.parametrize("specify_nitems_per_input", [True, False])
def test_NetCDFtoZarrSequentialRecipe_options(
daily_xarray_dataset,
netcdf_local_paths,
@@ -118,13 +129,13 @@ def test_NetCDFtoZarrSequentialRecipe_options(
# chunk validation
sequence_chunks = ds_target.chunks["time"]
- if target_chunks is None:
- target_chunks = {}
- seq_chunk_len = target_chunks.pop("time", None) or (items_per_file * inputs_per_chunk)
+ seq_chunk_len = target_chunks.get("time", None) or (items_per_file * inputs_per_chunk)
# we expect all chunks but the last to have the expected size
assert all([item == seq_chunk_len for item in sequence_chunks[:-1]])
for other_dim, chunk_len in target_chunks.items():
- all([item == chunk_len for item in ds_target.chunks[other_dim][:-1]])
+ if other_dim == "time":
+ continue
+ assert all([item == chunk_len for item in ds_target.chunks[other_dim][:-1]])
ds_target.load()
ds_expected = daily_xarray_dataset.compute()
@@ -187,8 +198,6 @@ def test_NetCDFtoZarrMultiVarSequentialRecipe(
_manually_execute_recipe(r)
ds_target = xr.open_zarr(tmp_target.get_mapper(), consolidated=True).compute()
- print(ds_target)
- print(daily_xarray_dataset)
assert ds_target.identical(daily_xarray_dataset)

MODIFY
tests/test_utils.py
tests/test_utils.py
@@ -30,6 +30,12 @@ def test_chunk_conflicts():
@pytest.mark.parametrize("conflicts", [{}, {0}, {0, 1}])
def test_locks(conflicts):
+
+ # first make sure the locks work without a cluster
+ with utils.lock_for_conflicts(conflicts):
+ # todo; how to actually test for concurrency! hard!
+ pass
+
# TOOD: move this into a fixture
with LocalCluster(n_workers=1, processes=False, threads_per_worker=1,) as cluster, Client(
cluster

wip; getting closer