prefect-client 2.19.4__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}

@@ -4,40 +4,39 @@ import os
 import signal
 import socket
 import sys
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import AsyncExitStack
-from functools import partial
-from typing import List, Optional, Type
+from contextvars import copy_context
+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.engine import emit_task_run_state_change_event, propose_state
 from prefect.exceptions import Abort, PrefectHTTPStatusError
 from prefect.logging.loggers import get_logger
 from prefect.results import ResultFactory
 from prefect.settings import (
     PREFECT_API_URL,
-    PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING,
     PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS,
 )
 from prefect.states import Pending
-from prefect.task_engine import submit_autonomous_task_run_to_engine
-from prefect.task_runners import (
-    BaseTaskRunner,
-    ConcurrentTaskRunner,
-)
+from prefect.task_engine import run_task_async, run_task_sync
 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
 
@@ -52,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.
@@ -64,18 +63,17 @@ class TaskServer:
     Args:
         - tasks: A list of tasks to serve. These tasks will be submitted to the engine
             when a scheduled task run is found.
-        - task_runner: The task runner to use for executing the tasks. Defaults to
-            `ConcurrentTaskRunner`.
+        - limit: The maximum number of tasks that can be run concurrently. Defaults to 10.
+            Pass `None` to remove the limit.
     """
 
     def __init__(
         self,
         *tasks: Task,
-        task_runner: Optional[Type[BaseTaskRunner]] = None,
+        limit: Optional[int] = 10,
     ):
-        self.tasks: List[Task] = tasks
+        self.tasks: List[Task] = list(tasks)
 
-        self.task_runner: BaseTaskRunner = task_runner or ConcurrentTaskRunner()
         self.started: bool = False
         self.stopping: bool = False
 
@@ -84,10 +82,12 @@ 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()
+        self._executor = ThreadPoolExecutor()
+        self._limiter = anyio.CapacityLimiter(limit) if limit else None
 
     @property
     def _client_id(self) -> str:
@@ -95,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))
@@ -105,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:
@@ -126,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(
@@ -148,8 +148,10 @@ class TaskServer:
             keys=[task.task_key for task in self.tasks],
             client_id=self._client_id,
         ):
+            if self._limiter:
+                await self._limiter.acquire_on_behalf_of(task_run.id)
             logger.info(f"Received task run: {task_run.id} - {task_run.name}")
-            await self._submit_scheduled_task_run(task_run)
+            self._runs_task_group.start_soon(self._submit_scheduled_task_run, task_run)
 
     async def _submit_scheduled_task_run(self, task_run: TaskRun):
         logger.debug(
@@ -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
 
@@ -171,12 +173,17 @@ class TaskServer:
         # state_details. If there is no parameters_id, then the task was created
         # without parameters.
         parameters = {}
+        wait_for = []
+        run_context = None
         if should_try_to_read_parameters(task, task_run):
             parameters_id = task_run.state.state_details.task_parameters_id
             task.persist_result = True
             factory = await ResultFactory.from_autonomous_task(task)
             try:
-                parameters = await factory.read_parameters(parameters_id)
+                run_data = await factory.read_parameters(parameters_id)
+                parameters = run_data.get("parameters", {})
+                wait_for = run_data.get("wait_for", [])
+                run_context = run_data.get("context", None)
             except Exception as exc:
                 logger.exception(
                     f"Failed to read parameters for task run {task_run.id!r}",
@@ -194,9 +201,11 @@ class TaskServer:
         )
 
         try:
+            new_state = Pending()
+            new_state.state_details.deferred = True
             state = await propose_state(
                 client=get_client(),  # TODO prove that we cannot use self._client here
-                state=Pending(),
+                state=new_state,
                 task_run_id=task_run.id,
             )
         except Abort as exc:
@@ -225,44 +234,61 @@ class TaskServer:
             validated_state=state,
         )
 
-        self._runs_task_group.start_soon(
-            partial(
-                submit_autonomous_task_run_to_engine,
+        if task.isasync:
+            await run_task_async(
                 task=task,
+                task_run_id=task_run.id,
                 task_run=task_run,
                 parameters=parameters,
-                task_runner=self.task_runner,
-                client=self._client,
+                wait_for=wait_for,
+                return_type="state",
+                context=run_context,
             )
-        )
+        else:
+            context = copy_context()
+            future = self._executor.submit(
+                context.run,
+                run_task_sync,
+                task=task,
+                task_run_id=task_run.id,
+                task_run=task_run,
+                parameters=parameters,
+                wait_for=wait_for,
+                return_type="state",
+                context=run_context,
+            )
+            await asyncio.wrap_future(future)
+        if self._limiter:
+            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()
 
         await self._exit_stack.enter_async_context(self._client)
-        await self._exit_stack.enter_async_context(self.task_runner.start())
-        await self._runs_task_group.__aenter__()
+        await self._exit_stack.enter_async_context(self._runs_task_group)
+        self._exit_stack.enter_context(self._executor)
 
         self.started = True
         return self
 
     async def __aexit__(self, *exc_info):
-        logger.debug("Stopping task server...")
+        logger.debug("Stopping task worker...")
         self.started = False
-        await self._runs_task_group.__aexit__(*exc_info)
         await self._exit_stack.__aexit__(*exc_info)
 
 
 @sync_compatible
-async def serve(*tasks: Task, task_runner: Optional[Type[BaseTaskRunner]] = None):
+async def serve(*tasks: Task, limit: Optional[int] = 10):
     """Serve the provided tasks so that their runs may be submitted to and executed.
     in the engine. Tasks do not need to be within a flow run context to be submitted.
     You must `.submit` the same task object that you pass to `serve`.
@@ -270,13 +296,13 @@ async def serve(*tasks: Task, task_runner: Optional[Type[BaseTaskRunner]] = None
     Args:
         - tasks: A list of tasks to serve. When a scheduled task run is found for a
             given task, the task run will be submitted to the engine for execution.
-        - task_runner: The task runner to use for executing the tasks. Defaults to
-            `ConcurrentTaskRunner`.
+        - limit: The maximum number of tasks that can be run concurrently. Defaults to 10.
+            Pass `None` to remove the limit.
 
     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):
@@ -291,18 +317,21 @@ async def serve(*tasks: Task, task_runner: Optional[Type[BaseTaskRunner]] = None
             serve(say, yell)
         ```
     """
-    if not PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING.value():
-        raise RuntimeError(
-            "To enable task scheduling, set PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING"
-            " to True."
-        )
+    task_worker = TaskWorker(*tasks, limit=limit)
 
-    task_server = TaskServer(*tasks, task_runner=task_runner)
     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...")