prefect-client 2.20.4__py3-none-any.whl → 3.0.0__py3-none-any.whl

prefect/concurrency/services.py

@@ -2,39 +2,53 @@ import asyncio
 import concurrent.futures
 from contextlib import asynccontextmanager
 from typing import (
+    TYPE_CHECKING,
+    AsyncGenerator,
     FrozenSet,
     Optional,
     Tuple,
 )
 
 import httpx
-from prefect._vendor.starlette import status
+from starlette import status
 
-from prefect import get_client
 from prefect._internal.concurrency import logger
 from prefect._internal.concurrency.services import QueueService
-from prefect.client.orchestration import PrefectClient
+from prefect.client.orchestration import get_client
 from prefect.utilities.timeout import timeout_async
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
 
 class ConcurrencySlotAcquisitionService(QueueService):
     def __init__(self, concurrency_limit_names: FrozenSet[str]):
         super().__init__(concurrency_limit_names)
-        self._client: PrefectClient
+        self._client: "PrefectClient"
         self.concurrency_limit_names = sorted(list(concurrency_limit_names))
 
     @asynccontextmanager
-    async def _lifespan(self):
+    async def _lifespan(self) -> AsyncGenerator[None, None]:
         async with get_client() as client:
             self._client = client
             yield
 
     async def _handle(
-        self, item: Tuple[int, str, Optional[float], concurrent.futures.Future]
-    ):
-        occupy, mode, timeout_seconds, future = item
+        self,
+        item: Tuple[
+            int,
+            str,
+            Optional[float],
+            concurrent.futures.Future,
+            Optional[bool],
+            Optional[int],
+        ],
+    ) -> None:
+        occupy, mode, timeout_seconds, future, create_if_missing, max_retries = item
         try:
-            response = await self.acquire_slots(occupy, mode, timeout_seconds)
+            response = await self.acquire_slots(
+                occupy, mode, timeout_seconds, create_if_missing, max_retries
+            )
         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
@@ -45,27 +59,41 @@ class ConcurrencySlotAcquisitionService(QueueService):
             future.set_result(response)
 
     async def acquire_slots(
-        self, slots: int, mode: str, timeout_seconds: Optional[float] = None
-    ):
-        with timeout_async(timeout_seconds):
+        self,
+        slots: int,
+        mode: str,
+        timeout_seconds: Optional[float] = None,
+        create_if_missing: Optional[bool] = False,
+        max_retries: Optional[int] = None,
+    ) -> 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 (
                         isinstance(exc, httpx.HTTPStatusError)
                         and exc.response.status_code == status.HTTP_423_LOCKED
                     ):
+                        if max_retries is not None and max_retries <= 0:
+                            raise exc
                         retry_after = float(exc.response.headers["Retry-After"])
                         await asyncio.sleep(retry_after)
+                        if max_retries is not None:
+                            max_retries -= 1
                     else:
                         raise exc
                 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], Optional[int]]
+    ) -> concurrent.futures.Future:
         with self._lock:
             if self._stopped:
                 raise RuntimeError("Cannot put items in a stopped service instance.")
@@ -73,7 +101,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, max_retries = item
+            self._queue.put_nowait(
+                (occupy, mode, timeout_seconds, future, create_if_missing, max_retries)
+            )
 
         return future

prefect/concurrency/sync.py

@@ -1,5 +1,15 @@
 from contextlib import contextmanager
-from typing import List, Optional, Union, cast
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Generator,
+    List,
+    Optional,
+    TypeVar,
+    Union,
+    cast,
+)
 
 import pendulum
 
@@ -22,13 +32,17 @@ from .events import (
     _emit_concurrency_release_events,
 )
 
