External Plugin Service (grpc)

.github/workflows/pythonpublish.yml

@@ -142,6 +142,48 @@ jobs:
           cache-from: type=gha
           cache-to: type=gha,mode=max
 
+  build-and-push-external-plugin-service-images:
+    runs-on: ubuntu-latest
+    needs: deploy
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: "0"
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v1
+      - name: Set up Docker Buildx
+        id: buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Login to GitHub Container Registry
+        if: ${{ github.event_name == 'release' }}
+        uses: docker/login-action@v1
+        with:
+          registry: ghcr.io
+          username: "${{ secrets.FLYTE_BOT_USERNAME }}"
+          password: "${{ secrets.FLYTE_BOT_PAT }}"
+      - name: Prepare External Plugin Service Image Names
+        id: external-plugin-service-names
+        uses: docker/metadata-action@v3
+        with:
+          images: |
+            ghcr.io/${{ github.repository_owner }}/flytekit
+          tags: |
+            external-plugin-service-latest
+            external-plugin-service-${{ github.sha }}
+            external-plugin-service-${{ needs.deploy.outputs.version }}
+      - name: Push External Plugin Service Image to GitHub Registry
+        uses: docker/build-push-action@v2
+        with:
+          context: "."
+          platforms: linux/arm64, linux/amd64
+          push: ${{ github.event_name == 'release' }}
+          tags: ${{ steps.external-plugin-service-names.outputs.tags }}
+          build-args: |
+            VERSION=${{ needs.deploy.outputs.version }}
+          file: ./Dockerfile
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+
   build-and-push-spark-images:
     runs-on: ubuntu-latest
     needs: deploy

Dockerfile.external-plugin-service

@@ -0,0 +1,10 @@
+FROM python:3.9-slim-buster
+
+MAINTAINER Flyte Team <[email protected]>
+LABEL org.opencontainers.image.source=https://github.com/flyteorg/flytekit
+
+ARG VERSION
+RUN pip install -U flytekit==$VERSION \
+	flytekitplugins-bigquery==$VERSION \
+
+CMD pyflyte serve --port 80

doc-requirements.txt

@@ -244,7 +244,7 @@ flask==2.2.3
     # via mlflow
 flatbuffers==23.1.21
     # via tensorflow
-flyteidl==1.3.12
+flyteidl==1.3.16
     # via flytekit
 fonttools==4.38.0
     # via matplotlib

flytekit/clients/friendly.py

@@ -1007,7 +1007,7 @@ def get_upload_signed_url(
 
     def get_download_signed_url(
         self, native_url: str, expires_in: datetime.timedelta = None
-    ) -> _data_proxy_pb2.CreateUploadLocationResponse:
+    ) -> _data_proxy_pb2.CreateDownloadLocationRequest:
         expires_in_pb = None
         if expires_in:
             expires_in_pb = Duration()

flytekit/clis/sdk_in_container/pyflyte.py

@@ -15,6 +15,7 @@
 from flytekit.clis.sdk_in_container.register import register
 from flytekit.clis.sdk_in_container.run import run
 from flytekit.clis.sdk_in_container.serialize import serialize
+from flytekit.clis.sdk_in_container.serve import serve
 from flytekit.configuration.internal import LocalSDK
 from flytekit.exceptions.base import FlyteException
 from flytekit.exceptions.user import FlyteInvalidInputException
@@ -134,6 +135,7 @@ def main(ctx, pkgs: typing.List[str], config: str, verbose: bool):
 main.add_command(run)
 main.add_command(register)
 main.add_command(backfill)
+main.add_command(serve)
 main.add_command(build)
 main.add_command(launchplan)
 main.epilog

flytekit/clis/sdk_in_container/serve.py

@@ -0,0 +1,46 @@
+from concurrent import futures
+
+import click
+import grpc
+from flyteidl.service.external_plugin_service_pb2_grpc import add_ExternalPluginServiceServicer_to_server
+
+from flytekit.extend.backend.external_plugin_service import BackendPluginServer
+
+_serve_help = """Start a grpc server for the external plugin service."""
+
+
+@click.command("serve", help=_serve_help)
+@click.option(
+    "--port",
+    default="80",
+    is_flag=False,
+    type=int,
+    help="Grpc port for the external plugin service",
+)
+@click.option(
+    "--worker",
+    default="10",
+    is_flag=False,
+    type=int,
+    help="Number of workers for the grpc server",
+)
+@click.option(
+    "--timeout",
+    default=None,
+    is_flag=False,
+    type=int,
+    help="It will wait for the specified number of seconds before shutting down grpc server. It should only be used "
+    "for testing.",
+)
+@click.pass_context
+def serve(_: click.Context, port, worker, timeout):
+    """
+    Start a grpc server for the external plugin service.
+    """
+    click.secho("Starting the external plugin service...", fg="blue")
+    server = grpc.server(futures.ThreadPoolExecutor(max_workers=worker))
+    add_ExternalPluginServiceServicer_to_server(BackendPluginServer(), server)
+
+    server.add_insecure_port(f"[::]:{port}")
+    server.start()
+    server.wait_for_termination(timeout=timeout)

