prefect-client 2.20.2__py3-none-any.whl → 3.0.0__py3-none-any.whl

prefect/utilities/asyncutils.py

@@ -3,15 +3,13 @@ Utilities for interoperability with async functions and workers from various con
 """
 
 import asyncio
-import ctypes
 import inspect
 import threading
 import warnings
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import asynccontextmanager
-from contextvars import copy_context
+from contextvars import ContextVar, copy_context
 from functools import partial, wraps
-from threading import Thread
 from typing import (
     Any,
     Awaitable,
@@ -20,25 +18,31 @@ from typing import (
     Dict,
     List,
     Optional,
-    Type,
     TypeVar,
     Union,
     cast,
+    overload,
 )
 from uuid import UUID, uuid4
 
 import anyio
 import anyio.abc
+import anyio.from_thread
 import anyio.to_thread
 import sniffio
-from anyio.from_thread import start_blocking_portal
 from typing_extensions import Literal, ParamSpec, TypeGuard
 
+from prefect._internal.concurrency.api import _cast_to_call, from_sync
+from prefect._internal.concurrency.threads import (
+    get_run_sync_loop,
+    in_run_sync_loop,
+)
 from prefect.logging import get_logger
 
 T = TypeVar("T")
 P = ParamSpec("P")
 R = TypeVar("R")
+F = TypeVar("F", bound=Callable[..., Any])
 Async = Literal[True]
 Sync = Literal[False]
 A = TypeVar("A", Async, Sync, covariant=True)
@@ -48,6 +52,14 @@ EVENT_LOOP_GC_REFS = {}
 
 PREFECT_THREAD_LIMITER: Optional[anyio.CapacityLimiter] = None
 
+RUNNING_IN_RUN_SYNC_LOOP_FLAG = ContextVar("running_in_run_sync_loop", default=False)
+RUNNING_ASYNC_FLAG = ContextVar("run_async", default=False)
+BACKGROUND_TASKS: set[asyncio.Task] = set()
+background_task_lock = threading.Lock()
+
+# Thread-local storage to keep track of worker thread state
+_thread_local = threading.local()
+
 logger = get_logger()
 
 
@@ -84,12 +96,47 @@ def is_async_gen_fn(func):
     return inspect.isasyncgenfunction(func)
 
 
-def run_sync(coroutine: Coroutine[Any, Any, T]) -> T:
+def create_task(coroutine: Coroutine) -> asyncio.Task:
     """
-    Runs a coroutine from a synchronous context. A thread will be spawned
-    to run the event loop if necessary, which allows coroutines to run in
-    environments like Jupyter notebooks where the event loop runs on the main
-    thread.
+    Replacement for asyncio.create_task that will ensure that tasks aren't
+    garbage collected before they complete. Allows for "fire and forget"
+    behavior in which tasks can be created and the application can move on.
+    Tasks can also be awaited normally.
+
+    See https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task
+    for details (and essentially this implementation)
+    """
+
+    task = asyncio.create_task(coroutine)
+
+    # Add task to the set. This creates a strong reference.
+    # Take a lock because this might be done from multiple threads.
+    with background_task_lock:
+        BACKGROUND_TASKS.add(task)
+
+    # To prevent keeping references to finished tasks forever,
+    # make each task remove its own reference from the set after
+    # completion:
+    task.add_done_callback(BACKGROUND_TASKS.discard)
+
+    return task
+
+
+def _run_sync_in_new_thread(coroutine: Coroutine[Any, Any, T]) -> T:
+    """
+    Note: this is an OLD implementation of `run_coro_as_sync` which liberally created
+    new threads and new loops. This works, but prevents sharing any objects
+    across coroutines, in particular httpx clients, which are very expensive to
+    instantiate.
+
+    This is here for historical purposes and can be removed if/when it is no
+    longer needed for reference.
+
+    ---
+
+    Runs a coroutine from a synchronous context. A thread will be spawned to run
+    the event loop if necessary, which allows coroutines to run in environments
+    like Jupyter notebooks where the event loop runs on the main thread.
 
     Args:
         coroutine: The coroutine to run.
@@ -98,15 +145,25 @@ def run_sync(coroutine: Coroutine[Any, Any, T]) -> T:
         The return value of the coroutine.
 
     Example:
-        Basic usage:
-        ```python
-        async def my_async_function(x: int) -> int:
+        Basic usage: ```python async def my_async_function(x: int) -> int:
             return x + 1
 
