prefect-client 3.2.6__tar.gz → 3.2.8__tar.gz

{prefect_client-3.2.6 → prefect_client-3.2.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prefect-client
-Version: 3.2.6
+Version: 3.2.8
 Summary: Workflow orchestration and management.
 Project-URL: Changelog, https://github.com/PrefectHQ/prefect/releases
 Project-URL: Documentation, https://docs.prefect.io
@@ -57,7 +57,7 @@ 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.29.0,>=0.14.0
-Requires-Dist: websockets<14.0,>=10.4
+Requires-Dist: websockets<16.0,>=13.0
 Provides-Extra: notifications
 Requires-Dist: apprise<2.0.0,>=1.1.0; extra == 'notifications'
 Description-Content-Type: text/markdown

{prefect_client-3.2.6 → prefect_client-3.2.8}/pyproject.toml

@@ -61,7 +61,7 @@ dependencies = [
     "typing_extensions>=4.5.0,<5.0.0",
     "ujson>=5.8.0,<6.0.0",
     "uvicorn>=0.14.0,!=0.29.0",
-    "websockets>=10.4,<14.0",
+    "websockets>=13.0,<16.0",
 ]
 [project.urls]
 Changelog = "https://github.com/PrefectHQ/prefect/releases"

prefect_client-3.2.8/src/prefect/_build_info.py

@@ -0,0 +1,5 @@
+# Generated by versioningit
+__version__ = "3.2.8"
+__build_date__ = "2025-02-28 00:11:08.300212+00:00"
+__git_commit__ = "2c0860d5dcd6156fc3fae75d53f7f48f4bd1a522"
+__dirty__ = False

{prefect_client-3.2.6 → prefect_client-3.2.8}/src/prefect/_experimental/bundles.py

@@ -6,9 +6,12 @@ import gzip
 import multiprocessing
 import multiprocessing.context
 import os
+import subprocess
+import sys
 from typing import Any, TypedDict
 
 import cloudpickle
+import uv
 
 from prefect.client.schemas.objects import FlowRun
 from prefect.context import SettingsContext, get_settings_context, serialize_context
@@ -17,6 +20,7 @@ from prefect.flow_engine import run_flow
 from prefect.flows import Flow
 from prefect.settings.context import get_current_settings
 from prefect.settings.models.root import Settings
+from prefect.utilities.slugify import slugify
 
 
 class SerializedBundle(TypedDict):
@@ -28,6 +32,7 @@ class SerializedBundle(TypedDict):
     function: str
     context: str
     flow_run: dict[str, Any]
+    dependencies: str
 
 
 def _serialize_bundle_object(obj: Any) -> str:
@@ -66,6 +71,17 @@ def create_bundle_for_flow_run(
         "function": _serialize_bundle_object(flow),
         "context": _serialize_bundle_object(context),
         "flow_run": flow_run.model_dump(mode="json"),
+        "dependencies": subprocess.check_output(
+            [
+                uv.find_uv_bin(),
+                "pip",
+                "freeze",
+                # Exclude editable installs because we won't be able to install them in the execution environment
+                "--exclude-editable",
+            ]
+        )
+        .decode()
+        .strip(),
     }
 
 
@@ -129,6 +145,14 @@ def execute_bundle_in_subprocess(
 
     ctx = multiprocessing.get_context("spawn")
 
+    # Install dependencies if necessary
+    if dependencies := bundle.get("dependencies"):
+        subprocess.check_call(
+            [uv.find_uv_bin(), "pip", "install", *dependencies.split("\n")],
+            # Copy the current environment to ensure we install into the correct venv
+            env=os.environ,
+        )
+
     process = ctx.Process(
         target=_extract_and_run_flow,
         kwargs={
@@ -141,3 +165,52 @@ def execute_bundle_in_subprocess(
     process.start()
 
     return process
+
+
+def convert_step_to_command(step: dict[str, Any], key: str) -> list[str]:
+    """
+    Converts a bundle upload or execution step to a command.
+
+    Args:
+        step: The step to convert.
+        key: The key to use for the remote file when downloading or uploading.
+
+    Returns:
+        A list of strings representing the command to run the step.
+    """
+    # Start with uv run
+    command = ["uv", "run"]
+
+    step_keys = list(step.keys())
+
+    if len(step_keys) != 1:
+        raise ValueError("Expected exactly one function in step")
+
+    function_fqn = step_keys[0]
+    function_args = step[function_fqn]
+
+    # Add the `--with` argument to handle dependencies for running the step
+    requires: list[str] | str = function_args.get("requires", [])
+    if isinstance(requires, str):
+        requires = [requires]
+    if requires:
+        command.extend(["--with", ",".join(requires)])
+
+    # Add the `--python` argument to handle the Python version for running the step
+    python_version = sys.version_info
+    command.extend(["--python", f"{python_version.major}.{python_version.minor}"])
+
+    # Add the `-m` argument to defined the function to run
+    command.extend(["-m", function_fqn])
+
+    # Add any arguments with values defined in the step
+    for arg_name, arg_value in function_args.items():
+        if arg_name == "requires":
+            continue
+
+        command.extend([f"--{slugify(arg_name)}", arg_value])
+
+    # Add the `--key` argument to specify the remote file name
+    command.extend(["--key", key])
+
+    return command

prefect_client-3.2.8/src/prefect/_waiters.py

@@ -0,0 +1,254 @@
+from __future__ import annotations
+
+import asyncio
+import atexit
+import threading
+import uuid
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+)
+
+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 import get_logger
+
+if TYPE_CHECKING:
+    import logging
+
+
+class FlowRunWaiter:
+    """
+    A service used for waiting for a flow run to finish.
+
+    This service listens for flow run events and provides a way to wait for a specific
+    flow run to finish. This is useful for waiting for a flow 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 flow run events.
+
+    The service can be used to wait for a flow run to finish by calling
+    `FlowRunWaiter.wait_for_flow_run` with the flow run ID to wait for. The method
+    will return when the flow 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 flow
+    from prefect.flow_engine import run_flow_async
+    from prefect.flow_runs import FlowRunWaiter
+
+
+    @flow
+    async def test_flow():
+        await asyncio.sleep(5)
+        print("Done!")
+
+
+    async def main():
+        flow_run_id = uuid4()
+        asyncio.create_flow(run_flow_async(flow=test_flow, flow_run_id=flow_run_id))
+
+        await FlowRunWaiter.wait_for_flow_run(flow_run_id)
+        print("Flow run finished")
+
+
+    if __name__ == "__main__":
+        asyncio.run(main())
+    ```
+    """
+
+    _instance: Self | None = None
+    _instance_lock = threading.Lock()
+
+    def __init__(self):
+        self.logger: "logging.Logger" = get_logger("FlowRunWaiter")
+        self._consumer_task: asyncio.Task[None] | None = None
+        self._observed_completed_flow_runs: TTLCache[uuid.UUID, bool] = TTLCache(
+            maxsize=10000, ttl=600
+        )
+        self._completion_events: dict[uuid.UUID, asyncio.Event] = {}
+        self._completion_callbacks: dict[uuid.UUID, Callable[[], None]] = {}
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._observed_completed_flow_runs_lock = threading.Lock()
+        self._completion_events_lock = threading.Lock()
+        self._started = False
+
+    def start(self) -> None:
+        """
+        Start the FlowRunWaiter service.
+        """
+        if self._started:
+            return
+        self.logger.debug("Starting FlowRunWaiter")
+        loop_thread = get_global_loop()
+
+        if not asyncio.get_running_loop() == loop_thread.loop:
+            raise RuntimeError("FlowRunWaiter must run on the global loop thread.")
+
+        self._loop = loop_thread.loop
+        if TYPE_CHECKING:
+            assert self._loop is not None
+
+        consumer_started = asyncio.Event()
+        self._consumer_task = self._loop.create_task(
+            self._consume_events(consumer_started)
+        )
+        asyncio.run_coroutine_threadsafe(consumer_started.wait(), self._loop)
+
+        loop_thread.add_shutdown_call(create_call(self.stop))
+        atexit.register(self.stop)
+        self._started = True
+
+    async def _consume_events(self, consumer_started: asyncio.Event):
+        async with get_events_subscriber(
+            filter=EventFilter(
+                event=EventNameFilter(
+                    name=[
+                        f"prefect.flow-run.{state.name.title()}"
+                        for state in TERMINAL_STATES
+                    ],
+                )
+            )
+        ) as subscriber:
+            consumer_started.set()
+            async for event in subscriber:
+                try:
+                    self.logger.debug(
+                        f"Received event: {event.resource['prefect.resource.id']}"
+                    )
+                    flow_run_id = uuid.UUID(
+                        event.resource["prefect.resource.id"].replace(
+                            "prefect.flow-run.", ""
+                        )
+                    )
+
+                    with self._observed_completed_flow_runs_lock:
+                        # Cache the flow run ID for a short period of time to avoid
+                        # unnecessary waits
+                        self._observed_completed_flow_runs[flow_run_id] = True
+                    with self._completion_events_lock:
+                        # Set the event for the flow run ID if it is in the cache
+                        # so the waiter can wake up the waiting coroutine
+                        if flow_run_id in self._completion_events:
+                            self._completion_events[flow_run_id].set()
+                        if flow_run_id in self._completion_callbacks:
+                            self._completion_callbacks[flow_run_id]()
+                except Exception as exc:
+                    self.logger.error(f"Error processing event: {exc}")
+
+    def stop(self) -> None:
+        """
+        Stop the FlowRunWaiter service.
+        """
+        self.logger.debug("Stopping FlowRunWaiter")
+        if self._consumer_task:
+            self._consumer_task.cancel()
+            self._consumer_task = None
+        self.__class__._instance = None
+        self._started = False
+
+    @classmethod
+    async def wait_for_flow_run(
+        cls, flow_run_id: uuid.UUID, timeout: float | None = None
+    ) -> None:
+        """
+        Wait for a flow 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:
+            flow_run_id: The ID of the flow run to wait for.
+            timeout: The maximum time to wait for the flow run to
+                finish. Defaults to None.
+        """
+        instance = cls.instance()
+        with instance._observed_completed_flow_runs_lock:
+            if flow_run_id in instance._observed_completed_flow_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 flow run ID so the consumer can set it
+            # when the event is received
+            instance._completion_events[flow_run_id] = finished_event
+
+        try:
+            # Now check one more time whether the flow run arrived before we start to
+            # wait on it, in case it came in while we were setting up the event above.
+            with instance._observed_completed_flow_runs_lock:
+                if flow_run_id in instance._observed_completed_flow_runs:
+                    return
+
+            with anyio.move_on_after(delay=timeout):
+                await from_async.wait_for_call_in_loop_thread(
+                    create_call(finished_event.wait)
+                )
+        finally:
+            with instance._completion_events_lock:
+                # Remove the event from the cache after it has been waited on
+                instance._completion_events.pop(flow_run_id, None)
+
+    @classmethod
+    def add_done_callback(
+        cls, flow_run_id: uuid.UUID, callback: Callable[[], None]
+    ) -> None:
+        """
+        Add a callback to be called when a flow run finishes.
+
+        Args:
+            flow_run_id: The ID of the flow run to wait for.
+            callback: The callback to call when the flow run finishes.
+        """
+        instance = cls.instance()
+        with instance._observed_completed_flow_runs_lock:
+            if flow_run_id in instance._observed_completed_flow_runs:
+                callback()
+                return
+
+        with instance._completion_events_lock:
+            # Cache the event for the flow run ID so the consumer can set it
+            # when the event is received
+            instance._completion_callbacks[flow_run_id] = callback
+
+    @classmethod
+    def instance(cls) -> Self:
+        """
+        Get the singleton instance of FlowRunWaiter.
+        """
+        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_client-3.2.6 → prefect_client-3.2.8}/src/prefect/client/orchestration/__init__.py

@@ -81,6 +81,7 @@ from prefect.client.orchestration._blocks_types.client import (
 
 import prefect
 import prefect.exceptions
+from prefect.logging.loggers import get_run_logger
 import prefect.settings
 import prefect.states
 from prefect.client.constants import SERVER_API_VERSION
@@ -1182,8 +1183,17 @@ class PrefectClient(
         if api_version.major != client_version.major:
             raise RuntimeError(
                 f"Found incompatible versions: client: {client_version}, server: {api_version}. "
-                f"Major versions must match."
+                "Major versions must match."
             )
+        if api_version < client_version:
+            warning_message = (
+                "Your Prefect server is running an older version of Prefect than your client which may result in unexpected behavior. "
+                f"Please upgrade your Prefect server from version {api_version} to version {client_version} or higher."
+            )
+            try:
+                get_run_logger().warning(warning_message)
+            except prefect.context.MissingContextError:
+                self.logger.warning(warning_message)
 
     async def __aenter__(self) -> Self:
         """
@@ -1523,8 +1533,17 @@ class SyncPrefectClient(
         if api_version.major != client_version.major:
             raise RuntimeError(
                 f"Found incompatible versions: client: {client_version}, server: {api_version}. "
-                f"Major versions must match."
+                "Major versions must match."
             )
+        if api_version < client_version:
+            warning_message = (
+                "Your Prefect server is running an older version of Prefect than your client which may result in unexpected behavior. "
+                f"Please upgrade your Prefect server from version {api_version} to version {client_version} or higher."
+            )
+            try:
+                get_run_logger().warning(warning_message)
+            except prefect.context.MissingContextError:
+                self.logger.warning(warning_message)
 
     def set_task_run_name(self, task_run_id: UUID, name: str) -> httpx.Response:
         task_run_data = TaskRunUpdate(name=name)

{prefect_client-3.2.6 → prefect_client-3.2.8}/src/prefect/client/subscriptions.py

@@ -5,6 +5,7 @@ from typing import Any, Generic, Optional, TypeVar
 
 import orjson
 import websockets
+import websockets.asyncio.client
 import websockets.exceptions
 from starlette.status import WS_1008_POLICY_VIOLATION
 from typing_extensions import Self
@@ -45,7 +46,7 @@ class Subscription(Generic[S]):
         return self
 
     @property
-    def websocket(self) -> websockets.WebSocketClientProtocol:
+    def websocket(self) -> websockets.asyncio.client.ClientConnection:
         if not self._websocket:
             raise RuntimeError("Subscription is not connected")
         return self._websocket

{prefect_client-3.2.6 → prefect_client-3.2.8}/src/prefect/events/clients.py

@@ -27,12 +27,12 @@ from prometheus_client import Counter
 from python_socks.async_.asyncio import Proxy
 from typing_extensions import Self
 from websockets import Subprotocol
+from websockets.asyncio.client import ClientConnection, connect
 from websockets.exceptions import (
     ConnectionClosed,
     ConnectionClosedError,
     ConnectionClosedOK,
 )
-from websockets.legacy.client import Connect, WebSocketClientProtocol
 
 from prefect.events import Event
 from prefect.logging import get_logger
@@ -91,7 +91,7 @@ def events_out_socket_from_api_url(url: str) -> str:
     return http_to_ws(url) + "/events/out"
 
 
-class WebsocketProxyConnect(Connect):
+class WebsocketProxyConnect(connect):
     def __init__(self: Self, uri: str, **kwargs: Any):
         # super() is intentionally deferred to the _proxy_connect method
         # to allow for the socket to be established first
@@ -130,7 +130,7 @@ class WebsocketProxyConnect(Connect):
             ctx.verify_mode = ssl.CERT_NONE
             self._kwargs.setdefault("ssl", ctx)
 
-    async def _proxy_connect(self: Self) -> WebSocketClientProtocol:
+    async def _proxy_connect(self: Self) -> ClientConnection:
         if self._proxy:
             sock = await self._proxy.connect(
                 dest_host=self._host,
@@ -142,7 +142,7 @@ class WebsocketProxyConnect(Connect):
         proto = await self.__await_impl__()
         return proto
 
-    def __await__(self: Self) -> Generator[Any, None, WebSocketClientProtocol]:
+    def __await__(self: Self) -> Generator[Any, None, ClientConnection]:
         return self._proxy_connect().__await__()
 
 
@@ -311,7 +311,7 @@ def _get_api_url_and_key(
 class PrefectEventsClient(EventsClient):
     """A Prefect Events client that streams events to a Prefect server"""
 
-    _websocket: Optional[WebSocketClientProtocol]
+    _websocket: Optional[ClientConnection]
     _unconfirmed_events: List[Event]
 
     def __init__(
@@ -403,19 +403,11 @@ class PrefectEventsClient(EventsClient):
             await self.emit(event)
         logger.debug("Finished resending unconfirmed events.")
 
-    async def _checkpoint(self, event: Event) -> None:
+    async def _checkpoint(self) -> None:
         assert self._websocket
 
-        self._unconfirmed_events.append(event)
-
         unconfirmed_count = len(self._unconfirmed_events)
 
-        logger.debug(
-            "Added event id=%s to unconfirmed events list. "
-            "There are now %s unconfirmed events.",
-            event.id,
-            unconfirmed_count,
-        )
         if unconfirmed_count < self._checkpoint_every:
             return
 
@@ -433,6 +425,16 @@ class PrefectEventsClient(EventsClient):
 
     async def _emit(self, event: Event) -> None:
         self._log_debug("Emitting event id=%s.", event.id)
+
+        self._unconfirmed_events.append(event)
+
+        logger.debug(
+            "Added event id=%s to unconfirmed events list. "
+            "There are now %s unconfirmed events.",
+            event.id,
+            len(self._unconfirmed_events),
+        )
+
         for i in range(self._reconnection_attempts + 1):
             self._log_debug("Emit reconnection attempt %s.", i)
             try:
@@ -450,7 +452,7 @@ class PrefectEventsClient(EventsClient):
                 self._log_debug("Sending event id=%s.", event.id)
                 await self._websocket.send(event.model_dump_json())
                 self._log_debug("Checkpointing event id=%s.", event.id)
-                await self._checkpoint(event)
+                await self._checkpoint()
 
                 return
             except ConnectionClosed:
@@ -537,7 +539,7 @@ class PrefectCloudEventsClient(PrefectEventsClient):
         )
         self._connect = websocket_connect(
             self._events_socket_url,
-            extra_headers={"Authorization": f"bearer {api_key}"},
+            additional_headers={"Authorization": f"bearer {api_key}"},
         )
 
 
@@ -562,7 +564,7 @@ class PrefectEventSubscriber:
 
     """
 
-    _websocket: Optional[WebSocketClientProtocol]
+    _websocket: Optional[ClientConnection]
     _filter: "EventFilter"
     _seen_events: MutableMapping[UUID, bool]
 

{prefect_client-3.2.6 → prefect_client-3.2.8}/src/prefect/events/schemas/automations.py

@@ -84,7 +84,7 @@ class Trigger(PrefectBaseModel, abc.ABC, extra="ignore"):  # type: ignore[call-a
                 getattr(self, "name", None)
                 or f"Automation for deployment {self._deployment_id}"
             ),
-            description="",
+            description=getattr(self, "description", ""),
             enabled=getattr(self, "enabled", True),
             trigger=trigger,
             actions=self.actions(),