prefect-client 2.14.20__py3-none-any.whl → 2.15.0__py3-none-any.whl

prefect/task_engine.py

@@ -1,3 +1,4 @@
+import threading
 from contextlib import AsyncExitStack
 from typing import (
     Any,
@@ -8,7 +9,7 @@ from typing import (
 )
 
 import anyio
-from anyio import start_blocking_portal
+import greenback
 from typing_extensions import Literal
 
 from prefect._internal.concurrency.api import create_call, from_async, from_sync
@@ -21,7 +22,7 @@ from prefect.engine import (
 )
 from prefect.futures import PrefectFuture
 from prefect.results import ResultFactory
-from prefect.task_runners import BaseTaskRunner, SequentialTaskRunner
+from prefect.task_runners import BaseTaskRunner
 from prefect.tasks import Task
 from prefect.utilities.asyncutils import sync_compatible
 
@@ -32,28 +33,28 @@ EngineReturnType = Literal["future", "state", "result"]
 async def submit_autonomous_task_to_engine(
     task: Task,
     task_run: TaskRun,
+    task_runner: Type[BaseTaskRunner],
     parameters: Optional[Dict] = None,
     wait_for: Optional[Iterable[PrefectFuture]] = None,
     mapped: bool = False,
     return_type: EngineReturnType = "future",
-    task_runner: Optional[Type[BaseTaskRunner]] = None,
+    client=None,
 ) -> Any:
-    parameters = parameters or {}
     async with AsyncExitStack() as stack:
+        if not task_runner._started:
+            task_runner_ctx = await stack.enter_async_context(task_runner.start())
+        else:
+            task_runner_ctx = task_runner
+        parameters = parameters or {}
         with EngineContext(
             flow=None,
             flow_run=None,
             autonomous_task_run=task_run,
-            task_runner=await stack.enter_async_context(
-                (task_runner if task_runner else SequentialTaskRunner()).start()
-            ),
-            client=await stack.enter_async_context(get_client()),
+            task_runner=task_runner_ctx,
+            client=client or await stack.enter_async_context(get_client()),
             parameters=parameters,
             result_factory=await ResultFactory.from_task(task),
             background_tasks=await stack.enter_async_context(anyio.create_task_group()),
-            sync_portal=(
-                stack.enter_context(start_blocking_portal()) if task.isasync else None
-            ),
         ) as flow_run_context:
             begin_run = create_call(
                 begin_task_map if mapped else get_task_call_return_value,
@@ -63,8 +64,10 @@ async def submit_autonomous_task_to_engine(
                 wait_for=wait_for,
                 return_type=return_type,
                 task_runner=task_runner,
+                user_thread=threading.current_thread(),
             )
             if task.isasync:
                 return await from_async.wait_for_call_in_loop_thread(begin_run)
             else:
+                await greenback.ensure_portal()
                 return from_sync.wait_for_call_in_loop_thread(begin_run)

prefect/task_server.py

@@ -1,58 +1,67 @@
 import asyncio
 import signal
 import sys
+from contextlib import AsyncExitStack
 from functools import partial
-from typing import Iterable, Optional
+from typing import Optional, Type
 
 import anyio
 import anyio.abc
-import pendulum
 
 from prefect import Task, get_client
 from prefect._internal.concurrency.api import create_call, from_sync
 from prefect.client.schemas.objects import TaskRun
 from prefect.client.subscriptions import Subscription
+from prefect.engine import propose_state
 from prefect.logging.loggers import get_logger
 from prefect.results import ResultFactory
 from prefect.settings import (
     PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING,
     PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS,
 )
+from prefect.states import Pending
 from prefect.task_engine import submit_autonomous_task_to_engine
+from prefect.task_runners import BaseTaskRunner, ConcurrentTaskRunner
 from prefect.utilities.asyncutils import asyncnullcontext, sync_compatible
-from prefect.utilities.collections import distinct
 from prefect.utilities.processutils import _register_signal
 
 logger = get_logger("task_server")
 
 
+class StopTaskServer(Exception):
+    """Raised when the task server is stopped."""
+
+    pass
+
+
 class TaskServer:
-    """This class is responsible for serving tasks that may be executed autonomously
-    (i.e., without a parent flow run).
+    """This class is responsible for serving tasks that may be executed autonomously by a
+    task runner in the engine.
 
-    When `start()` is called, the task server will subscribe to the task run scheduling
-    topic and poll for scheduled task runs. When a scheduled task run is found, it
-    will submit the task run to the engine for execution, using `submit_autonomous_task_to_engine`
-    to construct a minimal `EngineContext` for the task run.
+    When `start()` is called, the task server 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.
 
     Args:
         - tasks: A list of tasks to serve. These tasks will be submitted to the engine
             when a scheduled task run is found.
-        - tags: A list of tags to apply to the task server. Defaults to `["autonomous"]`.
+        - task_runner: The task runner to use for executing the tasks. Defaults to
+            `ConcurrentTaskRunner`.
     """
 
     def __init__(
         self,
         *tasks: Task,
-        tags: Optional[Iterable[str]] = None,
+        task_runner: Optional[Type[BaseTaskRunner]] = None,
     ):
         self.tasks: list[Task] = tasks
-        self.tags: Iterable[str] = tags or ["autonomous"]
-        self.last_polled: Optional[pendulum.DateTime] = None
-        self.started = False
-        self.stopping = False
+        self.task_runner: Type[BaseTaskRunner] = task_runner or ConcurrentTaskRunner()
+        self.started: bool = False
+        self.stopping: bool = False
 
         self._client = get_client()
+        self._exit_stack = AsyncExitStack()
 
         if not asyncio.get_event_loop().is_running():
             raise RuntimeError(
@@ -89,18 +98,21 @@ class TaskServer:
                 " calling .start()"
             )
 
-        logger.info("Stopping task server...")
         self.started = False
         self.stopping = True
 
+        raise StopTaskServer
+
     async def _subscribe_to_task_scheduling(self):
-        subscription = Subscription(TaskRun, "/task_runs/subscriptions/scheduled")
-        logger.debug(f"Created: {subscription}")
-        async for task_run in subscription:
+        async for task_run in Subscription(
+            TaskRun,
+            "/task_runs/subscriptions/scheduled",
+            [task.task_key for task in self.tasks],
+        ):
             logger.info(f"Received task run: {task_run.id} - {task_run.name}")
-            await self._submit_pending_task_run(task_run)
+            await self._submit_scheduled_task_run(task_run)
 
-    async def _submit_pending_task_run(self, task_run: TaskRun):
+    async def _submit_scheduled_task_run(self, task_run: TaskRun):
         logger.debug(
             f"Found task run: {task_run.name!r} in state: {task_run.state.name!r}"
         )
@@ -142,7 +154,18 @@ class TaskServer:
             f"Submitting run {task_run.name!r} of task {task.name!r} to engine"
         )
 
-        task_run.tags = distinct(task_run.tags + list(self.tags))
+        state = await propose_state(
+            client=get_client(),  # TODO prove that we cannot use self._client here
+            state=Pending(),
+            task_run_id=task_run.id,
+        )
+
+        if not state.is_pending():
+            logger.warning(
+                f"Aborted task run {task_run.id!r} -"
+                f" server returned a non-pending state {state.type.value!r}."
+                " Task run may have already begun execution."
+            )
 
         self._runs_task_group.start_soon(
             partial(
@@ -150,14 +173,17 @@ class TaskServer:
                 task=task,
                 task_run=task_run,
                 parameters=parameters,
+                task_runner=self.task_runner,
+                client=self._client,
             )
         )
 
     async def __aenter__(self):
         logger.debug("Starting task server...")
-        self._client = get_client()
-        await self._client.__aenter__()
-        await self._runs_task_group.__aenter__()
+
+        self._client = await self._exit_stack.enter_async_context(get_client())
+        await self._exit_stack.enter_async_context(self._runs_task_group)
+        await self._exit_stack.enter_async_context(self.task_runner.start())
 
         self.started = True
         return self
@@ -165,20 +191,21 @@ class TaskServer:
     async def __aexit__(self, *exc_info):
         logger.debug("Stopping task server...")
         self.started = False
-        if self._runs_task_group:
-            await self._runs_task_group.__aexit__(*exc_info)
-        if self._client:
-            await self._client.__aexit__(*exc_info)
+
+        await self._exit_stack.__aexit__(*exc_info)
 
 
 @sync_compatible
-async def serve(*tasks: Task, tags: Optional[Iterable[str]] = None):
-    """Serve the provided tasks so that they may be executed autonomously.
+async def serve(*tasks: Task, task_runner: Optional[Type[BaseTaskRunner]] = None):
+    """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`.
 
     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.
-        - tags: A list of tags to apply to the task server. Defaults to `["autonomous"]`.
+        - task_runner: The task runner to use for executing the tasks. Defaults to
+            `ConcurrentTaskRunner`.
 
     Example:
         ```python
@@ -193,7 +220,7 @@ async def serve(*tasks: Task, tags: Optional[Iterable[str]] = None):
         def yell(message: str):
             print(message.upper())
 
-        # starts a long-lived process that listens scheduled runs of these tasks
+        # starts a long-lived process that listens for scheduled runs of these tasks
         if __name__ == "__main__":
             serve(say, yell)
         ```
@@ -204,5 +231,12 @@ async def serve(*tasks: Task, tags: Optional[Iterable[str]] = None):
             " to True."
         )
 
-    task_server = TaskServer(*tasks, tags=tags)
-    await task_server.start()
+    task_server = TaskServer(*tasks, task_runner=task_runner)
+    try:
+        await task_server.start()
+
+    except StopTaskServer:
+        logger.info("Task server stopped.")
+
+    except asyncio.CancelledError:
+        logger.info("Task server interrupted, stopping...")

prefect/utilities/asyncutils.py

@@ -28,6 +28,8 @@ import anyio.abc
 import sniffio
 from typing_extensions import Literal, ParamSpec, TypeGuard
 
+from prefect.logging import get_logger
+
 T = TypeVar("T")
 P = ParamSpec("P")
 R = TypeVar("R")
@@ -40,6 +42,8 @@ EVENT_LOOP_GC_REFS = {}
 
 PREFECT_THREAD_LIMITER: Optional[anyio.CapacityLimiter] = None
 
+logger = get_logger()
+
 
 def get_thread_limiter():
     global PREFECT_THREAD_LIMITER
@@ -51,7 +55,7 @@ def get_thread_limiter():
 
 
 def is_async_fn(
-    func: Union[Callable[P, R], Callable[P, Awaitable[R]]]
+    func: Union[Callable[P, R], Callable[P, Awaitable[R]]],
 ) -> TypeGuard[Callable[P, Awaitable[R]]]:
     """
     Returns `True` if a function returns a coroutine.
@@ -334,7 +338,13 @@ async def add_event_loop_shutdown_callback(coroutine_fn: Callable[[], Awaitable]
     EVENT_LOOP_GC_REFS[key] = on_shutdown(key)
 
     # Begin iterating so it will be cleaned up as an incomplete generator
-    await EVENT_LOOP_GC_REFS[key].__anext__()
+    try:
+        await EVENT_LOOP_GC_REFS[key].__anext__()
+    # There is a poorly understood edge case we've seen in CI where the key is
+    # removed from the dict before we begin generator iteration.
+    except KeyError:
+        logger.warn("The event loop shutdown callback was not properly registered. ")
+        pass
 
 
 class GatherIncomplete(RuntimeError):

prefect/utilities/collections.py

@@ -106,7 +106,7 @@ def dict_to_flatdict(
 
 
 def flatdict_to_dict(
-    dct: Dict[Tuple[KT, ...], VT]
+    dct: Dict[Tuple[KT, ...], VT],
 ) -> Dict[KT, Union[VT, Dict[KT, VT]]]:
     """Converts a flattened dictionary back to a nested dictionary.
 

prefect/utilities/filesystem.py

@@ -3,6 +3,7 @@ Utilities for working with file systems
 """
 import os
 import pathlib
+import threading
 from contextlib import contextmanager
 from pathlib import Path, PureWindowsPath
 from typing import Union
@@ -51,6 +52,9 @@ def filter_files(
     return included_files
 
 
+chdir_lock = threading.Lock()
+
+
 @contextmanager
 def tmpchdir(path: str):
     """
@@ -62,11 +66,12 @@ def tmpchdir(path: str):
 
     owd = os.getcwd()
 
-    try:
-        os.chdir(path)
-        yield path
-    finally:
-        os.chdir(owd)
+    with chdir_lock:
+        try:
+            os.chdir(path)
+            yield path
+        finally:
+            os.chdir(owd)
 
 
 def filename(path: str) -> str:

{prefect_client-2.14.20.dist-info → prefect_client-2.15.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 2.14.20
+Version: 2.15.0
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.
@@ -31,6 +31,7 @@ Requires-Dist: coolname <3.0.0,>=1.0.4
 Requires-Dist: croniter <3.0.0,>=1.0.12
 Requires-Dist: fsspec >=2022.5.0
 Requires-Dist: graphviz >=0.20.1
+Requires-Dist: greenback >=1.2.0
 Requires-Dist: griffe >=0.20.0
 Requires-Dist: httpcore <2.0.0,>=0.15.0
 Requires-Dist: httpx[http2] !=0.23.2,>=0.23
@@ -46,12 +47,13 @@ Requires-Dist: pyyaml <7.0.0,>=5.4.1
 Requires-Dist: rich <14.0,>=11.0
 Requires-Dist: ruamel.yaml >=0.17.0
 Requires-Dist: sniffio <2.0.0,>=1.3.0
-Requires-Dist: starlette <0.33.0,>=0.27.0
 Requires-Dist: toml >=0.10.0
 Requires-Dist: typing-extensions <5.0.0,>=4.5.0
 Requires-Dist: ujson <6.0.0,>=5.8.0
 Requires-Dist: uvicorn >=0.14.0
 Requires-Dist: websockets <13.0,>=10.4
+Requires-Dist: itsdangerous
+Requires-Dist: python-multipart >=0.0.7
 Requires-Dist: importlib-metadata >=4.4 ; python_version < "3.10"
 Requires-Dist: pendulum <3.0 ; python_version < "3.12"
 Requires-Dist: pendulum <4,>=3.0.0 ; python_version >= "3.12"