prefect-client 3.0.0rc1__py3-none-any.whl → 3.0.0rc2__py3-none-any.whl

prefect/task_runs.py

@@ -0,0 +1,203 @@
+import asyncio
+import atexit
+import threading
+import uuid
+from typing import Dict, Optional
+
+import anyio
+from cachetools import TTLCache
+from typing_extensions import Self
+
+from prefect._internal.concurrency.api import create_call, from_async, from_sync
+from prefect._internal.concurrency.threads import get_global_loop
+from prefect.client.schemas.objects import TERMINAL_STATES
+from prefect.events.clients import get_events_subscriber
+from prefect.events.filters import EventFilter, EventNameFilter
+from prefect.logging.loggers import get_logger
+
+
+class TaskRunWaiter:
+    """
+    A service used for waiting for a task run to finish.
+
+    This service listens for task run events and provides a way to wait for a specific
+    task run to finish. This is useful for waiting for a task run to finish before
+    continuing execution.
+
+    The service is a singleton and must be started before use. The service will
+    automatically start when the first instance is created. A single websocket
+    connection is used to listen for task run events.
+
+    The service can be used to wait for a task run to finish by calling
+    `TaskRunWaiter.wait_for_task_run` with the task run ID to wait for. The method
+    will return when the task run has finished or the timeout has elapsed.
+
+    The service will automatically stop when the Python process exits or when the
+    global loop thread is stopped.
+
+    Example:
+    ```python
+    import asyncio
+    from uuid import uuid4
+
+    from prefect import task
+    from prefect.task_engine import run_task_async
+    from prefect.task_runs import TaskRunWaiter
+
+
+    @task
+    async def test_task():
+        await asyncio.sleep(5)
+        print("Done!")
+
+
+    async def main():
+        task_run_id = uuid4()
+        asyncio.create_task(run_task_async(task=test_task, task_run_id=task_run_id))
+
+        await TaskRunWaiter.wait_for_task_run(task_run_id)
+        print("Task run finished")
+
+
+    if __name__ == "__main__":
+        asyncio.run(main())
+    ```
+    """
+
+    _instance: Optional[Self] = None
+    _instance_lock = threading.Lock()
+
+    def __init__(self):
+        self.logger = get_logger("TaskRunWaiter")
+        self._consumer_task: Optional[asyncio.Task] = None
+        self._observed_completed_task_runs: TTLCache[uuid.UUID, bool] = TTLCache(
+            maxsize=100, ttl=60
+        )
+        self._completion_events: Dict[uuid.UUID, asyncio.Event] = {}
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._observed_completed_task_runs_lock = threading.Lock()
+        self._completion_events_lock = threading.Lock()
+        self._started = False
+
+    def start(self):
+        """
+        Start the TaskRunWaiter service.
+        """
+        if self._started:
+            return
+        self.logger.info("Starting TaskRunWaiter")
+        loop_thread = get_global_loop()
+
+        if not asyncio.get_running_loop() == loop_thread._loop:
+            raise RuntimeError("TaskRunWaiter must run on the global loop thread.")
+
+        self._loop = loop_thread._loop
+        self._consumer_task = self._loop.create_task(self._consume_events())
+
+        loop_thread.add_shutdown_call(create_call(self.stop))
+        atexit.register(self.stop)
+        self._started = True
+
+    async def _consume_events(self):
+        async with get_events_subscriber(
+            filter=EventFilter(
+                event=EventNameFilter(
+                    name=[
+                        f"prefect.task-run.{state.name.title()}"
+                        for state in TERMINAL_STATES
+                    ],
+                )
+            )
+        ) as subscriber:
+            async for event in subscriber:
+                try:
+                    self.logger.info(
+                        f"Received event: {event.resource['prefect.resource.id']}"
+                    )
+                    task_run_id = uuid.UUID(
+                        event.resource["prefect.resource.id"].replace(
+                            "prefect.task-run.", ""
+                        )
+                    )
+                    with self._observed_completed_task_runs_lock:
+                        # Cache the task run ID for a short period of time to avoid
+                        # unnecessary waits
+                        self._observed_completed_task_runs[task_run_id] = True
+                    with self._completion_events_lock:
+                        # Set the event for the task run ID if it is in the cache
+                        # so the waiter can wake up the waiting coroutine
+                        if task_run_id in self._completion_events:
+                            self._completion_events[task_run_id].set()
+                except Exception as exc:
+                    self.logger.error(f"Error processing event: {exc}")
+
+    def stop(self):
+        """
+        Stop the TaskRunWaiter service.
+        """
+        self.logger.debug("Stopping TaskRunWaiter")
+        if self._consumer_task:
+            self._consumer_task.cancel()
+            self._consumer_task = None
+        self.__class__._instance = None
+        self._started = False
+
+    @classmethod
+    async def wait_for_task_run(
+        cls, task_run_id: uuid.UUID, timeout: Optional[float] = None
+    ):
+        """
+        Wait for a task run to finish.
+
+        Note this relies on a websocket connection to receive events from the server
+        and will not work with an ephemeral server.
+
+        Args:
+            task_run_id: The ID of the task run to wait for.
+            timeout: The maximum time to wait for the task run to
+                finish. Defaults to None.
+        """
+        instance = cls.instance()
+        with instance._observed_completed_task_runs_lock:
+            if task_run_id in instance._observed_completed_task_runs:
+                return
+
+        # Need to create event in loop thread to ensure it can be set
+        # from the loop thread
+        finished_event = await from_async.wait_for_call_in_loop_thread(
+            create_call(asyncio.Event)
+        )
+        with instance._completion_events_lock:
+            # Cache the event for the task run ID so the consumer can set it
+            # when the event is received
+            instance._completion_events[task_run_id] = finished_event
+
+        with anyio.move_on_after(delay=timeout):
+            await from_async.wait_for_call_in_loop_thread(
+                create_call(finished_event.wait)
+            )
+
+        with instance._completion_events_lock:
+            # Remove the event from the cache after it has been waited on
+            instance._completion_events.pop(task_run_id, None)
+
+    @classmethod
+    def instance(cls):
+        """
+        Get the singleton instance of TaskRunWaiter.
+        """
+        with cls._instance_lock:
+            if cls._instance is None:
+                cls._instance = cls._new_instance()
+            return cls._instance
+
+    @classmethod
+    def _new_instance(cls):
+        instance = cls()
+
+        if threading.get_ident() == get_global_loop().thread.ident:
+            instance.start()
+        else:
+            from_sync.call_soon_in_loop_thread(create_call(instance.start)).result()
+
+        return instance