flytekit/extend/backend/base_plugin.py

@@ -0,0 +1,66 @@
+import typing
+from abc import ABC, abstractmethod
+
+import grpc
+from flyteidl.core.tasks_pb2 import TaskTemplate
+from flyteidl.service.external_plugin_service_pb2 import (
+    RETRYABLE_FAILURE,
+    RUNNING,
+    SUCCEEDED,
+    State,
+    TaskCreateResponse,
+    TaskDeleteResponse,
+    TaskGetResponse,
+)
+
+from flytekit.models.literals import LiteralMap
+
+
+class BackendPluginBase(ABC):
+    def __init__(self, task_type: str):
+        self._task_type = task_type
+
+    @property
+    def task_type(self) -> str:
+        return self._task_type
+
+    @abstractmethod
+    def create(
+        self,
+        context: grpc.ServicerContext,
+        output_prefix: str,
+        task_template: TaskTemplate,
+        inputs: typing.Optional[LiteralMap] = None,
+    ) -> TaskCreateResponse:
+        pass
+
+    @abstractmethod
+    def get(self, context: grpc.ServicerContext, job_id: str) -> TaskGetResponse:
+        pass
+
+    @abstractmethod
+    def delete(self, context: grpc.ServicerContext, job_id: str) -> TaskDeleteResponse:
+        pass
+
+
+class BackendPluginRegistry(object):
+    _REGISTRY: typing.Dict[str, BackendPluginBase] = {}
+
+    @staticmethod
+    def register(plugin: BackendPluginBase):
+        BackendPluginRegistry._REGISTRY[plugin.task_type] = plugin
+
+    @staticmethod
+    def get_plugin(task_type: str) -> BackendPluginBase:
+        return BackendPluginRegistry._REGISTRY[task_type]
+
+
+def convert_to_flyte_state(state: str) -> State:
+    state = state.lower()
+    if state in ["failed"]:
+        return RETRYABLE_FAILURE
+    elif state in ["done", "succeeded"]:
+        return SUCCEEDED
+    elif state in ["running"]:
+        return RUNNING
+    raise ValueError(f"Unrecognized state: {state}")

flytekit/extend/backend/external_plugin_service.py

@@ -0,0 +1,42 @@
+import grpc
+from flyteidl.service.external_plugin_service_pb2 import (
+    TaskCreateRequest,
+    TaskCreateResponse,
+    TaskDeleteRequest,
+    TaskDeleteResponse,
+    TaskGetRequest,
+    TaskGetResponse,
+)
+from flyteidl.service.external_plugin_service_pb2_grpc import ExternalPluginServiceServicer
+
+from flytekit.extend.backend import model
+from flytekit.extend.backend.base_plugin import BackendPluginRegistry
+
+
+class BackendPluginServer(ExternalPluginServiceServicer):
+    def CreateTask(self, request: TaskCreateRequest, context: grpc.ServicerContext) -> TaskCreateResponse:
+        try:
+            req = model.TaskCreateRequest.from_flyte_idl(request)
+            plugin = BackendPluginRegistry.get_plugin(req.template.type)
+            return plugin.create(
+                context=context, inputs=req.inputs, output_prefix=req.output_prefix, task_template=req.template
+            )
+        except Exception as e:
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(f"failed to create task with error {e}")
+
+    def GetTask(self, request: TaskGetRequest, context: grpc.ServicerContext) -> TaskGetResponse:
+        try:
+            plugin = BackendPluginRegistry.get_plugin(request.task_type)
+            return plugin.get(context=context, job_id=request.job_id)
+        except Exception as e:
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(f"failed to get task with error {e}")
+
+    def DeleteTask(self, request: TaskDeleteRequest, context: grpc.ServicerContext) -> TaskDeleteResponse:
+        try:
+            plugin = BackendPluginRegistry.get_plugin(request.task_type)
+            return plugin.delete(context=context, job_id=request.job_id)
+        except Exception as e:
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(f"failed to delete task with error {e}")

flytekit/extend/backend/model.py

@@ -0,0 +1,40 @@
+from typing import Optional
+
+from flyteidl.service import external_plugin_service_pb2
+
+from flytekit.models import common, task
+from flytekit.models.literals import LiteralMap
+
+
+class TaskCreateRequest(common.FlyteIdlEntity):
+    def __init__(self, output_prefix: str, template: task.TaskTemplate, inputs: Optional[LiteralMap] = None):
+        self._output_prefix = output_prefix
+        self._template = template
+        self._inputs = inputs
+
+    @property
+    def output_prefix(self) -> str:
+        return self._output_prefix
+
+    @property
+    def template(self) -> task.TaskTemplate:
+        return self._template
+
+    @property
+    def inputs(self) -> Optional[LiteralMap]:
+        return self._inputs
+
+    def to_flyte_idl(self) -> external_plugin_service_pb2.TaskCreateRequest:
+        return external_plugin_service_pb2.TaskCreateRequest(
+            output_prefix=self.output_prefix,
+            template=self.template.to_flyte_idl(),
+            inputs=self.inputs.to_flyte_idl(),
+        )
+
+    @classmethod
+    def from_flyte_idl(cls, proto):
+        return cls(
+            output_prefix=proto.output_prefix,
+            template=task.TaskTemplate.from_flyte_idl(proto.template),
+            inputs=LiteralMap.from_flyte_idl(proto.inputs) if proto.inputs is not None else None,
+        )

