Add agent ctrl-c handler to call the delete function. (Reupload) (#1782)

* reupload the agent-ctrl-c handler and grpc dependency Signed-off-by: Future Outlier <[email protected]> * image config annotation Signed-off-by: Future Outlier <[email protected]> * Trigger RTD build Signed-off-by: Future Outlier <[email protected]> * update grpc dependency (correct version) Signed-off-by: Future Outlier <[email protected]> --------- Signed-off-by: Future Outlier <[email protected]> Co-authored-by: Future Outlier <[email protected]>
flyteorg · Aug 9, 2023 · 6291301 · 6291301
1 parent 1777f4b
commit 6291301
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 3 deletions.
diff --git a/doc-requirements.txt b/doc-requirements.txt
@@ -332,7 +332,7 @@ great-expectations==0.17.6
     # via -r doc-requirements.in
 greenlet==2.0.2
     # via sqlalchemy
-grpcio==1.56.2
+grpcio==1.51.1
     # via
     #   -r doc-requirements.in
     #   flytekit
@@ -342,7 +342,7 @@ grpcio==1.56.2
     #   ray
     #   tensorboard
     #   tensorflow
-grpcio-status==1.56.2
+grpcio-status==1.51.1
     # via
     #   flytekit
     #   google-api-core

diff --git a/flytekit/configuration/__init__.py b/flytekit/configuration/__init__.py
@@ -228,11 +228,14 @@ def look_up_image_info(name: str, tag: str, optional_tag: bool = False) -> Image
 @dataclass(init=True, repr=True, eq=True, frozen=True)
 class ImageConfig(object):
     """
+    We recommend you to use ImageConfig.auto(img_name=None) to create an ImageConfig.
+    For example, ImageConfig.auto(img_name=""ghcr.io/flyteorg/flytecookbook:v1.0.0"") will create an ImageConfig.
+
     ImageConfig holds available images which can be used at registration time. A default image can be specified
     along with optional additional images. Each image in the config must have a unique name.
 
     Attributes:
-        default_image (str): The default image to be used as a container for task serialization.
+        default_image (Optional[Image]): The default image to be used as a container for task serialization.
         images (List[Image]): Optional, additional images which can be used in task container definitions.
     """
 

diff --git a/flytekit/extend/backend/base_agent.py b/flytekit/extend/backend/base_agent.py
@@ -1,7 +1,11 @@
+import signal
+import sys
 import time
 import typing
 from abc import ABC, abstractmethod
 from collections import OrderedDict
+from functools import partial
+from types import FrameType
 
 import grpc
 from flyteidl.admin.agent_pb2 import (
@@ -148,6 +152,7 @@ def execute(self, **kwargs) -> typing.Any:
         output_prefix = ctx.file_access.get_random_local_directory()
         cp_entity = get_serializable(m, settings=SerializationSettings(ImageConfig()), entity=entity)
         res = agent.create(dummy_context, output_prefix, cp_entity.template, inputs)
+        signal.signal(signal.SIGINT, partial(self.signal_handler, agent, dummy_context, res.resource_meta))
         state = RUNNING
         metadata = res.resource_meta
         progress = Progress(transient=True)
@@ -164,3 +169,14 @@ def execute(self, **kwargs) -> typing.Any:
             raise Exception(f"Failed to run the task {entity.name}")
 
         return LiteralMap.from_flyte_idl(res.resource.outputs)
+
+    def signal_handler(
+        self,
+        agent: AgentBase,
+        context: grpc.ServicerContext,
+        resource_meta: bytes,
+        signum: int,
+        frame: FrameType,
+    ) -> typing.Any:
+        agent.delete(context, resource_meta)
+        sys.exit(1)