use get_event_loop in eager coroutines

Signed-off-by: Niels Bantilan <[email protected]>

flytekit/clis/sdk_in_container/run.py

@@ -632,7 +632,7 @@ def _run(*args, **kwargs):
                     output = entity(**inputs)
                     if inspect.iscoroutine(output):
                         # TODO: make eager mode workflows run with local-mode
-                        output = asyncio.run(output)
+                        output = asyncio.get_event_loop().run_until_complete(output)
                     click.echo(output)
                     return
 

flytekit/core/base_task.py

@@ -754,7 +754,7 @@ def dispatch_execute(
                     elif exec_ctx.execution_state.mode == ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION:
                         # If executed inside of a workflow being executed locally, then run the coroutine to get the
                         # actual results.
-                        return asyncio.run(
+                        return asyncio.get_event_loop().run_until_complete(
                             self._async_execute(
                                 native_inputs,
                                 native_outputs,

flytekit/core/workflow.py

@@ -325,7 +325,7 @@ def local_execute(self, ctx: FlyteContext, **kwargs) -> Union[Tuple[Promise], Pr
 
         if inspect.iscoroutine(function_outputs):
             # handle coroutines for eager workflows
-            function_outputs = asyncio.run(function_outputs)
+            function_outputs = asyncio.get_event_loop().run_until_complete(function_outputs)
 
         # First handle the empty return case.
         # A workflow function may return a task that doesn't return anything

flytekit/experimental/eager_function.py

@@ -369,7 +369,7 @@ def node_cleanup(sig, frame, loop, async_stack: AsyncStack):
     for node in async_stack.call_stack:
         terminations.append(node.async_entity.terminate())
     results = asyncio.gather(*terminations)
-    results = asyncio.run(results)
+    results = loop.run_until_complete(results)
     logger.debug(f"Successfully terminated subtasks {results}")
     loop.close()
 

tests/flytekit/unit/experimental/test_eager_workflows.py

@@ -286,6 +286,7 @@ async def eager_wf_flyte_directory() -> str:
 @mock.patch("flytekit.core.utils.load_proto_from_file")
 @mock.patch("flytekit.core.data_persistence.FileAccessProvider.get_data")
 @mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data")
+@mock.patch("flytekit.core.utils.write_proto_to_file")
 def test_eager_workflow_dispatch_preserves_event_loop(*args):
     """Test that event loop is preserved when dispatching eager workflows."""