Fix plugin regressions (flyteorg#688)

* fix pandera regression Signed-off-by: Niels Bantilan <[email protected]> * install plugin with pip Signed-off-by: Niels Bantilan <[email protected]> * fix pandera plugin tests Signed-off-by: Niels Bantilan <[email protected]> * wip Signed-off-by: Niels Bantilan <[email protected]> * wip Signed-off-by: Niels Bantilan <[email protected]> * wip Signed-off-by: Niels Bantilan <[email protected]> * wip Signed-off-by: Niels Bantilan <[email protected]> * wip Signed-off-by: Niels Bantilan <[email protected]> * wip Signed-off-by: Niels Bantilan <[email protected]> * add spark flytekit plugin to papermill test_requires Signed-off-by: Niels Bantilan <[email protected]> * wip Signed-off-by: Niels Bantilan <[email protected]> * wip Signed-off-by: Niels Bantilan <[email protected]> * wip Signed-off-by: Niels Bantilan <[email protected]> * wip Signed-off-by: Niels Bantilan <[email protected]> * add sqlalchemy to great expectations plugin Signed-off-by: Niels Bantilan <[email protected]> * wip Signed-off-by: Niels Bantilan <[email protected]> * wip Signed-off-by: Niels Bantilan <[email protected]> * plugins plugins plugins! Signed-off-by: Niels Bantilan <[email protected]> * lint Signed-off-by: Niels Bantilan <[email protected]>
AdrianoKF · Oct 11, 2021 · eacdc04 · eacdc04
1 parent fe02c81
commit eacdc04
Show file tree

Hide file tree

Showing 11 changed files with 118 additions and 46 deletions.
diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml
@@ -53,7 +53,21 @@ jobs:
     strategy:
       matrix:
         python-version: [3.8]
-        plugin-names: ["flytekit-aws-athena", "flytekit-aws-sagemaker", "flytekit-data-fsspec", "flytekit-dolt", "flytekit-greatexpectations", "flytekit-hive", "flytekit-k8s-pod", "flytekit-kf-pytorch", "flytekit-kf-tensorflow", "flytekit-papermill", "flytekit-spark", "flytekit-sqlalchemy", "flytekit-pandera", "flytekit-snowflake"]
+        plugin-names:
+          - flytekit-aws-athena
+          - flytekit-aws-sagemaker
+          - flytekit-data-fsspec
+          - flytekit-dolt
+          - flytekit-greatexpectations
+          - flytekit-hive
+          - flytekit-k8s-pod
+          - flytekit-kf-pytorch
+          - flytekit-kf-tensorflow
+          - flytekit-papermill
+          - flytekit-spark
+          - flytekit-sqlalchemy
+          - flytekit-pandera
+          - flytekit-snowflake
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python ${{ matrix.python-version }}
@@ -72,7 +86,8 @@ jobs:
         python -m pip install --upgrade pip setuptools wheel
         make setup
         cd plugins/${{ matrix.plugin-names }}
-        pip install -r requirements.txt
+        # install in develop mode so setuptools doesn't override flytekit on current branch
+        pip install -e .
         pip freeze
     - name: Test with coverage
       run: |

diff --git a/plugins/Makefile b/plugins/Makefile
@@ -1,6 +1,6 @@
 .PHONY: test
 test:
-	pytest tests
+	find . -maxdepth 1 -type d | grep 'flytekit-' | xargs -L1 pytest
 
 .PHONY: build_all_plugins
 build_all_plugins:

diff --git a/plugins/flytekit-aws-sagemaker/tests/test_hpo.py b/plugins/flytekit-aws-sagemaker/tests/test_hpo.py
@@ -21,7 +21,8 @@
 )
 from flytekit.models.sagemaker.parameter_ranges import IntegerParameterRange, ParameterRangeOneOf
 from flytekit.models.sagemaker.training_job import AlgorithmName, AlgorithmSpecification, TrainingJobResourceConfig
-from plugins.tests.awssagemaker.test_training import _get_reg_settings
+
+from .test_training import _get_reg_settings
 
 
 def test_hpo_for_builtin():

diff --git a/plugins/flytekit-greatexpectations/setup.py b/plugins/flytekit-greatexpectations/setup.py
@@ -4,7 +4,7 @@
 
 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
 
-plugin_requires = ["flytekit>=0.21.0,<1.0.0", "great-expectations>=0.13.30"]
+plugin_requires = ["flytekit>=0.21.0,<1.0.0", "great-expectations>=0.13.30", "sqlalchemy>=1.4.23"]
 
 __version__ = "0.0.0+develop"
 

diff --git a/plugins/flytekit-pandera/flytekitplugins/pandera/schema.py b/plugins/flytekit-pandera/flytekitplugins/pandera/schema.py
@@ -11,6 +11,8 @@
 from flytekit.types.schema import FlyteSchema, PandasSchemaWriter, SchemaFormat, SchemaOpenMode
 from flytekit.types.schema.types import FlyteSchemaTransformer
 
