Chatgpt Agent

flytekit/__init__.py

@@ -245,6 +245,8 @@
     StructuredDatasetType,
 )
 
+from flytekit.extend.backend.task_executor import TaskExecutor  # isort:skip. This is for circular import avoidance.
+
 __version__ = "0.0.0+develop"
 
 

flytekit/core/base_task.py

@@ -105,6 +105,7 @@ class TaskMetadata(object):
     retries: int = 0
     timeout: Optional[Union[datetime.timedelta, int]] = None
     pod_template_name: Optional[str] = None
+    runtime_flavor: str = "python"
 
     def __post_init__(self):
         if self.timeout:
@@ -130,7 +131,9 @@ def to_taskmetadata_model(self) -> _task_model.TaskMetadata:
         return _task_model.TaskMetadata(
             discoverable=self.cache,
             runtime=_task_model.RuntimeMetadata(
-                _task_model.RuntimeMetadata.RuntimeType.FLYTE_SDK, __version__, "python"
+                _task_model.RuntimeMetadata.RuntimeType.FLYTE_SDK,
+                __version__,
+                self.runtime_flavor,
             ),
             timeout=self.timeout,
             retries=self.retry_strategy,
@@ -170,12 +173,13 @@ def __init__(
         task_type_version=0,
         security_ctx: Optional[SecurityContext] = None,
         docs: Optional[Documentation] = None,
+        runtime_flavor: str = "python",
         **kwargs,
     ):
         self._task_type = task_type
         self._name = name
         self._interface = interface
-        self._metadata = metadata if metadata else TaskMetadata()
+        self._metadata = metadata if metadata else TaskMetadata(runtime_flavor=runtime_flavor)
         self._task_type_version = task_type_version
         self._security_ctx = security_ctx
         self._docs = docs
@@ -419,6 +423,7 @@ def __init__(
         environment: Optional[Dict[str, str]] = None,
         disable_deck: Optional[bool] = None,
         enable_deck: Optional[bool] = None,
+        runtime_flavor: str = "python",
         **kwargs,
     ):
         """
@@ -434,11 +439,13 @@ def __init__(
                 execution of the task. Supplied as a dictionary of key/value pairs
             disable_deck (bool): (deprecated) If true, this task will not output deck html file
             enable_deck (bool): If true, this task will output deck html file
+            runtime_flavor (str): default is "python", we can set it to "sync_plugin" for flytepropeller to execute sync plugin task
         """
         super().__init__(
             task_type=task_type,
             name=name,
             interface=transform_interface_to_typed_interface(interface),
+            runtime_flavor=runtime_flavor,
             **kwargs,
         )
         self._python_interface = interface if interface else Interface()

flytekit/core/external_api_task.py

@@ -0,0 +1,74 @@
+import collections
+import inspect
+from abc import abstractmethod
+from typing import Any, Dict, Optional, TypeVar
+
+import jsonpickle
+from flyteidl.admin.agent_pb2 import DoTaskResponse
+from typing_extensions import get_type_hints
+
+from flytekit.configuration import SerializationSettings
+from flytekit.core.base_task import PythonTask
+from flytekit.core.interface import Interface
+from flytekit.extend.backend.base_agent import AsyncAgentExecutorMixin
+
+T = TypeVar("T")
+TASK_MODULE = "task_module"
+TASK_NAME = "task_name"
+TASK_CONFIG_PKL = "task_config_pkl"
+TASK_TYPE = "api_task"
+USE_SYNC_PLUGIN = "sync_plugin"  # Indicates that the sync plugin in FlytePropeller should be used to run this task
+
+
+class ExternalApiTask(AsyncAgentExecutorMixin, PythonTask):
+    """
+    Base class for all external API tasks. External API tasks are tasks that are designed to run until they receive a
+    response from an external service. When the response is received, the task will complete. External API tasks are
+    designed to be run by the flyte agent.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        config: Optional[T] = None,
+        task_type: str = TASK_TYPE,
+        return_type: Optional[Any] = None,
+        **kwargs,
+    ):
+        type_hints = get_type_hints(self.do, include_extras=True)
+        signature = inspect.signature(self.do)
+        inputs = collections.OrderedDict()
+        outputs = collections.OrderedDict({"o0": return_type}) if return_type else collections.OrderedDict()
+
+        for k, _ in signature.parameters.items():  # type: ignore
+            annotation = type_hints.get(k, None)
+            inputs[k] = annotation
+
+        super().__init__(
+            task_type=task_type,
+            name=name,
+            task_config=config,
+            interface=Interface(inputs=inputs, outputs=outputs),
+            runtime_flavor=USE_SYNC_PLUGIN,
+            **kwargs,
+        )
+
+        self._config = config
+
+    @abstractmethod
+    async def do(self, **kwargs) -> DoTaskResponse:
+        """
+        Initiate an HTTP request to an external service such as OpenAI or Vertex AI and retrieve the response.
+        """
+        raise NotImplementedError
+
+    def get_custom(self, settings: SerializationSettings = None) -> Dict[str, Any]:
+        cfg = {
+            TASK_MODULE: type(self).__module__,
+            TASK_NAME: type(self).__name__,
+        }
+
+        if self._config is not None:
+            cfg[TASK_CONFIG_PKL] = jsonpickle.encode(self._config)
+
+        return cfg

flytekit/extend/backend/agent_service.py

@@ -7,6 +7,8 @@
     CreateTaskResponse,
     DeleteTaskRequest,
     DeleteTaskResponse,
+    DoTaskRequest,
+    DoTaskResponse,
     GetTaskRequest,
     GetTaskResponse,
 )
@@ -23,6 +25,7 @@
 create_operation = "create"
 get_operation = "get"
 delete_operation = "delete"
+do_operation = "do"
 
 # Follow the naming convention. https://prometheus.io/docs/practices/naming/
 request_success_count = Counter(
@@ -44,7 +47,7 @@
 def agent_exception_handler(func):
     async def wrapper(
         self,
-        request: typing.Union[CreateTaskRequest, GetTaskRequest, DeleteTaskRequest],
+        request: typing.Union[CreateTaskRequest, GetTaskRequest, DeleteTaskRequest, DoTaskRequest],
         context: grpc.ServicerContext,
         *args,
         **kwargs,
@@ -60,6 +63,11 @@ async def wrapper(
         elif isinstance(request, DeleteTaskRequest):
             task_type = request.task_type
             operation = delete_operation
+        elif isinstance(request, DoTaskRequest):
+            task_type = request.template.type
+            operation = do_operation
+            if request.inputs:
+                input_literal_size.labels(task_type=task_type).observe(request.inputs.ByteSize())
         else:
             context.set_code(grpc.StatusCode.UNIMPLEMENTED)
             context.set_details("Method not implemented!")
@@ -122,3 +130,13 @@ async def DeleteTask(self, request: DeleteTaskRequest, context: grpc.ServicerCon
         if agent.asynchronous:
             return await agent.async_delete(context=context, resource_meta=request.resource_meta)
         return await asyncio.get_running_loop().run_in_executor(None, agent.delete, context, request.resource_meta)
+
+    @agent_exception_handler
+    async def DoTask(self, request: DoTaskRequest, context: grpc.ServicerContext) -> DoTaskResponse:
+        tmp = TaskTemplate.from_flyte_idl(request.template)
+        inputs = LiteralMap.from_flyte_idl(request.inputs) if request.inputs else None
+        agent = AgentRegistry.get_agent(tmp.type)
+        logger.info(f"{tmp.type} agent start doing the job")
+        if agent.asynchronous:
+            return await agent.async_do(context=context, inputs=inputs, task_template=tmp)
+        return await asyncio.get_running_loop().run_in_executor(None, agent.do, context, "", inputs, tmp)

flytekit/extend/backend/base_agent.py

@@ -6,7 +6,7 @@
 from abc import ABC
 from collections import OrderedDict
 from functools import partial
-from types import FrameType
+from types import FrameType, coroutine
 
 import grpc
 from flyteidl.admin.agent_pb2 import (
@@ -16,6 +16,7 @@
     SUCCEEDED,
     CreateTaskResponse,
     DeleteTaskResponse,
+    DoTaskResponse,
     GetTaskResponse,
     State,
 )
@@ -81,6 +82,17 @@ def delete(self, context: grpc.ServicerContext, resource_meta: bytes) -> DeleteT
         """
         raise NotImplementedError
 
+    def do(
+        self,
+        context: grpc.ServicerContext,
+        task_template: TaskTemplate,
+        inputs: typing.Optional[LiteralMap] = None,
+    ) -> DoTaskResponse:
+        """
+        Return the result of executing a task. It should return error code if the task execution failed.
+        """
+        raise NotImplementedError
+
     async def async_create(
         self,
         context: grpc.ServicerContext,
@@ -107,6 +119,17 @@ async def async_delete(self, context: grpc.ServicerContext, resource_meta: bytes
         """
         raise NotImplementedError
 
+    async def async_do(
+        self,
+        context: grpc.ServicerContext,
+        task_template: TaskTemplate,
+        inputs: typing.Optional[LiteralMap] = None,
+    ) -> DoTaskResponse:
+        """
+        Return the result of executing a task. It should return error code if the task execution failed.
+        """
+        raise NotImplementedError
+
 
 class AgentRegistry(object):
     """
@@ -135,7 +158,8 @@ def convert_to_flyte_state(state: str) -> State:
     Convert the state from the agent to the state in flyte.
     """
     state = state.lower()
-    if state in ["failed", "timedout", "canceled"]:
+    # timedout is weird but correct, you can refer here: https://github.com/databricks/databricks-sdk-py/pull/407
+    if state in ["failed", "timeout", "timedout", "canceled"]:
         return RETRYABLE_FAILURE
     elif state in ["done", "succeeded", "success"]:
         return SUCCEEDED
@@ -155,25 +179,38 @@ def get_agent_secret(secret_key: str) -> str:
     return flytekit.current_context().secrets.get(secret_key)
 
 
+def _get_grpc_context() -> grpc.ServicerContext:
+    from unittest.mock import MagicMock
+
+    grpc_ctx = MagicMock(spec=grpc.ServicerContext)
+    return grpc_ctx
+
+
 class AsyncAgentExecutorMixin:
     """
     This mixin class is used to run the agent task locally, and it's only used for local execution.
     Task should inherit from this class if the task can be run in the agent.
     """
 
-    _is_canceled = None
-    _agent = None
-    _entity = None
+    _clean_up_task: coroutine = None
+    _agent: AgentBase = None
+    _entity: PythonTask = None
+    _ctx: FlyteContext = FlyteContext.current_context()
+    _grpc_ctx: grpc.ServicerContext = _get_grpc_context()
 
     def execute(self, **kwargs) -> typing.Any:
+        from flytekit.extend.backend.task_executor import TaskExecutor
         from flytekit.tools.translator import get_serializable
 
         self._entity = typing.cast(PythonTask, self)
         task_template = get_serializable(OrderedDict(), SerializationSettings(ImageConfig()), self._entity).template
         self._agent = AgentRegistry.get_agent(task_template.type)
 
-        res = asyncio.run(self._create(task_template, kwargs))
-        res = asyncio.run(self._get(resource_meta=res.resource_meta))
+        if isinstance(self._agent, TaskExecutor):
+            res = asyncio.run(self._do(task_template, kwargs))
+        else:
+            res = asyncio.run(self._create(task_template, kwargs))
+            res = asyncio.run(self._get(resource_meta=res.resource_meta))
 
         if res.resource.state != SUCCEEDED:
             raise FlyteUserException(f"Failed to run the task {self._entity.name}")
@@ -183,53 +220,51 @@ def execute(self, **kwargs) -> typing.Any:
     async def _create(
         self, task_template: TaskTemplate, inputs: typing.Dict[str, typing.Any] = None
     ) -> CreateTaskResponse:
-        ctx = FlyteContext.current_context()
-        grpc_ctx = _get_grpc_context()
-
-        # Convert python inputs to literals
-        literals = {}
-        for k, v in inputs.items():
-            literals[k] = TypeEngine.to_literal(ctx, v, type(v), self._entity.interface.inputs[k].type)
-        inputs = LiteralMap(literals) if literals else None
-        output_prefix = ctx.file_access.get_random_local_directory()
+        inputs = self.get_input_literal_map(inputs)
+        output_prefix = self._ctx.file_access.get_random_local_directory()
 
         if self._agent.asynchronous:
-            res = await self._agent.async_create(grpc_ctx, output_prefix, task_template, inputs)
+            res = await self._agent.async_create(self._grpc_ctx, output_prefix, task_template, inputs)
         else:
-            res = self._agent.create(grpc_ctx, output_prefix, task_template, inputs)
+            res = self._agent.create(self._grpc_ctx, output_prefix, task_template, inputs)
 
         signal.signal(signal.SIGINT, partial(self.signal_handler, res.resource_meta))  # type: ignore
         return res
 
     async def _get(self, resource_meta: bytes) -> GetTaskResponse:
         state = RUNNING
-        grpc_ctx = _get_grpc_context()
-
         while not is_terminal_state(state):
             time.sleep(1)
             if self._agent.asynchronous:
-                res = await self._agent.async_get(grpc_ctx, resource_meta)
-                if self._is_canceled:
-                    await self._is_canceled
+                res = await self._agent.async_get(self._grpc_ctx, resource_meta)
+                if self._clean_up_task:
+                    await self._clean_up_task
                     sys.exit(1)
             else:
-                res = self._agent.get(grpc_ctx, resource_meta)
+                res = self._agent.get(self._grpc_ctx, resource_meta)
             state = res.resource.state
             logger.info(f"Task state: {state}")
         return res
 
+    async def _do(self, task_template: TaskTemplate, inputs: typing.Dict[str, typing.Any] = None):
+        inputs = self.get_input_literal_map(inputs)
+        if self._agent.asynchronous:
+            res = await self._agent.async_do(self._grpc_ctx, task_template, inputs)
+        else:
+            res = self._agent.do(self._grpc_ctx, task_template, inputs)
+        return res
+
     def signal_handler(self, resource_meta: bytes, signum: int, frame: FrameType) -> typing.Any:
-        grpc_ctx = _get_grpc_context()
         if self._agent.asynchronous:
-            if self._is_canceled is None:
-                self._is_canceled = asyncio.create_task(self._agent.async_delete(grpc_ctx, resource_meta))
+            if self._clean_up_task is None:
+                self._clean_up_task = asyncio.create_task(self._agent.async_delete(self._grpc_ctx, resource_meta))
         else:
-            self._agent.delete(grpc_ctx, resource_meta)
+            self._agent.delete(self._grpc_ctx, resource_meta)
             sys.exit(1)
 
-
-def _get_grpc_context():
-    from unittest.mock import MagicMock
-
-    grpc_ctx = MagicMock(spec=grpc.ServicerContext)
-    return grpc_ctx
+    def get_input_literal_map(self, inputs: typing.Dict[str, typing.Any] = None) -> typing.Optional[LiteralMap]:
+        # Convert python inputs to literals
+        literals = {}
+        for k, v in inputs.items():
+            literals[k] = TypeEngine.to_literal(self._ctx, v, type(v), self._entity.interface.inputs[k].type)
+        return LiteralMap(literals) if literals else None