flytekitplugins-xarray-zarr 1.16.0b5__tar.gz

flytekitplugins_xarray_zarr-1.16.0b5/PKG-INFO

@@ -0,0 +1,29 @@
+Metadata-Version: 2.4
+Name: flytekitplugins-xarray-zarr
+Version: 1.16.0b5
+Summary: Xarray Zarr plugin for flytekit
+Author: flyteorg
+Author-email: admin@flyte.org
+License: apache2
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.9
+Requires-Dist: dask[distributed]>=2022.10.2
+Requires-Dist: flytekit<2.0.0,>=1.3.0b2
+Requires-Dist: xarray
+Requires-Dist: zarr
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

flytekitplugins_xarray_zarr-1.16.0b5/README.md

@@ -0,0 +1,48 @@
+# Flytekit Xarray Zarr Plugin
+The Xarray Zarr plugin adds support to persist xarray datasets and dataarrays to zarr between tasks. If a dask cluster is present (see flytekitplugins-dask), it will attempt to connect to the distributed client before we call `.to_zarr(url)` call. This prevents the need to explicitly connect to a distributed client within the task.
+
+If deck is enabled, we also render the datasets/dataarrays to html.
+
+To install the plugin, run the following command:
+
+```bash
+pip install flytekitplugins-xarray-zarr
+```
+
+## Example
+
+```python
+import dask.array as da
+import xarray as xr
+from flytekit import task, workflow
+from flytekitplugins.dask import Dask, WorkerGroup
+
+
+@task(
+    task_config=Dask(workers=WorkerGroup(number_of_workers=6)),
+    enable_deck=True,
+)
+def generate_xarray_task() -> xr.Dataset:
+    return xr.Dataset(
+        {
+            "variable": (
+                ("time", "x", "y"),
+                da.random.uniform(size=(1024, 1024, 1024)),
+            )
+        },
+    )
+
+
+@task(
+    task_config=Dask(workers=WorkerGroup(number_of_workers=6)),
+    enable_deck=True,
+)
+def preprocess_xarray_task(ds: xr.Dataset) -> xr.Dataset:
+    return ds * 2
+
+
+@workflow
+def xarray_workflow() -> xr.Dataset:
+    ds = generate_xarray_task()
+    return preprocess_xarray_task(ds=ds)
+```

flytekitplugins_xarray_zarr-1.16.0b5/flytekitplugins/xarray/__init__.py

@@ -0,0 +1,14 @@
+"""
+.. currentmodule:: flytekitplugins.xarray
+
+This package contains things that are useful when extending Flytekit.
+
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+
+   XarrayDaZarrTypeTransformer
+   XarrayZarrTypeTransformer
+"""
+
+from .xarray_transformers import XarrayDaZarrTypeTransformer, XarrayZarrTypeTransformer

flytekitplugins_xarray_zarr-1.16.0b5/flytekitplugins/xarray/xarray_transformers.py

