prefect-client 3.0.0rc12__py3-none-any.whl → 3.0.0rc14__py3-none-any.whl

prefect/concurrency/asyncio.py

@@ -2,6 +2,7 @@ import asyncio
 from contextlib import asynccontextmanager
 from typing import AsyncGenerator, List, Literal, Optional, Union, cast
 
+import anyio
 import httpx
 import pendulum
 
@@ -14,6 +15,7 @@ except ImportError:
 from prefect.client.orchestration import get_client
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
 
+from .context import ConcurrencyContext
 from .events import (
     _emit_concurrency_acquisition_events,
     _emit_concurrency_release_events,
@@ -34,6 +36,7 @@ async def concurrency(
     names: Union[str, List[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
+    create_if_missing: Optional[bool] = True,
 ) -> AsyncGenerator[None, None]:
     """A context manager that acquires and releases concurrency slots from the
     given concurrency limits.
@@ -43,6 +46,7 @@ async def concurrency(
         occupy: The number of slots to acquire and hold from each limit.
         timeout_seconds: The number of seconds to wait for the slots to be acquired before
             raising a `TimeoutError`. A timeout of `None` will wait indefinitely.
+        create_if_missing: Whether to create the concurrency limits if they do not exist.
 
     Raises:
         TimeoutError: If the slots are not acquired within the given timeout.
@@ -60,9 +64,17 @@ async def concurrency(
         await resource_heavy()
     ```
     """
+    if not names:
+        yield
+        return
+
     names = names if isinstance(names, list) else [names]
+
     limits = await _acquire_concurrency_slots(
-        names, occupy, timeout_seconds=timeout_seconds
+        names,
+        occupy,
+        timeout_seconds=timeout_seconds,
+        create_if_missing=create_if_missing,
     )
     acquisition_time = pendulum.now("UTC")
     emitted_events = _emit_concurrency_acquisition_events(limits, occupy)
@@ -71,13 +83,28 @@ async def concurrency(
         yield
     finally:
         occupancy_period = cast(Interval, (pendulum.now("UTC") - acquisition_time))
-        await _release_concurrency_slots(
-            names, occupy, occupancy_period.total_seconds()
-        )
+        try:
+            await _release_concurrency_slots(
+                names, occupy, occupancy_period.total_seconds()
+            )
+        except anyio.get_cancelled_exc_class():
+            # The task was cancelled before it could release the slots. Add the
+            # slots to the cleanup list so they can be released when the
+            # concurrency context is exited.
+            if ctx := ConcurrencyContext.get():
+                ctx.cleanup_slots.append(
+                    (names, occupy, occupancy_period.total_seconds())
+                )
+
         _emit_concurrency_release_events(limits, occupy, emitted_events)
 
 
-async def rate_limit(names: Union[str, List[str]], occupy: int = 1) -> None:
+async def rate_limit(
+    names: Union[str, List[str]],
+    occupy: int = 1,
+    timeout_seconds: Optional[float] = None,
+    create_if_missing: Optional[bool] = True,
+) -> None:
     """Block execution until an `occupy` number of slots of the concurrency
     limits given in `names` are acquired. Requires that all given concurrency
     limits have a slot decay.
@@ -85,9 +112,22 @@ async def rate_limit(names: Union[str, List[str]], occupy: int = 1) -> None:
     Args:
         names: The names of the concurrency limits to acquire slots from.
         occupy: The number of slots to acquire and hold from each limit.
+        timeout_seconds: The number of seconds to wait for the slots to be acquired before
+            raising a `TimeoutError`. A timeout of `None` will wait indefinitely.
+        create_if_missing: Whether to create the concurrency limits if they do not exist.
     """
+    if not names:
+        return
+
     names = names if isinstance(names, list) else [names]
-    limits = await _acquire_concurrency_slots(names, occupy, mode="rate_limit")
+
+    limits = await _acquire_concurrency_slots(
+        names,
+        occupy,
+        mode="rate_limit",
+        timeout_seconds=timeout_seconds,
+        create_if_missing=create_if_missing,
+    )
     _emit_concurrency_acquisition_events(limits, occupy)
 
 
@@ -96,9 +136,10 @@ async def _acquire_concurrency_slots(
     slots: int,
     mode: Union[Literal["concurrency"], Literal["rate_limit"]] = "concurrency",
     timeout_seconds: Optional[float] = None,
+    create_if_missing: Optional[bool] = True,
 ) -> List[MinimalConcurrencyLimitResponse]:
     service = ConcurrencySlotAcquisitionService.instance(frozenset(names))
-    future = service.send((slots, mode, timeout_seconds))
+    future = service.send((slots, mode, timeout_seconds, create_if_missing))
     response_or_exception = await asyncio.wrap_future(future)
 
     if isinstance(response_or_exception, Exception):

prefect/concurrency/context.py

@@ -0,0 +1,24 @@
+from contextvars import ContextVar
+from typing import List, Tuple
+
+from prefect.client.orchestration import get_client
+from prefect.context import ContextModel, Field
+
+
+class ConcurrencyContext(ContextModel):
+    __var__: ContextVar = ContextVar("concurrency")
+
+    # Track the slots that have been acquired but were not able to be released
+    # due to cancellation or some other error. These slots are released when
+    # the context manager exits.
+    cleanup_slots: List[Tuple[List[str], int, float]] = Field(default_factory=list)
+
+    def __exit__(self, *exc_info):
+        if self.cleanup_slots:
+            with get_client(sync_client=True) as client:
+                for names, occupy, occupancy_seconds in self.cleanup_slots:
+                    client.release_concurrency_slots(
+                        names=names, slots=occupy, occupancy_seconds=occupancy_seconds
+                    )
+
+        return super().__exit__(*exc_info)

prefect/concurrency/services.py

@@ -34,11 +34,16 @@ class ConcurrencySlotAcquisitionService(QueueService):
             yield
 
     async def _handle(
-        self, item: Tuple[int, str, Optional[float], concurrent.futures.Future]
+        self,
+        item: Tuple[
+            int, str, Optional[float], concurrent.futures.Future, Optional[bool]
+        ],
     ) -> None:
-        occupy, mode, timeout_seconds, future = item
+        occupy, mode, timeout_seconds, future, create_if_missing = item
         try:
-            response = await self.acquire_slots(occupy, mode, timeout_seconds)
+            response = await self.acquire_slots(
+                occupy, mode, timeout_seconds, create_if_missing
+            )
         except Exception as exc:
             # If the request to the increment endpoint fails in a non-standard
             # way, we need to set the future's result so that the caller can
@@ -49,13 +54,20 @@ class ConcurrencySlotAcquisitionService(QueueService):
             future.set_result(response)
 
     async def acquire_slots(
-        self, slots: int, mode: str, timeout_seconds: Optional[float] = None
+        self,
+        slots: int,
+        mode: str,
+        timeout_seconds: Optional[float] = None,
+        create_if_missing: Optional[bool] = False,
     ) -> httpx.Response:
         with timeout_async(seconds=timeout_seconds):
             while True:
                 try:
                     response = await self._client.increment_concurrency_slots(
-                        names=self.concurrency_limit_names, slots=slots, mode=mode
+                        names=self.concurrency_limit_names,
+                        slots=slots,
+                        mode=mode,
+                        create_if_missing=create_if_missing,
                     )
                 except Exception as exc:
                     if (
@@ -69,7 +81,9 @@ class ConcurrencySlotAcquisitionService(QueueService):
                 else:
                     return response
 
-    def send(self, item: Tuple[int, str, Optional[float]]) -> concurrent.futures.Future:
+    def send(
+        self, item: Tuple[int, str, Optional[float], Optional[bool]]
+    ) -> concurrent.futures.Future:
         with self._lock:
             if self._stopped:
                 raise RuntimeError("Cannot put items in a stopped service instance.")
@@ -77,7 +91,9 @@ class ConcurrencySlotAcquisitionService(QueueService):
             logger.debug("Service %r enqueuing item %r", self, item)
             future: concurrent.futures.Future = concurrent.futures.Future()
 
-            occupy, mode, timeout_seconds = item
-            self._queue.put_nowait((occupy, mode, timeout_seconds, future))
+            occupy, mode, timeout_seconds, create_if_missing = item
+            self._queue.put_nowait(
+                (occupy, mode, timeout_seconds, future, create_if_missing)
+            )
 
         return future

prefect/concurrency/sync.py

@@ -40,6 +40,7 @@ def concurrency(
     names: Union[str, List[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
+    create_if_missing: Optional[bool] = True,
 ) -> Generator[None, None, None]:
     """A context manager that acquires and releases concurrency slots from the
     given concurrency limits.
@@ -49,6 +50,7 @@ def concurrency(
         occupy: The number of slots to acquire and hold from each limit.
         timeout_seconds: The number of seconds to wait for the slots to be acquired before
             raising a `TimeoutError`. A timeout of `None` will wait indefinitely.
+        create_if_missing: Whether to create the concurrency limits if they do not exist.
 
     Raises:
         TimeoutError: If the slots are not acquired within the given timeout.
@@ -66,10 +68,18 @@ def concurrency(
         resource_heavy()
     ```
     """
+    if not names:
+        yield
+        return
+
     names = names if isinstance(names, list) else [names]
 
     limits: List[MinimalConcurrencyLimitResponse] = _call_async_function_from_sync(
-        _acquire_concurrency_slots, names, occupy, timeout_seconds=timeout_seconds
+        _acquire_concurrency_slots,
+        names,
+        occupy,
+        timeout_seconds=timeout_seconds,
+        create_if_missing=create_if_missing,
     )
     acquisition_time = pendulum.now("UTC")
     emitted_events = _emit_concurrency_acquisition_events(limits, occupy)
@@ -87,7 +97,12 @@ def concurrency(
         _emit_concurrency_release_events(limits, occupy, emitted_events)
 
 
-def rate_limit(names: Union[str, List[str]], occupy: int = 1) -> None:
+def rate_limit(
+    names: Union[str, List[str]],
+    occupy: int = 1,
+    timeout_seconds: Optional[float] = None,
+    create_if_missing: Optional[bool] = True,
+) -> None:
     """Block execution until an `occupy` number of slots of the concurrency
     limits given in `names` are acquired. Requires that all given concurrency
     limits have a slot decay.
@@ -95,10 +110,22 @@ def rate_limit(names: Union[str, List[str]], occupy: int = 1) -> None:
     Args:
         names: The names of the concurrency limits to acquire slots from.
         occupy: The number of slots to acquire and hold from each limit.
+        timeout_seconds: The number of seconds to wait for the slots to be acquired before
+            raising a `TimeoutError`. A timeout of `None` will wait indefinitely.
+        create_if_missing: Whether to create the concurrency limits if they do not exist.
     """
+    if not names:
+        return
+
     names = names if isinstance(names, list) else [names]
+
     limits = _call_async_function_from_sync(
-        _acquire_concurrency_slots, names, occupy, mode="rate_limit"
+        _acquire_concurrency_slots,
+        names,
+        occupy,
+        mode="rate_limit",
+        timeout_seconds=timeout_seconds,
+        create_if_missing=create_if_missing,
     )
     _emit_concurrency_acquisition_events(limits, occupy)
 

prefect/context.py

@@ -10,12 +10,13 @@ import os
 import sys
 import warnings
 import weakref
-from contextlib import ExitStack, contextmanager
+from contextlib import ExitStack, asynccontextmanager, contextmanager
 from contextvars import ContextVar, Token
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
     Any,
+    AsyncGenerator,
     Dict,
     Generator,
     Mapping,
@@ -44,6 +45,7 @@ from prefect.settings import PREFECT_HOME, Profile, Settings
 from prefect.states import State
 from prefect.task_runners import TaskRunner
 from prefect.utilities.asyncutils import run_coro_as_sync
+from prefect.utilities.services import start_client_metrics_server
 
 T = TypeVar("T")
 
@@ -177,36 +179,34 @@ class ContextModel(BaseModel):
         return self.model_dump(exclude_unset=True)
 
 
-class ClientContext(ContextModel):
+class SyncClientContext(ContextModel):
     """
-    A context for managing the Prefect client instances.
+    A context for managing the sync Prefect client instances.
 
     Clients were formerly tracked on the TaskRunContext and FlowRunContext, but
     having two separate places and the addition of both sync and async clients
     made it difficult to manage. This context is intended to be the single
-    source for clients.
+    source for sync clients.
 
-    The client creates both sync and async clients, which can either be read
-    directly from the context object OR loaded with get_client, inject_client,
-    or other Prefect utilities.
+    The client creates a sync client, which can either be read directly from
+    the context object OR loaded with get_client, inject_client, or other
+    Prefect utilities.
 
-    with ClientContext.get_or_create() as ctx:
+    with SyncClientContext.get_or_create() as ctx:
         c1 = get_client(sync_client=True)
         c2 = get_client(sync_client=True)
         assert c1 is c2
-        assert c1 is ctx.sync_client
+        assert c1 is ctx.client
     """
 
-    __var__ = ContextVar("clients")
-    sync_client: SyncPrefectClient
-    async_client: PrefectClient
+    __var__ = ContextVar("sync-client-context")
+    client: SyncPrefectClient
     _httpx_settings: Optional[dict[str, Any]] = PrivateAttr(None)
     _context_stack: int = PrivateAttr(0)
 
     def __init__(self, httpx_settings: Optional[dict[str, Any]] = None):
         super().__init__(
-            sync_client=get_client(sync_client=True, httpx_settings=httpx_settings),
-            async_client=get_client(sync_client=False, httpx_settings=httpx_settings),
+            client=get_client(sync_client=True, httpx_settings=httpx_settings),
         )
         self._httpx_settings = httpx_settings
         self._context_stack = 0
@@ -214,8 +214,7 @@ class ClientContext(ContextModel):
     def __enter__(self):
         self._context_stack += 1
         if self._context_stack == 1:
-            self.sync_client.__enter__()
-            run_coro_as_sync(self.async_client.__aenter__())
+            self.client.__enter__()
             return super().__enter__()
         else:
             return self
@@ -223,18 +222,74 @@ class ClientContext(ContextModel):
     def __exit__(self, *exc_info):
         self._context_stack -= 1
         if self._context_stack == 0:
-            self.sync_client.__exit__(*exc_info)
-            run_coro_as_sync(self.async_client.__aexit__(*exc_info))
+            self.client.__exit__(*exc_info)
             return super().__exit__(*exc_info)
 
     @classmethod
     @contextmanager
-    def get_or_create(cls) -> Generator["ClientContext", None, None]:
-        ctx = ClientContext.get()
+    def get_or_create(cls) -> Generator["SyncClientContext", None, None]:
+        ctx = SyncClientContext.get()
         if ctx:
             yield ctx
         else:
-            with ClientContext() as ctx:
+            with SyncClientContext() as ctx:
+                yield ctx
+
+
+class AsyncClientContext(ContextModel):
+    """
+    A context for managing the async Prefect client instances.
+
+    Clients were formerly tracked on the TaskRunContext and FlowRunContext, but
+    having two separate places and the addition of both sync and async clients
+    made it difficult to manage. This context is intended to be the single
+    source for async clients.
+
+    The client creates an async client, which can either be read directly from
+    the context object OR loaded with get_client, inject_client, or other
+    Prefect utilities.
+
+    with AsyncClientContext.get_or_create() as ctx:
+        c1 = get_client(sync_client=False)
+        c2 = get_client(sync_client=False)
+        assert c1 is c2
+        assert c1 is ctx.client
+    """
+
+    __var__ = ContextVar("async-client-context")
+    client: PrefectClient
+    _httpx_settings: Optional[dict[str, Any]] = PrivateAttr(None)
+    _context_stack: int = PrivateAttr(0)
+
+    def __init__(self, httpx_settings: Optional[dict[str, Any]] = None):
+        super().__init__(
+            client=get_client(sync_client=False, httpx_settings=httpx_settings),
+        )
+        self._httpx_settings = httpx_settings
+        self._context_stack = 0
+
+    async def __aenter__(self):
+        self._context_stack += 1
+        if self._context_stack == 1:
+            await self.client.__aenter__()
+            return super().__enter__()
+        else:
+            return self
+
+    async def __aexit__(self, *exc_info):
+        self._context_stack -= 1
+        if self._context_stack == 0:
+            await self.client.__aexit__(*exc_info)
+            return super().__exit__(*exc_info)
+
+    @classmethod
+    @asynccontextmanager
+    async def get_or_create(cls) -> AsyncGenerator[Self, None]:
+        ctx = cls.get()
+        if ctx:
+            yield ctx
+        else:
+            with cls() as ctx:
                 yield ctx
 
 
@@ -248,6 +303,11 @@ class RunContext(ContextModel):
         client: The Prefect client instance being used for API communication
     """
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        start_client_metrics_server()
+
     start_time: DateTime = Field(default_factory=lambda: pendulum.now("UTC"))
     input_keyset: Optional[Dict[str, Dict[str, str]]] = None
     client: Union[PrefectClient, SyncPrefectClient]
@@ -300,7 +360,7 @@ class EngineContext(RunContext):
         default_factory=weakref.WeakValueDictionary
     )
 
-    # Events worker to emit events to Prefect Cloud
+    # Events worker to emit events
     events: Optional[EventsWorker] = None
 
     __var__: ContextVar = ContextVar("flow_run")
@@ -601,7 +661,7 @@ def root_settings_context():
             ),
             file=sys.stderr,
         )
-        active_name = "default"
+        active_name = "ephemeral"
 
     with use_profile(
         profiles[active_name],

prefect/events/clients.py

@@ -19,6 +19,7 @@ import httpx
 import orjson
 import pendulum
 from cachetools import TTLCache
+from prometheus_client import Counter
 from typing_extensions import Self
 from websockets import Subprotocol
 from websockets.client import WebSocketClientProtocol, connect
@@ -36,6 +37,30 @@ from prefect.settings import PREFECT_API_KEY, PREFECT_API_URL, PREFECT_CLOUD_API
 if TYPE_CHECKING:
     from prefect.events.filters import EventFilter
 
+EVENTS_EMITTED = Counter(
+    "prefect_events_emitted",
+    "The number of events emitted by Prefect event clients",
+    labelnames=["client"],
+)
+EVENTS_OBSERVED = Counter(
+    "prefect_events_observed",
+    "The number of events observed by Prefect event subscribers",
+    labelnames=["client"],
+)
+EVENT_WEBSOCKET_CONNECTIONS = Counter(
+    "prefect_event_websocket_connections",
+    (
+        "The number of times Prefect event clients have connected to an event stream, "
+        "broken down by direction (in/out) and connection (initial/reconnect)"
+    ),
+    labelnames=["client", "direction", "connection"],
+)
+EVENT_WEBSOCKET_CHECKPOINTS = Counter(
+    "prefect_event_websocket_checkpoints",
+    "The number of checkpoints performed by Prefect event clients",
+    labelnames=["client"],
+)
+
 logger = get_logger(__name__)
 
 
@@ -82,6 +107,10 @@ def get_events_subscriber(
 class EventsClient(abc.ABC):
     """The abstract interface for all Prefect Events clients"""
 
+    @property
+    def client_name(self) -> str:
+        return self.__class__.__name__
+
     async def emit(self, event: Event) -> None:
         """Emit a single event"""
         if not hasattr(self, "_in_context"):
@@ -89,7 +118,11 @@ class EventsClient(abc.ABC):
                 "Events may only be emitted while this client is being used as a "
                 "context manager"
             )
-        return await self._emit(event)
+
+        try:
+            return await self._emit(event)
+        finally:
+            EVENTS_EMITTED.labels(self.client_name).inc()
 
     @abc.abstractmethod
     async def _emit(self, event: Event) -> None:  # pragma: no cover
@@ -140,6 +173,11 @@ class AssertingEventsClient(EventsClient):
         cls.last = None
         cls.all = []
 
+    def pop_events(self) -> List[Event]:
+        events = self.events
+        self.events = []
+        return events
+
     async def _emit(self, event: Event) -> None:
         self.events.append(event)
 
@@ -294,6 +332,8 @@ class PrefectEventsClient(EventsClient):
         # don't clear the list, just the ones that we are sure of.
         self._unconfirmed_events = self._unconfirmed_events[unconfirmed_count:]
 
+        EVENT_WEBSOCKET_CHECKPOINTS.labels(self.client_name).inc()
+
     async def _emit(self, event: Event) -> None:
         for i in range(self._reconnection_attempts + 1):
             try:
@@ -421,10 +461,17 @@ class PrefectEventSubscriber:
         if self._reconnection_attempts < 0:
             raise ValueError("reconnection_attempts must be a non-negative integer")
 
+    @property
+    def client_name(self) -> str:
+        return self.__class__.__name__
+
     async def __aenter__(self) -> Self:
         # Don't handle any errors in the initial connection, because these are most
         # likely a permission or configuration issue that should propagate
-        await self._reconnect()
+        try:
+            await self._reconnect()
+        finally:
+            EVENT_WEBSOCKET_CONNECTIONS.labels(self.client_name, "out", "initial")
         return self
 
     async def _reconnect(self) -> None:
@@ -498,7 +545,12 @@ class PrefectEventSubscriber:
                 # Otherwise, after the first time through this loop, we're recovering
                 # from a ConnectionClosed, so reconnect now.
                 if not self._websocket or i > 0:
-                    await self._reconnect()
+                    try:
+                        await self._reconnect()
+                    finally:
+                        EVENT_WEBSOCKET_CONNECTIONS.labels(
+                            self.client_name, "out", "reconnect"
+                        )
                     assert self._websocket
 
                 while True:
@@ -509,7 +561,10 @@ class PrefectEventSubscriber:
                         continue
                     self._seen_events[event.id] = True
 
-                    return event
+                    try:
+                        return event
+                    finally:
+                        EVENTS_OBSERVED.labels(self.client_name).inc()
             except ConnectionClosedOK:
                 logger.debug('Connection closed with "OK" status')
                 raise StopAsyncIteration

prefect/events/worker.py

@@ -17,7 +17,6 @@ from .clients import (
     EventsClient,
     NullEventsClient,
     PrefectCloudEventsClient,
-    PrefectEphemeralEventsClient,
     PrefectEventsClient,
 )
 from .related import related_resources_from_run_context
@@ -97,7 +96,15 @@ class EventsWorker(QueueService[Event]):
             elif should_emit_events_to_running_server():
                 client_type = PrefectEventsClient
             elif should_emit_events_to_ephemeral_server():
-                client_type = PrefectEphemeralEventsClient
+                # create an ephemeral API if none was provided
+                from prefect.server.api.server import SubprocessASGIServer
+
+                server = SubprocessASGIServer()
+                server.start()
+                assert server.server_process is not None, "Server process did not start"
+
+                client_kwargs = {"api_url": server.api_url}
+                client_type = PrefectEventsClient
             else:
                 client_type = NullEventsClient
 

prefect/flow_engine.py

@@ -29,7 +29,8 @@ from prefect.client.orchestration import SyncPrefectClient, get_client
 from prefect.client.schemas import FlowRun, TaskRun
 from prefect.client.schemas.filters import FlowRunFilter
 from prefect.client.schemas.sorting import FlowRunSort
-from prefect.context import ClientContext, FlowRunContext, TagsContext
+from prefect.concurrency.context import ConcurrencyContext
+from prefect.context import FlowRunContext, SyncClientContext, TagsContext
 from prefect.exceptions import (
     Abort,
     Pause,
@@ -505,6 +506,8 @@ class FlowRunEngine(Generic[P, R]):
                     task_runner=task_runner,
                 )
             )
+            stack.enter_context(ConcurrencyContext())
+
             # set the logger to the flow run logger
             self.logger = flow_run_logger(flow_run=self.flow_run, flow=self.flow)
 
@@ -529,8 +532,8 @@ class FlowRunEngine(Generic[P, R]):
         """
         Enters a client context and creates a flow run if needed.
         """
-        with ClientContext.get_or_create() as client_ctx:
-            self._client = client_ctx.sync_client
+        with SyncClientContext.get_or_create() as client_ctx:
+            self._client = client_ctx.client
             self._is_started = True
 
             if not self.flow_run: