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

prefect/task_runners.py

@@ -1,365 +1,461 @@
-"""
-Interface and implementations of various task runners.
-
-[Task Runners](/concepts/task-runners/) in Prefect are responsible for managing the execution of Prefect task runs. Generally speaking, users are not expected to interact with task runners outside of configuring and initializing them for a flow.
-
-Example:
-    ```
-    >>> from prefect import flow, task
-    >>> from prefect.task_runners import SequentialTaskRunner
-    >>> from typing import List
-    >>>
-    >>> @task
-    >>> def say_hello(name):
-    ...     print(f"hello {name}")
-    >>>
-    >>> @task
-    >>> def say_goodbye(name):
-    ...     print(f"goodbye {name}")
-    >>>
-    >>> @flow(task_runner=SequentialTaskRunner())
-    >>> def greetings(names: List[str]):
-    ...     for name in names:
-    ...         say_hello(name)
-    ...         say_goodbye(name)
-    >>>
-    >>> greetings(["arthur", "trillian", "ford", "marvin"])
-    hello arthur
-    goodbye arthur
-    hello trillian
-    goodbye trillian
-    hello ford
-    goodbye ford
-    hello marvin
-    goodbye marvin
-    ```
-
-    Switching to a `DaskTaskRunner`:
-    ```
-    >>> from prefect_dask.task_runners import DaskTaskRunner
-    >>> flow.task_runner = DaskTaskRunner()
-    >>> greetings(["arthur", "trillian", "ford", "marvin"])
-    hello arthur
-    goodbye arthur
-    hello trillian
-    hello ford
-    goodbye marvin
-    hello marvin
-    goodbye ford
-    goodbye trillian
-    ```
-
-For usage details, see the [Task Runners](/concepts/task-runners/) documentation.
-"""
 import abc
-from contextlib import AsyncExitStack, asynccontextmanager
+import asyncio
+import sys
+import threading
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from contextvars import copy_context
 from typing import (
     TYPE_CHECKING,
     Any,
-    AsyncIterator,
-    Awaitable,
-    Callable,
+    Coroutine,
     Dict,
+    Generic,
+    Iterable,
+    List,
     Optional,
     Set,
-    TypeVar,
+    overload,
 )
-from uuid import UUID
 
-import anyio
+from typing_extensions import ParamSpec, Self, TypeVar
 
-from prefect._internal.concurrency.primitives import Event
-from prefect.client.schemas.objects import State
-from prefect.logging import get_logger
-from prefect.states import exception_to_crashed_state
-from prefect.utilities.collections import AutoEnum
+from prefect.client.schemas.objects import TaskRunInput
+from prefect.exceptions import MappingLengthMismatch, MappingMissingIterable
+from prefect.futures import (
+    PrefectConcurrentFuture,
+    PrefectDistributedFuture,
+    PrefectFuture,
+    PrefectFutureList,
+)
+from prefect.logging.loggers import get_logger, get_run_logger
+from prefect.utilities.annotations import allow_failure, quote, unmapped
+from prefect.utilities.callables import (
+    collapse_variadic_parameters,
+    explode_variadic_parameter,
+    get_parameter_defaults,
+)
+from prefect.utilities.collections import isiterable
 
 if TYPE_CHECKING:
-    import anyio.abc
+    from prefect.tasks import Task
 
-
-T = TypeVar("T", bound="BaseTaskRunner")
+P = ParamSpec("P")
+T = TypeVar("T")
 R = TypeVar("R")
+F = TypeVar("F", bound=PrefectFuture, default=PrefectConcurrentFuture)
 
 
-class TaskConcurrencyType(AutoEnum):
-    SEQUENTIAL = AutoEnum.auto()
-    CONCURRENT = AutoEnum.auto()
-    PARALLEL = AutoEnum.auto()
-
+class TaskRunner(abc.ABC, Generic[F]):
+    """
+    Abstract base class for task runners.
 
-CONCURRENCY_MESSAGES = {
-    TaskConcurrencyType.SEQUENTIAL: "sequentially",
-    TaskConcurrencyType.CONCURRENT: "concurrently",
-    TaskConcurrencyType.PARALLEL: "in parallel",
-}
+    A task runner is responsible for submitting tasks to the task run engine running
+    in an execution environment. Submitted tasks are non-blocking and return a future
+    object that can be used to wait for the task to complete and retrieve the result.
 
+    Task runners are context managers and should be used in a `with` block to ensure
+    proper cleanup of resources.
+    """
 
