Annotations - separate out plugins into top level folder

flytekit/annotated/type_engine.py

@@ -198,7 +198,7 @@ def register(cls, transformer: TypeTransformer):
  if transformer.python_type in cls._REGISTRY:
  existing = cls._REGISTRY[transformer.python_type]
  raise ValueError(
- f"Transformer f{existing.name} for type{transformer.python_type} is already registered."
+ f"Transformer {existing.name} for type {transformer.python_type} is already registered."
  f" Cannot override with {transformer.name}"
  )
  cls._REGISTRY[transformer.python_type] = transformer

plugins/README.md

@@ -0,0 +1,3 @@
+
+Read for background:
+https://medium.com/@jherreras/python-microlibs-5be9461ad979

plugins/hive/flytekitplugins/hive/__init__.py

@@ -0,0 +1 @@
+from .task import HiveSelectTask, HiveTask, HiveConfig

plugins/hive/setup.py

@@ -0,0 +1,32 @@
+from setuptools import setup
+
+PLUGIN_NAME = "hive"
+
+microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
+
+plugin_requires = ["flytekit==0.16.0a2", "hmsclient>=0.0.1,<1.0.0"]
+
+setup(
+ name=microlib_name,
+ version="0.1.0",
+ author="flyteorg",
+ author_email="admin@flyte.org",
+ description="Your microlib descriton",
+ namespace_packages=["flytekitplugins"],
+ packages=[f"flytekitplugins.{PLUGIN_NAME}"],
+ install_requires=plugin_requires,
+ license="apache2",
+ python_requires=">=3.7",
+ classifiers=[
+ "Intended Audience :: Science/Research",
+ "Intended Audience :: Developers",
+ "License :: OSI Approved :: Apache Software License",
+ "Programming Language :: Python :: 3.7",
+ "Programming Language :: Python :: 3.8",
+ "Topic :: Scientific/Engineering",
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
+ "Topic :: Software Development",
+ "Topic :: Software Development :: Libraries",
+ "Topic :: Software Development :: Libraries :: Python Modules",
+ ],
+)

tests/flytekit/unit/taskplugins/hive/test_hive_task.py → plugins/hive/tests/test_hive_task.py

@@ -1,13 +1,13 @@
 import pandas
 import pytest
+from flytekitplugins.hive.task import HiveConfig, HiveSelectTask, HiveTask
 
 from flytekit.annotated import context_manager
 from flytekit.annotated.base_task import kwtypes
 from flytekit.annotated.context_manager import Image, ImageConfig
 from flytekit.annotated.testing import task_mock
 from flytekit.annotated.workflow import workflow
 from flytekit.common.translator import get_serializable
-from flytekit.taskplugins.hive.task import HiveConfig, HiveSelectTask, HiveTask
 from flytekit.types.schema import FlyteSchema
 
 

flytekit/taskplugins/notebook/task.py → plugins/notebook/flytekitplugins/notebook/task.py

@@ -47,7 +47,7 @@ class NotebookTask(PythonInstanceTask[T]):
 
  ...
  # cell begin
- from flytekit.taskplugins.notebook import record_outputs
+ from flytekit.plugins.notebook import record_outputs
 
  record_outputs(x=val_x, y=val_y)
  #cell end

plugins/notebook/setup.py

@@ -0,0 +1,32 @@
+from setuptools import setup
+
+PLUGIN_NAME = "notebook"
+
+microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
+
+plugin_requires = ["flytekit==0.16.0a2", "papermill>=1.2.0", "nbconvert>=6.0.7", "ipykernel>=5.0.0"]
+
+setup(
+ name=microlib_name,
+ version="0.1.0",
+ author="flyteorg",
+ author_email="admin@flyte.org",
+ description="Your microlib descriton",
+ namespace_packages=["flytekitplugins"],
+ packages=[f"flytekitplugins.{PLUGIN_NAME}"],
+ install_requires=plugin_requires,
+ license="apache2",
+ python_requires=">=3.7",
+ classifiers=[
+ "Intended Audience :: Science/Research",
+ "Intended Audience :: Developers",
+ "License :: OSI Approved :: Apache Software License",
+ "Programming Language :: Python :: 3.7",
+ "Programming Language :: Python :: 3.8",
+ "Topic :: Scientific/Engineering",
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
+ "Topic :: Software Development",
+ "Topic :: Software Development :: Libraries",
+ "Topic :: Software Development :: Libraries :: Python Modules",
+ ],
+)

