prefect-client 3.1.10__py3-none-any.whl → 3.1.12__py3-none-any.whl

prefect/concurrency/asyncio.py

@@ -1,42 +1,25 @@
-import asyncio
+from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator, List, Literal, Optional, Union, cast
+from typing import Optional, Union
 
 import anyio
-import httpx
 import pendulum
 
-from prefect._internal.compatibility.deprecated import deprecated_parameter
-
-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 prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
-from prefect.logging.loggers import get_run_logger
-
-from .context import ConcurrencyContext
-from .events import (
-    _emit_concurrency_acquisition_events,
-    _emit_concurrency_release_events,
+from ._asyncio import (
+    AcquireConcurrencySlotTimeoutError as AcquireConcurrencySlotTimeoutError,
 )
-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."""
+from ._asyncio import ConcurrencySlotAcquisitionError as ConcurrencySlotAcquisitionError
+from ._asyncio import aacquire_concurrency_slots, arelease_concurrency_slots
+from ._events import (
+    emit_concurrency_acquisition_events,
+    emit_concurrency_release_events,
+)
+from .context import ConcurrencyContext
 
 
 @asynccontextmanager
 async def concurrency(
-    names: Union[str, List[str]],
+    names: Union[str, list[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
     max_retries: Optional[int] = None,
@@ -78,7 +61,7 @@ async def concurrency(
 
     names = names if isinstance(names, list) else [names]
 
-    limits = await _aacquire_concurrency_slots(
+    limits = await aacquire_concurrency_slots(
         names,
         occupy,
         timeout_seconds=timeout_seconds,
@@ -87,14 +70,14 @@ async def concurrency(
         strict=strict,
     )
     acquisition_time = pendulum.now("UTC")
-    emitted_events = _emit_concurrency_acquisition_events(limits, occupy)
+    emitted_events = emit_concurrency_acquisition_events(limits, occupy)
 
     try:
         yield
     finally:
-        occupancy_period = cast(Interval, (pendulum.now("UTC") - acquisition_time))
+        occupancy_period = pendulum.now("UTC") - acquisition_time
         try:
-            await _arelease_concurrency_slots(
+            await arelease_concurrency_slots(
                 names, occupy, occupancy_period.total_seconds()
             )
         except anyio.get_cancelled_exc_class():
@@ -106,11 +89,11 @@ async def concurrency(
                     (names, occupy, occupancy_period.total_seconds())
                 )
 
-        _emit_concurrency_release_events(limits, occupy, emitted_events)
+        emit_concurrency_release_events(limits, occupy, emitted_events)
 
 
 async def rate_limit(
-    names: Union[str, List[str]],
+    names: Union[str, list[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
     create_if_missing: Optional[bool] = None,
@@ -137,7 +120,7 @@ async def rate_limit(
 
     names = names if isinstance(names, list) else [names]
 
-    limits = await _aacquire_concurrency_slots(
+    limits = await aacquire_concurrency_slots(
         names,
         occupy,
         mode="rate_limit",
@@ -145,71 +128,4 @@ async def rate_limit(
         create_if_missing=create_if_missing,
         strict=strict,
     )
-    _emit_concurrency_acquisition_events(limits, occupy)
-
-
-@deprecated_parameter(
-    name="create_if_missing",
-    start_date="Sep 2024",
-    end_date="Oct 2024",
-    when=lambda x: x is not None,
-    help="Limits must be explicitly created before acquiring concurrency slots; see `strict` if you want to enforce this behavior.",
-)
-async def _aacquire_concurrency_slots(
-    names: List[str],
-    slots: int,
-    mode: Literal["concurrency", "rate_limit"] = "concurrency",
-    timeout_seconds: Optional[float] = None,
-    create_if_missing: Optional[bool] = None,
-    max_retries: Optional[int] = None,
-    strict: bool = False,
-) -> List[MinimalConcurrencyLimitResponse]:
-    service = ConcurrencySlotAcquisitionService.instance(frozenset(names))
-    future = service.send(
-        (slots, mode, timeout_seconds, create_if_missing, max_retries)
-    )
-    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 slots timed out after {timeout_seconds} second(s)"
-            ) from response_or_exception
-
-        raise ConcurrencySlotAcquisitionError(
-            f"Unable to acquire concurrency slots on {names!r}"
-        ) from response_or_exception
-
-    retval = _response_to_minimal_concurrency_limit_response(response_or_exception)
-
-    if strict and not retval:
-        raise ConcurrencySlotAcquisitionError(
-            f"Concurrency limits {names!r} must be created before acquiring slots"
-        )
-    elif not retval:
-        try:
-            logger = get_run_logger()
-            logger.warning(
-                f"Concurrency limits {names!r} do not exist - skipping acquisition."
-            )
-        except Exception:
-            pass
-    return retval
-
-
-async def _arelease_concurrency_slots(
-    names: List[str], slots: int, occupancy_seconds: float
-) -> List[MinimalConcurrencyLimitResponse]:
-    async with get_client() as client:
-        response = await client.release_concurrency_slots(
-            names=names, slots=slots, occupancy_seconds=occupancy_seconds
-        )
-        return _response_to_minimal_concurrency_limit_response(response)
-
-
-def _response_to_minimal_concurrency_limit_response(
-    response: httpx.Response,
-) -> List[MinimalConcurrencyLimitResponse]:
-    return [
-        MinimalConcurrencyLimitResponse.model_validate(obj_) for obj_ in response.json()
-    ]
+    emit_concurrency_acquisition_events(limits, occupy)

prefect/concurrency/context.py

@@ -1,19 +1,21 @@
 from contextvars import ContextVar
-from typing import List, Tuple
+from typing import Any, ClassVar
+
+from typing_extensions import Self
 
 from prefect.client.orchestration import get_client
 from prefect.context import ContextModel, Field
 
 
 class ConcurrencyContext(ContextModel):
-    __var__: ContextVar = ContextVar("concurrency")
+    __var__: ClassVar[ContextVar[Self]] = 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)
+    cleanup_slots: list[tuple[list[str], int, float]] = Field(default_factory=list)
 
-    def __exit__(self, *exc_info):
+    def __exit__(self, *exc_info: Any) -> None:
         if self.cleanup_slots:
             with get_client(sync_client=True) as client:
                 for names, occupy, occupancy_seconds in self.cleanup_slots:

prefect/concurrency/services.py

@@ -1,31 +1,30 @@
 import asyncio
-import concurrent.futures
+from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from typing import (
-    TYPE_CHECKING,
-    AsyncGenerator,
-    FrozenSet,
-    Optional,
-    Tuple,
-)
+from typing import TYPE_CHECKING, Optional
 
 import httpx
 from starlette import status
+from typing_extensions import TypeAlias, Unpack
 
 from prefect._internal.concurrency import logger
-from prefect._internal.concurrency.services import QueueService
+from prefect._internal.concurrency.services import FutureQueueService
 from prefect.client.orchestration import get_client
 from prefect.utilities.timeout import timeout_async
 
 if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient
 
+_Item: TypeAlias = tuple[int, str, Optional[float], Optional[bool], Optional[int]]
 
-class ConcurrencySlotAcquisitionService(QueueService):
-    def __init__(self, concurrency_limit_names: FrozenSet[str]):
+
+class ConcurrencySlotAcquisitionService(
+    FutureQueueService[Unpack[_Item], httpx.Response]
+):
+    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))
+        self._client: PrefectClient
+        self.concurrency_limit_names: list[str] = sorted(list(concurrency_limit_names))
 
     @asynccontextmanager
     async def _lifespan(self) -> AsyncGenerator[None, None]:
@@ -33,32 +32,7 @@ class ConcurrencySlotAcquisitionService(QueueService):
             self._client = client
             yield
 
-    async def _handle(
-        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, 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
-            # handle the exception and then re-raise.
-            future.set_result(exc)
-            raise exc
-        else:
-            future.set_result(response)
-
-    async def acquire_slots(
+    async def acquire(
         self,
         slots: int,
         mode: str,
@@ -69,44 +43,22 @@ class ConcurrencySlotAcquisitionService(QueueService):
         with timeout_async(seconds=timeout_seconds):
             while True:
                 try:
-                    response = await self._client.increment_concurrency_slots(
+                    return await self._client.increment_concurrency_slots(
                         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"])
-                        logger.debug(
-                            f"Unable to acquire concurrency slot. Retrying in {retry_after} second(s)."
-                        )
-                        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], Optional[bool], Optional[int]]
-    ) -> concurrent.futures.Future:
-        with self._lock:
-            if self._stopped:
-                raise RuntimeError("Cannot put items in a stopped service instance.")
+                except httpx.HTTPStatusError as exc:
+                    if not exc.response.status_code == status.HTTP_423_LOCKED:
+                        raise
 
-            logger.debug("Service %r enqueuing item %r", self, item)
-            future: concurrent.futures.Future = concurrent.futures.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
+                    if max_retries is not None and max_retries <= 0:
+                        raise exc
+                    retry_after = float(exc.response.headers["Retry-After"])
+                    logger.debug(
+                        f"Unable to acquire concurrency slot. Retrying in {retry_after} second(s)."
+                    )
+                    await asyncio.sleep(retry_after)
+                    if max_retries is not None:
+                        max_retries -= 1

prefect/concurrency/sync.py

@@ -1,71 +1,54 @@
+from collections.abc import Generator
 from contextlib import contextmanager
-from typing import (
-    Generator,
-    List,
-    Optional,
-    TypeVar,
-    Union,
-    cast,
-)
+from typing import Optional, TypeVar, Union
 
 import pendulum
 from typing_extensions import Literal
 
-from prefect.utilities.asyncutils import run_coro_as_sync
-
-try:
-    from pendulum import Interval
-except ImportError:
-    # pendulum < 3
-    from pendulum.period import Period as Interval  # type: ignore
-
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
+from prefect.utilities.asyncutils import run_coro_as_sync
 
-from .asyncio import (
-    _aacquire_concurrency_slots,
-    _arelease_concurrency_slots,
+from ._asyncio import (
+    aacquire_concurrency_slots,
+    arelease_concurrency_slots,
 )
-from .events import (
-    _emit_concurrency_acquisition_events,
-    _emit_concurrency_release_events,
+from ._events import (
+    emit_concurrency_acquisition_events,
+    emit_concurrency_release_events,
 )
 
 T = TypeVar("T")
 
 
 def _release_concurrency_slots(
-    names: List[str], slots: int, occupancy_seconds: float
-) -> List[MinimalConcurrencyLimitResponse]:
+    names: list[str], slots: int, occupancy_seconds: float
+) -> list[MinimalConcurrencyLimitResponse]:
     result = run_coro_as_sync(
-        _arelease_concurrency_slots(names, slots, occupancy_seconds)
+        arelease_concurrency_slots(names, slots, occupancy_seconds)
     )
-    if result is None:
-        raise RuntimeError("Failed to release concurrency slots")
     return result
 
 
 def _acquire_concurrency_slots(
-    names: List[str],
+    names: list[str],
     slots: int,
     mode: Literal["concurrency", "rate_limit"] = "concurrency",
     timeout_seconds: Optional[float] = None,
     create_if_missing: Optional[bool] = None,
     max_retries: Optional[int] = None,
     strict: bool = False,
-) -> List[MinimalConcurrencyLimitResponse]:
+) -> list[MinimalConcurrencyLimitResponse]:
     result = run_coro_as_sync(
-        _aacquire_concurrency_slots(
+        aacquire_concurrency_slots(
             names, slots, mode, timeout_seconds, create_if_missing, max_retries, strict
         )
     )
-    if result is None:
-        raise RuntimeError("Failed to acquire concurrency slots")
     return result
 
 
 @contextmanager
 def concurrency(
-    names: Union[str, List[str]],
+    names: Union[str, list[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
     max_retries: Optional[int] = None,
@@ -107,7 +90,7 @@ def concurrency(
 
     names = names if isinstance(names, list) else [names]
 
-    limits: List[MinimalConcurrencyLimitResponse] = _acquire_concurrency_slots(
+    limits: list[MinimalConcurrencyLimitResponse] = _acquire_concurrency_slots(
         names,
         occupy,
         timeout_seconds=timeout_seconds,
@@ -116,22 +99,18 @@ def concurrency(
         max_retries=max_retries,
     )
     acquisition_time = pendulum.now("UTC")
-    emitted_events = _emit_concurrency_acquisition_events(limits, occupy)
+    emitted_events = emit_concurrency_acquisition_events(limits, occupy)
 
     try:
         yield
     finally:
-        occupancy_period = cast(Interval, pendulum.now("UTC") - acquisition_time)
-        _release_concurrency_slots(
-            names,
-            occupy,
-            occupancy_period.total_seconds(),
-        )
-        _emit_concurrency_release_events(limits, occupy, emitted_events)
+        occupancy_period = pendulum.now("UTC") - acquisition_time
+        _release_concurrency_slots(names, occupy, occupancy_period.total_seconds())
+        emit_concurrency_release_events(limits, occupy, emitted_events)
 
 
 def rate_limit(
-    names: Union[str, List[str]],
+    names: Union[str, list[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
     create_if_missing: Optional[bool] = None,
@@ -166,4 +145,4 @@ def rate_limit(
         create_if_missing=create_if_missing,
         strict=strict,
     )
-    _emit_concurrency_acquisition_events(limits, occupy)
+    emit_concurrency_acquisition_events(limits, occupy)

prefect/concurrency/v1/_asyncio.py

@@ -0,0 +1,63 @@
+import asyncio
+from typing import Optional
+from uuid import UUID
+
+import httpx
+
+from prefect.client.orchestration import get_client
+from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
+from prefect.utilities.asyncutils import sync_compatible
+
+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."""
+
+
+@sync_compatible
+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))
+    try:
+        response = await asyncio.wrap_future(future)
+    except TimeoutError as timeout:
+        raise AcquireConcurrencySlotTimeoutError(
+            f"Attempt to acquire concurrency limits timed out after {timeout_seconds} second(s)"
+        ) from timeout
+    except Exception as exc:
+        raise ConcurrencySlotAcquisitionError(
+            f"Unable to acquire concurrency limits {names!r}"
+        ) from exc
+    else:
+        return _response_to_concurrency_limit_response(response)
+
+
+@sync_compatible
+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: list[MinimalConcurrencyLimitResponse] = response.json() or []
+    return [
+        MinimalConcurrencyLimitResponse.model_validate(limit) for limit in data if data
+    ]

