Annotations - separate out plugins into top level folder

Makefile

@@ -35,6 +35,7 @@ lint: ## Run linters
 test: lint ## Run tests
  pytest tests/flytekit/unit -k "not test_task"
  pytest tests/scripts
+ pytest plugins/tests
  shellcheck **/*.sh
 
 requirements-spark3.txt: export CUSTOM_COMPILE_COMMAND := make requirements-spark3.txt
@@ -56,4 +57,4 @@ requirements: requirements.txt dev-requirements.txt requirements-spark3.txt ## C
 .PHONY: coverage
 coverage:
  coverage run -m pytest tests/flytekit/unit/annotated
- coverage report -m --include="flytekit/annotated/*,flytekit/types/*,flytekit/taskplugins*"
+ coverage report -m --include="flytekit/annotated/*,flytekit/types/*"

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

flytekit/taskplugins/sqlite3/task.py → flytekit/extras/sqlite3/task.py

@@ -6,10 +6,11 @@
 import typing
 from dataclasses import dataclass
 
+import pandas as pd
+
 from flytekit import FlyteContext, kwtypes
 from flytekit.annotated.base_sql_task import SQLTask
 from flytekit.annotated.python_function_task import PythonInstanceTask
-from flytekit.plugins import pandas as pd
 from flytekit.types.schema import FlyteSchema
 
 

flytekit/plugins/__init__.py

@@ -1,3 +1,7 @@
+"""
+This file is for old style plugins - for new plugins that work with the Python native-typed Flytekit, please
+refer to the plugin specific directory underneath the plugins folder at the top level of this repository.
+"""
 from flytekit.tools import lazy_loader as _lazy_loader
 
 pyspark = _lazy_loader.lazy_load_module("pyspark") # type: _lazy_loader._LazyLoadModule

plugins/Makefile

@@ -0,0 +1,7 @@
+.PHONY: test
+test:
+ pytest tests
+
+.PHONY: install-all-dev
+install-all-dev:
+ pip install -e .

plugins/README.md

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

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

@@ -4,6 +4,7 @@
 
 from flyteidl.plugins.sagemaker import hyperparameter_tuning_job_pb2 as _pb2_hpo_job
 from flyteidl.plugins.sagemaker import parameter_ranges_pb2 as _pb2_params
+from flytekitplugins.aws.training import SagemakerBuiltinAlgorithmsTask, SagemakerCustomTrainingTask
 from google.protobuf import json_format
 from google.protobuf.json_format import MessageToDict
 
@@ -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/aws/flytekitplugins/aws/training.py

@@ -3,6 +3,7 @@
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, TypeVar
 
+from flytekitplugins.aws.distributed_training import DistributedTrainingContext
 from google.protobuf.json_format import MessageToDict
 
 import flytekit
@@ -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/aws/setup.py

@@ -0,0 +1,32 @@
+from setuptools import setup
+
+PLUGIN_NAME = "aws"
+
+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="AWS Plugins for flytekit",
+ 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",
+ ],
+)

plugins/hive/flytekitplugins/hive/__init__.py

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

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="This package holds Hive plugins for flytekit",
+ 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",
+ ],
+)

plugins/kfpytorch/setup.py

@@ -0,0 +1,32 @@
+from setuptools import setup
+
+PLUGIN_NAME = "kfpytorch"
+
+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="K8s based Pytorch plugin for Flytekit",
+ 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/kfpytorch/tests/test_pytorch_task.py

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

plugins/kftensorflow/setup.py

@@ -0,0 +1,33 @@
+from setuptools import setup
+
+PLUGIN_NAME = "kftensorflow"
+
+microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
+
+# TODO: Requirements are missing, add them back in later.
+plugin_requires = ["flytekit>=0.16.0a2"]
+
+setup(
+ name=microlib_name,
+ version="0.1.0",
+ author="flyteorg",
+ author_email="admin@flyte.org",
+ description="K8s based Tensorflow plugin for flytekit",
+ 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",
+ ],
+)

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

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

plugins/papermill/setup.py

@@ -0,0 +1,32 @@
+from setuptools import setup
+
+PLUGIN_NAME = "papermill"
+
+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="This is the flytekit papermill plugin",
+ 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",
+ ],
+)

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

@@ -2,18 +2,19 @@
 
 from flyteidl.core import tasks_pb2 as _core_task
 from google.protobuf.json_format import MessageToDict
+from k8s.io.api.core.v1.generated_pb2 import Container, EnvVar, PodSpec, ResourceRequirements
+from k8s.io.apimachinery.pkg.api.resource.generated_pb2 import Quantity
 
 from flytekit.annotated.context_manager import FlyteContext, SerializationSettings
 from flytekit.annotated.promise import Promise
 from flytekit.annotated.python_function_task import PythonFunctionTask
 from flytekit.annotated.task import TaskPlugins
 from flytekit.common.exceptions import user as _user_exceptions
 from flytekit.models import task as _task_models
-from flytekit.plugins import k8s as _lazy_k8s
 
 
 class Pod(object):
- def __init__(self, pod_spec: _lazy_k8s.io.api.core.v1.generated_pb2.PodSpec, primary_container_name: str):
+ def __init__(self, pod_spec: PodSpec, primary_container_name: str):
  if not pod_spec:
  raise _user_exceptions.FlyteValidationException("A pod spec cannot be undefined")
  if not primary_container_name:
@@ -23,7 +24,7 @@ def __init__(self, pod_spec: _lazy_k8s.io.api.core.v1.generated_pb2.PodSpec, pri
  self._primary_container_name = primary_container_name
 
  @property
- def pod_spec(self) -> _lazy_k8s.io.api.core.v1.generated_pb2.PodSpec:
+ def pod_spec(self) -> PodSpec:
  return self._pod_spec
 
  @property
@@ -46,9 +47,7 @@ def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
  break
  if not primary_exists:
  # insert a placeholder primary container if it is not defined in the pod spec.
- containers.extend(
- [_lazy_k8s.io.api.core.v1.generated_pb2.Container(name=self.task_config.primary_container_name)]
- )
+ containers.extend([Container(name=self.task_config.primary_container_name)])
 
  final_containers = []
  for container in containers:
@@ -65,26 +64,21 @@ def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
  del container.args[:]
  container.args.extend(sdk_default_container.args)
 
- resource_requirements = _lazy_k8s.io.api.core.v1.generated_pb2.ResourceRequirements()
+ resource_requirements = ResourceRequirements()
  for resource in sdk_default_container.resources.limits:
  resource_requirements.limits[
  _core_task.Resources.ResourceName.Name(resource.name).lower()
- ].CopyFrom(_lazy_k8s.io.apimachinery.pkg.api.resource.generated_pb2.Quantity(string=resource.value))
+ ].CopyFrom(Quantity(string=resource.value))
  for resource in sdk_default_container.resources.requests:
  resource_requirements.requests[
  _core_task.Resources.ResourceName.Name(resource.name).lower()
- ].CopyFrom(_lazy_k8s.io.apimachinery.pkg.api.resource.generated_pb2.Quantity(string=resource.value))
+ ].CopyFrom(Quantity(string=resource.value))
  if resource_requirements.ByteSize():
  # Important! Only copy over resource requirements if they are non-empty.
  container.resources.CopyFrom(resource_requirements)
 
  del container.env[:]
- container.env.extend(
- [
- _lazy_k8s.io.api.core.v1.generated_pb2.EnvVar(name=key, value=val)
- for key, val in sdk_default_container.env.items()
- ]
- )
+ container.env.extend([EnvVar(name=key, value=val) for key, val in sdk_default_container.env.items()])
 
  final_containers.append(container)
 

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", "k8s-proto>=0.0.3,<1.0.0"]
+
+setup(
+ name=microlib_name,
+ version="0.1.0",
+ author="flyteorg",
+ author_email="admin@flyte.org",
+ description="Flytekit plugin to support K8s Pod tasks",
+ 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",
+ ],
+)