-class BaseTaskRunner(metaclass=abc.ABCMeta):
-    def __init__(self) -> None:
+    def __init__(self):
         self.logger = get_logger(f"task_runner.{self.name}")
-        self._started: bool = False
-
-    @property
-    @abc.abstractmethod
-    def concurrency_type(self) -> TaskConcurrencyType:
-        pass  # noqa
+        self._started = False
 
     @property
     def name(self):
+        """The name of this task runner"""
         return type(self).__name__.lower().replace("taskrunner", "")
 
-    def duplicate(self):
-        """
-        Return a new task runner instance with the same options.
-        """
-        # The base class returns `NotImplemented` to indicate that this is not yet
-        # implemented by a given task runner.
-        return NotImplemented
-
-    def __eq__(self, other: object) -> bool:
-        """
-        Returns true if the task runners use the same options.
-        """
-        if type(other) == type(self) and (
-            # Compare public attributes for naive equality check
-            # Subclasses should implement this method with a check init option equality
-            {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
-            == {k: v for k, v in other.__dict__.items() if not k.startswith("_")}
-        ):
-            return True
-        else:
-            return NotImplemented
+    @abc.abstractmethod
+    def duplicate(self) -> Self:
+        """Return a new instance of this task runner with the same configuration."""
+        ...
 
     @abc.abstractmethod
-    async def submit(
+    def submit(
         self,
-        key: UUID,
-        call: Callable[..., Awaitable[State[R]]],
-    ) -> None:
+        task: "Task",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+    ) -> F:
         """
-        Submit a call for execution and return a `PrefectFuture` that can be used to
-        get the call result.
+        Submit a task to the task run engine.
 
         Args:
-            task_run: The task run being submitted.
-            task_key: A unique key for this orchestration run of the task. Can be used
-                for caching.
-            call: The function to be executed
-            run_kwargs: A dict of keyword arguments to pass to `call`
+            task: The task to submit.
+            parameters: The parameters to use when running the task.
+            wait_for: A list of futures that the task depends on.
 
         Returns:
-            A future representing the result of `call` execution
+            A future object that can be used to wait for the task to complete and
+            retrieve the result.
         """
-        raise NotImplementedError()
+        ...
 
-    @abc.abstractmethod
-    async def wait(self, key: UUID, timeout: float = None) -> Optional[State]:
+    def map(
+        self,
+        task: "Task",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+    ) -> PrefectFutureList[F]:
         """
-        Given a `PrefectFuture`, wait for its return state up to `timeout` seconds.
-        If it is not finished after the timeout expires, `None` should be returned.
+        Submit multiple tasks to the task run engine.
 
-        Implementers should be careful to ensure that this function never returns or
-        raises an exception.
-        """
-        raise NotImplementedError()
+        Args:
+            task: The task to submit.
+            parameters: The parameters to use when running the task.
+            wait_for: A list of futures that the task depends on.
 
-    @asynccontextmanager
-    async def start(
-        self: T,
-    ) -> AsyncIterator[T]:
+        Returns:
+            An iterable of future objects that can be used to wait for the tasks to
+            complete and retrieve the results.
         """
-        Start the task runner, preparing any resources necessary for task submission.
+        if not self._started:
+            raise RuntimeError(
+                "The task runner must be started before submitting work."
+            )
 
-        Children should implement `_start` to prepare and clean up resources.
+        from prefect.utilities.engine import (
+            collect_task_run_inputs_sync,
+            resolve_inputs_sync,
+        )
 
-        Yields:
-            The prepared task runner
-        """
-        if self._started:
-            raise RuntimeError("The task runner is already started!")
-
-        async with AsyncExitStack() as exit_stack:
-            self.logger.debug("Starting task runner...")
-            try:
-                await self._start(exit_stack)
-                self._started = True
-                yield self
-            finally:
-                self.logger.debug("Shutting down task runner...")
-                self._started = False
-
-    async def _start(self, exit_stack: AsyncExitStack) -> None:
-        """
-        Create any resources required for this task runner to submit work.
+        # We need to resolve some futures to map over their data, collect the upstream
+        # links beforehand to retain relationship tracking.
+        task_inputs = {
+            k: collect_task_run_inputs_sync(v, max_depth=0)
+            for k, v in parameters.items()
+        }
+
+        # Resolve the top-level parameters in order to get mappable data of a known length.
+        # Nested parameters will be resolved in each mapped child where their relationships
+        # will also be tracked.
+        parameters = resolve_inputs_sync(parameters, max_depth=0)
+
+        # Ensure that any parameters in kwargs are expanded before this check
+        parameters = explode_variadic_parameter(task.fn, 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}"
+            )
 
-        Cleanup of resources should be submitted to the `exit_stack`.
-        """
-        pass  # noqa
+        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}"
+            )
 
