prefect-client 3.2.7__py3-none-any.whl → 3.2.8__py3-none-any.whl

prefect/_build_info.py

@@ -1,5 +1,5 @@
 # Generated by versioningit
-__version__ = "3.2.7"
-__build_date__ = "2025-02-21 19:38:38.599741+00:00"
-__git_commit__ = "d4d9001e40eb378ec68a97f7e3bab07d89c7557c"
+__version__ = "3.2.8"
+__build_date__ = "2025-02-28 00:11:08.300212+00:00"
+__git_commit__ = "2c0860d5dcd6156fc3fae75d53f7f48f4bd1a522"
 __dirty__ = False

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/_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/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/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]