@@ -0,0 +1,102 @@
+import typing
+
+import dask.distributed as dd
+
+import xarray as xr
+from flytekit import (
+    Blob,
+    BlobMetadata,
+    BlobType,
+    FlyteContext,
+    Literal,
+    LiteralType,
+    Scalar,
+)
+from flytekit.extend import TypeEngine, TypeTransformer
+
+
+class XarrayZarrTypeTransformer(TypeTransformer[xr.Dataset]):
+    _TYPE_INFO = BlobType(format="binary", dimensionality=BlobType.BlobDimensionality.MULTIPART)
+
+    def __init__(self) -> None:
+        super().__init__(name="Xarray Dataset", t=xr.Dataset)
+
+    def get_literal_type(self, t: typing.Type[xr.Dataset]) -> LiteralType:
+        return LiteralType(blob=self._TYPE_INFO)
+
+    def to_literal(
+        self,
+        ctx: FlyteContext,
+        python_val: xr.Dataset,
+        python_type: typing.Type[xr.Dataset],
+        expected: LiteralType,
+    ) -> Literal:
+        remote_dir = ctx.file_access.get_random_remote_path("data.zarr")
+        # Opening with the dask client will attach the client eliminating the
+        # need for users to connect to the client if a task tasks a xr.Dataset
+        # type.
+        with dd.Client(timeout=120):
+            python_val.to_zarr(remote_dir, mode="w")
+        return Literal(scalar=Scalar(blob=Blob(uri=remote_dir, metadata=BlobMetadata(type=self._TYPE_INFO))))
+
+    def to_python_value(
+        self,
+        ctx: FlyteContext,
+        lv: Literal,
+        expected_python_type: typing.Type[xr.Dataset],
+    ) -> xr.Dataset:
+        return xr.open_zarr(lv.scalar.blob.uri)
+
+    def to_html(
+        self,
+        ctx: FlyteContext,
+        python_val: xr.Dataset,
+        expected_python_type: LiteralType,
+    ) -> str:
+        return python_val._repr_html_()
+
+
+class XarrayDaZarrTypeTransformer(TypeTransformer[xr.DataArray]):
+    _TYPE_INFO = BlobType(format="binary", dimensionality=BlobType.BlobDimensionality.MULTIPART)
+
+    def __init__(self) -> None:
+        super().__init__(name="Xarray DataArray", t=xr.DataArray)
+
+    def get_literal_type(self, t: typing.Type[xr.DataArray]) -> LiteralType:
+        return LiteralType(blob=self._TYPE_INFO)
+
+    def to_literal(
+        self,
+        ctx: FlyteContext,
+        python_val: xr.DataArray,
+        python_type: typing.Type[xr.DataArray],
+        expected: LiteralType,
+    ) -> Literal:
+        remote_dir = ctx.file_access.get_random_remote_path("data.zarr")
+        # Opening with the dask client will attach the client eliminating the
+        # need for users to connect to the client if a task tasks a xr.Dataset
+        # type.
+        with dd.Client(timeout=120):
+            python_val.to_zarr(remote_dir, mode="w")
+        return Literal(scalar=Scalar(blob=Blob(uri=remote_dir, metadata=BlobMetadata(type=self._TYPE_INFO))))
+
+    def to_python_value(
+        self,
+        ctx: FlyteContext,
+        lv: Literal,
+        expected_python_type: typing.Type[xr.DataArray],
+    ) -> xr.DataArray:
+        # xr.open_zarr always opens a dataset, so we take the first variable
+        return list(xr.open_zarr(lv.scalar.blob.uri).data_vars.values())[0]
+
+    def to_html(
+        self,
+        ctx: FlyteContext,
+        python_val: xr.DataArray,
+        expected_python_type: LiteralType,
+    ) -> str:
+        return python_val._repr_html_()
+
+
+TypeEngine.register(XarrayZarrTypeTransformer())
+TypeEngine.register(XarrayDaZarrTypeTransformer())

flytekitplugins_xarray_zarr-1.16.0b5/flytekitplugins_xarray_zarr.egg-info/PKG-INFO

@@ -0,0 +1,29 @@
+Metadata-Version: 2.4
+Name: flytekitplugins-xarray-zarr
+Version: 1.16.0b5
+Summary: Xarray Zarr plugin for flytekit
+Author: flyteorg
+Author-email: admin@flyte.org
+License: apache2
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.9
+Requires-Dist: dask[distributed]>=2022.10.2
+Requires-Dist: flytekit<2.0.0,>=1.3.0b2
+Requires-Dist: xarray
+Requires-Dist: zarr
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

flytekitplugins_xarray_zarr-1.16.0b5/flytekitplugins_xarray_zarr.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+README.md
+setup.py
+flytekitplugins/xarray/__init__.py
+flytekitplugins/xarray/xarray_transformers.py
+flytekitplugins_xarray_zarr.egg-info/PKG-INFO
+flytekitplugins_xarray_zarr.egg-info/SOURCES.txt
+flytekitplugins_xarray_zarr.egg-info/dependency_links.txt
+flytekitplugins_xarray_zarr.egg-info/entry_points.txt
+flytekitplugins_xarray_zarr.egg-info/namespace_packages.txt
+flytekitplugins_xarray_zarr.egg-info/requires.txt
+flytekitplugins_xarray_zarr.egg-info/top_level.txt
+tests/test_xarray_zarr_plugin.py