prefect/{task_server.py → task_worker.py}

@@ -11,10 +11,12 @@ from typing import List, Optional
 
 import anyio
 import anyio.abc
+from exceptiongroup import BaseExceptionGroup  # novermin
 from websockets.exceptions import InvalidStatusCode
 
-from prefect import Task, get_client
+from prefect import Task
 from prefect._internal.concurrency.api import create_call, from_sync
+from prefect.client.orchestration import get_client
 from prefect.client.schemas.objects import TaskRun
 from prefect.client.subscriptions import Subscription
 from prefect.exceptions import Abort, PrefectHTTPStatusError
@@ -30,11 +32,11 @@ from prefect.utilities.asyncutils import asyncnullcontext, sync_compatible
 from prefect.utilities.engine import emit_task_run_state_change_event, propose_state
 from prefect.utilities.processutils import _register_signal
 
-logger = get_logger("task_server")
+logger = get_logger("task_worker")
 
 
-class StopTaskServer(Exception):
-    """Raised when the task server is stopped."""
+class StopTaskWorker(Exception):
+    """Raised when the task worker is stopped."""
 
     pass
 
@@ -49,11 +51,11 @@ def should_try_to_read_parameters(task: Task, task_run: TaskRun) -> bool:
     return new_enough_state_details and task_accepts_parameters
 
 