tests/flytekit/unit/taskplugins/notebook/test_spark_notebook.py → plugins/notebook/tests/test_spark_notebook.py

@@ -1,8 +1,9 @@
 import os
 
+from plugins.notebook.flytekitplugins.notebook import NotebookTask
+from plugins.spark.flytekitplugins.spark import Spark
+
 from flytekit import kwtypes
-from flytekit.taskplugins.notebook import NotebookTask
-from flytekit.taskplugins.spark import Spark
 from flytekit.types.schema import FlyteSchema
 
 

tests/flytekit/unit/taskplugins/notebook/test_task.py → plugins/notebook/tests/test_task.py

@@ -1,10 +1,11 @@
 import datetime
 import os
 
+from plugins.notebook.flytekitplugins.notebook import NotebookTask
+
 from flytekit import kwtypes
-from flytekit.taskplugins.notebook import NotebookTask
 from flytekit.types.file import PythonNotebook
-from tests.flytekit.unit.taskplugins.notebook.testdata.datatype import X
+from .testdata.datatype import X
 
 
 def _get_nb_path(name: str, suffix: str = "", abs: bool = True, ext: str = ".ipynb") -> str:

tests/flytekit/unit/taskplugins/notebook/testdata/nb-complex.ipynb → plugins/notebook/tests/testdata/nb-complex.ipynb

@@ -21,7 +21,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "from tests.flytekit.unit.taskplugins.notebook.testdata.datatype import X\n",
+ "from plugins.notebook.tests.testdata.datatype import X\n",
  "\n",
  "h = h + \" world!\"\n",
  "x = X(x=n)"
@@ -46,7 +46,7 @@
  }
  ],
  "source": [
- "from flytekit.taskplugins.notebook import record_outputs\n",
+ "from flytekitplugins.notebook import record_outputs\n",
  "from flytekit.types.file import PythonNotebook\n",
  "\n",
  "record_outputs(h=h, w=PythonNotebook(w), x=x)"

tests/flytekit/unit/taskplugins/notebook/testdata/nb-multi.ipynb → plugins/notebook/tests/testdata/nb-multi.ipynb

@@ -57,7 +57,7 @@
  }
  ],
  "source": [
- "from flytekit.taskplugins.notebook import record_outputs\n",
+ "from flytekitplugins.notebook import record_outputs\n",
  "\n",
  "record_outputs(z=z, m=m, h=h, n=n)"
  ]

tests/flytekit/unit/taskplugins/notebook/testdata/nb-simple.ipynb → plugins/notebook/tests/testdata/nb-simple.ipynb

@@ -33,7 +33,7 @@
  },
  "outputs": [],
  "source": [
- "from flytekit.taskplugins.notebook import record_outputs\n",
+ "from flytekitplugins.notebook import record_outputs\n",
  "\n",
  "record_outputs(square=out)"
  ]

tests/flytekit/unit/taskplugins/notebook/testdata/nb-spark.ipynb → plugins/notebook/tests/testdata/nb-spark.ipynb

@@ -19,7 +19,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "from flytekit.taskplugins.spark import new_spark_session\n",
+ "from flytekitplugins.spark import new_spark_session\n",
  "session = new_spark_session(name=\"notebook\", conf=local_spark_conf)\n",
  "new_df = session.createDataFrame([(\"Bob\", 10)], [\"name\", \"age\"])"
  ]
@@ -43,7 +43,7 @@
  }
  ],
  "source": [
- "from flytekit.taskplugins.notebook import record_outputs\n",
+ "from flytekitplugins.notebook import record_outputs\n",
  "\n",
  "record_outputs(df=new_df)"
  ]

plugins/pod/setup.py

