got tests passing for XarrayZarrRecipe
MODIFY
pangeo_forge_recipes/recipes/xarray_zarr.py
pangeo_forge_recipes/recipes/xarray_zarr.py
@@ -417,7 +417,7 @@ def prepare_target(*, config: XarrayZarrRecipe) -> None:
if config.cache_metadata:
# if nitems_per_input is not constant, we need to cache this info
assert config.metadata_cache is not None # for mypy
- recipe_meta = {"input_sequence_lens": config.input_sequence_lens}
+ recipe_meta = {"input_sequence_lens": input_sequence_lens}
config.metadata_cache[_GLOBAL_METADATA_KEY] = recipe_meta

MODIFY
tests/recipe_tests/test_XarrayZarrRecipe.py
tests/recipe_tests/test_XarrayZarrRecipe.py
@@ -2,7 +2,6 @@ from contextlib import nullcontext as does_not_raise
from dataclasses import replace
from unittest.mock import patch
-import dask.core
import pytest
import xarray as xr
import zarr
@@ -378,122 +377,3 @@ def test_lock_timeout(netCDFtoZarr_recipe_sequential_only, execute_recipe_no_das
execute_recipe_no_dask(recipe)
assert p.call_args[1]["timeout"] == 1
-
-
-class MyRecipe(BaseRecipe):
- def __init__(self) -> None:
- super().__init__()
- self.cache = {}
- self.target = None
- self.finalized = False
- self.cache_inputs = True
-
- @property
- def prepare_target(self):
- def _():
- self.target = {}
-
- return _
-
- @property
- def cache_input(self):
- def _(input_key):
- self.cache[input_key] = input_key
-
- return _
-
- @property
- def store_chunk(self):
- def _(chunk_key):
- self.target[chunk_key] = self.cache[chunk_key]
-
- return _
-
- @property
- def finalize_target(self):
- def _():
- self.finalized = True
-
- return _
-
- def iter_inputs(self):
- return iter(range(4))
-
- def iter_chunks(self):
- return iter(range(4))
-
-
-def test_base_recipe():
- recipe = MyRecipe()
- recipe.to_function()()
- assert recipe.finalized
- assert recipe.target == {i: i for i in range(4)}
-
- import dask
-
- dask.config.set(scheduler="single-threaded")
- recipe = MyRecipe()
- recipe.to_dask().compute()
- assert recipe.finalized
- assert recipe.target == {i: i for i in range(4)}
-
- recipe = MyRecipe()
- recipe.to_prefect().run()
- assert recipe.finalized
- assert recipe.target == {i: i for i in range(4)}
-
-
-def _make_filename_for_memory_usage_test(time):
- import pandas as pd
-
- input_url_pattern = (
- "https://arthurhouhttps.pps.eosdis.nasa.gov/gpmdata/{yyyy}/{mm}/{dd}/"
- "imerg/3B-HHR.MS.MRG.3IMERG.{yyyymmdd}-S{sh}{sm}00-E{eh}{em}59.{MMMM}.V06B.HDF5"
- ).format(
- yyyy=time.strftime("%Y"),
- mm=time.strftime("%m"),
- dd=time.strftime("%d"),
- yyyymmdd=time.strftime("%Y%m%d"),
- sh=time.strftime("%H"),
- sm=time.strftime("%M"),
- eh=time.strftime("%H"),
- em=(time + pd.Timedelta("29 min")).strftime("%M"),
- MMMM=f"{(time.hour*60 + time.minute):04}",
- )
- return input_url_pattern
-
-
-def _simple_func(*args, **kwargs):
- return None
-
-
-@pytest.mark.timeout(90)
-@pytest.mark.filterwarnings("ignore:Large object")
-def test_memory_usage():
- # https://github.com/pangeo-forge/pangeo-forge-recipes/issues/151
- # Requires >4 GiB of memory to run.
- pd = pytest.importorskip("pandas")
- distributed = pytest.importorskip("distributed")
-
- dates = pd.date_range("2020-05-31T00:00:00", "2021-05-31T23:59:59", freq="30min")
- time_concat_dim = ConcatDim("time", dates, nitems_per_file=1)
- pattern = FilePattern(_make_filename_for_memory_usage_test, time_concat_dim,)
-
- recipe = XarrayZarrRecipe(
- pattern, xarray_open_kwargs={"group": "Grid", "decode_coords": "all"}, inputs_per_chunk=1,
- )
-
- obj = recipe.to_dask()
- dsk = obj.dask
-
- for k, v in dsk.items():
- if dask.core.istask(v):
- _, *args = v
- dsk[k] = (_simple_func,) + tuple(args)
-
- with dask.config.set(
- **{"distributed.worker.memory.pause": 0.95, "distributed.worker.memory.terminate": 0.9}
- ):
- with distributed.Client(n_workers=1, threads_per_worker=1, memory_limit="4G") as client:
- print("submitting")
- client.get(dsk, [obj.key])

ADD
tests/recipe_tests/test_graph_memory.py
+from contextlib import nullcontext as does_not_raise
+from dataclasses import replace
+from unittest.mock import patch
+
+import dask.core
+from dask.highlevelgraph import HighLevelGraph
+from dask.delayed import Delayed
+import pytest
+import xarray as xr
+import zarr
+
+# need to import this way (rather than use pytest.lazy_fixture) to make it work with dask
+from pytest_lazyfixture import lazy_fixture
+
+from pangeo_forge_recipes.patterns import ConcatDim, FilePattern
+from pangeo_forge_recipes.recipes.base import BaseRecipe
+from pangeo_forge_recipes.recipes.xarray_zarr import XarrayZarrRecipe
+
+
+def _make_filename_for_memory_usage_test(time):
+ import pandas as pd
+
+ input_url_pattern = (
+ "https://arthurhouhttps.pps.eosdis.nasa.gov/gpmdata/{yyyy}/{mm}/{dd}/"
+ "imerg/3B-HHR.MS.MRG.3IMERG.{yyyymmdd}-S{sh}{sm}00-E{eh}{em}59.{MMMM}.V06B.HDF5"
+ ).format(
+ yyyy=time.strftime("%Y"),
+ mm=time.strftime("%m"),
+ dd=time.strftime("%d"),
+ yyyymmdd=time.strftime("%Y%m%d"),
+ sh=time.strftime("%H"),
+ sm=time.strftime("%M"),
+ eh=time.strftime("%H"),
+ em=(time + pd.Timedelta("29 min")).strftime("%M"),
+ MMMM=f"{(time.hour*60 + time.minute):04}",
+ )
+ return input_url_pattern
+
+
+def _simple_func(*args, **kwargs):
+ return None
+
+
+@pytest.mark.timeout(90)
+@pytest.mark.filterwarnings("ignore:Large object")
+def test_memory_usage():
+ # https://github.com/pangeo-forge/pangeo-forge-recipes/issues/151
+ # Requires >4 GiB of memory to run.
+ pd = pytest.importorskip("pandas")
+ distributed = pytest.importorskip("distributed")
+
+ dates = pd.date_range("2020-05-31T00:00:00", "2021-05-31T23:59:59", freq="30min")
+ time_concat_dim = ConcatDim("time", dates, nitems_per_file=1)
+ pattern = FilePattern(_make_filename_for_memory_usage_test, time_concat_dim,)
+
+ recipe = XarrayZarrRecipe(
+ pattern, xarray_open_kwargs={"group": "Grid", "decode_coords": "all"}, inputs_per_chunk=1,
+ )
+
+ delayed = recipe.to_dask()
+
+ # in what follows, we rebuild the dask HighLevelGraph, replacing all of the
+ # actual functions with dummy functions that do nothin
+ dsk = delayed.dask
+ key = delayed.key
+ layers = dsk.layers
+ dependencies = dsk.dependencies
+ newlayers = {}
+ for layer_name, subgraph in layers.items():
+ new_subgraph = dict(subgraph)
+ for k, v in new_subgraph.items():
+ if dask.core.istask(v):
+ _, *args = v
+ new_subgraph[k] = (_simple_func,) + tuple(args)
+ newlayers[layer_name] = new_subgraph
+ new_hlg = HighLevelGraph(newlayers, dependencies)
+
+ with dask.config.set(
+ **{"distributed.worker.memory.pause": 0.95, "distributed.worker.memory.terminate": 0.9}
+ ):
+ with distributed.Client(n_workers=1, threads_per_worker=1, memory_limit="4G") as client:
+ print("submitting")
+ client.get(new_hlg, [key])

got tests passing for XarrayZarrRecipe