got the first test passing
DELETE
pangeo_forge_recipes/executors.py
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from typing import Dict, Callable, Any, Iterable, List, Optional
-from enum import Enum, auto
-
-
-class StageAnnotation(Enum):
- CONCURRENCY = auto()
- RETRIES = auto()
-
-
-Config = Any
-
-
-@dataclass(frozen=True)
-class Stage:
- function: Callable
- name: str
- mappable: Optional[Iterable] = None
- config: Optional[Config] = None
- annotations: Optional[Dict[StageAnnotation, Any]] = None
-
-
-Pipeline = List[Stage]
-
-
-def pipeline_to_function():
- pass

MODIFY
pangeo_forge_recipes/executors/base.py
pangeo_forge_recipes/executors/base.py
+import enum
from abc import ABC, abstractmethod
from dataclasses import dataclass
-from typing import Dict, Generic, Any, Iterable, Optional, Protocol, Hashable, Union, TypeVar
-import enum
-from typing import Protocol
+from typing import Any, Dict, Generic, Hashable, Iterable, Optional, Protocol, TypeVar, Union
Config = Any
# https://stackoverflow.com/questions/57837609/python-typing-signature-typing-callable-for-function-with-kwargs
class NoArgumentStageFunction(Protocol):
- def __call__(*, config: Optional[Config]=None) -> None: ...
+ def __call__(*, config: Optional[Config] = None) -> None:
+ ...
class SingleArgumentStageFunction(Protocol):
- def __call__(__a: Hashable, *, config: Optional[Config]=None) -> None: ...
+ def __call__(__a: Hashable, *, config: Optional[Config] = None) -> None:
+ ...
StageFunction = Union[NoArgumentStageFunction, SingleArgumentStageFunction]
@@ -22,8 +23,10 @@ class StageAnnotationType(enum.Enum):
CONCURRENCY = enum.auto()
RETRIES = enum.auto()
+
StageAnnotations = Dict[StageAnnotationType, Any]
+
@dataclass(frozen=True)
class Stage:
function: StageFunction
@@ -42,7 +45,6 @@ T = TypeVar("T")
class PipelineExecutor(Generic[T]):
-
@staticmethod
def compile(pipeline: Pipeline) -> T:
raise NotImplementedError

MODIFY
pangeo_forge_recipes/executors/dask.py
pangeo_forge_recipes/executors/dask.py
-from pprint import pprint
-from typing import Any, Callable, Dict, Optional
-
import dask
-import dask.bag as db
from dask.delayed import Delayed
-from dask.graph_manipulation import bind
from .base import Pipeline, PipelineExecutor
-# task = (apply, func, args, kwargs) # func(*args, **kwargs)
-
-
def wrap_map_task(function):
# dependencies are dummy args used to create dependence between stages

MODIFY
pangeo_forge_recipes/executors/function.py
pangeo_forge_recipes/executors/function.py
@@ -4,10 +4,8 @@ from .base import Pipeline, PipelineExecutor
class FunctionPipelineExecutor(PipelineExecutor[Callable]):
-
@staticmethod
def compile(pipeline: Pipeline):
-
def function():
for stage in pipeline.stages:
if stage.mappable is not None:

MODIFY
pangeo_forge_recipes/executors/prefect.py
pangeo_forge_recipes/executors/prefect.py
-from typing import Callable, Dict, Any, Optional
+from typing import Any, Callable, Dict, Optional
from prefect import Flow, task, unmapped
@@ -9,13 +9,13 @@ def annotations_to_task_kwargs(annotations: Optional[StageAnnotations]) -> Dict[
task_kwargs = {}
for ann_type in annotations or []:
if ann_type == StageAnnotationType.retries:
- task_kwargs['max_retries'] = annotations[ann_type]
+ task_kwargs["max_retries"] = annotations[ann_type]
if ann_type == StageAnnotationType.concurrency:
raise ValueError("Haven't figured out concurrency yet.")
return task_kwargs
-class PrefectPipelineExecutor(PipelineExecutor[Flow]):
+class PrefectPipelineExecutor(PipelineExecutor[Flow]):
@staticmethod
def compile(pipeline: Pipeline):
@@ -28,12 +28,11 @@ class PrefectPipelineExecutor(PipelineExecutor[Flow]):
stage_task_called = stage_task.map(
stage.mappable,
config=unmapped(pipeline.config),
- upstream_tasks=[unmapped(t) for t in upstream_tasks]
+ upstream_tasks=[unmapped(t) for t in upstream_tasks],
)
else:
stage_task_called = stage_task(
- config=pipeline.config,
- upstream_tasks=upstream_tasks
+ config=pipeline.config, upstream_tasks=upstream_tasks
)
upstream_tasks = [stage_task_called]

MODIFY
pangeo_forge_recipes/recipes/__init__.py
pangeo_forge_recipes/recipes/__init__.py
-from .reference_hdf_zarr import HDFReferenceRecipe
+# from .reference_hdf_zarr import HDFReferenceRecipe
from .xarray_zarr import XarrayZarrRecipe
-__all__ = ["XarrayZarrRecipe", "HDFReferenceRecipe"]
+__all__ = ["XarrayZarrRecipe"] # , "HDFReferenceRecipe"]

MODIFY
pangeo_forge_recipes/recipes/base.py
pangeo_forge_recipes/recipes/base.py
+from __future__ import annotations
+
from abc import ABC, abstractmethod
from dataclasses import dataclass, replace
-from typing import Callable, Hashable, Iterable
+from typing import Callable, ClassVar, Hashable, Iterable
+from ..executors.base import Pipeline
from ..patterns import FilePattern, prune_pattern
# How to manually execute a recipe: ###
@@ -29,173 +32,31 @@ from ..patterns import FilePattern, prune_pattern
# 6)
-class RecipeConfig(ABC):
- pass
-
-
+@dataclass
class BaseRecipe(ABC):
- """Base recipe class from which all other Recipes inherit.
- """
+ _compiler: ClassVar[RecipeCompiler]
- @property
- @abstractmethod
- def prepare_target(self) -> Callable[[RecipeConfig], None]:
- """Prepare the recipe for execution by initializing the target.
- Attribute that returns a callable function.
- """
- pass
-
- @abstractmethod
- def iter_inputs(self) -> Iterable[Hashable]:
- """Iterate over all inputs."""
- pass
-
- @property
- @abstractmethod
- def cache_input(self) -> Callable[[RecipeConfig, Hashable], None]:
- """Copy an input from its source location to the cache.
- Attribute that returns a callable function.
- """
- pass
-
- @abstractmethod
- def iter_chunks(self) -> Iterable[Hashable]:
- """Iterate over all target chunks."""
- pass
-
- @property
- @abstractmethod
- def store_chunk(self) -> Callable[[RecipeConfig, Hashable], None]:
- """Store a chunk of data in the target.
- Attribute that returns a callable function.
- """
- pass
+ def to_function(self):
+ from ..executors.function import FunctionPipelineExecutor
- @property
- @abstractmethod
- def finalize_target(self) -> Callable[[RecipeConfig], None]:
- """Final step to finish the recipe after data has been written.
- Attribute that returns a callable function.
- """
- pass
+ return FunctionPipelineExecutor.compile(self._compiler())
- def to_function(self) -> Callable[[], None]:
- """
- Translate the recipe to a Python function for execution.
- """
+ def to_dask(self):
+ from pangeo_forge_recipes.executors.dask import DaskPipelineExecutor
- def pipeline():
- # TODO: allow recipes to customize which stages to run
- for input_key in self.iter_inputs():
- self.cache_input(self.config, input_key)
- self.prepare_target(self.config)
- for chunk_key in self.iter_chunks():
- self.store_chunk(self.config, chunk_key)
- self.finalize_target(self.config)
+ return DaskPipelineExecutor.compile(self._compiler())
- return pipeline
+ def to_prefect(self):
+ from pangeo_forge_recipes.executors.prefect import PrefectPipelineExecutor
- def to_dask(self):
- """
- Translate the recipe to a dask.Delayed object for parallel execution.
- """
- # This manually builds a Dask task graph with each stage of the recipe.
- # We use a few "checkpoints" to ensure that downstream tasks depend
- # on upstream tasks being done before starting. We use a manual task
- # graph rather than dask.delayed to avoid some expensive tokenization
- # in dask.delayed
- import dask
- from dask.delayed import Delayed
-
- # TODO: HighlevelGraph layers for each of these mapped inputs.
- # Cache Input --------------------------------------------------------
- dsk = {}
- token = dask.base.tokenize(self)
-
- config_ = f"config-{token}"
- dsk[config_] = self.config
-
- # TODO: allow recipes to customize which stages to run
- for i, input_key in enumerate(self.iter_inputs()):
- dsk[(f"cache_input-{token}", i)] = (self.cache_input, config_, input_key)
-
- # Prepare Target ------------------------------------------------------
- dsk[f"checkpoint_0-{token}"] = (lambda *args: None, list(dsk))
- dsk[f"prepare_target-{token}"] = (
- _checkpoint,
- f"checkpoint_0-{token}",
- self.prepare_target,
- config_,
- )
-
- # Store Chunk --------------------------------------------------------
- keys = []
- for i, chunk_key in enumerate(self.iter_chunks()):
- k = (f"store_chunk-{token}", i)
- dsk[k] = (
- _checkpoint,
- f"prepare_target-{token}",
- self.store_chunk,
- config_,
- chunk_key,
- )
- keys.append(k)
-
- # Finalize Target -----------------------------------------------------
- dsk[f"checkpoint_1-{token}"] = (lambda *args: None, keys)
- key = f"finalize_target-{token}"
- dsk[key] = (_checkpoint, f"checkpoint_1-{token}", self.finalize_target, config_)
-
- return Delayed(key, dsk)
+ return PrefectPipelineExecutor.compile(self._compiler())
- def to_prefect(self):
- """Compile the recipe to a Prefect.Flow object."""
- from prefect import Flow, task, unmapped
-
- # TODO: allow recipes to customize which stages to run
- cache_input_task = task(self.cache_input, name="cache_input")
- prepare_target_task = task(self.prepare_target, name="prepare_target")
- store_chunk_task = task(self.store_chunk, name="store_chunk")
- finalize_target_task = task(self.finalize_target, name="finalize_target")
-
- with Flow("pangeo-forge-recipe") as flow:
- # we hope this will be interpreted as a Prefect Constant
- # https://docs.prefect.io/core/concepts/tasks.html#constants
- config = self.config
- cache_task = cache_input_task.map(
- config=unmapped(config), input_key=list(self.iter_inputs())
- )
- upstream_tasks = [cache_task]
- prepare_task = prepare_target_task(config=config, upstream_tasks=upstream_tasks)
- store_task = store_chunk_task.map(
- config=unmapped(config),
- chunk_key=list(self.iter_chunks()),
- upstream_tasks=[unmapped(prepare_task)],
- )
- _ = finalize_target_task(config=config, upstream_tasks=[store_task])
-
- return flow
-
- def __iter__(self):
- # TODO: allow recipes to customize which stages to run
- yield self.cache_input, self.iter_inputs()
- yield self.prepare_target, []
- yield self.store_chunk, self.iter_chunks()
- yield self.finalize_target, []
-
- # https://stackoverflow.com/questions/59986413/achieving-multiple-inheritance-using-python-dataclasses
- def __post_init__(self):
- # just intercept the __post_init__ calls so they
- # aren't relayed to `object`
- pass
-
-
-def _checkpoint(checkpoint, func, *args, **kwargs):
- return func(*args, **kwargs)
+
+RecipeCompiler = Callable[[BaseRecipe], Pipeline]
@dataclass
-class FilePatternRecipeMixin:
+class FilePatternMixin:
file_pattern: FilePattern
def copy_pruned(self, nkeep: int = 2):

MODIFY
pangeo_forge_recipes/recipes/xarray_zarr.py
pangeo_forge_recipes/recipes/xarray_zarr.py
@@ -19,10 +19,11 @@ import xarray as xr
import zarr
from ..chunk_grid import ChunkGrid
+from ..executors.base import Pipeline, Stage
from ..patterns import CombineOp, DimIndex, FilePattern, Index
from ..storage import AbstractTarget, CacheFSSpecTarget, MetadataTarget, file_opener
from ..utils import calc_subsets, fix_scalar_attr_encoding, lock_for_conflicts
-from .base import BaseRecipe, RecipeConfig
+from .base import BaseRecipe, FilePatternMixin
# use this filename to store global recipe metadata in the metadata_cache
# it will be written once (by prepare_target) and read many times (by store_chunk)
@@ -137,7 +138,7 @@ def chunk_position(chunk_key: ChunkKey) -> int:
return subset_factor * concat_idx + subset_idx
-def cache_input(config: XarrayZarrRecipeConfig, input_key: InputKey) -> None:
+def cache_input(input_key: InputKey, *, config: XarrayZarrRecipe) -> None:
if config.cache_inputs:
if config.file_pattern.is_opendap:
raise ValueError("Can't cache opendap inputs")
@@ -155,13 +156,13 @@ def cache_input(config: XarrayZarrRecipeConfig, input_key: InputKey) -> None:
if config.metadata_cache is None:
raise ValueError("metadata_cache is not set.")
logger.info(f"Caching metadata for input '{input_key!s}'")
- with open_input(config, input_key) as ds:
+ with open_input(input_key, config=config) as ds:
input_metadata = ds.to_dict(data=False)
config.metadata_cache[_input_metadata_fname(input_key)] = input_metadata
def region_and_conflicts_for_chunk(
- config: XarrayZarrRecipeConfig, chunk_key: ChunkKey
+ config: XarrayZarrRecipe, chunk_key: ChunkKey
) -> Tuple[Dict[str, slice], Dict[str, Set[int]]]:
# return a dict suitable to pass to xr.to_zarr(region=...)
# specifies where in the overall array to put this chunk's data
@@ -205,7 +206,7 @@ def region_and_conflicts_for_chunk(
@contextmanager
-def open_input(config: XarrayZarrRecipeConfig, input_key: InputKey) -> xr.Dataset:
+def open_input(input_key: InputKey, *, config: XarrayZarrRecipe) -> xr.Dataset:
fname = config.file_pattern[input_key]
logger.info(f"Opening input with Xarray {input_key!s}: '{fname}'")
@@ -259,14 +260,14 @@ def subset_dataset(ds: xr.Dataset, subset_spec: DimIndex) -> xr.Dataset:
@contextmanager
-def open_chunk(config: XarrayZarrRecipeConfig, chunk_key: ChunkKey) -> xr.Dataset:
+def open_chunk(chunk_key: ChunkKey, *, config: XarrayZarrRecipe) -> xr.Dataset:
logger.info(f"Opening inputs for chunk {chunk_key!s}")
ninputs = config.file_pattern.dims[config.file_pattern.concat_dims[0]]
inputs = inputs_for_chunk(chunk_key, config.inputs_per_chunk, ninputs)
# need to open an unknown number of contexts at the same time
with ExitStack() as stack:
- dsets = [stack.enter_context(open_input(config, input_key)) for input_key in inputs]
+ dsets = [stack.enter_context(open_input(input_key, config=config)) for input_key in inputs]
# subset before chunking; hopefully lazy
subset_dims = [dim_idx for dim_idx in chunk_key if dim_idx.operation == CombineOp.SUBSET]
@@ -341,7 +342,7 @@ def calculate_sequence_lens(
return np.atleast_1d(sequence_lens.squeeze()).tolist()
-def prepare_target(config: XarrayZarrRecipeConfig) -> None:
+def prepare_target(*, config: XarrayZarrRecipe) -> None:
try:
ds = open_target(config.target)
logger.info("Found an existing dataset in target")
@@ -365,7 +366,7 @@ def prepare_target(config: XarrayZarrRecipeConfig) -> None:
init_chunks = list(filter(filter_init_chunks, config.iter_chunks()))
for chunk_key in init_chunks:
- with open_chunk(config, chunk_key) as ds:
+ with open_chunk(chunk_key, config=config) as ds:
# ds is already chunked
# https://github.com/pydata/xarray/blob/5287c7b2546fc8848f539bb5ee66bb8d91d8496f/xarray/core/variable.py#L1069
@@ -420,11 +421,11 @@ def prepare_target(config: XarrayZarrRecipeConfig) -> None:
config.metadata_cache[_GLOBAL_METADATA_KEY] = recipe_meta
-def store_chunk(config: XarrayZarrRecipeConfig, chunk_key: ChunkKey) -> None:
+def store_chunk(chunk_key: ChunkKey, *, config: XarrayZarrRecipe) -> None:
if config.target is None:
raise ValueError("target has not been set.")
- with open_chunk(config, chunk_key) as ds_chunk:
+ with open_chunk(chunk_key, config=config) as ds_chunk:
# writing a region means that all the variables MUST have concat_dim
to_drop = [v for v in ds_chunk.variables if config.concat_dim not in ds_chunk[v].dims]
ds_chunk = ds_chunk.drop_vars(to_drop)
@@ -483,7 +484,7 @@ def _gather_coordinate_dimensions(group: zarr.Group) -> List[str]:
)
-def finalize_target(config: XarrayZarrRecipeConfig) -> None:
+def finalize_target(*, config: XarrayZarrRecipe) -> None:
if config.target is None:
raise ValueError("target has not been set.")
@@ -517,13 +518,23 @@ def finalize_target(config: XarrayZarrRecipeConfig) -> None:
zarr.consolidate_metadata(target_mapper)
+def xarray_zarr_recipe_compiler(recipe: XarrayZarrRecipe) -> Pipeline:
+ stages = [
+ Stage(name="cache_input", function=cache_input, mappable=list(recipe.iter_inputs())),
+ Stage(name="prepare_target", function=prepare_target),
+ Stage(name="store_chunk", function=store_chunk, mappable=list(recipe.iter_chunks())),
+ Stage(name="finalize_target", function=finalize_target),
+ ]
+ return Pipeline(stages=stages, config=recipe)
+
+
# Notes about dataclasses:
# - https://www.python.org/dev/peps/pep-0557/#inheritance
# - https://stackoverflow.com/questions/51575931/class-inheritance-in-python-3-7-dataclasses
@dataclass
-class XarrayZarrRecipeConfig(RecipeConfig):
+class XarrayZarrRecipe(BaseRecipe, FilePatternMixin):
"""This configuration represents a dataset composed of many individual NetCDF files.
This class uses Xarray to read and write data and writes its output to Zarr.
The organization of the source files is described by the ``file_pattern``.
@@ -567,7 +578,8 @@ class XarrayZarrRecipeConfig(RecipeConfig):
time dimension. Multiple dimensions are allowed.
"""
- file_pattern: FilePattern
+ _compiler = xarray_zarr_recipe_compiler
+
inputs_per_chunk: int = 1
target_chunks: Dict[str, int] = field(default_factory=dict)
target: Optional[AbstractTarget] = None
@@ -668,8 +680,7 @@ class XarrayZarrRecipeConfig(RecipeConfig):
raise ValueError("Inputs and chunks are inconsistent")
def iter_inputs(self) -> Iterator[InputKey]:
- for input_key in self.file_pattern:
- yield input_key
+ yield from self.file_pattern
def iter_chunks(self) -> Iterator[ChunkKey]:
for input_key in self.iter_inputs():
@@ -707,56 +718,3 @@ class XarrayZarrRecipeConfig(RecipeConfig):
"""Convenience function for users to introspect recipe."""
ninputs = self.file_pattern.dims[self.file_pattern.concat_dims[0]]
return inputs_for_chunk(chunk_key, self.inputs_per_chunk, ninputs)
-
-
-@dataclass
-class XarrayZarrRecipe(BaseRecipe):
-
- config: XarrayZarrRecipeConfig
-
- # Each stage of the recipe follows the same pattern:
- # 1. A top-level function, e.g. `prepare_target`, that does the actual work.
- # 2. A public property, e.g. `.prepare_target`, that calls the partially applied function
- # with the provided arguments if any (e.g. a chunk_key)
- # This ensures that the actual function objects shipped to and executed on
- # workers do not contain any references to the `recipe` object itself, which is complicated
- # to serialize.
-
- @property
- def prepare_target(self) -> Callable[[XarrayZarrRecipeConfig], None]:
- return prepare_target
-
- @property
- def cache_input(self) -> Callable[[XarrayZarrRecipeConfig, Hashable], None]:
- return cache_input
-
- @property
- def store_chunk(self) -> Callable[[XarrayZarrRecipeConfig, Hashable], None]:
- return store_chunk
-
- @property
- def finalize_target(self) -> Callable[[XarrayZarrRecipeConfig], None]:
- return finalize_target
-
- # ------------------------------------------------------------------------
- # Convenience methods
- @contextmanager
- def open_input(self, input_key):
- with open_input(self.config, input_key) as ds:
- yield ds
-
- @contextmanager
- def open_chunk(self, chunk_key):
- with open_chunk(self.config, chunk_key,) as ds:
- yield ds
-
- def iter_inputs(self) -> Iterator[InputKey]:
- for input_key in self.config.iter_inputs():
- yield input_key
-
- def iter_chunks(self) -> Iterator[ChunkKey]:
- for chunk_key in self.config.iter_chunks():
- yield chunk_key
-
- def inputs_for_chunk(self, chunk_key: ChunkKey) -> Sequence[InputKey]:
- return self.config(inputs_for_chunk)

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

MODIFY
tests/recipe_tests/test_XarrayZarrRecipe.py
tests/recipe_tests/test_XarrayZarrRecipe.py
@@ -12,10 +12,10 @@ 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, XarrayZarrRecipeConfig
+from pangeo_forge_recipes.recipes.xarray_zarr import XarrayZarrRecipe
-def make_netCDFtoZarr_recipe_config(
+def make_netCDFtoZarr_recipe(
file_pattern, xarray_dataset, target, cache, metadata_target, extra_kwargs=None
):
kwargs = dict(
@@ -23,7 +23,7 @@ def make_netCDFtoZarr_recipe_config(
)
if extra_kwargs:
kwargs.update(extra_kwargs)
- return XarrayZarrRecipeConfig(file_pattern=file_pattern, **kwargs), xarray_dataset, target
+ return XarrayZarrRecipe, file_pattern, kwargs, xarray_dataset, target
@pytest.fixture
@@ -34,7 +34,7 @@ def netCDFtoZarr_recipe_sequential_only(
tmp_cache,
tmp_metadata_target,
):
- return make_netCDFtoZarr_recipe_config(
+ return make_netCDFtoZarr_recipe(
netcdf_local_file_pattern_sequential,
daily_xarray_dataset,
tmp_target,
@@ -47,7 +47,7 @@ def netCDFtoZarr_recipe_sequential_only(
def netCDFtoZarr_recipe(
netcdf_local_file_pattern, daily_xarray_dataset, tmp_target, tmp_cache, tmp_metadata_target
):
- return make_netCDFtoZarr_recipe_config(
+ return make_netCDFtoZarr_recipe(
netcdf_local_file_pattern, daily_xarray_dataset, tmp_target, tmp_cache, tmp_metadata_target
)
@@ -56,7 +56,7 @@ def netCDFtoZarr_recipe(
def netCDFtoZarr_http_recipe(
netcdf_http_file_pattern, daily_xarray_dataset, tmp_target, tmp_cache, tmp_metadata_target
):
- return make_netCDFtoZarr_recipe_config(
+ return make_netCDFtoZarr_recipe(
netcdf_http_file_pattern, daily_xarray_dataset, tmp_target, tmp_cache, tmp_metadata_target
)
@@ -69,7 +69,7 @@ def netCDFtoZarr_http_recipe_sequential_1d(
tmp_cache,
tmp_metadata_target,
):
- return make_netCDFtoZarr_recipe_config(
+ return make_netCDFtoZarr_recipe(
netcdf_http_file_pattern_sequential_1d,
daily_xarray_dataset,
tmp_target,
@@ -88,7 +88,7 @@ def netCDFtoZarr_subset_recipe(
extra_kwargs = dict(subset_inputs={"time": 2})
- return make_netCDFtoZarr_recipe_config(
+ return make_netCDFtoZarr_recipe(
netcdf_local_file_pattern,
daily_xarray_dataset,
tmp_target,
@@ -112,18 +112,12 @@ recipes_no_subset = [
def test_recipe(recipe_fixture, execute_recipe):
"""The basic recipe test. Use this as a template for other tests."""
- config, ds_expected, target = recipe_fixture
- rec = XarrayZarrRecipe(config)
+ RecipeClass, file_pattern, kwargs, ds_expected, target = recipe_fixture
+ rec = RecipeClass(file_pattern, **kwargs)
execute_recipe(rec)
ds_actual = xr.open_zarr(target.get_mapper()).load()
xr.testing.assert_identical(ds_actual, ds_expected)
- with rec.open_input(next(rec.iter_inputs())):
- pass
-
- with rec.open_chunk(next(rec.iter_chunks())):
- pass
-
@pytest.mark.parametrize("recipe_fixture", all_recipes)
@pytest.mark.parametrize("nkeep", [1, 2])
@@ -358,7 +352,7 @@ def test_consolidate_dimension_coordinates_with_coordinateless_dimension(
tmp_cache,
tmp_metadata_target,
):
- RecipeClass, file_pattern, kwargs, ds_expected, target = make_netCDFtoZarr_recipe_config(
+ RecipeClass, file_pattern, kwargs, ds_expected, target = make_netCDFtoZarr_recipe(
netcdf_local_file_pattern_sequential_with_coordinateless_dimension,
daily_xarray_dataset_with_coordinateless_dimension,
tmp_target,

MODIFY
tests/test_pipelines.py
tests/test_pipelines.py
@@ -5,21 +5,25 @@ Test Pipline Executors
import pytest
from pytest_lazyfixture import lazy_fixture
-# from pangeo_forge_recipes.executors.dask import DaskPipelineExecutor
-from pangeo_forge_recipes.executors.prefect import PrefectPipelineExecutor
+from pangeo_forge_recipes.executors.base import Pipeline, Stage
+from pangeo_forge_recipes.executors.dask import DaskPipelineExecutor
from pangeo_forge_recipes.executors.function import FunctionPipelineExecutor
-from pangeo_forge_recipes.executors.base import Stage, Pipeline
+from pangeo_forge_recipes.executors.prefect import PrefectPipelineExecutor
@pytest.fixture
def pipeline_no_config(tmpdir_factory):
tmp = tmpdir_factory.mktemp("pipeline_data")
+
def func0(config=None):
tmp.join("func0.log").ensure(file=True)
assert not tmp.join("func1_a.log").check(file=True)
+
stage0 = Stage(function=func0, name="create_first_file")
+
def func1(arg, config=None):
tmp.join(f"func1_{arg}.log").ensure(file=True)
+
stage1 = Stage(function=func1, name="create_many_files", mappable=["a", "b", 3])
pipeline = Pipeline(stages=[stage0, stage1])
return pipeline, {}, tmp
@@ -28,14 +32,18 @@ def pipeline_no_config(tmpdir_factory):
@pytest.fixture
def pipeline_with_config(tmpdir_factory):
tmp = tmpdir_factory.mktemp("pipeline_data")
+
def func0(config=None):
- prefix = config['prefix']
+ prefix = config["prefix"]
tmp.join(f"{prefix}func0.log").ensure(file=True)
assert not tmp.join(f"{prefix}-func1_a.log").check(file=True)
+
stage0 = Stage(function=func0, name="create_first_file")
+
def func1(arg, config=None):
- prefix = config['prefix']
+ prefix = config["prefix"]
tmp.join(f"{prefix}func1_{arg}.log").ensure(file=True)
+
stage1 = Stage(function=func1, name="create_many_files", mappable=["a", "b", 3])
config = {"prefix": "special-"}
pipeline = Pipeline(stages=[stage0, stage1], config=config)
@@ -43,13 +51,11 @@ def pipeline_with_config(tmpdir_factory):
@pytest.mark.parametrize(
- "Executor", [FunctionPipelineExecutor, PrefectPipelineExecutor],
+ "Executor", [FunctionPipelineExecutor, PrefectPipelineExecutor, DaskPipelineExecutor],
)
@pytest.mark.parametrize(
- "pipeline_config_tmpdir", [
- lazy_fixture("pipeline_no_config"),
- lazy_fixture("pipeline_with_config"),
- ]
+ "pipeline_config_tmpdir",
+ [lazy_fixture("pipeline_no_config"), lazy_fixture("pipeline_with_config"),],
)
def test_pipeline(pipeline_config_tmpdir, Executor):
pipeline, config, tmpdir = pipeline_config_tmpdir
@@ -57,5 +63,10 @@ def test_pipeline(pipeline_config_tmpdir, Executor):
Executor.execute(plan)
prefix = config.get("prefix", "")
- for fname in [f"{prefix}func0.log", f"{prefix}func1_a.log", f"{prefix}func1_b.log", f"{prefix}func1_3.log"]:
+ for fname in [
+ f"{prefix}func0.log",
+ f"{prefix}func1_a.log",
+ f"{prefix}func1_b.log",
+ f"{prefix}func1_3.log",
+ ]:
assert tmpdir.join(fname).check(file=True)

got the first test passing