@@ -0,0 +1,32 @@
+from setuptools import setup
+
+PLUGIN_NAME = "pod"
+
+microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
+
+plugin_requires = ["flytekit==0.16.0a2"]
+
+setup(
+ name=microlib_name,
+ version="0.1.0",
+ author="flyteorg",
+ author_email="admin@flyte.org",
+ description="Your microlib descriton",
+ namespace_packages=["flytekitplugins"],
+ packages=[f"flytekitplugins.{PLUGIN_NAME}"],
+ install_requires=plugin_requires,
+ license="apache2",
+ python_requires=">=3.7",
+ classifiers=[
+ "Intended Audience :: Science/Research",
+ "Intended Audience :: Developers",
+ "License :: OSI Approved :: Apache Software License",
+ "Programming Language :: Python :: 3.7",
+ "Programming Language :: Python :: 3.8",
+ "Topic :: Scientific/Engineering",
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
+ "Topic :: Software Development",
+ "Topic :: Software Development :: Libraries",
+ "Topic :: Software Development :: Libraries :: Python Modules",
+ ],
+)

tests/flytekit/unit/taskplugins/pod/test_pod.py → plugins/pod/tests/test_pod.py

@@ -1,13 +1,13 @@
 from typing import List
 
 from k8s.io.api.core.v1 import generated_pb2
+from plugins.pod.flytekitplugins.pod.task import Pod, PodFunctionTask
 
 from flytekit.annotated import context_manager
 from flytekit.annotated.context_manager import ExecutionState, Image, ImageConfig, RegistrationSettings
 from flytekit.annotated.dynamic_workflow_task import dynamic
 from flytekit.annotated.resources import Resources
 from flytekit.annotated.task import task
-from flytekit.taskplugins.pod.task import Pod, PodFunctionTask
 
 
 def get_pod_spec():