plugins/flytekit-bigquery/flytekitplugins/bigquery/__init__.py

@@ -11,4 +11,5 @@
    BigQueryTask
 """
 
+from .backend_plugin import BigQueryPlugin
 from .task import BigQueryConfig, BigQueryTask

plugins/flytekit-bigquery/flytekitplugins/bigquery/backend_plugin.py

@@ -0,0 +1,93 @@
+import datetime
+from typing import Dict, Optional
+
+import grpc
+from flyteidl.service.external_plugin_service_pb2 import (
+    SUCCEEDED,
+    TaskCreateResponse,
+    TaskDeleteResponse,
+    TaskGetResponse,
+)
+from google.cloud import bigquery
+
+from flytekit import FlyteContextManager, StructuredDataset
+from flytekit.core.type_engine import TypeEngine
+from flytekit.extend.backend.base_plugin import BackendPluginBase, BackendPluginRegistry, convert_to_flyte_state
+from flytekit.models import literals
+from flytekit.models.literals import LiteralMap
+from flytekit.models.task import TaskTemplate
+from flytekit.models.types import LiteralType, StructuredDatasetType
+
+pythonTypeToBigQueryType: Dict[type, str] = {
+    # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#data_type_sizes
+    list: "ARRAY",
+    bool: "BOOL",
+    bytes: "BYTES",
+    datetime.datetime: "DATETIME",
+    float: "FLOAT64",
+    int: "INT64",
+    str: "STRING",
+}
+
+
+class BigQueryPlugin(BackendPluginBase):
+    def __init__(self):
+        super().__init__(task_type="bigquery_query_job_task")
+
+    def create(
+        self,
+        context: grpc.ServicerContext,
+        output_prefix: str,
+        task_template: TaskTemplate,
+        inputs: Optional[LiteralMap] = None,
+    ) -> TaskCreateResponse:
+        job_config = None
+        if inputs:
+            ctx = FlyteContextManager.current_context()
+            python_interface_inputs = {
+                name: TypeEngine.guess_python_type(lt.type) for name, lt in task_template.interface.inputs.items()
+            }
+            native_inputs = TypeEngine.literal_map_to_kwargs(ctx, inputs, python_interface_inputs)
+
+            job_config = bigquery.QueryJobConfig(
+                query_parameters=[
+                    bigquery.ScalarQueryParameter(name, pythonTypeToBigQueryType[python_interface_inputs[name]], val)
+                    for name, val in native_inputs.items()
+                ]
+            )
+
+        custom = task_template.custom
+        client = bigquery.Client(project=custom["ProjectID"], location=custom["Location"])
+        query_job = client.query(task_template.sql.statement, job_config=job_config)
+
+        return TaskCreateResponse(job_id=str(query_job.job_id))
+
+    def get(self, context: grpc.ServicerContext, job_id: str) -> TaskGetResponse:
+        client = bigquery.Client()
+        job = client.get_job(job_id)
+        cur_state = convert_to_flyte_state(str(job.state))
+        res = None
+
+        if cur_state == SUCCEEDED:
+            ctx = FlyteContextManager.current_context()
+            output_location = f"bq://{job.destination.project}:{job.destination.dataset_id}.{job.destination.table_id}"
+            res = literals.LiteralMap(
+                {
+                    "results": TypeEngine.to_literal(
+                        ctx,
+                        StructuredDataset(uri=output_location),
+                        StructuredDataset,
+                        LiteralType(structured_dataset_type=StructuredDatasetType(format="")),
+                    )
+                }
+            )
+
+        return TaskGetResponse(state=cur_state, outputs=res.to_flyte_idl())
+
+    def delete(self, context: grpc.ServicerContext, job_id: str) -> TaskDeleteResponse:
+        client = bigquery.Client()
+        client.cancel_job(job_id)
+        return TaskDeleteResponse()
+
+
+BackendPluginRegistry.register(BigQueryPlugin())

plugins/flytekit-bigquery/flytekitplugins/bigquery/task.py

@@ -81,3 +81,6 @@ def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
     def get_sql(self, settings: SerializationSettings) -> Optional[_task_model.Sql]:
         sql = _task_model.Sql(statement=self.query_template, dialect=_task_model.Sql.Dialect.ANSI)
         return sql
+
+    def execute(self, **kwargs) -> Any:
+        raise Exception("Cannot run a SQL Task natively, please mock.")