refactor pipelines and tasks to use functions defined outside the pipeline, introduce pipelines and tasks directories
DELETE
pangeo_forge/pipeline.py
-import os
-from abc import ABC, abstractmethod
-from typing import List
-
-import fsspec
-import zarr
-from prefect import Flow, task
-
-from .utils import chunked_iterable
-
-
-class AbstractPipeline(ABC):
- @property
- @abstractmethod
- def sources(self) -> List[str]:
- pass
-
- @property
- @abstractmethod
- def targets(self) -> List[str]:
- pass
-
- @abstractmethod
- def run(self) -> None:
- pass
-
-
-class XarrayPrefectPipelineMixin:
- @task
- def download(self, source_url):
- target_url = os.path.join(self.cache_location, str(hash(source_url)))
-
- # there is probably a better way to do caching!
- try:
- fsspec.open(target_url).open()
- return target_url
- except FileNotFoundError:
- pass
-
- with fsspec.open(source_url, mode="rb") as source:
- with fsspec.open(target_url, mode="wb") as target:
- target.write(source.read())
- return target_url
-
- @task
- def combine_and_write(self, sources, target, append_dim, first=True):
- import xarray as xr
-
- # while debugging this, I had itermittent fsspec / hdf5 read errors related to
- # "trying to read from a closed file"
- # but they seem to have gone away for now
- double_open_files = [fsspec.open(url).open() for url in sources]
- ds = xr.open_mfdataset(double_open_files, combine="nested", concat_dim=self.concat_dim)
- # by definition, this should be a contiguous chunk
- ds = ds.chunk({append_dim: len(sources)})
-
- if first:
- kwargs = dict(mode="w")
- else:
- kwargs = dict(mode="a", append_dim=append_dim)
-
- mapper = fsspec.get_mapper(target)
- ds.to_zarr(mapper, **kwargs)
-
- @task
- def consolidate_metadata(self):
- mapper = fsspec.get_mapper(self.targets)
- zarr.consolidate_metadata(mapper)
-
- @property
- def flow(self) -> Flow:
- with Flow("Pangeo-Forge") as flow:
-
- cached_sources = [self.download(k) for k in self.sources]
-
- first = True
- write_tasks = []
- for source_group in chunked_iterable(cached_sources, self.files_per_chunk):
- write_task = self.combine_and_write(
- source_group, self.targets, self.concat_dim, first=first
- )
- write_tasks.append(write_task)
- first = False
- cm = self.consolidate_metadata(self.targets)
-
- # create dependencies in imperative mode
- for n in range(1, len(write_tasks)):
- write_tasks[n].set_upstream(write_tasks[n - 1], flow=flow)
- cm.set_upstream(write_tasks[-1], flow=flow)
-
- return flow
-
- def run(self):
- self.flow.run()

ADD
pangeo_forge/pipelines/README.md
+# `pangeo_forge.pipelines`
+
+A collection of Pangeo-forge _Pipelines_.
+
+## Details
+
+- An abstract Pipeline class is defined in `base.py`.
+- A general file naming convention might be `{source}_{container}_{target}.py`. For example, a pipeline that downloads netCDF files from a remote http server and publishes a Zarr store would be `http_xarray_zarr.py`. The main class in each module would follow a similar convention, e.g. `HttpXarrayZarr`.

ADD
pangeo_forge/pipelines/__init__.py
pangeo_forge/pipelines/__init__.py
ADD
pangeo_forge/pipelines/base.py
+
+from abc import ABC, abstractmethod
+from typing import List
+
+from prefect import Flow
+
+
+class AbstractPipeline(ABC):
+ name = 'AbstractPipeline'
+
+ @property
+ @abstractmethod
+ def sources(self) -> List[str]:
+ pass
+
+ @property
+ @abstractmethod
+ def targets(self) -> List[str]:
+ pass
+
+ @abstractmethod
+ def flow(self) -> Flow:
+ pass

ADD
pangeo_forge/pipelines/http_xarray_zarr.py
+from prefect import Flow
+
+from ..utils import chunked_iterable
+
+from ..tasks.http import download
+from ..tasks.xarray import combine_and_write
+from ..tasks.zarr import consolidate_metadata
+
+class HttpXarrayZarrMixin:
+
+ @property
+ def flow(self):
+
+ if len(self.targets) == 1:
+ target = self.targets[0]
+ else:
+ raise ValueError('Zarr target requires self.targets be a length one list')
+
+ with Flow(self.name) as _flow:
+
+ cached_sources = [download(k, self.cache_location) for k in self.sources]
+
+ first = True
+ write_tasks = []
+ for source_group in chunked_iterable(cached_sources, self.files_per_chunk):
+ write_task = combine_and_write(
+ source_group, target, self.append_dim, self.concat_dim, first=first
+ )
+ write_tasks.append(write_task)
+ first = False
+ cm = consolidate_metadata(target)
+
+ # create dependencies in imperative mode
+ for n in range(1, len(write_tasks)):
+ write_tasks[n].set_upstream(write_tasks[n - 1], flow=_flow)
+ cm.set_upstream(write_tasks[-1], flow=_flow)
+
+ return _flow

ADD
pangeo_forge/tasks/README.md
+# `pangeo_forge.tasks`
+
+A collection of Prefect Tasks for Pangeo-Forge.
+
+## Details
+
+- This is simply a collection of Prefect Tasks. Much like https://docs.prefect.io/api/latest/#task-library or https://github.com/PrefectHQ/prefect/tree/master/src/prefect/tasks.
+- This should be considered low-level public API.

ADD
pangeo_forge/tasks/__init__.py
pangeo_forge/tasks/__init__.py
ADD
pangeo_forge/tasks/http.py
+import os
+
+import fsspec
+
+from prefect import task
+
+
+@task
+def download(source_url, cache_location):
+ '''
+ Download a remote file to a cache.
+
+ Parameters
+ ----------
+ source_url : str
+ Path or url to the source file.
+ cache_location : str
+ Path or url to the target location for the source file.
+
+ Returns
+ -------
+ target_url : str
+ Path or url in the form of `{cache_location}/hash({source_url})`.
+ '''
+ target_url = os.path.join(cache_location, str(hash(source_url)))
+
+ # there is probably a better way to do caching!
+ try:
+ fsspec.open(target_url).open()
+ return target_url
+ except FileNotFoundError:
+ pass
+
+ with fsspec.open(source_url, mode="rb") as source:
+ with fsspec.open(target_url, mode="wb") as target:
+ target.write(source.read())
+ return target_url
\ No newline at end of file

ADD
pangeo_forge/tasks/xarray.py
+import fsspec
+import xarray as xr
+
+from prefect import task
+
+@task
+def combine_and_write(sources, target, append_dim, concat_dim, first=True):
+ '''
+ Combine one or more source datasets into a single `Xarray.Dataset`, then
+ write them to a Zarr store.
+
+ Parameters
+ ----------
+ sources : list of str
+ Path or url to the source files.
+ target : str
+ Path or url to the target location of the Zarr store.
+ append_dim : str
+ Name of the dimension of which datasets should be appended during write.
+ concat_dim : str
+ Name of the dimension of which datasets should be concatenated during read.
+ first : bool
+ Flag to indicate if the target dataset should be expected to already exist.
+
+ Returns
+ -------
+ target_url : str
+ Path or url in the form of `{cache_location}/hash({source_url})`.
+ '''
+
+ double_open_files = [fsspec.open(url).open() for url in sources]
+ ds = xr.open_mfdataset(double_open_files, combine="nested", concat_dim=concat_dim)
+
+ # by definition, this should be a contiguous chunk
+ ds = ds.chunk({append_dim: len(sources)})
+
+ if first:
+ kwargs = dict(mode="w")
+ else:
+ kwargs = dict(mode="a", append_dim=append_dim)
+
+ mapper = fsspec.get_mapper(target)
+ ds.to_zarr(mapper, **kwargs)

ADD
pangeo_forge/tasks/zarr.py
+import fsspec
+import zarr
+
+from prefect import task
+
+
+@task
+def consolidate_metadata(target):
+ '''
+ Consolidate Zarr metadata
+
+ Parameters
+ ----------
+ target : str
+ Path or url of the Zarr store.
+ '''
+ mapper = fsspec.get_mapper(target)
+ zarr.consolidate_metadata(mapper)

MODIFY
setup.cfg
setup.cfg
@@ -12,4 +12,4 @@ line_length=100
[tool:pytest]
log_cli = True
-log_level = DEBUG
+log_level = INFO

MODIFY
tests/test_pipeline.py
tests/test_pipeline.py
@@ -10,7 +10,8 @@ import pytest
import xarray as xr
# classes tested here
-from pangeo_forge.pipeline import AbstractPipeline, XarrayPrefectPipelineMixin
+from pangeo_forge.pipelines.base import AbstractPipeline
+from pangeo_forge.pipelines.http_xarray_zarr import HttpXarrayZarrMixin
# where to run the http server
_PORT = "8080"
@@ -94,13 +95,22 @@ def test_fixture_http_files(daily_xarray_dataset, netcdf_http_server):
# a pipeline to load that data
-class Pipeline(XarrayPrefectPipelineMixin, AbstractPipeline):
+class MyPipeline(HttpXarrayZarrMixin, AbstractPipeline):
def __init__(
- self, name, cache_path, target_path, concat_dim, files_per_chunk, url_base, nfiles
+ self,
+ name,
+ cache_path,
+ target_path,
+ concat_dim,
+ append_dim,
+ files_per_chunk,
+ url_base,
+ nfiles,
):
self.name = name
self.cache_location = f"{cache_path}/{name}-cache/"
self.target_location = f"{target_path}/{name}.zarr"
+ self.append_dim = append_dim
self.concat_dim = concat_dim
self.files_per_chunk = files_per_chunk
@@ -121,20 +131,21 @@ class Pipeline(XarrayPrefectPipelineMixin, AbstractPipeline):
# a basic pipeline test
-
-
def test_pipeline(daily_xarray_dataset, netcdf_http_server, tmpdir):
name = "TEST_DATASET"
cache_dir = tmpdir.mkdir("cache")
target_dir = tmpdir.mkdir("target")
concat_dim = "time"
+ append_dim = "time"
files_per_chunk = 5
url_base, paths = netcdf_http_server
nfiles = len(paths)
- pipeline = Pipeline(name, cache_dir, target_dir, concat_dim, files_per_chunk, url_base, nfiles)
- pipeline.run()
+ pipeline = MyPipeline(
+ name, cache_dir, target_dir, concat_dim, append_dim, files_per_chunk, url_base, nfiles
+ )
+ pipeline.flow.run()
ds_test = xr.open_zarr(pipeline.targets[0])
assert ds_test.identical(daily_xarray_dataset)

refactor pipelines and tasks to use functions defined outside the pipeline, introduce pipelines and tasks directories