-class TaskServer:
+class TaskWorker:
     """This class is responsible for serving tasks that may be executed in the background
     by a task runner via the traditional engine machinery.
 
-    When `start()` is called, the task server will open a websocket connection to a
+    When `start()` is called, the task worker will open a websocket connection to a
     server-side queue of scheduled task runs. When a scheduled task run is found, the
     scheduled task run is submitted to the engine for execution with a minimal `EngineContext`
     so that the task run can be governed by orchestration rules.
@@ -70,7 +72,7 @@ class TaskServer:
         *tasks: Task,
         limit: Optional[int] = 10,
     ):
-        self.tasks: List[Task] = tasks
+        self.tasks: List[Task] = list(tasks)
 
         self.started: bool = False
         self.stopping: bool = False
@@ -80,7 +82,7 @@ class TaskServer:
 
         if not asyncio.get_event_loop().is_running():
             raise RuntimeError(
-                "TaskServer must be initialized within an async context."
+                "TaskWorker must be initialized within an async context."
             )
 
         self._runs_task_group: anyio.abc.TaskGroup = anyio.create_task_group()
@@ -93,7 +95,7 @@ class TaskServer:
 
     def handle_sigterm(self, signum, frame):
         """
-        Shuts down the task server when a SIGTERM is received.
+        Shuts down the task worker when a SIGTERM is received.
         """
         logger.info("SIGTERM received, initiating graceful shutdown...")
         from_sync.call_in_loop_thread(create_call(self.stop))
@@ -103,12 +105,12 @@ class TaskServer:
     @sync_compatible
     async def start(self) -> None:
         """
-        Starts a task server, which runs the tasks provided in the constructor.
+        Starts a task worker, which runs the tasks provided in the constructor.
         """
         _register_signal(signal.SIGTERM, self.handle_sigterm)
 
         async with asyncnullcontext() if self.started else self:
-            logger.info("Starting task server...")
+            logger.info("Starting task worker...")
             try:
                 await self._subscribe_to_task_scheduling()
             except InvalidStatusCode as exc:
@@ -124,17 +126,17 @@ class TaskServer:
 
     @sync_compatible
     async def stop(self):
-        """Stops the task server's polling cycle."""
+        """Stops the task worker's polling cycle."""
         if not self.started:
             raise RuntimeError(
-                "Task server has not yet started. Please start the task server by"
+                "Task worker has not yet started. Please start the task worker by"
                 " calling .start()"
             )
 
         self.started = False
         self.stopping = True
 
