wrote code but test doesn't work
MODIFY
pangeo_forge/recipe.py
pangeo_forge/recipe.py
@@ -4,6 +4,8 @@ A Pangeo Forge Recipe
import json
import logging
+import os
+import tempfile
from abc import ABC, abstractmethod
from contextlib import ExitStack, contextmanager
from dataclasses import dataclass, field
@@ -107,6 +109,7 @@ class BaseRecipe(ABC):
"""
pipeline = [] # type: MultiStagePipeline
+ if getattr(self, "cache_inputs", False): # TODO: formalize this contract
pipeline.append(Stage(self.cache_input, list(self.iter_inputs())))
pipeline.append(Stage(self.prepare_target))
pipeline.append(Stage(self.store_chunk, list(self.iter_chunks())))
@@ -144,6 +147,22 @@ def _copy_btw_filesystems(input_opener, output_opener, BLOCK_SIZE=10_000_000):
target.write(data)
+@contextmanager
+def _maybe_open_or_copy_to_local(opener, copy_to_local, orig_name):
+ _, suffix = os.path.splitext(orig_name)
+ if copy_to_local:
+ ntf = tempfile.NamedTemporaryFile(suffix=suffix)
+ tmp_name = ntf.name
+ logger.info(f"Copying {orig_name} to local file {tmp_name}")
+ target_opener = open(tmp_name, mode="wb")
+ _copy_btw_filesystems(opener, target_opener)
+ yield tmp_name
+ ntf.close() # cleans up the temporary file
+ else:
+ with opener as fp:
+ yield fp
+
+
# Notes about dataclasses:
# - https://www.python.org/dev/peps/pep-0557/#inheritance
# - https://stackoverflow.com/questions/51575931/class-inheritance-in-python-3-7-dataclasses
@@ -171,7 +190,7 @@ class NetCDFtoZarrRecipe(BaseRecipe):
:param input_cache: A location in which to cache temporary data.
:param metadata_cache: A location in which to cache metadata for inputs and chunks.
Required if ``nitems_per_input=None``.
- :param require_cache: Whether to allow opening inputs directly which have not
+ :param cache_inputs: Whether to allow opening inputs directly which have not
yet been cached. This could lead to very unstanble behavior if the inputs
live behind a slow network connection.
:param copy_input_to_local_file: Whether to copy the inputs to a temporary
@@ -198,8 +217,8 @@ class NetCDFtoZarrRecipe(BaseRecipe):
target: AbstractTarget = field(default_factory=UninitializedTarget)
input_cache: AbstractTarget = field(default_factory=UninitializedTarget)
metadata_cache: AbstractTarget = field(default_factory=UninitializedTarget)
- require_cache: bool = True
- copy_input_to_local_file = False
+ cache_inputs: bool = True
+ copy_input_to_local_file: bool = False
consolidate_zarr: bool = True
xarray_open_kwargs: dict = field(default_factory=dict)
xarray_concat_kwargs: dict = field(default_factory=dict)
@@ -367,21 +386,25 @@ class NetCDFtoZarrRecipe(BaseRecipe):
@contextmanager
def input_opener(self, fname: str):
try:
- with self.input_cache.open(fname, mode="rb") as f:
logger.info(f"Opening '{fname}' from cache")
- yield f
+ opener = self.input_cache.open(fname, mode="rb")
+ with _maybe_open_or_copy_to_local(opener, self.copy_input_to_local_file, fname) as fp:
+ yield fp
except (IOError, FileNotFoundError):
- if self.require_cache:
+ if self.cache_inputs:
raise FileNotFoundError(
f"You are trying to open input {fname}, but the file is "
"not cached yet. First call `cache_input` or set "
- "`require_cache=False`."
+ "`cache_inputs=False`."
)
else:
logger.info(f"No cache found. Opening input `{fname}` directly.")
# This will bypass the cache. May be slow.
- with fsspec.open(fname, mode="rb", **self.fsspec_open_kwargs) as f:
- yield f
+ opener = fsspec.open(fname, mode="rb", **self.fsspec_open_kwargs)
+ with _maybe_open_or_copy_to_local(
+ opener, self.copy_input_to_local_file, fname
+ ) as fp:
+ yield fp
@contextmanager
def open_input(self, input_key: Hashable):

MODIFY
tests/conftest.py
tests/conftest.py
@@ -257,6 +257,7 @@ def execute_recipe(request, dask_cluster):
if request.param == "manual":
def execute(r):
+ if r.cache_inputs:
for input_key in r.iter_inputs():
r.cache_input(input_key)
r.prepare_target()

MODIFY
tests/test_recipes.py
tests/test_recipes.py
@@ -23,6 +23,22 @@ def test_recipe(recipe_fixture, execute_recipe):
xr.testing.assert_identical(ds_actual, ds_expected)
+@pytest.mark.parametrize("cache_inputs", [True, False])
+@pytest.mark.parametrize("copy_input_to_local_file", [True, False])
+def test_recipe_caching_copying(
+ netCDFtoZarr_sequential_recipe, execute_recipe, cache_inputs, copy_input_to_local_file
+):
+ """The basic recipe test. Use this as a template for other tests."""
+
+ RecipeClass, kwargs, ds_expected, target = netCDFtoZarr_sequential_recipe
+ rec = RecipeClass(
+ **kwargs, cache_inputs=cache_inputs, copy_input_to_local_file=copy_input_to_local_file
+ )
+ execute_recipe(rec)
+ ds_actual = xr.open_zarr(target.get_mapper()).load()
+ xr.testing.assert_identical(ds_actual, ds_expected)
+
+
# function passed to preprocessing
def incr_date(ds, filename=""):
# add one day

wrote code but test doesn't work