-    def __str__(self) -> str:
-        return type(self).__name__
+        map_length = list(lengths)[0]
 
+        futures: List[PrefectFuture] = []
+        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()}
+            )
 
-class SequentialTaskRunner(BaseTaskRunner):
-    """
-    A simple task runner that executes calls as they are submitted.
+            # Add default values for parameters; these are skipped earlier since they should
+            # not be mapped over
+            for key, value in get_parameter_defaults(task.fn).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 = collapse_variadic_parameters(task.fn, call_parameters)
+
+            futures.append(
+                self.submit(
+                    task=task,
+                    parameters=call_parameters,
+                    wait_for=wait_for,
+                    dependencies=task_inputs,
+                )
+            )
 
-    If writing synchronous tasks, this runner will always execute tasks sequentially.
-    If writing async tasks, this runner will execute tasks sequentially unless grouped
-    using `anyio.create_task_group` or `asyncio.gather`.
-    """
+        return PrefectFutureList(futures)
 
-    def __init__(self) -> None:
-        super().__init__()
-        self._results: Dict[str, State] = {}
+    def __enter__(self):
+        if self._started:
+            raise RuntimeError("This task runner is already started")
 
-    @property
-    def concurrency_type(self) -> TaskConcurrencyType:
-        return TaskConcurrencyType.SEQUENTIAL
+        self.logger.debug("Starting task runner")
+        self._started = True
+        return self
 
-    def duplicate(self):
-        return type(self)()
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.logger.debug("Stopping task runner")
+        self._started = False
 
-    async def submit(
-        self,
-        key: UUID,
-        call: Callable[..., Awaitable[State[R]]],
-    ) -> None:
-        # Run the function immediately and store the result in memory
-        try:
-            result = await call()
-        except BaseException as exc:
-            result = await exception_to_crashed_state(exc)
 
-        self._results[key] = result
+class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
+    def __init__(self, max_workers: Optional[int] = None):
+        super().__init__()
+        self._executor: Optional[ThreadPoolExecutor] = None
+        self._max_workers = sys.maxsize if max_workers is None else max_workers
+        self._cancel_events: Dict[uuid.UUID, threading.Event] = {}
 
-    async def wait(self, key: UUID, timeout: float = None) -> Optional[State]:
-        return self._results[key]
+    def duplicate(self) -> "ThreadPoolTaskRunner":
+        return type(self)(max_workers=self._max_workers)
 
+    @overload
+    def submit(
+        self,
+        task: "Task[P, Coroutine[Any, Any, R]]",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+    ) -> PrefectConcurrentFuture[R]:
+        ...
+
+    @overload
+    def submit(
+        self,
+        task: "Task[Any, R]",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+    ) -> PrefectConcurrentFuture[R]:
+        ...
+
+    def submit(
+        self,
+        task: "Task",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+    ):
+        """
+        Submit a task to the task run engine running in a separate thread.
 
-class ConcurrentTaskRunner(BaseTaskRunner):
-    """
-    A concurrent task runner that allows tasks to switch when blocking on IO.
-    Synchronous tasks will be submitted to a thread pool maintained by `anyio`.
-
-    Example:
-        ```
-        Using a thread for concurrency:
-        >>> from prefect import flow
-        >>> from prefect.task_runners import ConcurrentTaskRunner
-        >>> @flow(task_runner=ConcurrentTaskRunner)
-        >>> def my_flow():
-        >>>     ...
-        ```
-    """
+        Args:
+            task: The task to submit.
+            parameters: The parameters to use when running the task.
+            wait_for: A list of futures that the task depends on.
 
-    def __init__(self):
-        # TODO: Consider adding `max_workers` support using anyio capacity limiters
+        Returns:
+            A future object that can be used to wait for the task to complete and
+            retrieve the result.
+        """
+        if not self._started or self._executor is None:
+            raise RuntimeError("Task runner is not started")
 