prefect/concurrency/v1/{events.py → _events.py}

@@ -1,18 +1,18 @@
-from typing import Dict, List, Literal, Optional, Union
+from typing import 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(
+def emit_concurrency_event(
     phase: Union[Literal["acquired"], Literal["released"]],
     primary_limit: MinimalConcurrencyLimitResponse,
-    related_limits: List[MinimalConcurrencyLimitResponse],
+    related_limits: list[MinimalConcurrencyLimitResponse],
     task_run_id: UUID,
     follows: Union[Event, None] = None,
 ) -> Union[Event, None]:
-    resource: Dict[str, str] = {
+    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),
@@ -38,24 +38,22 @@ def _emit_concurrency_event(
     )
 
 
-def _emit_concurrency_acquisition_events(
-    limits: List[MinimalConcurrencyLimitResponse],
+def emit_concurrency_acquisition_events(
+    limits: list[MinimalConcurrencyLimitResponse],
     task_run_id: UUID,
-) -> Dict[UUID, Optional[Event]]:
-    events = {}
+) -> dict[UUID, Optional[Event]]:
+    events: dict[UUID, Optional[Event]] = {}
     for limit in limits:
-        event = _emit_concurrency_event("acquired", limit, limits, task_run_id)
+        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]],
+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]
-        )
+        emit_concurrency_event("released", limit, limits, task_run_id, events[limit.id])