Add support for the new Presto task (#89)

* AVSW-41872: Extend the Flyte SDK with a handler for the new task

* Changs to presto object type

* coding session 1

* remaining changes & cleanup

* PR feedback

* prefix implicits with __

* small stuff

* move the schema_type call into constructor?

* override add_inputs

* reverting schema feedback

* bump version

* changes

Co-authored-by: lu4nm3 <lmedina@lyft.com>
Co-authored-by: Luis Medina <3936213+lu4nm3@users.noreply.github.com>
Co-authored-by: Yee Hing Tong <ytong@lyft.com>

Author: 4 people

flytekit/__init__.py

@@ -1,4 +1,4 @@
 from __future__ import absolute_import
 import flytekit.plugins
 
-__version__ = '0.6.0b5'
+__version__ = '0.6.0b6'

flytekit/common/constants.py

@@ -20,7 +20,7 @@ class SdkTaskType(object):
     HIVE_JOB = "hive"
     SIDECAR_TASK = "sidecar"
     SENSOR_TASK = "sensor-task"
-
+    PRESTO_TASK = "presto"
 
 GLOBAL_INPUT_NODE_ID = ''
 

flytekit/common/tasks/presto_task.py

@@ -0,0 +1,158 @@
+from __future__ import absolute_import
+
+from google.protobuf.json_format import MessageToDict as _MessageToDict
+from flytekit import __version__
+
+from flytekit.common import constants as _constants
+from flytekit.common.tasks import task as _base_task
+from flytekit.models import (
+    interface as _interface_model
+)
+from flytekit.models import literals as _literals, types as _types, \
+    task as _task_model
+
+from flytekit.common import interface as _interface
+import datetime as _datetime
+from flytekit.models import presto as _presto_models
+from flytekit.common.exceptions.user import \
+    FlyteValueException as _FlyteValueException
+from flytekit.common.exceptions import scopes as _exception_scopes
+
+
+class SdkPrestoTask(_base_task.SdkTask):
+    """
+    This class includes the logic for building a task that executes as a Presto task.
+    """
+
+    def __init__(
+            self,
+            statement,
+            output_schema,
+            routing_group=None,
+            catalog=None,
+            schema=None,
+            task_inputs=None,
+            interruptible=False,
+            discoverable=False,
+            discovery_version=None,
+            retries=1,
+            timeout=None,
+    ):
+        """
+        :param Text statement: Presto query specification
+        :param flytekit.common.types.schema.Schema output_schema: Schema that represents that data queried from Presto
+        :param Text routing_group: The routing group that a Presto query should be sent to for the given environment
+        :param Text catalog: The catalog to set for the given Presto query
+        :param Text schema: The schema to set for the given Presto query
+        :param dict[Text,flytekit.common.types.base_sdk_types.FlyteSdkType] task_inputs: Optional inputs to the Presto task
+        :param bool discoverable:
+        :param Text discovery_version: String describing the version for task discovery purposes
+        :param int retries: Number of retries to attempt
+        :param datetime.timedelta timeout:
+        """
+
+        # Set as class fields which are used down below to configure implicit
+        # parameters
+        self._routing_group = routing_group or ""
+        self._catalog = catalog or ""
+        self._schema = schema or ""
+
+        metadata = _task_model.TaskMetadata(
+            discoverable,
+            # This needs to have the proper version reflected in it
+            _task_model.RuntimeMetadata(
+                _task_model.RuntimeMetadata.RuntimeType.FLYTE_SDK, __version__,
+                "python"),
+            timeout or _datetime.timedelta(seconds=0),
+            _literals.RetryStrategy(retries),
+            interruptible,
+            discovery_version,
+            "This is deprecated!"
+        )
+
+        presto_query = _presto_models.PrestoQuery(
+            routing_group=routing_group or "",
+            catalog=catalog or "",
+            schema=schema or "",
+            statement=statement
+        )
+
+        # Here we set the routing_group, catalog, and schema as implicit
+        # parameters for caching purposes
+        i = _interface.TypedInterface(
+            {
+                "__implicit_routing_group": _interface_model.Variable(
+                    type=_types.LiteralType(simple=_types.SimpleType.STRING),
+                    description="The routing group set as an implicit input"
+                ),
+                "__implicit_catalog": _interface_model.Variable(
+                    type=_types.LiteralType(simple=_types.SimpleType.STRING),
+                    description="The catalog set as an implicit input"
+                ),
+                "__implicit_schema": _interface_model.Variable(
+                    type=_types.LiteralType(simple=_types.SimpleType.STRING),
+                    description="The schema set as an implicit input"
+                )
+            },
+            {
+                # Set the schema for the Presto query as an output
+                "results": _interface_model.Variable(
+                    type=_types.LiteralType(schema=output_schema.schema_type),
+                    description="The schema for the Presto query"
+                )
+            })
+
+        super(SdkPrestoTask, self).__init__(
+            _constants.SdkTaskType.PRESTO_TASK,
+            metadata,
+            i,
+            _MessageToDict(presto_query.to_flyte_idl()),
+        )
+
+        # Set user provided inputs
+        task_inputs(self)
+
+    # Override method in order to set the implicit inputs
+    def __call__(self, *args, **kwargs):
+        kwargs["__implicit_routing_group"] = self.routing_group
+        kwargs["__implicit_catalog"] = self.catalog
+        kwargs["__implicit_schema"] = self.schema
+
+        return super(SdkPrestoTask, self).__call__(
+            *args, **kwargs
+        )
+
+    @_exception_scopes.system_entry_point
+    def add_inputs(self, inputs):
+        """
+        Adds the inputs to this task.  This can be called multiple times, but it will fail if an input with a given
+        name is added more than once, a name collides with an output, or if the name doesn't exist as an arg name in
+        the wrapped function.
+        :param dict[Text, flytekit.models.interface.Variable] inputs: names and variables
+        """
+        self._validate_inputs(inputs)
+        self.interface.inputs.update(inputs)
+
+    @property
+    def routing_group(self):
+        """
+        The routing group that a Presto query should be sent to for the given environment
+        :rtype: Text
+        """
+        return self._routing_group
+
+    @property
+    def catalog(self):
+        """
+        The catalog to set for the given Presto query
+        :rtype: Text
+        """
+        return self._catalog
+
+    @property
+    def schema(self):
+        """
+        The schema to set for the given Presto query
+        :rtype: Text
+        """
+        return self._schema

flytekit/common/tasks/sdk_runnable.py

@@ -235,6 +235,17 @@ def __init__(
     _banned_inputs = {}
     _banned_outputs = {}
 
+    @_exception_scopes.system_entry_point
+    def add_inputs(self, inputs):
+        """
+        Adds the inputs to this task.  This can be called multiple times, but it will fail if an input with a given
+        name is added more than once, a name collides with an output, or if the name doesn't exist as an arg name in
+        the wrapped function.
+        :param dict[Text, flytekit.models.interface.Variable] inputs: names and variables
+        """
+        self._validate_inputs(inputs)
+        self.interface.inputs.update(inputs)
+        
     @classmethod
     def promote_from_model(cls, base_model):
         # TODO: If the task exists in this container, we should be able to retrieve it.

flytekit/common/tasks/task.py

@@ -6,12 +6,13 @@
 
 from flytekit.common import interface as _interfaces, nodes as _nodes, sdk_bases as _sdk_bases
 from flytekit.common.core import identifier as _identifier
-from flytekit.common.exceptions import user as _user_exceptions, scopes as _exception_scopes
+from flytekit.common.exceptions import scopes as _exception_scopes
 from flytekit.common.mixins import registerable as _registerable, hash as _hash_mixin
 from flytekit.configuration import internal as _internal_config
 from flytekit.engines import loader as _engine_loader
 from flytekit.models import common as _common_model, task as _task_model
 from flytekit.models.core import workflow as _workflow_model, identifier as _identifier_model
+from flytekit.common.exceptions import user as _user_exceptions
 
 
 class SdkTask(
@@ -189,14 +190,7 @@ def validate(self):
 
     @_exception_scopes.system_entry_point
     def add_inputs(self, inputs):
-        """
-        Adds the inputs to this task.  This can be called multiple times, but it will fail if an input with a given
-        name is added more than once, a name collides with an output, or if the name doesn't exist as an arg name in
-        the wrapped function.
-        :param dict[Text, flytekit.models.interface.Variable] inputs: names and variables
-        """
-        self._validate_inputs(inputs)
-        self.interface.inputs.update(inputs)
+        raise _user_exceptions.FlyteUserException("You can not add inputs to this task")
 
     @_exception_scopes.system_entry_point
     def add_outputs(self, outputs):

flytekit/models/presto.py

@@ -0,0 +1,77 @@
+from __future__ import absolute_import
+
+## Todo - change this to qubole_presto once Luis's PR get's merged
+# from flyteidl.plugins import qubole_presto as _qubole
+from flyteidl.plugins import presto_pb2 as _presto
+
+from flytekit.models import common as _common
+
+
+class PrestoQuery(_common.FlyteIdlEntity):
+    def __init__(self, routing_group, catalog, schema, statement):
+        """
+        Initializes a new PrestoQuery.
+
+        :param string routing_group:
+        :param string catalog:
+        :param string schema:
+        :param string statement:
+
+        """
+        self._routing_group = routing_group
+        self._catalog = catalog
+        self._schema = schema
+        self._statement = statement
+
+    @property
+    def routing_group(self):
+        """
+        The query string.
+        :rtype: str
+        """
+        return self._routing_group
+
+    @property
+    def catalog(self):
+        """
+        :rtype: int
+        """
+        return self._catalog
+
+    @property
+    def schema(self):
+        """
+        :rtype: int
+        """
+        return self._schema
+
+    @property
+    def statement(self):
+        """
+        :rtype: int
+        """
+        return self._statement
+
+    def to_flyte_idl(self):
+        """
+        :rtype: _presto.PrestoQuery 
+        """
+        return _presto.PrestoQuery(
+            routing_group=self._routing_group,
+            catalog=self._catalog,
+            schema=self._schema,
+            statement=self._statement
+        )
+
+    @classmethod
+    def from_flyte_idl(cls, pb2_object):
+        """
+        :param _presto.PrestoQuery pb2_object:
+        :return: PrestoQuery
+        """
+        return cls(
+            routing_group=pb2_object.routing_group,
+            catalog=pb2_object.catalog,
+            schema=pb2_object.schema,
+            statement=pb2_object.statement
+        )

flytekit/sdk/tasks.py

@@ -7,6 +7,7 @@
 from flytekit.common.exceptions import user as _user_exceptions
 from flytekit.common.tasks import sdk_runnable as _sdk_runnable_tasks, sdk_dynamic as _sdk_dynamic, \
     spark_task as _sdk_spark_tasks, hive_task as _sdk_hive_tasks, sidecar_task as _sdk_sidecar_tasks
+from flytekit.common.tasks import task as _task
 from flytekit.common.types import helpers as _type_helpers
 from flytekit.models import interface as _interface_model
 
@@ -35,7 +36,7 @@ def my_task(wf_params, in1, in2, out1, out2):
     """
 
     def apply_inputs_wrapper(task):
-        if not isinstance(task, _sdk_runnable_tasks.SdkRunnableTask):
+        if not isinstance(task, _task.SdkTask):
             additional_msg = \
                 "Inputs can only be applied to a task. Did you forget the task decorator on method '{}.{}'?".format(
                     task.__module__,

setup.py

@@ -29,7 +29,7 @@
         ]
     },
     install_requires=[
-        "flyteidl>=0.17.8,<1.0.0",
+        "flyteidl>=0.17.9,<1.0.0",
         "click>=6.6,<8.0",
         "croniter>=0.3.20,<4.0.0",
         "deprecation>=2.0,<3.0",

tests/flytekit/common/workflows/presto.py

@@ -0,0 +1,27 @@
+from __future__ import absolute_import
+
+from flytekit.sdk.tasks import inputs
+from flytekit.sdk.types import Types
+from flytekit.sdk.workflow import workflow_class, Input, Output
+from flytekit.common.tasks.presto_task import SdkPrestoTask
+
+schema = Types.Schema([("a", Types.String), ("b", Types.Integer)])
+
+presto_task = SdkPrestoTask(
+    task_inputs=inputs(ds=Types.String, rg=Types.String),
+    statement="SELECT * FROM hive.city.fact_airport_sessions WHERE ds = '{{ .Inputs.ds}}' LIMIT 10",
+    output_schema=schema,
+    routing_group="{{ .Inputs.rg }}",
+    # catalog="hive",
+    # schema="city",
+)
+
+
+@workflow_class()
+class PrestoWorkflow(object):
+    ds = Input(Types.String, required=True, help="Test string with no default")
+    # routing_group = Input(Types.String, required=True, help="Test string with no default")
+
+    p_task = presto_task(ds=ds, rg='etl')
+
+    output_a = Output(p_task.outputs.results, sdk_type=schema)