prefect-client 2.19.4__py3-none-any.whl → 3.0.0rc2__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,7 +18,6 @@ from typing import (
     Dict,
     List,
     Optional,
-    Type,
     TypeVar,
     Union,
     cast,
@@ -29,9 +26,16 @@ from uuid import UUID, uuid4
 
 import anyio
 import anyio.abc
+import anyio.from_thread
+import anyio.to_thread
 import sniffio
 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")
@@ -46,6 +50,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()
 
 
@@ -82,12 +94,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:
+    """
+    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)
     """
-    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.
+
+    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.
@@ -96,15 +143,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()
@@ -113,102 +170,99 @@ 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,
+) -> R:
     """
-    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, cancellable=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():
+        """
+        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)
+
+    # 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)
+        return call.result()
 
 
-async def run_sync_in_interruptible_worker_thread(
+async def run_sync_in_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`.
-    """
-
-    class NotSet:
-        pass
+    Runs a sync function in a new worker thread so that the main thread's event loop
+    is not blocked.
 
-    thread: Thread = None
-    result = NotSet
-    event = asyncio.Event()
-    loop = asyncio.get_running_loop()
+    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.
 
-    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)
+    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)
+    result = await anyio.to_thread.run_sync(
+        call_with_mark, call, abandon_on_cancel=True, limiter=get_thread_limiter()
+    )
+    return result
 
-    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
-
-    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,
-                cancellable=True,
-                limiter=get_thread_limiter(),
-            )
-        )
 
-    assert result is not NotSet
-    return result
+def call_with_mark(call):
+    mark_as_worker_thread()
+    return call()
 
 
 def run_async_from_worker_thread(
@@ -226,13 +280,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_module
-    except AttributeError:
-        return False
-    else:
-        return True
+    return getattr(_thread_local, "is_worker_thread", False)
 
 
 def in_async_main_thread() -> bool:
@@ -245,7 +298,7 @@ def in_async_main_thread() -> bool:
         return not in_async_worker_thread()
 
 
-def sync_compatible(async_fn: T) -> T:
+def sync_compatible(async_fn: T, force_sync: bool = False) -> T:
     """
     Converts an async function into a dual async and sync function.
 
@@ -261,47 +314,53 @@ def sync_compatible(async_fn: T) -> T:
     """
 
     @wraps(async_fn)
-    def coroutine_wrapper(*args, **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:
-            return async_fn(*args, **kwargs)
-
-        global_thread_portal = get_global_loop()
-        current_thread = threading.current_thread()
-        current_call = get_current_call()
-        current_loop = get_running_loop()
+    def coroutine_wrapper(*args, _sync: bool = None, **kwargs):
+        from prefect.context import MissingContextError, get_run_context
+        from prefect.settings import (
+            PREFECT_EXPERIMENTAL_DISABLE_SYNC_COMPAT,
+        )
 
-        if current_thread.ident == global_thread_portal.thread.ident:
-            logger.debug(f"{async_fn} --> return coroutine for internal await")
-            # In the prefect async context; return the coro for us to await
+        if PREFECT_EXPERIMENTAL_DISABLE_SYNC_COMPAT or _sync is False:
             return async_fn(*args, **kwargs)
-        elif in_async_main_thread() and (
-            not current_call or is_async_fn(current_call.fn)
-        ):
-            # 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.

prefect/utilities/callables.py

@@ -5,36 +5,33 @@ Utilities for working with Python callables.
 import ast
 import importlib.util
 import inspect
+import warnings
 from functools import partial
 from pathlib import Path
 from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
 
 import cloudpickle
-
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-from prefect._internal.pydantic.v1_schema import has_v1_type_as_param
-
-if HAS_PYDANTIC_V2:
-    import pydantic.v1 as pydantic
-
-    from prefect._internal.pydantic.v2_schema import (
-        create_v2_schema,
-        process_v2_params,
-    )
-else:
-    import pydantic
-
+import pydantic
 from griffe.dataclasses import Docstring
 from griffe.docstrings.dataclasses import DocstringSectionKind
 from griffe.docstrings.parsers import Parser, parse
 from typing_extensions import Literal
 
+from prefect._internal.pydantic.v1_schema import has_v1_type_as_param
+from prefect._internal.pydantic.v2_schema import (
+    create_v2_schema,
+    process_v2_params,
+)
 from prefect.exceptions import (
+    MappingLengthMismatch,
+    MappingMissingIterable,
     ParameterBindError,
     ReservedArgumentError,
     SignatureMismatchError,
 )
 from prefect.logging.loggers import disable_logger, get_logger
+from prefect.utilities.annotations import allow_failure, quote, unmapped
+from prefect.utilities.collections import isiterable
 from prefect.utilities.importtools import safe_load_namespace
 
 logger = get_logger(__name__)
@@ -230,15 +227,14 @@ class ParameterSchema(pydantic.BaseModel):
     title: Literal["Parameters"] = "Parameters"
     type: Literal["object"] = "object"
     properties: Dict[str, Any] = pydantic.Field(default_factory=dict)
-    required: List[str] = None
-    definitions: Optional[Dict[str, Any]] = None
+    required: List[str] = pydantic.Field(default_factory=list)
+    definitions: Dict[str, Any] = pydantic.Field(default_factory=dict)
 
-    def dict(self, *args, **kwargs):
-        """Exclude `None` fields by default to comply with
-        the OpenAPI spec.
-        """
-        kwargs.setdefault("exclude_none", True)
-        return super().dict(*args, **kwargs)
+    def model_dump_for_openapi(self) -> Dict[str, Any]:
+        result = self.model_dump(mode="python", exclude_none=True)
+        if "required" in result and not result["required"]:
+            del result["required"]
+        return result
 
 
 def parameter_docstrings(docstring: Optional[str]) -> Dict[str, str]:
@@ -286,21 +282,31 @@ def process_v1_params(
         name = param.name
 
     type_ = Any if param.annotation is inspect._empty else param.annotation
-    field = pydantic.Field(
-        default=... if param.default is param.empty else param.default,
-        title=param.name,
-        description=docstrings.get(param.name, None),
-        alias=aliases.get(name),
-        position=position,
-    )
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", category=pydantic.warnings.PydanticDeprecatedSince20
+        )
+        field = pydantic.Field(
+            default=... if param.default is param.empty else param.default,
+            title=param.name,
+            description=docstrings.get(param.name, None),
+            alias=aliases.get(name),
+            position=position,
+        )
     return name, type_, field
 
 
 def create_v1_schema(name_: str, model_cfg, **model_fields):
-    model: "pydantic.BaseModel" = pydantic.create_model(
-        name_, __config__=model_cfg, **model_fields
-    )
-    return model.schema(by_alias=True)
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", category=pydantic.warnings.PydanticDeprecatedSince20
+        )
+
+        model: "pydantic.BaseModel" = pydantic.create_model(
+            name_, __config__=model_cfg, **model_fields
+        )
+        return model.schema(by_alias=True)
 
 
 def parameter_schema(fn: Callable) -> ParameterSchema:
@@ -381,16 +387,20 @@ def generate_parameter_schema(
     model_fields = {}
     aliases = {}
 
-    class ModelConfig:
-        arbitrary_types_allowed = True
-
-    if HAS_PYDANTIC_V2 and not has_v1_type_as_param(signature):
+    if not has_v1_type_as_param(signature):
         create_schema = create_v2_schema
         process_params = process_v2_params
+
+        config = pydantic.ConfigDict(arbitrary_types_allowed=True)
     else:
         create_schema = create_v1_schema
         process_params = process_v1_params
 
+        class ModelConfig:
+            arbitrary_types_allowed = True
+
+        config = ModelConfig
+
     for position, param in enumerate(signature.parameters.values()):
         name, type_, field = process_params(
             param, position=position, docstrings=docstrings, aliases=aliases
@@ -398,16 +408,14 @@ def generate_parameter_schema(
         # Generate a Pydantic model at each step so we can check if this parameter
         # type supports schema generation
         try:
-            create_schema(
-                "CheckParameter", model_cfg=ModelConfig, **{name: (type_, field)}
-            )
+            create_schema("CheckParameter", model_cfg=config, **{name: (type_, field)})
         except (ValueError, TypeError):
             # This field's type is not valid for schema creation, update it to `Any`
             type_ = Any
         model_fields[name] = (type_, field)
 
     # Generate the final model and schema
-    schema = create_schema("Parameters", model_cfg=ModelConfig, **model_fields)
+    schema = create_schema("Parameters", model_cfg=config, **model_fields)
     return ParameterSchema(**schema)
 
 
@@ -621,3 +629,75 @@ def _get_docstring_from_source(source_code: str, func_name: str) -> Optional[str
     ):
         return func_def.body[0].value.value
     return None
+
+
+def expand_mapping_parameters(
+    func: Callable, parameters: Dict[str, Any]
+) -> List[Dict[str, Any]]:
+    """
+    Generates a list of call parameters to be used for individual calls in a mapping
+    operation.
+
+    Args:
+        func: The function to be called
+        parameters: A dictionary of parameters with iterables to be mapped over
+
+    Returns:
+        List: A list of dictionaries to be used as parameters for each
+            call in the mapping operation
+    """
+    # Ensure that any parameters in kwargs are expanded before this check
+    parameters = explode_variadic_parameter(func, parameters)
+
+    iterable_parameters = {}
+    static_parameters = {}
+    annotated_parameters = {}
+    for key, val in parameters.items():
+        if isinstance(val, (allow_failure, quote)):
+            # Unwrap annotated parameters to determine if they are iterable
+            annotated_parameters[key] = val
+            val = val.unwrap()
+
+        if isinstance(val, unmapped):
+            static_parameters[key] = val.value
+        elif isiterable(val):
+            iterable_parameters[key] = list(val)
+        else:
+            static_parameters[key] = val
+
+    if not len(iterable_parameters):
+        raise MappingMissingIterable(
+            "No iterable parameters were received. Parameters for map must "
+            f"include at least one iterable. Parameters: {parameters}"
+        )
+
+    iterable_parameter_lengths = {
+        key: len(val) for key, val in iterable_parameters.items()
+    }
+    lengths = set(iterable_parameter_lengths.values())
+    if len(lengths) > 1:
+        raise MappingLengthMismatch(
+            "Received iterable parameters with different lengths. Parameters for map"
+            f" must all be the same length. Got lengths: {iterable_parameter_lengths}"
+        )
+
+    map_length = list(lengths)[0]
+
+    call_parameters_list = []
+    for i in range(map_length):
+        call_parameters = {key: value[i] for key, value in iterable_parameters.items()}
+        call_parameters.update({key: value for key, value in static_parameters.items()})
+
+        # Add default values for parameters; these are skipped earlier since they should
+        # not be mapped over
+        for key, value in get_parameter_defaults(func).items():
+            call_parameters.setdefault(key, value)
+
+        # Re-apply annotations to each key again
+        for key, annotation in annotated_parameters.items():
+            call_parameters[key] = annotation.rewrap(call_parameters[key])
+
+        # Collapse any previously exploded kwargs
+        call_parameters_list.append(collapse_variadic_parameters(func, call_parameters))
+
+    return call_parameters_list