-        # Runtime attributes
-        self._task_group: anyio.abc.TaskGroup = None
-        self._result_events: Dict[UUID, Event] = {}
-        self._results: Dict[UUID, Any] = {}
-        self._keys: Set[UUID] = set()
+        from prefect.context import FlowRunContext
+        from prefect.task_engine import run_task_async, run_task_sync
 
-        super().__init__()
+        task_run_id = uuid.uuid4()
+        cancel_event = threading.Event()
+        self._cancel_events[task_run_id] = cancel_event
+        context = copy_context()
 
-    @property
-    def concurrency_type(self) -> TaskConcurrencyType:
-        return TaskConcurrencyType.CONCURRENT
+        flow_run_ctx = FlowRunContext.get()
+        if flow_run_ctx:
+            get_run_logger(flow_run_ctx).info(
+                f"Submitting task {task.name} to thread pool executor..."
+            )
+        else:
+            self.logger.info(f"Submitting task {task.name} to thread pool executor...")
+
+        submit_kwargs = dict(
+            task=task,
+            task_run_id=task_run_id,
+            parameters=parameters,
+            wait_for=wait_for,
+            return_type="state",
+            dependencies=dependencies,
+            context=dict(cancel_event=cancel_event),
+        )
 
-    def duplicate(self):
-        return type(self)()
+        if task.isasync:
+            # TODO: Explore possibly using a long-lived thread with an event loop
+            # for better performance
+            future = self._executor.submit(
+                context.run,
+                asyncio.run,
+                run_task_async(**submit_kwargs),
+            )
+        else:
+            future = self._executor.submit(
+                context.run,
+                run_task_sync,
+                **submit_kwargs,
+            )
+        prefect_future = PrefectConcurrentFuture(
+            task_run_id=task_run_id, wrapped_future=future
+        )
+        return prefect_future
 
-    async def submit(
+    @overload
+    def map(
         self,
-        key: UUID,
-        call: Callable[[], Awaitable[State[R]]],
-    ) -> None:
-        if not self._started:
-            raise RuntimeError(
-                "The task runner must be started before submitting work."
-            )
+        task: "Task[P, Coroutine[Any, Any, R]]",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+    ) -> PrefectFutureList[PrefectConcurrentFuture[R]]:
+        ...
+
+    @overload
+    def map(
+        self,
+        task: "Task[Any, R]",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+    ) -> PrefectFutureList[PrefectConcurrentFuture[R]]:
+        ...
 
-        if not self._task_group:
-            raise RuntimeError(
-                "The concurrent task runner cannot be used to submit work after "
-                "serialization."
-            )
+    def map(
+        self,
+        task: "Task",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+    ):
+        return super().map(task, parameters, wait_for)
 
-        # Create an event to set on completion
-        self._result_events[key] = Event()
+    def cancel_all(self):
+        for event in self._cancel_events.values():
+            event.set()
+            self.logger.debug("Set cancel event")
 
-        # Rely on the event loop for concurrency
-        self._task_group.start_soon(self._run_and_store_result, key, call)
+        if self._executor is not None:
+            self._executor.shutdown(cancel_futures=True)
+            self._executor = None
 
-    async def wait(
-        self,
-        key: UUID,
-        timeout: float = None,
-    ) -> Optional[State]:
-        if not self._task_group:
-            raise RuntimeError(
-                "The concurrent task runner cannot be used to wait for work after "
-                "serialization."
-            )
+    def __enter__(self):
+        super().__enter__()
+        self._executor = ThreadPoolExecutor(max_workers=self._max_workers)
+        return self
 
