Add async delete function in base_agent (#1800)

Signed-off-by: Future Outlier <[email protected]> Co-authored-by: Future Outlier <[email protected]>
flyteorg · Aug 19, 2023 · b49867a · b49867a
1 parent 6acc97b
commit b49867a
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/flytekit/extend/backend/base_agent.py b/flytekit/extend/backend/base_agent.py
@@ -182,7 +182,7 @@ def execute(self, **kwargs) -> typing.Any:
             res = asyncio.run(agent.async_create(dummy_context, output_prefix, cp_entity.template, inputs))
         else:
             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))
+        signal.signal(signal.SIGINT, partial(self.signal_handler, agent, dummy_context, res.resource_meta))
         state = RUNNING
         metadata = res.resource_meta
         progress = Progress(transient=True)
@@ -211,5 +211,8 @@ def signal_handler(
         signum: int,
         frame: FrameType,
     ) -> typing.Any:
-        agent.delete(context, resource_meta)
+        if agent.asynchronous:
+            asyncio.run(agent.async_delete(context, resource_meta))
+        else:
+            agent.delete(context, resource_meta)
         sys.exit(1)