-        run_sync(my_async_function(1))
-        ```
+        run_sync(my_async_function(1)) ```
     """
+
     # ensure context variables are properly copied to the async frame
+    async def context_local_wrapper():
+        """
+        Wrapper that is submitted using copy_context().run to ensure
+        the RUNNING_ASYNC_FLAG mutations are tightly scoped to this coroutine's frame.
+        """
+        token = RUNNING_ASYNC_FLAG.set(True)
+        try:
+            result = await coroutine
+        finally:
+            RUNNING_ASYNC_FLAG.reset(token)
+        return result
+
     context = copy_context()
     try:
         loop = asyncio.get_running_loop()
@@ -115,102 +172,111 @@ def run_sync(coroutine: Coroutine[Any, Any, T]) -> T:
 
     if loop and loop.is_running():
         with ThreadPoolExecutor() as executor:
-            future = executor.submit(context.run, asyncio.run, coroutine)
-            return cast(T, future.result())
+            future = executor.submit(context.run, asyncio.run, context_local_wrapper())
+            result = cast(T, future.result())
     else:
-        return context.run(asyncio.run, coroutine)
+        result = context.run(asyncio.run, context_local_wrapper())
+    return result
 
 
-async def run_sync_in_worker_thread(
-    __fn: Callable[..., T], *args: Any, **kwargs: Any
-) -> T:
+def run_coro_as_sync(
+    coroutine: Awaitable[R],
+    force_new_thread: bool = False,
+    wait_for_result: bool = True,
+) -> Union[R, None]:
     """
-    Runs a sync function in a new worker thread so that the main thread's event loop
-    is not blocked
+    Runs a coroutine from a synchronous context, as if it were a synchronous
+    function.
 
-    Unlike the anyio function, this defaults to a cancellable thread and does not allow
-    passing arguments to the anyio function so users can pass kwargs to their function.
+    The coroutine is scheduled to run in the "run sync" event loop, which is
+    running in its own thread and is started the first time it is needed. This
+    allows us to share objects like async httpx clients among all coroutines
+    running in the loop.
 
-    Note that cancellation of threads will not result in interrupted computation, the
-    thread may continue running — the outcome will just be ignored.
-    """
-    call = partial(__fn, *args, **kwargs)
-    return await anyio.to_thread.run_sync(
-        call, abandon_on_cancel=True, limiter=get_thread_limiter()
-    )
+    If run_sync is called from within the run_sync loop, it will run the
+    coroutine in a new thread, because otherwise a deadlock would occur. Note
+    that this behavior should not appear anywhere in the Prefect codebase or in
+    user code.
 
+    Args:
+        coroutine (Awaitable): The coroutine to be run as a synchronous function.
+        force_new_thread (bool, optional): If True, the coroutine will always be run in a new thread.
+            Defaults to False.
+        wait_for_result (bool, optional): If True, the function will wait for the coroutine to complete
+            and return the result. If False, the function will submit the coroutine to the "run sync"
+            event loop and return immediately, where it will eventually be run. Defaults to True.
 
-def raise_async_exception_in_thread(thread: Thread, exc_type: Type[BaseException]):
+    Returns:
+        The result of the coroutine if wait_for_result is True, otherwise None.
     """
-    Raise an exception in a thread asynchronously.
 
-    This will not interrupt long-running system calls like `sleep` or `wait`.
-    """
-    ret = ctypes.pythonapi.PyThreadState_SetAsyncExc(
-        ctypes.c_long(thread.ident), ctypes.py_object(exc_type)
-    )
-    if ret == 0:
-        raise ValueError("Thread not found.")
+    async def coroutine_wrapper() -> Union[R, None]:
+        """
+        Set flags so that children (and grandchildren...) of this task know they are running in a new
+        thread and do not try to run on the run_sync thread, which would cause a
+        deadlock.
+        """
+        token1 = RUNNING_IN_RUN_SYNC_LOOP_FLAG.set(True)
+        token2 = RUNNING_ASYNC_FLAG.set(True)
+        try:
+            # use `asyncio.create_task` because it copies context variables automatically
+            task = create_task(coroutine)
+            if wait_for_result:
+                return await task
+        finally:
+            RUNNING_IN_RUN_SYNC_LOOP_FLAG.reset(token1)
+            RUNNING_ASYNC_FLAG.reset(token2)
 
+    # if we are already in the run_sync loop, or a descendent of a coroutine
+    # that is running in the run_sync loop, we need to run this coroutine in a
+    # new thread
+    if in_run_sync_loop() or RUNNING_IN_RUN_SYNC_LOOP_FLAG.get() or force_new_thread:
+        return from_sync.call_in_new_thread(coroutine_wrapper)
 
-async def run_sync_in_interruptible_worker_thread(
-    __fn: Callable[..., T], *args: Any, **kwargs: Any
-) -> T:
-    """
-    Runs a sync function in a new interruptible worker thread so that the main
-    thread's event loop is not blocked
-
-    Unlike the anyio function, this performs best-effort cancellation of the
-    thread using the C API. Cancellation will not interrupt system calls like
-    `sleep`.
-    """
+    # otherwise, we can run the coroutine in the run_sync loop
+    # and wait for the result
+    else:
+        call = _cast_to_call(coroutine_wrapper)
+        runner = get_run_sync_loop()
+        runner.submit(call)
+        try:
+            return call.result()
+        except KeyboardInterrupt:
+            call.cancel()
 
-    class NotSet:
-        pass
+            logger.debug("Coroutine cancelled due to KeyboardInterrupt.")
+            raise
 
-    thread: Thread = None
-    result = NotSet
-    event = asyncio.Event()
-    loop = asyncio.get_running_loop()
 
-    def capture_worker_thread_and_result():
-        # Captures the worker thread that AnyIO is using to execute the function so
-        # the main thread can perform actions on it
-        nonlocal thread, result
-        try:
-            thread = threading.current_thread()
-            result = __fn(*args, **kwargs)
-        except BaseException as exc:
-            result = exc
-            raise
-        finally:
-            loop.call_soon_threadsafe(event.set)
+async def run_sync_in_worker_thread(
+    __fn: Callable[..., T], *args: Any, **kwargs: Any
+) -> T:
+    """
+    Runs a sync function in a new worker thread so that the main thread's event loop
+    is not blocked.
 