-        raise StopTaskServer
+        raise StopTaskWorker
 
     async def _subscribe_to_task_scheduling(self):
         logger.info(
@@ -159,11 +161,11 @@ class TaskServer:
         task = next((t for t in self.tasks if t.task_key == task_run.task_key), None)
 
         if not task:
-            if PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS.value():
+            if PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS:
                 logger.warning(
-                    f"Task {task_run.name!r} not found in task server registry."
+                    f"Task {task_run.name!r} not found in task worker registry."
                 )
-                await self._client._client.delete(f"/task_runs/{task_run.id}")
+                await self._client._client.delete(f"/task_runs/{task_run.id}")  # type: ignore
 
             return
 
@@ -260,14 +262,14 @@ class TaskServer:
             self._limiter.release_on_behalf_of(task_run.id)
 
     async def execute_task_run(self, task_run: TaskRun):
-        """Execute a task run in the task server."""
+        """Execute a task run in the task worker."""
         async with self if not self.started else asyncnullcontext():
             if self._limiter:
                 await self._limiter.acquire_on_behalf_of(task_run.id)
             await self._submit_scheduled_task_run(task_run)
 
     async def __aenter__(self):
-        logger.debug("Starting task server...")
+        logger.debug("Starting task worker...")
 
         if self._client._closed:
             self._client = get_client()
@@ -280,7 +282,7 @@ class TaskServer:
         return self
 
     async def __aexit__(self, *exc_info):
-        logger.debug("Stopping task server...")
+        logger.debug("Stopping task worker...")
         self.started = False
         await self._exit_stack.__aexit__(*exc_info)
 
@@ -300,7 +302,7 @@ async def serve(*tasks: Task, limit: Optional[int] = 10):
     Example:
         ```python
         from prefect import task
-        from prefect.task_server import serve
+        from prefect.task_worker import serve
 
         @task(log_prints=True)
         def say(message: str):
@@ -315,13 +317,21 @@ async def serve(*tasks: Task, limit: Optional[int] = 10):
             serve(say, yell)
         ```
     """
-    task_server = TaskServer(*tasks, limit=limit)
+    task_worker = TaskWorker(*tasks, limit=limit)
 
     try:
-        await task_server.start()
+        await task_worker.start()
+
+    except BaseExceptionGroup as exc:  # novermin
+        exceptions = exc.exceptions
+        n_exceptions = len(exceptions)
+        logger.error(
+            f"Task worker stopped with {n_exceptions} exception{'s' if n_exceptions != 1 else ''}:"
+            f"\n" + "\n".join(str(e) for e in exceptions)
+        )
 
-    except StopTaskServer:
-        logger.info("Task server stopped.")
+    except StopTaskWorker:
+        logger.info("Task worker stopped.")
 
-    except asyncio.CancelledError:
-        logger.info("Task server interrupted, stopping...")
+    except (asyncio.CancelledError, KeyboardInterrupt):
+        logger.info("Task worker interrupted, stopping...")

prefect/tasks.py

@@ -43,6 +43,7 @@ from prefect.context import (
 )
 from prefect.futures import PrefectDistributedFuture, PrefectFuture
 from prefect.logging.loggers import get_logger
+from prefect.records.cache_policies import DEFAULT, CachePolicy
 from prefect.results import ResultFactory, ResultSerializer, ResultStorage
 from prefect.settings import (
     PREFECT_TASK_DEFAULT_RETRIES,
@@ -62,7 +63,6 @@ from prefect.utilities.importtools import to_qualified_name
 if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient
     from prefect.context import TaskRunContext
-    from prefect.task_runners import BaseTaskRunner
     from prefect.transactions import Transaction
 
 T = TypeVar("T")  # Generic type var for capturing the inner return type of async funcs
@@ -145,6 +145,7 @@ class Task(Generic[P, R]):
             tags are combined with any tags defined by a `prefect.tags` context at
             task runtime.
         version: An optional string specifying the version of this task definition
+        cache_policy: A cache policy that determines the level of caching for this task
         cache_key_fn: An optional callable that, given the task run context and call
             parameters, generates a string key; if the key matches a previous completed
             state, that state result will be restored instead of running the task again.
@@ -204,6 +205,7 @@ class Task(Generic[P, R]):
         description: Optional[str] = None,
         tags: Optional[Iterable[str]] = None,
         version: Optional[str] = None,
+        cache_policy: Optional[CachePolicy] = NotSet,
         cache_key_fn: Optional[
             Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
         ] = None,
@@ -303,10 +305,23 @@ class Task(Generic[P, R]):
 
             self.task_key = f"{self.fn.__qualname__}-{task_origin_hash}"
 
+        # TODO: warn of precedence of cache policies and cache key fn if both provided?
+        if cache_key_fn:
+            cache_policy = CachePolicy.from_cache_key_fn(cache_key_fn)
+
+        # TODO: manage expiration and cache refresh
         self.cache_key_fn = cache_key_fn
         self.cache_expiration = cache_expiration
         self.refresh_cache = refresh_cache
 
+        if cache_policy is NotSet and result_storage_key is None:
+            self.cache_policy = DEFAULT
+        elif result_storage_key:
+            # TODO: handle this situation with double storage
+            self.cache_policy = None
+        else:
+            self.cache_policy = cache_policy
+
         # TaskRunPolicy settings
         # TODO: We can instantiate a `TaskRunPolicy` and add Pydantic bound checks to
         #       validate that the user passes positive numbers here
@@ -358,6 +373,7 @@ class Task(Generic[P, R]):
         name: str = None,
         description: str = None,
         tags: Iterable[str] = None,
+        cache_policy: CachePolicy = NotSet,
         cache_key_fn: Callable[
             ["TaskRunContext", Dict[str, Any]], Optional[str]
         ] = None,
@@ -469,6 +485,9 @@ class Task(Generic[P, R]):
             name=name or self.name,
             description=description or self.description,
             tags=tags or copy(self.tags),
+            cache_policy=cache_policy
+            if cache_policy is not NotSet
+            else self.cache_policy,
             cache_key_fn=cache_key_fn or self.cache_key_fn,
             cache_expiration=cache_expiration or self.cache_expiration,
             task_run_name=task_run_name,
@@ -582,7 +601,7 @@ class Task(Generic[P, R]):
             else:
                 state = Pending()
 
-            # store parameters for background tasks so that task servers
+            # store parameters for background tasks so that task worker
             # can retrieve them at runtime
             if deferred and (parameters or wait_for):
                 parameters_id = uuid4()
@@ -755,8 +774,6 @@ class Task(Generic[P, R]):
         """
         Submit a run of the task to the engine.
 
-        If writing an async task, this call must be awaited.
-
         Will create a new task run in the backing API and submit the task to the flow's
         task runner. This call only blocks execution while the task is being submitted,
         once it is submitted, the flow function will continue executing.
@@ -849,7 +866,11 @@ class Task(Generic[P, R]):
         flow_run_context = FlowRunContext.get()
 
         if not flow_run_context:
-            raise ValueError("Task.submit() must be called within a flow")
+            raise RuntimeError(
+                "Unable to determine task runner to use for submission. If you are"
+                " submitting a task outside of a flow, please use `.delay`"
+                " to submit the task run for deferred execution."
+            )
 
         task_viz_tracker = get_task_viz_tracker()
         if task_viz_tracker:
@@ -897,6 +918,7 @@ class Task(Generic[P, R]):
         *args: Any,
         return_state: bool = False,
         wait_for: Optional[Iterable[PrefectFuture]] = None,
+        deferred: bool = False,
         **kwargs: Any,
     ):
         """
@@ -1010,6 +1032,7 @@ class Task(Generic[P, R]):
             [[11, 21], [12, 22], [13, 23]]
         """
 
