first pass at working executors
ADD
pangeo_forge_recipes/executors/__init__.py
pangeo_forge_recipes/executors/__init__.py
ADD
pangeo_forge_recipes/executors/base.py
+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
+
+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: ...
+
+
+class SingleArgumentStageFunction(Protocol):
+ def __call__(__a: Hashable, *, config: Optional[Config]=None) -> None: ...
+
+
+StageFunction = Union[NoArgumentStageFunction, SingleArgumentStageFunction]
+
+
+class StageAnnotationType(enum.Enum):
+ CONCURRENCY = enum.auto()
+ RETRIES = enum.auto()
+
+StageAnnotations = Dict[StageAnnotationType, Any]
+
+@dataclass(frozen=True)
+class Stage:
+ function: StageFunction
+ name: str
+ mappable: Optional[Iterable] = None
+ annotations: Optional[StageAnnotations] = None
+
+
+@dataclass(frozen=True)
+class Pipeline:
+ stages: Iterable[Stage]
+ config: Optional[Config] = None
+
+
+T = TypeVar("T")
+
+
+class PipelineExecutor(Generic[T]):
+
+ @staticmethod
+ def compile(pipeline: Pipeline) -> T:
+ raise NotImplementedError
+
+ @staticmethod
+ def execute(plan: T):
+ raise NotImplementedError

ADD
pangeo_forge_recipes/executors/dask.py
pangeo_forge_recipes/executors/dask.py
ADD
pangeo_forge_recipes/executors/function.py
+from typing import Callable
+
+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:
+ for m in stage.mappable:
+ stage.function(m, config=pipeline.config)
+ else:
+ stage.function(config=pipeline.config)
+
+ return function
+
+ @staticmethod
+ def execute(func: Callable) -> None:
+ func()

ADD
pangeo_forge_recipes/executors/prefect.py
+from typing import Callable, Dict, Any, Optional
+
+from prefect import Flow, task, unmapped
+
+from .base import Pipeline, PipelineExecutor, StageAnnotations, StageAnnotationType
+
+
+def annotations_to_task_kwargs(annotations: Optional[StageAnnotations]) -> Dict[str, Any]:
+ task_kwargs = {}
+ for ann_type in annotations or []:
+ if ann_type == StageAnnotationType.retries:
+ 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]):
+
+ @staticmethod
+ def compile(pipeline: Pipeline):
+
+ with Flow("pangeo-forge-recipe") as flow:
+ upstream_tasks = None
+ for stage in pipeline.stages:
+ task_kwargs = annotations_to_task_kwargs(stage.annotations)
+ stage_task = task(stage.function, name=stage.name, **task_kwargs)
+ if stage.mappable is not None:
+ stage_task_called = stage_task.map(
+ stage.mappable,
+ config=unmapped(pipeline.config),
+ upstream_tasks=[unmapped(t) for t in upstream_tasks]
+ )
+ else:
+ stage_task_called = stage_task(
+ config=pipeline.config,
+ upstream_tasks=upstream_tasks
+ )
+ upstream_tasks = [stage_task_called]
+
+ return flow
+
+ @staticmethod
+ def execute(flow: Flow):
+ return flow.run()

ADD
tests/test_pipelines.py
+"""
+Test Pipline Executors
+"""
+
+import pytest
+
+# from pangeo_forge_recipes.executors.dask import DaskPipelineExecutor
+from pangeo_forge_recipes.executors.prefect import PrefectPipelineExecutor
+from pangeo_forge_recipes.executors.function import FunctionPipelineExecutor
+from pangeo_forge_recipes.executors.base import Stage, Pipeline
+
+
+@pytest.fixture
+def example_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])
+
+ # MultiStagePipeline
+ pipeline = Pipeline(stages=[stage0, stage1])
+ return pipeline, tmp
+
+
+@pytest.mark.parametrize(
+ "Executor", [FunctionPipelineExecutor, PrefectPipelineExecutor]
+)
+def test_pipeline(example_pipeline_no_config, Executor):
+ pipeline, tmpdir = example_pipeline_no_config
+ plan = Executor.compile(pipeline)
+ Executor.execute(plan)
+ for fname in ["func0.log", "func1_a.log", "func1_b.log", "func1_3.log"]:
+ assert tmpdir.join(fname).check(file=True)

first pass at working executors