ChatGPT Agent V2

.github/workflows/monodocs_build.yml

@@ -38,7 +38,7 @@ jobs:
         working-directory: ${{ github.workspace }}/flyte
         run: |
           conda activate monodocs-env
-          pip install -e ./flyteidl
+          pip install git+https://github.com/flyteorg/flyte.git@b5c1d99c7f85764775c9924d5b98948e4d8478d0#subdirectory=flyteidl
           conda info
           conda list
           conda config --show-sources

.github/workflows/pythonbuild.yml

@@ -210,6 +210,7 @@ jobs:
           # onnx-tensorflow needs a version of tensorflow that does not work with protobuf>4.
           # The issue is being tracked on the tensorflow side in https://github.com/tensorflow/tensorflow/issues/53234#issuecomment-1330111693
           # flytekit-onnx-tensorflow
+          - flytekit-chatgpt
           - flytekit-pandera
           - flytekit-papermill
           - flytekit-polars
@@ -275,6 +276,7 @@ jobs:
           pip freeze
       - name: Test with coverage
         run: |
+          pip install git+https://github.com/flyteorg/flyte.git@b5c1d99c7f85764775c9924d5b98948e4d8478d0#subdirectory=flyteidl
           cd plugins/${{ matrix.plugin-names }}
           # onnx plugins does not support protobuf>4 yet (in fact it is tensorflow that
           # does not support that yet). More details in https://github.com/onnx/onnx/issues/4239.

Dockerfile.agent

@@ -14,6 +14,7 @@ RUN pip install --no-cache-dir -U flytekit==$VERSION \
   flytekitplugins-mmcloud==$VERSION \
   flytekitplugins-spark==$VERSION \
   flytekitplugins-snowflake==$VERSION \
+  flytekitplugins-chatgpt==$VERSION \
   && apt-get clean autoclean \
   && apt-get autoremove --yes \
   && rm -rf /var/lib/{apt,dpkg,cache,log}/ \

Dockerfile.dev

@@ -29,8 +29,7 @@ COPY . /flytekit
 # 4. Create a non-root user 'flytekit' and set appropriate permissions for directories.
 RUN apt-get update && apt-get install build-essential vim libmagic1 git -y \
     && pip install --no-cache-dir -U --pre \
-        flyteidl \
-        -e /flytekit \
+        git+https://github.com/flyteorg/flyte.git@a46196d42e80f055efc9a013cb95c20342c72622#subdirectory=flyteidl \
         -e /flytekit/plugins/flytekit-k8s-pod \
         -e /flytekit/plugins/flytekit-deck-standard \
         -e /flytekit/plugins/flytekit-flyteinteractive \

Makefile

@@ -24,9 +24,9 @@ update_boilerplate:
 
 .PHONY: setup
 setup: install-piptools ## Install requirements
+	pip install git+https://github.com/flyteorg/flyte.git@b5c1d99c7f85764775c9924d5b98948e4d8478d0#subdirectory=flyteidl
 	pip install --pre -r dev-requirements.in
 
-
 .PHONY: fmt
 fmt:
 	pre-commit run ruff --all-files || true

flytekit/clis/sdk_in_container/serve.py

@@ -4,6 +4,7 @@
 from flyteidl.service.agent_pb2_grpc import (
     add_AgentMetadataServiceServicer_to_server,
     add_AsyncAgentServiceServicer_to_server,
+    add_SyncAgentServiceServicer_to_server,
 )
 from grpc import aio
 
@@ -52,7 +53,7 @@ def agent(_: click.Context, port, worker, timeout):
 
 async def _start_grpc_server(port: int, worker: int, timeout: int):
     click.secho("Starting up the server to expose the prometheus metrics...", fg="blue")
-    from flytekit.extend.backend.agent_service import AgentMetadataService, AsyncAgentService
+    from flytekit.extend.backend.agent_service import AgentMetadataService, AsyncAgentService, SyncAgentService
 
     try:
         from prometheus_client import start_http_server
@@ -64,6 +65,7 @@ async def _start_grpc_server(port: int, worker: int, timeout: int):
     server = aio.server(futures.ThreadPoolExecutor(max_workers=worker))
 
     add_AsyncAgentServiceServicer_to_server(AsyncAgentService(), server)
+    add_SyncAgentServiceServicer_to_server(SyncAgentService(), server)
     add_AgentMetadataServiceServicer_to_server(AgentMetadataService(), server)
 
     server.add_insecure_port(f"[::]:{port}")

flytekit/core/type_engine.py

@@ -16,6 +16,7 @@
 from typing import Dict, List, NamedTuple, Optional, Type, cast
 
 from dataclasses_json import DataClassJsonMixin, dataclass_json
+from flyteidl.core import literals_pb2
 from google.protobuf import json_format as _json_format
 from google.protobuf import struct_pb2 as _struct
 from google.protobuf.json_format import MessageToDict as _MessageToDict