+        from prefect.task_runners import TaskRunner
         from prefect.utilities.visualization import (
             VisualizationUnsupportedError,
             get_task_viz_tracker,
@@ -1026,22 +1049,22 @@ class Task(Generic[P, R]):
                 "`task.map()` is not currently supported by `flow.visualize()`"
             )
 
-        if not flow_run_context:
-            # TODO: Should we split out background task mapping into a separate method
-            # like we do for the `submit`/`apply_async` split?
+        if deferred:
             parameters_list = expand_mapping_parameters(self.fn, parameters)
-            # TODO: Make this non-blocking once we can return a list of futures
-            # instead of a list of task runs
-            return [
-                run_coro_as_sync(self.create_run(parameters=parameters, deferred=True))
+            futures = [
+                self.apply_async(kwargs=parameters, wait_for=wait_for)
                 for parameters in parameters_list
             ]
-
-        from prefect.task_runners import TaskRunner
-
-        task_runner = flow_run_context.task_runner
-        assert isinstance(task_runner, TaskRunner)
-        futures = task_runner.map(self, parameters, wait_for)
+        elif task_runner := getattr(flow_run_context, "task_runner", None):
+            assert isinstance(task_runner, TaskRunner)
+            futures = task_runner.map(self, parameters, wait_for)
+        else:
+            raise RuntimeError(
+                "Unable to determine task runner to use for mapped task runs. If"
+                " you are mapping a task outside of a flow, please provide"
+                " `deferred=True` to submit the mapped task runs for deferred"
+                " execution."
+            )
         if return_state:
             states = []
             for future in futures:
@@ -1059,7 +1082,7 @@ class Task(Generic[P, R]):
         dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     ) -> PrefectDistributedFuture:
         """
-        Create a pending task run for a task server to execute.
+        Create a pending task run for a task worker to execute.
 
         Args:
             args: Arguments to run the task with
@@ -1181,7 +1204,7 @@ class Task(Generic[P, R]):
         """
         return self.apply_async(args=args, kwargs=kwargs)
 