+T = TypeVar("T")
+
 
 @contextmanager
 def concurrency(
     names: Union[str, List[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
-):
+    create_if_missing: bool = True,
+    max_retries: Optional[int] = None,
+) -> Generator[None, None, None]:
     """A context manager that acquires and releases concurrency slots from the
     given concurrency limits.
 
@@ -37,6 +51,8 @@ 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.
+        max_retries: The maximum number of retries to acquire the concurrency slots.
 
     Raises:
         TimeoutError: If the slots are not acquired within the given timeout.
@@ -54,10 +70,19 @@ 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,
+        max_retries=max_retries,
     )
     acquisition_time = pendulum.now("UTC")
     emitted_events = _emit_concurrency_acquisition_events(limits, occupy)
@@ -75,7 +100,12 @@ def concurrency(
         _emit_concurrency_release_events(limits, occupy, emitted_events)
 
 
-def rate_limit(names: Union[str, List[str]], occupy: int = 1):
+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.
@@ -83,19 +113,33 @@ def rate_limit(names: Union[str, List[str]], occupy: int = 1):
     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)
 
 
-def _call_async_function_from_sync(fn, *args, **kwargs):
+def _call_async_function_from_sync(
+    fn: Callable[..., Awaitable[T]], *args: Any, **kwargs: Any
+) -> T:
     loop = get_running_loop()
     call = create_call(fn, *args, **kwargs)
 
     if loop is not None:
         return from_sync.call_soon_in_loop_thread(call).result()
     else:
-        return call()
+        return call()  # type: ignore [return-value]

prefect/concurrency/v1/asyncio.py

@@ -0,0 +1,143 @@
+import asyncio
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator, List, Optional, Union, cast
+from uuid import UUID
+
+import anyio
+import httpx
+import pendulum
+
+from ...client.schemas.responses import MinimalConcurrencyLimitResponse
+
+try:
+    from pendulum import Interval
+except ImportError:
+    # pendulum < 3
+    from pendulum.period import Period as Interval  # type: ignore
+
+from prefect.client.orchestration import get_client
+
+from .context import ConcurrencyContext
+from .events import (
+    _emit_concurrency_acquisition_events,
+    _emit_concurrency_release_events,
+)
+from .services import ConcurrencySlotAcquisitionService
+
+
+class ConcurrencySlotAcquisitionError(Exception):
+    """Raised when an unhandlable occurs while acquiring concurrency slots."""
+
+
+class AcquireConcurrencySlotTimeoutError(TimeoutError):
+    """Raised when acquiring a concurrency slot times out."""
+
+
+@asynccontextmanager
+async def concurrency(
+    names: Union[str, List[str]],
+    task_run_id: UUID,
+    timeout_seconds: Optional[float] = None,
+) -> AsyncGenerator[None, None]:
+    """A context manager that acquires and releases concurrency slots from the
+    given concurrency limits.
+
+    Args:
+        names: The names of the concurrency limits to acquire slots from.
+        task_run_id: The name of the task_run_id that is incrementing the slots.
+        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.
+
+    Raises:
+        TimeoutError: If the slots are not acquired within the given timeout.
+
+    Example:
+    A simple example of using the async `concurrency` context manager:
+    ```python
+    from prefect.concurrency.v1.asyncio import concurrency
+
+    async def resource_heavy():
+        async with concurrency("test", task_run_id):
+            print("Resource heavy task")
+
+    async def main():
+        await resource_heavy()
+    ```
+    """
+    if not names:
+        yield
+        return
+
+    names_normalized: List[str] = names if isinstance(names, list) else [names]
+
+    limits = await _acquire_concurrency_slots(
+        names_normalized,
+        task_run_id=task_run_id,
+        timeout_seconds=timeout_seconds,
+    )
+    acquisition_time = pendulum.now("UTC")
+    emitted_events = _emit_concurrency_acquisition_events(limits, task_run_id)
+
+    try:
+        yield
+    finally:
+        occupancy_period = cast(Interval, (pendulum.now("UTC") - acquisition_time))
+        try:
+            await _release_concurrency_slots(
+                names_normalized, task_run_id, 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_normalized, occupancy_period.total_seconds(), task_run_id)
+                )
+
+        _emit_concurrency_release_events(limits, emitted_events, task_run_id)
+
+
+async def _acquire_concurrency_slots(
+    names: List[str],
+    task_run_id: UUID,
+    timeout_seconds: Optional[float] = None,
+) -> List[MinimalConcurrencyLimitResponse]:
+    service = ConcurrencySlotAcquisitionService.instance(frozenset(names))
+    future = service.send((task_run_id, timeout_seconds))
+    response_or_exception = await asyncio.wrap_future(future)
+
+    if isinstance(response_or_exception, Exception):
+        if isinstance(response_or_exception, TimeoutError):
+            raise AcquireConcurrencySlotTimeoutError(
+                f"Attempt to acquire concurrency limits timed out after {timeout_seconds} second(s)"
+            ) from response_or_exception
+
+        raise ConcurrencySlotAcquisitionError(
+            f"Unable to acquire concurrency limits {names!r}"
+        ) from response_or_exception
+
+    return _response_to_concurrency_limit_response(response_or_exception)
+
+
+async def _release_concurrency_slots(
+    names: List[str],
+    task_run_id: UUID,
+    occupancy_seconds: float,
+) -> List[MinimalConcurrencyLimitResponse]:
+    async with get_client() as client:
+        response = await client.decrement_v1_concurrency_slots(
+            names=names,
+            task_run_id=task_run_id,
+            occupancy_seconds=occupancy_seconds,
+        )
+        return _response_to_concurrency_limit_response(response)
+
+
+def _response_to_concurrency_limit_response(
+    response: httpx.Response,
+) -> List[MinimalConcurrencyLimitResponse]:
+    data = response.json() or []
+    return [
+        MinimalConcurrencyLimitResponse.model_validate(limit) for limit in data if data
+    ]

prefect/concurrency/v1/context.py

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

prefect/concurrency/v1/events.py

@@ -0,0 +1,61 @@
+from typing import Dict, List, Literal, Optional, Union
+from uuid import UUID
+
+from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
+from prefect.events import Event, RelatedResource, emit_event
+
+
+def _emit_concurrency_event(
+    phase: Union[Literal["acquired"], Literal["released"]],
+    primary_limit: MinimalConcurrencyLimitResponse,
+    related_limits: List[MinimalConcurrencyLimitResponse],
+    task_run_id: UUID,
+    follows: Union[Event, None] = None,
+) -> Union[Event, None]:
+    resource: Dict[str, str] = {
+        "prefect.resource.id": f"prefect.concurrency-limit.v1.{primary_limit.id}",
+        "prefect.resource.name": primary_limit.name,
+        "limit": str(primary_limit.limit),
+        "task_run_id": str(task_run_id),
+    }
+
+    related = [
+        RelatedResource.model_validate(
+            {
+                "prefect.resource.id": f"prefect.concurrency-limit.v1.{limit.id}",
+                "prefect.resource.role": "concurrency-limit",
+            }
+        )
+        for limit in related_limits
+        if limit.id != primary_limit.id
+    ]
+
+    return emit_event(
+        f"prefect.concurrency-limit.v1.{phase}",
+        resource=resource,
+        related=related,
+        follows=follows,
+    )
+
+
+def _emit_concurrency_acquisition_events(
+    limits: List[MinimalConcurrencyLimitResponse],
+    task_run_id: UUID,
+) -> Dict[UUID, Optional[Event]]:
+    events = {}
+    for limit in limits:
+        event = _emit_concurrency_event("acquired", limit, limits, task_run_id)
+        events[limit.id] = event
+
+    return events
+
+
+def _emit_concurrency_release_events(
+    limits: List[MinimalConcurrencyLimitResponse],
+    events: Dict[UUID, Optional[Event]],
+    task_run_id: UUID,
+) -> None:
+    for limit in limits:
+        _emit_concurrency_event(
+            "released", limit, limits, task_run_id, events[limit.id]
+        )

prefect/concurrency/v1/services.py

@@ -0,0 +1,116 @@
+import asyncio
+import concurrent.futures
+from contextlib import asynccontextmanager
+from json import JSONDecodeError
+from typing import (
+    TYPE_CHECKING,
+    AsyncGenerator,
+    FrozenSet,
+    Optional,
+    Tuple,
+)
+from uuid import UUID
+
+import httpx
+from starlette import status
+
+from prefect._internal.concurrency import logger
+from prefect._internal.concurrency.services import QueueService
+from prefect.client.orchestration import get_client
+from prefect.utilities.timeout import timeout_async
+
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
+
+class ConcurrencySlotAcquisitionServiceError(Exception):
+    """Raised when an error occurs while acquiring concurrency slots."""
+
+
+class ConcurrencySlotAcquisitionService(QueueService):
+    def __init__(self, concurrency_limit_names: FrozenSet[str]):
+        super().__init__(concurrency_limit_names)
+        self._client: "PrefectClient"
+        self.concurrency_limit_names = sorted(list(concurrency_limit_names))
+
+    @asynccontextmanager
+    async def _lifespan(self) -> AsyncGenerator[None, None]:
+        async with get_client() as client:
+            self._client = client
+            yield
+
+    async def _handle(
+        self,
+        item: Tuple[
+            UUID,
+            concurrent.futures.Future,
+            Optional[float],
+        ],
+    ) -> None:
+        task_run_id, future, timeout_seconds = item
+        try:
+            response = await self.acquire_slots(task_run_id, timeout_seconds)
+        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
+            # handle the exception and then re-raise.
+            future.set_result(exc)
+            raise exc
+        else:
+            future.set_result(response)
+
+    async def acquire_slots(
+        self,
+        task_run_id: UUID,
+        timeout_seconds: Optional[float] = None,
+    ) -> httpx.Response:
+        with timeout_async(seconds=timeout_seconds):
+            while True:
+                try:
+                    response = await self._client.increment_v1_concurrency_slots(
+                        task_run_id=task_run_id,
+                        names=self.concurrency_limit_names,
+                    )
+                except Exception as exc:
+                    if (
+                        isinstance(exc, httpx.HTTPStatusError)
+                        and exc.response.status_code == status.HTTP_423_LOCKED
+                    ):
+                        retry_after = exc.response.headers.get("Retry-After")
+                        if retry_after:
+                            retry_after = float(retry_after)
+                            await asyncio.sleep(retry_after)
+                        else:
+                            # We received a 423 but no Retry-After header. This
+                            # should indicate that the server told us to abort
+                            # because the concurrency limit is set to 0, i.e.
+                            # effectively disabled.
+                            try:
+                                reason = exc.response.json()["detail"]
+                            except (JSONDecodeError, KeyError):
+                                logger.error(
+                                    "Failed to parse response from concurrency limit 423 Locked response: %s",
+                                    exc.response.content,
+                                )
+                                reason = "Concurrency limit is locked (server did not specify the reason)"
+                            raise ConcurrencySlotAcquisitionServiceError(
+                                reason
+                            ) from exc
+
+                    else:
+                        raise exc  # type: ignore
+                else:
+                    return response
+
+    def send(self, item: Tuple[UUID, Optional[float]]) -> concurrent.futures.Future:
+        with self._lock:
+            if self._stopped:
+                raise RuntimeError("Cannot put items in a stopped service instance.")
+
+            logger.debug("Service %r enqueuing item %r", self, item)
+            future: concurrent.futures.Future = concurrent.futures.Future()
+
+            task_run_id, timeout_seconds = item
+            self._queue.put_nowait((task_run_id, future, timeout_seconds))
+
+        return future

prefect/concurrency/v1/sync.py

@@ -0,0 +1,92 @@
+from contextlib import contextmanager
+from typing import (
+    Generator,
+    List,
+    Optional,
+    TypeVar,
+    Union,
+    cast,
+)
+from uuid import UUID
+
+import pendulum
+
+from ...client.schemas.responses import MinimalConcurrencyLimitResponse
+from ..sync import _call_async_function_from_sync
+
+try:
+    from pendulum import Interval
+except ImportError:
+    # pendulum < 3
+    from pendulum.period import Period as Interval  # type: ignore
+
+from .asyncio import (
+    _acquire_concurrency_slots,
+    _release_concurrency_slots,
+)
+from .events import (
+    _emit_concurrency_acquisition_events,
+    _emit_concurrency_release_events,
+)
+
+T = TypeVar("T")
+
+
+@contextmanager
+def concurrency(
+    names: Union[str, List[str]],
+    task_run_id: UUID,
+    timeout_seconds: Optional[float] = None,
+) -> Generator[None, None, None]:
+    """
+    A context manager that acquires and releases concurrency slots from the
+    given concurrency limits.
+
+    Args:
+        names: The names of the concurrency limits to acquire.
+        task_run_id: The task run ID acquiring the limits.
+        timeout_seconds: The number of seconds to wait to acquire the limits before
+            raising a `TimeoutError`. A timeout of `None` will wait indefinitely.
+
+    Raises:
+        TimeoutError: If the limits are not acquired within the given timeout.
+
+    Example:
+    A simple example of using the sync `concurrency` context manager:
+    ```python
+    from prefect.concurrency.v1.sync import concurrency
+
+    def resource_heavy():
+        with concurrency("test"):
+            print("Resource heavy task")
+
+    def main():
+        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,
+        timeout_seconds=timeout_seconds,
+        task_run_id=task_run_id,
+    )
+    acquisition_time = pendulum.now("UTC")
+    emitted_events = _emit_concurrency_acquisition_events(limits, task_run_id)
+
+    try:
+        yield
+    finally:
+        occupancy_period = cast(Interval, pendulum.now("UTC") - acquisition_time)
+        _call_async_function_from_sync(
+            _release_concurrency_slots,
+            names,
+            task_run_id,
+            occupancy_period.total_seconds(),
+        )
+        _emit_concurrency_release_events(limits, emitted_events, task_run_id)