-        return await self._get_run_result(key, timeout)
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.cancel_all()
+        if self._executor is not None:
+            self._executor.shutdown(cancel_futures=True)
+            self._executor = None
+        super().__exit__(exc_type, exc_value, traceback)
 
-    async def _run_and_store_result(
-        self, key: UUID, call: Callable[[], Awaitable[State[R]]]
-    ):
-        """
-        Simple utility to store the orchestration result in memory on completion
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, ThreadPoolTaskRunner):
+            return False
+        return self._max_workers == value._max_workers
 
-        Since this run is occurring on the main thread, we capture exceptions to prevent
-        task crashes from crashing the flow run.
-        """
-        try:
-            result = await call()
-        except BaseException as exc:
-            result = await exception_to_crashed_state(exc)
 
-        self._results[key] = result
-        self._result_events[key].set()
+# Here, we alias ConcurrentTaskRunner to ThreadPoolTaskRunner for backwards compatibility
+ConcurrentTaskRunner = ThreadPoolTaskRunner
 
-    async def _get_run_result(
-        self, key: UUID, timeout: float = None
-    ) -> Optional[State]:
-        """
-        Block until the run result has been populated.
-        """
-        result = None  # retval on timeout
 
-        # Note we do not use `asyncio.wrap_future` and instead use an `Event` to avoid
-        # stdlib behavior where the wrapped future is cancelled if the parent future is
-        # cancelled (as it would be during a timeout here)
-        with anyio.move_on_after(timeout):
-            await self._result_events[key].wait()
-            result = self._results[key]
+class PrefectTaskRunner(TaskRunner[PrefectDistributedFuture]):
+    def __init__(self):
+        super().__init__()
 
-        return result  # timeout reached
+    def duplicate(self) -> "PrefectTaskRunner":
+        return type(self)()
 
-    async def _start(self, exit_stack: AsyncExitStack):
+    @overload
+    def submit(
+        self,
+        task: "Task[P, Coroutine[Any, Any, R]]",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+    ) -> PrefectDistributedFuture[R]:
+        ...
+
+    @overload
+    def submit(
+        self,
+        task: "Task[Any, R]",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+    ) -> PrefectDistributedFuture[R]:
+        ...
+
+    def submit(
+        self,
+        task: "Task",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+    ):
         """
-        Start the process pool
+        Submit a task to the task run engine running in a separate thread.
+
+        Args:
+            task: The task to submit.
+            parameters: The parameters to use when running the task.
+            wait_for: A list of futures that the task depends on.
+
+        Returns:
+            A future object that can be used to wait for the task to complete and
+            retrieve the result.
         """
-        self._task_group = await exit_stack.enter_async_context(
-            anyio.create_task_group()
+        if not self._started:
+            raise RuntimeError("Task runner is not started")
+        from prefect.context import FlowRunContext
+
+        flow_run_ctx = FlowRunContext.get()
+        if flow_run_ctx:
+            get_run_logger(flow_run_ctx).info(
+                f"Submitting task {task.name} to for execution by a Prefect task worker..."
+            )
+        else:
+            self.logger.info(
+                f"Submitting task {task.name} to for execution by a Prefect task worker..."
+            )
+
+        return task.apply_async(
+            kwargs=parameters, wait_for=wait_for, dependencies=dependencies
         )
 
-    def __getstate__(self):
-        """
-        Allow the `ConcurrentTaskRunner` to be serialized by dropping the task group.
-        """
-        data = self.__dict__.copy()
-        data.update({k: None for k in {"_task_group"}})
-        return data
+    @overload
+    def map(
+        self,
+        task: "Task[P, Coroutine[Any, Any, R]]",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+    ) -> PrefectFutureList[PrefectDistributedFuture[R]]:
+        ...
+
+    @overload
+    def map(
+        self,
+        task: "Task[Any, R]",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+    ) -> PrefectFutureList[PrefectDistributedFuture[R]]:
+        ...
 
-    def __setstate__(self, data: dict):
-        """
-        When deserialized, we will no longer have a reference to the task group.
-        """
-        self.__dict__.update(data)
-        self._task_group = None
+    def map(
+        self,
+        task: "Task",
+        parameters: Dict[str, Any],
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+    ):
+        return super().map(task, parameters, wait_for)