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

prefect/task_runners.py

@@ -202,12 +202,13 @@ class TaskRunner(abc.ABC, Generic[F]):
 
 
 class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
-    def __init__(self):
+    def __init__(self, max_workers: Optional[int] = None):
         super().__init__()
         self._executor: Optional[ThreadPoolExecutor] = None
+        self._max_workers = max_workers
 
     def duplicate(self) -> "ThreadPoolTaskRunner":
-        return type(self)()
+        return type(self)(max_workers=self._max_workers)
 
     def submit(
         self,
@@ -278,7 +279,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
 
     def __enter__(self):
         super().__enter__()
-        self._executor = ThreadPoolExecutor()
+        self._executor = ThreadPoolExecutor(max_workers=self._max_workers)
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
@@ -287,6 +288,15 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
             self._executor = None
         super().__exit__(exc_type, exc_value, traceback)
 
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, ThreadPoolTaskRunner):
+            return False
+        return self._max_workers == value._max_workers
+
+
+# Here, we alias ConcurrentTaskRunner to ThreadPoolTaskRunner for backwards compatibility
+ConcurrentTaskRunner = ThreadPoolTaskRunner
+
 
 class PrefectTaskRunner(TaskRunner[PrefectDistributedFuture]):
     def __init__(self):
@@ -321,11 +331,11 @@ class PrefectTaskRunner(TaskRunner[PrefectDistributedFuture]):
         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 server..."
+                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 server..."
+                f"Submitting task {task.name} to for execution by a Prefect task worker..."
             )
 
         return task.apply_async(

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=10000, ttl=600
+        )
+        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.debug("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.debug(
+                        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,11 +82,11 @@ 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._executor = ThreadPoolExecutor(max_workers=limit if limit else None)
         self._limiter = anyio.CapacityLimiter(limit) if limit else None
 
     @property
@@ -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,18 +105,18 @@ 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:
                 if exc.status_code == 403:
                     logger.error(
-                        "Could not establish a connection to the `/task_runs/subscriptions/scheduled`"
+                        "403: Could not establish a connection to the `/task_runs/subscriptions/scheduled`"
                         f" endpoint found at:\n\n {PREFECT_API_URL.value()}"
                         "\n\nPlease double-check the values of your"
                         " `PREFECT_API_URL` and `PREFECT_API_KEY` environment variables."
@@ -124,32 +126,54 @@ 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(
-            f"Subscribing to tasks: {' | '.join(t.task_key.split('.')[-1] for t in self.tasks)}"
+        base_url = PREFECT_API_URL.value()
+        if base_url is None:
+            raise ValueError(
+                "`PREFECT_API_URL` must be set to use the task worker. "
+                "Task workers are not compatible with the ephemeral API."
+            )
+        task_keys_repr = " | ".join(
+            t.task_key.split(".")[-1].split("-")[0] for t in self.tasks
         )
+        logger.info(f"Subscribing to runs of task(s): {task_keys_repr}")
         async for task_run in Subscription(
             model=TaskRun,
             path="/task_runs/subscriptions/scheduled",
             keys=[task.task_key for task in self.tasks],
             client_id=self._client_id,
+            base_url=base_url,
         ):
+            logger.info(f"Received task run: {task_run.id} - {task_run.name}")
             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}")
-            self._runs_task_group.start_soon(self._submit_scheduled_task_run, task_run)
+            self._runs_task_group.start_soon(
+                self._safe_submit_scheduled_task_run, task_run
+            )
+
+    async def _safe_submit_scheduled_task_run(self, task_run: TaskRun):
+        try:
+            await self._submit_scheduled_task_run(task_run)
+        except BaseException as exc:
+            logger.exception(
+                f"Failed to submit task run {task_run.id!r}",
+                exc_info=exc,
+            )
+        finally:
+            if self._limiter:
+                self._limiter.release_on_behalf_of(task_run.id)
 
     async def _submit_scheduled_task_run(self, task_run: TaskRun):
         logger.debug(
@@ -159,11 +183,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
 
@@ -256,18 +280,16 @@ class TaskServer:
                 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)
+            await self._safe_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 +302,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 +322,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 +337,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...")