prefect-client 3.0.0rc13__py3-none-any.whl → 3.0.0rc15__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")
 
@@ -89,11 +91,12 @@ def hydrated_context(
             client = client or get_client(sync_client=True)
             if flow_run_context := serialized_context.get("flow_run_context"):
                 flow = flow_run_context["flow"]
+                task_runner = stack.enter_context(flow.task_runner.duplicate())
                 flow_run_context = FlowRunContext(
                     **flow_run_context,
                     client=client,
                     result_factory=run_coro_as_sync(ResultFactory.from_flow(flow)),
-                    task_runner=flow.task_runner.duplicate(),
+                    task_runner=task_runner,
                     detached=True,
                 )
                 stack.enter_context(flow_run_context)
@@ -177,36 +180,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 +215,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 +223,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 +304,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 +361,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 +662,7 @@ def root_settings_context():
             ),
             file=sys.stderr,
         )
-        active_name = "default"
+        active_name = "ephemeral"
 
     with use_profile(
         profiles[active_name],