diff --git a/examples/advanced/local_tasks.py b/examples/advanced/local_tasks.py index 6a27d13c8..0e98049f8 100644 --- a/examples/advanced/local_tasks.py +++ b/examples/advanced/local_tasks.py @@ -101,7 +101,7 @@ def noio_task(): @env.task async def parallel_main_no_io(q: str) -> int: print("Starting parallel_main_no_io", flush=True) - noio_task() + await noio_task.aio() await input_trace("hello world", "blah", 42) a = await output_trace() await noio_trace() diff --git a/examples/basics/dir_download_sync_repro.py b/examples/basics/dir_download_sync_repro.py index fbbf4c58c..cdc4fcab9 100644 --- a/examples/basics/dir_download_sync_repro.py +++ b/examples/basics/dir_download_sync_repro.py @@ -49,7 +49,7 @@ def download_directory_sync(d: Dir) -> list[str]: async def main() -> list[str]: remote_dir = await create_remote_directory() # Pass the Dir as an input to a sync task, which downloads it via download_sync(). - return download_directory_sync(d=remote_dir) + return await download_directory_sync.aio(d=remote_dir) if __name__ == "__main__": diff --git a/src/flyte/_task.py b/src/flyte/_task.py index 43d650f12..fe4d7e129 100644 --- a/src/flyte/_task.py +++ b/src/flyte/_task.py @@ -24,7 +24,12 @@ ) from flyte._pod import PodTemplate -from flyte.errors import RuntimeSystemError, RuntimeUserError, TraceDoesNotAllowNestedTasksError +from flyte.errors import ( + RuntimeSystemError, + RuntimeUserError, + SyncTaskInAsyncContextError, + TraceDoesNotAllowNestedTasksError, +) from ._cache import Cache, CacheRequest from ._context import internal_ctx @@ -363,6 +368,18 @@ def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] raise RuntimeSystemError("BadContext", "Controller is not initialized.") if self._call_as_synchronous: + from ._utils.asyncify import is_sync_execution_loop + + try: + running_loop = asyncio.get_running_loop() + except RuntimeError: + running_loop = None + if running_loop is not None and not is_sync_execution_loop(running_loop): + raise SyncTaskInAsyncContextError( + f"Synchronous task {self.name} was called from an async task. This blocks the event " + f"loop and will hang if the control plane becomes unreachable. " + f"Use `await {self.name}.aio(...)` instead." + ) fut = controller.submit_sync(self, *args, **kwargs) x = fut.result(None) return x diff --git a/src/flyte/_utils/asyncify.py b/src/flyte/_utils/asyncify.py index 33408045b..4d31de51f 100644 --- a/src/flyte/_utils/asyncify.py +++ b/src/flyte/_utils/asyncify.py @@ -14,6 +14,16 @@ T = TypeVar("T") P = ParamSpec("P") +SYNC_EXECUTION_LOOP_ATTR = "_flyte_sync_execution_loop" + + +def is_sync_execution_loop(loop: asyncio.AbstractEventLoop) -> bool: + """ + Returns True if the loop was created by `run_sync_with_loop` to host a single synchronous task body. Such loops + are private to that one function, so blocking them is safe. + """ + return loop.__dict__.get(SYNC_EXECUTION_LOOP_ATTR, False) + async def run_sync_with_loop( func: Callable[P, T], @@ -67,6 +77,7 @@ def _sync_thread_loop_runner() -> None: nonlocal execute_loop try: execute_loop = asyncio.new_event_loop() + execute_loop.__dict__[SYNC_EXECUTION_LOOP_ATTR] = True asyncio.set_event_loop(execute_loop) logger.debug(f"Created event loop for function '{func_name}' in thread '{full_thread_name}'") execute_loop_created.set() diff --git a/src/flyte/errors.py b/src/flyte/errors.py index 207317f4b..fb9bf4731 100644 --- a/src/flyte/errors.py +++ b/src/flyte/errors.py @@ -325,6 +325,17 @@ def __init__(self, message: str): super().__init__("TraceDoesNotAllowNestedTasksError", message) +class SyncTaskInAsyncContextError(RuntimeUserError): + """ + This error is raised when a synchronous task is called directly from an async task. The synchronous call path + blocks the event loop, which prevents the controller from reporting failures and can deadlock the run. Use + `await task.aio(...)` instead. + """ + + def __init__(self, message: str): + super().__init__("SyncTaskInAsyncContextError", message) + + class InvalidPackageError(RuntimeUserError): """Raised when an invalid system package is detected during image build.""" diff --git a/tests/flyte/local_controller/test_nested_hello_aio.py b/tests/flyte/local_controller/test_nested_hello_aio.py index 4f55a873d..0f9cf8aa7 100644 --- a/tests/flyte/local_controller/test_nested_hello_aio.py +++ b/tests/flyte/local_controller/test_nested_hello_aio.py @@ -28,7 +28,7 @@ async def say_hello_nested(data: str = "default string") -> str: squared = await asyncio.gather(*coros) - return say_hello(data=data, lt=squared) + return await say_hello.aio(data=data, lt=squared) def test_run_local_controller(): diff --git a/tests/flyte/test_sync_tasks.py b/tests/flyte/test_sync_tasks.py index 02bb99473..4d080e16c 100644 --- a/tests/flyte/test_sync_tasks.py +++ b/tests/flyte/test_sync_tasks.py @@ -1,6 +1,9 @@ from typing import List +import pytest + import flyte +from flyte.errors import RuntimeUserError, SyncTaskInAsyncContextError env = flyte.TaskEnvironment(name="test") @@ -31,3 +34,29 @@ def test_parent_action_local(): flyte.init() result = flyte.run(sync_parent_task, 3) assert result.outputs()[0] == ["Hello, world 0!", "Hello, world 1!", "Hello, world 2!"] + + +@env.task +async def async_parent_calling_sync_child(i: int) -> str: + return sync_task1(str(i)) + + +@env.task +async def async_parent_awaiting_sync_child(i: int) -> str: + return await sync_task1.aio(str(i)) + + +def test_sync_task_in_async_context_error_is_user_error(): + assert issubclass(SyncTaskInAsyncContextError, RuntimeUserError) + + +def test_sync_child_from_async_parent_raises(): + flyte.init() + with pytest.raises(RuntimeUserError, match="aio"): + flyte.run(async_parent_calling_sync_child, 1).outputs() + + +def test_sync_child_from_async_parent_with_aio_succeeds(): + flyte.init() + result = flyte.run(async_parent_awaiting_sync_child, 1) + assert result.outputs()[0] == "Hello, world 1!"