-    def serve(self, task_runner: Optional["BaseTaskRunner"] = None) -> "Task":
+    def serve(self) -> "Task":
         """Serve the task using the provided task runner. This method is used to
         establish a websocket connection with the Prefect server and listen for
         submitted task runs to execute.
@@ -1198,9 +1221,9 @@ class Task(Generic[P, R]):
 
             >>> my_task.serve()
         """
-        from prefect.task_server import serve
+        from prefect.task_worker import serve
 
-        serve(self, task_runner=task_runner)
+        serve(self)
 
 
 @overload
@@ -1215,6 +1238,7 @@ def task(
     description: str = None,
     tags: Iterable[str] = None,
     version: str = None,
+    cache_policy: CachePolicy = NotSet,
     cache_key_fn: Callable[["TaskRunContext", Dict[str, Any]], Optional[str]] = None,
     cache_expiration: datetime.timedelta = None,
     task_run_name: Optional[Union[Callable[[], str], str]] = None,
@@ -1249,6 +1273,7 @@ def task(
     description: str = None,
     tags: Iterable[str] = None,
     version: str = None,
+    cache_policy: CachePolicy = NotSet,
     cache_key_fn: Callable[["TaskRunContext", Dict[str, Any]], Optional[str]] = None,
     cache_expiration: datetime.timedelta = None,
     task_run_name: Optional[Union[Callable[[], str], str]] = None,
@@ -1391,6 +1416,7 @@ def task(
                 description=description,
                 tags=tags,
                 version=version,
+                cache_policy=cache_policy,
                 cache_key_fn=cache_key_fn,
                 cache_expiration=cache_expiration,
                 task_run_name=task_run_name,
@@ -1420,6 +1446,7 @@ def task(
                 description=description,
                 tags=tags,
                 version=version,
+                cache_policy=cache_policy,
                 cache_key_fn=cache_key_fn,
                 cache_expiration=cache_expiration,
                 task_run_name=task_run_name,

prefect/transactions.py

@@ -52,6 +52,7 @@ class Transaction(ContextModel):
     on_rollback_hooks: List[Callable[["Transaction"], None]] = Field(
         default_factory=list
     )
+    overwrite: bool = False
     _staged_value: Any = None
     __var__ = ContextVar("transaction")
 
@@ -122,7 +123,7 @@ class Transaction(ContextModel):
     def begin(self):
         # currently we only support READ_COMMITTED isolation
         # i.e., no locking behavior
-        if self.store and self.store.exists(key=self.key):
+        if not self.overwrite and self.store and self.store.exists(key=self.key):
             self.state = TransactionState.COMMITTED
 
     def read(self) -> dict:
@@ -215,6 +216,9 @@ def transaction(
     key: Optional[str] = None,
     store: Optional[RecordStore] = None,
     commit_mode: CommitMode = CommitMode.LAZY,
+    overwrite: bool = False,
 ) -> Generator[Transaction, None, None]:
-    with Transaction(key=key, store=store, commit_mode=commit_mode) as txn:
+    with Transaction(
+        key=key, store=store, commit_mode=commit_mode, overwrite=overwrite
+    ) as txn:
         yield txn