@@ -1210,6 +1211,16 @@ def dict_to_literal_map(
                 raise user_exceptions.FlyteTypeException(type(v), python_type, received_value=v)
         return LiteralMap(literal_map)
 
+    @classmethod
+    def dict_to_literal_map_pb(
+        cls,
+        ctx: FlyteContext,
+        d: typing.Dict[str, typing.Any],
+        type_hints: Optional[typing.Dict[str, type]] = None,
+    ) -> Optional[literals_pb2.LiteralMap]:
+        literal_map = cls.dict_to_literal_map(ctx, d, type_hints)
+        return literal_map.to_flyte_idl()
+
     @classmethod
     def get_available_transformers(cls) -> typing.KeysView[Type]:
         """

flytekit/extend/backend/agent_service.py

@@ -1,31 +1,42 @@
 import typing
+from http import HTTPStatus
 
 import grpc
 from flyteidl.admin.agent_pb2 import (
     CreateTaskRequest,
     CreateTaskResponse,
     DeleteTaskRequest,
     DeleteTaskResponse,
+    ExecuteTaskSyncRequest,
+    ExecuteTaskSyncResponse,
+    ExecuteTaskSyncResponseHeader,
     GetAgentRequest,
     GetAgentResponse,
     GetTaskRequest,
     GetTaskResponse,
     ListAgentsRequest,
     ListAgentsResponse,
+    Resource,
+)
+from flyteidl.service.agent_pb2_grpc import (
+    AgentMetadataServiceServicer,
+    AsyncAgentServiceServicer,
+    SyncAgentServiceServicer,
 )
-from flyteidl.service.agent_pb2_grpc import AgentMetadataServiceServicer, AsyncAgentServiceServicer
 from prometheus_client import Counter, Summary
 
-from flytekit import logger
+from flytekit import FlyteContext, logger
+from flytekit.core.type_engine import TypeEngine
 from flytekit.exceptions.system import FlyteAgentNotFound
-from flytekit.extend.backend.base_agent import AgentRegistry, mirror_async_methods
+from flytekit.extend.backend.base_agent import AgentRegistry, SyncAgentBase, mirror_async_methods
 from flytekit.models.literals import LiteralMap
 from flytekit.models.task import TaskTemplate
 
 metric_prefix = "flyte_agent_"
 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(
@@ -46,7 +57,24 @@
 input_literal_size = Summary(f"{metric_prefix}input_literal_bytes", "Size of input literal", ["task_type"])
 
 
-def agent_exception_handler(func: typing.Callable):
+def _handle_exception(e: Exception, context: grpc.ServicerContext, task_type: str, operation: str):
+    if isinstance(e, FlyteAgentNotFound):
+        error_message = f"Cannot find agent for task type: {task_type}."
+        logger.error(error_message)
+        context.set_code(grpc.StatusCode.NOT_FOUND)
+        context.set_details(error_message)
+        request_failure_count.labels(task_type=task_type, operation=operation, error_code=HTTPStatus.NOT_FOUND).inc()
+    else:
+        error_message = f"failed to {operation} {task_type} task with error: {e}."
+        logger.error(error_message)
+        context.set_code(grpc.StatusCode.INTERNAL)
+        context.set_details(error_message)
+        request_failure_count.labels(
+            task_type=task_type, operation=operation, error_code=HTTPStatus.INTERNAL_SERVER_ERROR
+        ).inc()
+
+
+def record_agent_metrics(func: typing.Callable):
     async def wrapper(
         self,
         request: typing.Union[CreateTaskRequest, GetTaskRequest, DeleteTaskRequest],
@@ -60,10 +88,10 @@
             if request.inputs:
                 input_literal_size.labels(task_type=task_type).observe(request.inputs.ByteSize())
         elif isinstance(request, GetTaskRequest):
-            task_type = request.task_type
+            task_type = request.task_type or request.task_category.name
             operation = get_operation
         elif isinstance(request, DeleteTaskRequest):
-            task_type = request.task_type
+            task_type = request.task_category or request.task_category.name
             operation = delete_operation
         else:
             context.set_code(grpc.StatusCode.UNIMPLEMENTED)
@@ -75,51 +103,90 @@
                 res = await func(self, request, context, *args, **kwargs)
             request_success_count.labels(task_type=task_type, operation=operation).inc()
             return res
-        except FlyteAgentNotFound:
-            error_message = f"Cannot find agent for task type: {task_type}."
-            logger.error(error_message)
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            context.set_details(error_message)
-            request_failure_count.labels(task_type=task_type, operation=operation, error_code="404").inc()
         except Exception as e:
-            error_message = f"failed to {operation} {task_type} task with error {e}."
-            logger.error(error_message)
-            context.set_code(grpc.StatusCode.INTERNAL)
-            context.set_details(error_message)
-            request_failure_count.labels(task_type=task_type, operation=operation, error_code="500").inc()
+            _handle_exception(e, context, task_type, operation)
 
     return wrapper
 
 
 class AsyncAgentService(AsyncAgentServiceServicer):
-    @agent_exception_handler
+    @record_agent_metrics
     async def CreateTask(self, request: CreateTaskRequest, context: grpc.ServicerContext) -> CreateTaskResponse:
-        tmp = TaskTemplate.from_flyte_idl(request.template)
+        template = TaskTemplate.from_flyte_idl(request.template)
         inputs = LiteralMap.from_flyte_idl(request.inputs) if request.inputs else None
-        agent = AgentRegistry.get_agent(tmp.type)
+        agent = AgentRegistry.get_agent(template.type, template.task_type_version)
 
-        logger.info(f"{tmp.type} agent start creating the job")
-        return await mirror_async_methods(
-            agent.create, output_prefix=request.output_prefix, task_template=tmp, inputs=inputs
-        )
+        logger.info(f"{agent.name} start creating the job")
+        resource_mata = await mirror_async_methods(agent.create, task_template=template, inputs=inputs)
+        return CreateTaskResponse(resource_meta=resource_mata.encode())
 
-    @agent_exception_handler
+    @record_agent_metrics
     async def GetTask(self, request: GetTaskRequest, context: grpc.ServicerContext) -> GetTaskResponse:
-        agent = AgentRegistry.get_agent(request.task_type)
-        logger.info(f"{agent.task_type} agent start checking the status of the job")
-        return await mirror_async_methods(agent.get, resource_meta=request.resource_meta)
+        if request.task_category:
+            agent = AgentRegistry.get_agent(request.task_category.name, request.task_category.version)
+        else:
+            agent = AgentRegistry.get_agent(request.task_type)
+        logger.info(f"{agent.name} start checking the status of the job")
+        res = await mirror_async_methods(agent.get, resource_meta=agent.metadata_type.decode(request.resource_meta))
+
+        if res.outputs is None:
+            outputs = None
+        elif isinstance(res.outputs, LiteralMap):
+            outputs = res.outputs.to_flyte_idl()
+        else:
+            ctx = FlyteContext.current_context()
+            outputs = TypeEngine.dict_to_literal_map_pb(ctx, res.outputs)
+        return GetTaskResponse(
+            resource=Resource(phase=res.phase, log_links=res.log_links, message=res.message, outputs=outputs)
+        )
 
-    @agent_exception_handler
+    @record_agent_metrics
     async def DeleteTask(self, request: DeleteTaskRequest, context: grpc.ServicerContext) -> DeleteTaskResponse:
-        agent = AgentRegistry.get_agent(request.task_type)
-        logger.info(f"{agent.task_type} agent start deleting the job")
-        return await mirror_async_methods(agent.delete, resource_meta=request.resource_meta)
+        if request.task_category:
+            agent = AgentRegistry.get_agent(request.task_category.name, request.task_category.version)
+        else:
+            agent = AgentRegistry.get_agent(request.task_type)
+        logger.info(f"{agent.name} start deleting the job")
+        return await mirror_async_methods(agent.delete, resource_meta=agent.metadata_type.decode(request.resource_meta))
+
+
+class SyncAgentService(SyncAgentServiceServicer):
+    async def ExecuteTaskSync(
+        self, request_iterator: typing.AsyncIterator[ExecuteTaskSyncRequest], context: grpc.ServicerContext
+    ) -> typing.AsyncIterator[ExecuteTaskSyncResponse]:
+        request = await request_iterator.__anext__()
+        template = TaskTemplate.from_flyte_idl(request.header.template)
+        task_type = template.type
+        try:
+            with request_latency.labels(task_type=task_type, operation=do_operation).time():
+                agent = AgentRegistry.get_agent(task_type, template.task_type_version)
+                if not isinstance(agent, SyncAgentBase):
+                    raise ValueError(f"[{agent.name}] does not support sync execution")
+
+                request = await request_iterator.__anext__()
+                literal_map = LiteralMap.from_flyte_idl(request.inputs) if request.inputs else None
+                res = await mirror_async_methods(agent.do, task_template=template, inputs=literal_map)
+
+                if res.outputs is None:
+                    outputs = None
+                elif isinstance(res.outputs, LiteralMap):
+                    outputs = res.outputs.to_flyte_idl()
+                else:
+                    ctx = FlyteContext.current_context()
+                    outputs = TypeEngine.dict_to_literal_map_pb(ctx, res.outputs)
+
+                header = ExecuteTaskSyncResponseHeader(
+                    resource=Resource(phase=res.phase, log_links=res.log_links, message=res.message, outputs=outputs)
+                )
+                yield ExecuteTaskSyncResponse(header=header)
+            request_success_count.labels(task_type=task_type, operation=do_operation).inc()
+        except Exception as e:
+            _handle_exception(e, context, template.type, do_operation)
 
 
 class AgentMetadataService(AgentMetadataServiceServicer):
     async def GetAgent(self, request: GetAgentRequest, context: grpc.ServicerContext) -> GetAgentResponse:
-        return GetAgentResponse(agent=AgentRegistry._METADATA[request.name])
+        return GetAgentResponse(agent=AgentRegistry.METADATA[request.name])
 
     async def ListAgents(self, request: ListAgentsRequest, context: grpc.ServicerContext) -> ListAgentsResponse:
-        agents = [agent for agent in AgentRegistry._METADATA.values()]
-        return ListAgentsResponse(agents=agents)
+        return ListAgentsResponse(agents=AgentRegistry.list_agents())