Naive implementation of Beam executor.
ADD
pangeo_forge_recipes/executors/beam.py
+from typing import Optional, Iterable
+
+import apache_beam as beam
+
+from .base import Pipeline, PipelineExecutor, Config, StageFunction
+
+
+def exec_ordered(last: int,
+ *,
+ current: int,
+ function: StageFunction,
+ config: Config,
+ mappable: Optional[Iterable]) -> int:
+ """Execute a Pipeline Stage, ensuring execution order."""
+ assert (last + 1) == current, 'stages are executing out of order!'
+
+ if mappable is not None:
+ for m in mappable:
+ function(m, config=config)
+ else:
+ function(config=config)
+
+ return current
+
+
+class BeamPipelineExecutor(PipelineExecutor[beam.PTransform]):
+ @staticmethod
+ def compile(pipeline: Pipeline) -> beam.PTransform:
+ pcoll = beam.Create([-1])
+ for step, stage in enumerate(pipeline.stages):
+ pcoll |= stage.name >> beam.Map(exec_ordered,
+ current=step,
+ function=stage.function,
+ mappable=stage.mappable,
+ config=pipeline.config)
+
+ return pcoll
+
+ @staticmethod
+ def execute(plan: beam.PTransform, *args, **kwargs):
+ with beam.Pipeline(*args, **kwargs) as p:
+ p | plan

MODIFY
tests/test_pipelines.py
tests/test_pipelines.py
@@ -9,6 +9,7 @@ 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.prefect import PrefectPipelineExecutor
+from pangeo_forge_recipes.executors.beam import BeamPipelineExecutor
@pytest.fixture
@@ -51,7 +52,7 @@ def pipeline_with_config(tmpdir_factory):
@pytest.mark.parametrize(
- "Executor", [FunctionPipelineExecutor, PrefectPipelineExecutor, DaskPipelineExecutor],
+ "Executor", [FunctionPipelineExecutor, PrefectPipelineExecutor, DaskPipelineExecutor, BeamPipelineExecutor],
)
@pytest.mark.parametrize(
"pipeline_config_tmpdir",
@@ -69,4 +70,4 @@ def test_pipeline(pipeline_config_tmpdir, Executor):
f"{prefix}func1_b.log",
f"{prefix}func1_3.log",
]:
- assert tmpdir.join(fname).check(file=True)
+ assert tmpdir.join(fname).check(file=True), f'File not found in temporary directory: {fname}.'

Naive implementation of Beam executor.