flytekitplugins_xarray_zarr-1.16.0b5/flytekitplugins_xarray_zarr.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

flytekitplugins_xarray_zarr-1.16.0b5/flytekitplugins_xarray_zarr.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[flytekit.plugins]
+xarray = flytekitplugins.xarray

flytekitplugins_xarray_zarr-1.16.0b5/flytekitplugins_xarray_zarr.egg-info/namespace_packages.txt

@@ -0,0 +1 @@
+flytekitplugins

flytekitplugins_xarray_zarr-1.16.0b5/flytekitplugins_xarray_zarr.egg-info/requires.txt

@@ -0,0 +1,4 @@
+dask[distributed]>=2022.10.2
+flytekit<2.0.0,>=1.3.0b2
+xarray
+zarr

flytekitplugins_xarray_zarr-1.16.0b5/flytekitplugins_xarray_zarr.egg-info/top_level.txt

@@ -0,0 +1 @@
+flytekitplugins

flytekitplugins_xarray_zarr-1.16.0b5/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

flytekitplugins_xarray_zarr-1.16.0b5/setup.py

@@ -0,0 +1,40 @@
+from setuptools import setup
+
+PLUGIN_NAME = "xarray"
+
+microlib_name = f"flytekitplugins-{PLUGIN_NAME}-zarr"
+
+plugin_requires = [
+    "dask[distributed]>=2022.10.2",
+    "flytekit>=1.3.0b2,<2.0.0",
+    "xarray",
+    "zarr",
+]
+
+__version__ = "1.16.0b5"
+
+setup(
+    name=microlib_name,
+    version=__version__,
+    author="flyteorg",
+    author_email="admin@flyte.org",
+    description="Xarray Zarr plugin for flytekit",
+    namespace_packages=["flytekitplugins"],
+    packages=[f"flytekitplugins.{PLUGIN_NAME}"],
+    install_requires=plugin_requires,
+    license="apache2",
+    python_requires=">=3.9",
+    classifiers=[
+        "Intended Audience :: Science/Research",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Topic :: Scientific/Engineering",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Software Development",
+        "Topic :: Software Development :: Libraries",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+    ],
+    entry_points={"flytekit.plugins": [f"{PLUGIN_NAME}=flytekitplugins.{PLUGIN_NAME}"]},
+)

flytekitplugins_xarray_zarr-1.16.0b5/tests/test_xarray_zarr_plugin.py

@@ -0,0 +1,48 @@
+from flytekit import task, workflow
+import numpy as np
+import dask.array as da
+import xarray as xr
+
+
+def _sample_dataset() -> xr.Dataset:
+    return xr.Dataset(
+        {"test": (("x", "y"), da.random.uniform(size=(32, 32)))},
+    )
+
+
+def test_xarray_zarr_dataarray_plugin():
+
+    @task
+    def _generate_xarray() -> xr.DataArray:
+        return _sample_dataset()["test"]
+
+    @task
+    def _consume_xarray(ds: xr.DataArray) -> xr.DataArray:
+        return ds
+
+    @workflow
+    def _xarray_wf() -> xr.DataArray:
+        ds = _generate_xarray()
+        return _consume_xarray(ds=ds)
+
+    array = _xarray_wf()
+    assert isinstance(array, xr.DataArray)
+
+
+def test_xarray_zarr_dataset_plugin():
+
+    @task
+    def _generate_xarray() -> xr.Dataset:
+        return _sample_dataset()
+
+    @task
+    def _consume_xarray(ds: xr.Dataset) -> xr.Dataset:
+        return ds
+
+    @workflow
+    def _xarray_wf() -> xr.Dataset:
+        ds = _generate_xarray()
+        return _consume_xarray(ds=ds)
+
+    array = _xarray_wf()
+    assert isinstance(array, xr.Dataset)