-    async def send_interrupt_to_thread():
-        # This task waits until the result is returned from the thread, if cancellation
-        # occurs during that time, we will raise the exception in the thread as well
-        try:
-            await event.wait()
-        except anyio.get_cancelled_exc_class():
-            # NOTE: We could send a SIGINT here which allow us to interrupt system
-            # calls but the interrupt bubbles from the child thread into the main thread
-            # and there is not a clear way to prevent it.
-            raise_async_exception_in_thread(thread, anyio.get_cancelled_exc_class())
-            raise
+    Unlike the anyio function, this defaults to a cancellable thread and does not allow
+    passing arguments to the anyio function so users can pass kwargs to their function.
 
-    async with anyio.create_task_group() as tg:
-        tg.start_soon(send_interrupt_to_thread)
-        tg.start_soon(
-            partial(
-                anyio.to_thread.run_sync,
-                capture_worker_thread_and_result,
-                abandon_on_cancel=True,
-                limiter=get_thread_limiter(),
-            )
+    Note that cancellation of threads will not result in interrupted computation, the
+    thread may continue running — the outcome will just be ignored.
+    """
+    # When running a sync function in a worker thread, we set this flag so that
+    # any root sync compatible functions will run as sync functions
+    token = RUNNING_ASYNC_FLAG.set(False)
+    try:
+        call = partial(__fn, *args, **kwargs)
+        result = await anyio.to_thread.run_sync(
+            call_with_mark, call, abandon_on_cancel=True, limiter=get_thread_limiter()
         )
+        return result
+    finally:
+        RUNNING_ASYNC_FLAG.reset(token)
 
-    assert result is not NotSet
-    return result
+
+def call_with_mark(call):
+    mark_as_worker_thread()
+    return call()
 
 
 def run_async_from_worker_thread(
@@ -228,13 +294,12 @@ def run_async_in_new_loop(__fn: Callable[..., Awaitable[T]], *args: Any, **kwarg
     return anyio.run(partial(__fn, *args, **kwargs))
 
 
+def mark_as_worker_thread():
+    _thread_local.is_worker_thread = True
+
+
 def in_async_worker_thread() -> bool:
-    try:
-        anyio.from_thread.threadlocals.current_async_backend
-    except AttributeError:
-        return False
-    else:
-        return True
+    return getattr(_thread_local, "is_worker_thread", False)
 
 
 def in_async_main_thread() -> bool:
@@ -247,7 +312,23 @@ def in_async_main_thread() -> bool:
         return not in_async_worker_thread()
 
 
-def sync_compatible(async_fn: T) -> T:
+@overload
+def sync_compatible(
+    async_fn: Callable[..., Coroutine[Any, Any, R]],
+) -> Callable[..., R]:
+    ...
+
+
+@overload
+def sync_compatible(
+    async_fn: Callable[..., Coroutine[Any, Any, R]],
+) -> Callable[..., Coroutine[Any, Any, R]]:
+    ...
+
+
+def sync_compatible(
+    async_fn: Callable[..., Coroutine[Any, Any, R]],
+) -> Callable[..., Union[R, Coroutine[Any, Any, R]]]:
     """
     Converts an async function into a dual async and sync function.
 
@@ -263,55 +344,53 @@ def sync_compatible(async_fn: T) -> T:
     """
 
     @wraps(async_fn)
-    def coroutine_wrapper(*args, _sync: Optional[bool] = None, **kwargs):
-        from prefect._internal.concurrency.api import create_call, from_sync
-        from prefect._internal.concurrency.calls import get_current_call, logger
-        from prefect._internal.concurrency.event_loop import get_running_loop
-        from prefect._internal.concurrency.threads import get_global_loop
-        from prefect.settings import PREFECT_EXPERIMENTAL_DISABLE_SYNC_COMPAT
-
-        if PREFECT_EXPERIMENTAL_DISABLE_SYNC_COMPAT or _sync is False:
-            return async_fn(*args, **kwargs)
+    def coroutine_wrapper(
+        *args: Any, _sync: Optional[bool] = None, **kwargs: Any
+    ) -> Union[R, Coroutine[Any, Any, R]]:
+        from prefect.context import MissingContextError, get_run_context
 
-        global_thread_portal = get_global_loop()
-        current_thread = threading.current_thread()
-        current_call = get_current_call()
-        current_loop = get_running_loop()
-
-        if (
-            current_thread.ident == global_thread_portal.thread.ident
-            and _sync is not True
-        ):
-            logger.debug(f"{async_fn} --> return coroutine for internal await")
-            # In the prefect async context; return the coro for us to await
+        if _sync is False:
             return async_fn(*args, **kwargs)
-        elif (
-            in_async_main_thread()
-            and (not current_call or is_async_fn(current_call.fn))
-            and _sync is not True
-        ):
-            # In the main async context; return the coro for them to await
-            logger.debug(f"{async_fn} --> return coroutine for user await")
-            return async_fn(*args, **kwargs)
-        elif in_async_worker_thread():
-            # In a sync context but we can access the event loop thread; send the async
-            # call to the parent
-            return run_async_from_worker_thread(async_fn, *args, **kwargs)
-        elif current_loop is not None:
-            logger.debug(f"{async_fn} --> run async in global loop portal")
-            # An event loop is already present but we are in a sync context, run the
-            # call in Prefect's event loop thread
-            return from_sync.call_soon_in_loop_thread(
-                create_call(async_fn, *args, **kwargs)
-            ).result()
+
+        is_async = True
+
+        # if _sync is set, we do as we're told
+        # otherwise, we make some determinations
+        if _sync is None:
+            try:
+                run_ctx = get_run_context()
+                parent_obj = getattr(run_ctx, "task", None)
+                if not parent_obj:
+                    parent_obj = getattr(run_ctx, "flow", None)
+                is_async = getattr(parent_obj, "isasync", True)
+            except MissingContextError:
+                # not in an execution context, make best effort to
+                # decide whether to syncify
+                try:
+                    asyncio.get_running_loop()
+                    is_async = True
+                except RuntimeError:
+                    is_async = False
+
+        async def ctx_call():
+            """
+            Wrapper that is submitted using copy_context().run to ensure
+            mutations of RUNNING_ASYNC_FLAG are tightly scoped to this coroutine's frame.
+            """
+            token = RUNNING_ASYNC_FLAG.set(True)
+            try:
+                result = await async_fn(*args, **kwargs)
+            finally:
+                RUNNING_ASYNC_FLAG.reset(token)
+            return result
+
+        if _sync is True:
+            return run_coro_as_sync(ctx_call())
+        elif _sync is False or RUNNING_ASYNC_FLAG.get() or is_async:
+            return ctx_call()
         else:
-            logger.debug(f"{async_fn} --> run async in new loop")
-            # Run in a new event loop, but use a `Call` for nested context detection
-            call = create_call(async_fn, *args, **kwargs)
-            return call()
+            return run_coro_as_sync(ctx_call())
 
-    # TODO: This is breaking type hints on the callable... mypy is behind the curve
-    #       on argument annotations. We can still fix this for editors though.
     if is_async_fn(async_fn):
         wrapper = coroutine_wrapper
     elif is_async_gen_fn(async_fn):
@@ -319,13 +398,13 @@ def sync_compatible(async_fn: T) -> T:
     else:
         raise TypeError("The decorated function must be async.")
 
-    wrapper.aio = async_fn
+    wrapper.aio = async_fn  # type: ignore
     return wrapper
 
 
 @asynccontextmanager
-async def asyncnullcontext():
-    yield
+async def asyncnullcontext(value=None, *args, **kwargs):
+    yield value
 
 
 def sync(__async_fn: Callable[P, Awaitable[T]], *args: P.args, **kwargs: P.kwargs) -> T:
@@ -340,7 +419,7 @@ def sync(__async_fn: Callable[P, Awaitable[T]], *args: P.args, **kwargs: P.kwarg
             "`sync` called from an asynchronous context; "
             "you should `await` the async function directly instead."
         )
-        with start_blocking_portal() as portal:
+        with anyio.start_blocking_portal() as portal:
             return portal.call(partial(__async_fn, *args, **kwargs))
     elif in_async_worker_thread():
         # In a sync context but we can access the event loop thread; send the async