+T = typing.TypeVar("T")
+
 
 class PanderaTransformer(TypeTransformer[pandera.typing.DataFrame]):
     _SUPPORTED_TYPES: typing.Dict[
@@ -53,6 +55,10 @@ def _get_schema_type(self, t: Type[pandera.typing.DataFrame]) -> SchemaType:
     def get_literal_type(self, t: Type[pandera.typing.DataFrame]) -> LiteralType:
         return LiteralType(schema=self._get_schema_type(t))
 
+    def assert_type(self, t: Type[T], v: T):
+        if not hasattr(t, "__origin__") and not isinstance(v, (t, pandas.DataFrame)):
+            raise TypeError(f"Type of Val '{v}' is not an instance of {t}")
+
     def to_literal(
         self,
         ctx: FlyteContext,

diff --git a/plugins/flytekit-pandera/tests/test_plugin.py b/plugins/flytekit-pandera/tests/test_plugin.py
@@ -42,7 +42,7 @@ def my_wf() -> pandera.typing.DataFrame[OutSchema]:
     # raise error when defining workflow using invalid data
     invalid_df = pandas.DataFrame({"col1": [1, 2, 3], "col2": list("abc")})
 
-    with pytest.raises(pandera.errors.SchemaError):
+    with pytest.raises(AssertionError):
 
         @workflow
         def invalid_wf() -> pandera.typing.DataFrame[OutSchema]:
@@ -65,7 +65,7 @@ def transform2_noop(df: pandera.typing.DataFrame[IntermediateSchema]) -> pandera
     def wf_invalid_output(df: pandera.typing.DataFrame[InSchema]) -> pandera.typing.DataFrame[OutSchema]:
         return transform2_noop(df=transform1(df=df))
 
-    with pytest.raises(pandera.errors.SchemaError, match="column 'col4' not in dataframe"):
+    with pytest.raises(TypeError, match="^Failed to convert return value"):
         wf_invalid_output(df=valid_df)
 
 

diff --git a/plugins/flytekit-papermill/setup.py b/plugins/flytekit-papermill/setup.py
@@ -4,7 +4,13 @@
 
 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
 
-plugin_requires = ["flytekit>=0.16.0b0,<1.0.0", "papermill>=1.2.0", "nbconvert>=6.0.7", "ipykernel>=5.0.0"]
+plugin_requires = [
+    "flytekit>=0.16.0b0,<1.0.0",
+    "flytekitplugins-spark>=0.16.0b0,<1.0.0",
+    "papermill>=1.2.0",
+    "nbconvert>=6.0.7",
+    "ipykernel>=5.0.0",
+]
 
 __version__ = "0.0.0+develop"
 

diff --git a/plugins/flytekit-papermill/tests/test_task.py b/plugins/flytekit-papermill/tests/test_task.py
@@ -5,7 +5,8 @@
 
 from flytekit import kwtypes
 from flytekit.types.file import PythonNotebook
-from plugins.tests.papermilltests.testdata.datatype import X
+
+from .testdata.datatype import X
 
 
 def _get_nb_path(name: str, suffix: str = "", abs: bool = True, ext: str = ".ipynb") -> str:

diff --git a/plugins/flytekit-papermill/tests/testdata/nb-complex.ipynb b/plugins/flytekit-papermill/tests/testdata/nb-complex.ipynb
@@ -3,63 +3,104 @@
   {
    "cell_type": "code",
    "execution_count": 12,
-   "metadata": {
-    "tags": [
-     "parameters"
-    ]
-   },
-   "outputs": [],
    "source": [
     "h = \"hello\"\n",
     "w = \"./nb-simple.ipynb\"\n",
     "n = 5"
-   ]
+   ],
+   "outputs": [],
+   "metadata": {
+    "tags": [
+     "parameters"
+    ]
+   }
   },
   {
    "cell_type": "code",
    "execution_count": 13,
-   "metadata": {},
-   "outputs": [],
    "source": [
     "# We're importing from the root folder because tests aren't released as part of the package, so there's nothing to\n",
-    "# import from in the actual package.\n",
-    "from plugins.tests.papermilltests.testdata.datatype import X\n",
+    "# import from in the actual package. Assume that we're in the root of the papermill plugin project.\n",
+    "from tests.testdata.datatype import X\n",
     "\n",
     "h = h + \" world!\"\n",
     "x = X(x=n)"
-   ]
+   ],
+   "outputs": [],
+   "metadata": {}
   },
   {
    "cell_type": "code",
    "execution_count": 14,
-   "metadata": {
-    "tags": [
-     "outputs"
-    ]
-   },
+   "source": [
+    "from flytekitplugins.papermill  import record_outputs\n",
+    "from flytekit.types.file import PythonNotebook\n",
+    "\n",
+    "record_outputs(h=h, w=PythonNotebook(w), x=x)"
+   ],
    "outputs": [
     {
+     "output_type": "execute_result",
      "data": {
-      "text/plain": "literals {\n  key: \"h\"\n  value {\n    scalar {\n      primitive {\n        string_value: \"hello world!\"\n      }\n    }\n  }\n}\nliterals {\n  key: \"w\"\n  value {\n    scalar {\n      blob {\n        metadata {\n          type {\n            format: \"ipynb\"\n          }\n        }\n        uri: \"/tmp/flyte/20210112_113703/mock_remote/b96d99e5767adf99ea71fa4d6f3374b8/nb-simple.ipynb\"\n      }\n    }\n  }\n}\nliterals {\n  key: \"x\"\n  value {\n    scalar {\n      generic {\n        fields {\n          key: \"x\"\n          value {\n            number_value: 5.0\n          }\n        }\n      }\n    }\n  }\n}"
+      "text/plain": [
+       "literals {\n",
+       "  key: \"h\"\n",
+       "  value {\n",
+       "    scalar {\n",
+       "      primitive {\n",
+       "        string_value: \"hello world!\"\n",
+       "      }\n",
+       "    }\n",
+       "  }\n",
+       "}\n",
+       "literals {\n",
+       "  key: \"w\"\n",
+       "  value {\n",
+       "    scalar {\n",
+       "      blob {\n",
+       "        metadata {\n",
+       "          type {\n",
+       "            format: \"ipynb\"\n",
+       "          }\n",
+       "        }\n",
+       "        uri: \"/tmp/flyte/20210112_113703/mock_remote/b96d99e5767adf99ea71fa4d6f3374b8/nb-simple.ipynb\"\n",
+       "      }\n",
+       "    }\n",
+       "  }\n",
+       "}\n",
+       "literals {\n",
+       "  key: \"x\"\n",
+       "  value {\n",
+       "    scalar {\n",
+       "      generic {\n",
+       "        fields {\n",
+       "          key: \"x\"\n",
+       "          value {\n",
+       "            number_value: 5.0\n",
+       "          }\n",
+       "        }\n",
+       "      }\n",
+       "    }\n",
+       "  }\n",
+       "}"
+      ]
      },
-     "execution_count": 14,
      "metadata": {},
-     "output_type": "execute_result"
+     "execution_count": 14
     }
    ],
-   "source": [
-    "from flytekitplugins.papermill  import record_outputs\n",
-    "from flytekit.types.file import PythonNotebook\n",
-    "\n",
-    "record_outputs(h=h, w=PythonNotebook(w), x=x)"
-   ]
+   "metadata": {
+    "tags": [
+     "outputs"
+    ]
+   }
   },
   {
    "cell_type": "code",
    "execution_count": 14,
-   "metadata": {},
+   "source": [],
    "outputs": [],
-   "source": []
+   "metadata": {}
   }
  ],
  "metadata": {

diff --git a/plugins/flytekit-sqlalchemy/tests/test_sql_tracker.py b/plugins/flytekit-sqlalchemy/tests/test_sql_tracker.py
@@ -3,7 +3,8 @@
 from flytekit.common.translator import get_serializable
 from flytekit.core import context_manager
 from flytekit.core.context_manager import Image, ImageConfig
-from plugins.tests.sqlalchemy.test_task import tk as not_tk
+
+from .test_task import tk as not_tk
 
 
 def test_sql_lhs():

diff --git a/plugins/setup.py b/plugins/setup.py
@@ -7,19 +7,20 @@
 
 PACKAGE_NAME = "flytekitplugins-parent"
 SOURCES = {
+    "flytekitplugins-athena": "flytekit-aws-athena",
+    "flytekitplugins-awssagemaker": "flytekit-aws-sagemaker",
+    "flytekitplugins-fsspec": "flytekit-data-fsspec",
+    "flytekitplugins-dolt": "flytekit-dolt",
+    "flytekitplugins-great_expectations": "flytekit-greatexpectations",
     "flytekitplugins-hive": "flytekit-hive",
-    "flytekitplugins-papermill": "flytekit-papermill",
-    "flytekitplugins-spark": "flytekit-spark",
     "flytekitplugins-pod": "flytekit-k8s-pod",
     "flytekitplugins-kfpytorch": "flytekit-kf-pytorch",
-    "flytekitplugins-awssagemaker": "flytekit-aws-sagemaker",
     "flytekitplugins-kftensorflow": "flytekit-kf-tensorflow",
     "flytekitplugins-pandera": "flytekit-pandera",
-    "flytekitplugins-dolt": "flytekit-dolt",
-    "flytekitplugins-sqlalchemy": "flytekit-sqlalchemy",
-    "flytekitplugins-athena": "flytekit-aws-athena",
-    "flytekitplugins-great_expectations": "flytekit-greatexpectations",
+    "flytekitplugins-papermill": "flytekit-papermill",
     "flytekitplugins-snowflake": "flytekit-snowflake",
+    "flytekitplugins-spark": "flytekit-spark",
+    "flytekitplugins-sqlalchemy": "flytekit-sqlalchemy",
 }
 
 
@@ -64,7 +65,7 @@ def run(self):
     version="0.1.0",
     author="flyteorg",
     author_email="[email protected]",
-    description="This is a fake package to help install all the plugins",
+    description="This is a microlib package to help install all the plugins",
     license="apache2",
     classifiers=["Private :: Do Not Upload to pypi server"],
     install_requires=[],