@@ -48,7 +48,7 @@ def simple_pod_task(i: int):
  assert primary_container["args"] == [
  "pyflyte-execute",
  "--task-module",
- "tests.flytekit.unit.taskplugins.pod.test_pod",
+ "tests.test_pod",
  "--task-name",
  "simple_pod_task",
  "--inputs",

plugins/pytorch/setup.py

@@ -0,0 +1,32 @@
+from setuptools import setup
+
+PLUGIN_NAME = "pytorch"
+
+microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
+
+plugin_requires = ["flytekit==0.16.0a2"]
+
+setup(
+ name=microlib_name,
+ version="0.1.0",
+ author="flyteorg",
+ author_email="admin@flyte.org",
+ description="Your microlib descriton",
+ namespace_packages=["flytekitplugins"],
+ packages=[f"flytekitplugins.{PLUGIN_NAME}"],
+ install_requires=plugin_requires,
+ license="apache2",
+ python_requires=">=3.7",
+ classifiers=[
+ "Intended Audience :: Science/Research",
+ "Intended Audience :: Developers",
+ "License :: OSI Approved :: Apache Software License",
+ "Programming Language :: Python :: 3.7",
+ "Programming Language :: Python :: 3.8",
+ "Topic :: Scientific/Engineering",
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
+ "Topic :: Software Development",
+ "Topic :: Software Development :: Libraries",
+ "Topic :: Software Development :: Libraries :: Python Modules",
+ ],
+)

tests/flytekit/unit/taskplugins/pytorch/test_pytorch_task.py → plugins/pytorch/tests/test_pytorch_task.py

@@ -1,7 +1,8 @@
+from plugins.pytorch.flytekitplugins.pytorch.task import PyTorch
+
 from flytekit import task
 from flytekit.annotated.context_manager import Image, ImageConfig, RegistrationSettings
 from flytekit.annotated.resources import Resources
-from flytekit.taskplugins.pytorch.task import PyTorch
 
 
 def test_pytorch_task():

flytekit/taskplugins/sagemaker/hpo.py → plugins/sagemaker/flytekitplugins/sagemaker/hpo.py

@@ -6,6 +6,7 @@
 from flyteidl.plugins.sagemaker import parameter_ranges_pb2 as _pb2_params
 from google.protobuf import json_format
 from google.protobuf.json_format import MessageToDict
+from flytekitplugins.sagemaker.training import SagemakerBuiltinAlgorithmsTask, SagemakerCustomTrainingTask
 
 from flytekit import FlyteContext
 from flytekit.annotated.base_task import PythonTask
@@ -17,7 +18,6 @@
 from flytekit.models.sagemaker import parameter_ranges as _params
 from flytekit.models.sagemaker import training_job as _training_job_model
 from flytekit.models.types import LiteralType
-from flytekit.taskplugins.sagemaker.training import SagemakerBuiltinAlgorithmsTask, SagemakerCustomTrainingTask
 
 
 @dataclass

flytekit/taskplugins/sagemaker/training.py → plugins/sagemaker/flytekitplugins/sagemaker/training.py

@@ -4,6 +4,7 @@
 from typing import Any, Callable, Dict, TypeVar
 
 from google.protobuf.json_format import MessageToDict
+from flytekitplugins.sagemaker.distributed_training import DistributedTrainingContext
 
 import flytekit
 from flytekit.annotated.base_task import PythonTask, kwtypes
@@ -13,7 +14,6 @@
 from flytekit.annotated.task import TaskPlugins
 from flytekit.common.tasks.sdk_runnable import ExecutionParameters
 from flytekit.models.sagemaker import training_job as _training_job_models
-from flytekit.taskplugins.sagemaker.distributed_training import DistributedTrainingContext
 from flytekit.types.directory.types import FlyteDirectory
 from flytekit.types.file import FlyteFile
 

plugins/sagemaker/setup.py

@@ -0,0 +1,32 @@
+from setuptools import setup
+
+PLUGIN_NAME = "sagemaker"
+
+microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
+
+plugin_requires = ["flytekit==0.16.0a2", "sagemaker-training>=3.6.2,<4.0.0"]
+
+setup(
+ name=microlib_name,
+ version="0.1.0",
+ author="flyteorg",
+ author_email="admin@flyte.org",
+ description="Your microlib descriton",
+ namespace_packages=["flytekitplugins"],
+ packages=[f"flytekitplugins.{PLUGIN_NAME}"],
+ install_requires=plugin_requires,
+ license="apache2",
+ python_requires=">=3.7",
+ classifiers=[
+ "Intended Audience :: Science/Research",
+ "Intended Audience :: Developers",
+ "License :: OSI Approved :: Apache Software License",
+ "Programming Language :: Python :: 3.7",
+ "Programming Language :: Python :: 3.8",
+ "Topic :: Scientific/Engineering",
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
+ "Topic :: Software Development",
+ "Topic :: Software Development :: Libraries",
+ "Topic :: Software Development :: Libraries :: Python Modules",
+ ],
+)

tests/flytekit/unit/taskplugins/sagemaker/test_hpo.py → plugins/sagemaker/tests/test_hpo.py

@@ -1,24 +1,26 @@
 import pytest
+from flytekitplugins.sagemaker.hpo import HPOJob, SagemakerHPOTask
+from flytekitplugins.sagemaker.hpo import HPOTuningJobConfigTransformer, ParameterRangesTransformer
+from flytekitplugins.sagemaker.training import SagemakerBuiltinAlgorithmsTask, SagemakerTrainingJobConfig
+from plugins.sagemaker.tests.test_training import _get_reg_settings
 
 from flytekit import FlyteContext
 from flytekit.common.types.primitives import Generic
-from flytekit.taskplugins.sagemaker import (
- AlgorithmName,
- AlgorithmSpecification,
- HPOJob,
+from flytekit.models.sagemaker.hpo_job import (
  HyperparameterTuningJobConfig,
  HyperparameterTuningObjective,
  HyperparameterTuningObjectiveType,
+ TrainingJobEarlyStoppingType,
+)
+from flytekit.models.sagemaker.parameter_ranges import (
  IntegerParameterRange,
  ParameterRangeOneOf,
- SagemakerBuiltinAlgorithmsTask,
- SagemakerHPOTask,
- SagemakerTrainingJobConfig,
- TrainingJobEarlyStoppingType,
+)
+from flytekit.models.sagemaker.training_job import (
+ AlgorithmName,
+ AlgorithmSpecification,
  TrainingJobResourceConfig,
 )
-from flytekit.taskplugins.sagemaker.hpo import HPOTuningJobConfigTransformer, ParameterRangesTransformer
-from tests.flytekit.unit.taskplugins.sagemaker.test_training import _get_reg_settings
 
 
 def test_hpo_for_builtin():