prefect-client 3.1.9__py3-none-any.whl → 3.1.11__py3-none-any.whl

prefect/concurrency/v1/asyncio.py

@@ -1,42 +1,30 @@
-import asyncio
+from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator, List, Optional, Union, cast
+from typing import TYPE_CHECKING, Optional, Union
 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 prefect.utilities.asyncutils import sync_compatible
-
-from .context import ConcurrencyContext
-from .events import (
-    _emit_concurrency_acquisition_events,
-    _emit_concurrency_release_events,
+from prefect.concurrency.v1._asyncio import (
+    acquire_concurrency_slots,
+    release_concurrency_slots,
 )
-from .services import ConcurrencySlotAcquisitionService
-
-
-class ConcurrencySlotAcquisitionError(Exception):
-    """Raised when an unhandlable occurs while acquiring concurrency slots."""
-
+from prefect.concurrency.v1._events import (
+    emit_concurrency_acquisition_events,
+    emit_concurrency_release_events,
+)
+from prefect.concurrency.v1.context import ConcurrencyContext
 
-class AcquireConcurrencySlotTimeoutError(TimeoutError):
-    """Raised when acquiring a concurrency slot times out."""
+from ._asyncio import (
+    AcquireConcurrencySlotTimeoutError as AcquireConcurrencySlotTimeoutError,
+)
+from ._asyncio import ConcurrencySlotAcquisitionError as ConcurrencySlotAcquisitionError
 
 
 @asynccontextmanager
 async def concurrency(
-    names: Union[str, List[str]],
+    names: Union[str, list[str]],
     task_run_id: UUID,
     timeout_seconds: Optional[float] = None,
 ) -> AsyncGenerator[None, None]:
@@ -69,24 +57,30 @@ async def concurrency(
         yield
         return
 
-    names_normalized: List[str] = names if isinstance(names, list) else [names]
+    names_normalized: list[str] = names if isinstance(names, list) else [names]
 
-    limits = await _acquire_concurrency_slots(
+    acquire_slots = acquire_concurrency_slots(
         names_normalized,
         task_run_id=task_run_id,
         timeout_seconds=timeout_seconds,
     )
+    if TYPE_CHECKING:
+        assert not isinstance(acquire_slots, list)
+    limits = await acquire_slots
     acquisition_time = pendulum.now("UTC")
-    emitted_events = _emit_concurrency_acquisition_events(limits, task_run_id)
+    emitted_events = emit_concurrency_acquisition_events(limits, task_run_id)
 
     try:
         yield
     finally:
-        occupancy_period = cast(Interval, (pendulum.now("UTC") - acquisition_time))
+        occupancy_period = pendulum.now("UTC") - acquisition_time
         try:
-            await _release_concurrency_slots(
+            release_slots = release_concurrency_slots(
                 names_normalized, task_run_id, occupancy_period.total_seconds()
             )
+            if TYPE_CHECKING:
+                assert not isinstance(release_slots, list)
+            await release_slots
         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
@@ -96,51 +90,4 @@ async def concurrency(
                     (names_normalized, occupancy_period.total_seconds(), task_run_id)
                 )
 
-        _emit_concurrency_release_events(limits, emitted_events, task_run_id)
-
-
-@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))
-    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)
-
-
-@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 = response.json() or []
-    return [
-        MinimalConcurrencyLimitResponse.model_validate(limit) for limit in data if data
-    ]
+        emit_concurrency_release_events(limits, emitted_events, task_run_id)

prefect/concurrency/v1/context.py

@@ -1,20 +1,22 @@
 from contextvars import ContextVar
-from typing import List, Tuple
+from typing import Any, ClassVar
 from uuid import UUID
 
+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_v1")
+    __var__: ClassVar[ContextVar[Self]] = 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)
+    cleanup_slots: list[tuple[list[str], float, UUID]] = 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, occupancy_seconds, task_run_id in self.cleanup_slots:

prefect/concurrency/v1/services.py

@@ -1,21 +1,16 @@
 import asyncio
-import concurrent.futures
+from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
 from json import JSONDecodeError
-from typing import (
-    TYPE_CHECKING,
-    AsyncGenerator,
-    FrozenSet,
-    Optional,
-    Tuple,
-)
+from typing import TYPE_CHECKING, Optional
 from uuid import UUID
 
 import httpx
 from starlette import status
+from typing_extensions import 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
 
@@ -27,11 +22,13 @@ class ConcurrencySlotAcquisitionServiceError(Exception):
     """Raised when an error occurs while acquiring concurrency slots."""
 
 
-class ConcurrencySlotAcquisitionService(QueueService):
-    def __init__(self, concurrency_limit_names: FrozenSet[str]):
+class ConcurrencySlotAcquisitionService(
+    FutureQueueService[Unpack[tuple[UUID, Optional[float]]], httpx.Response]
+):
+    def __init__(self, concurrency_limit_names: frozenset[str]) -> None:
         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]:
@@ -39,78 +36,35 @@ class ConcurrencySlotAcquisitionService(QueueService):
             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,
+    async def acquire(
+        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(
+                    return 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
+                except httpx.HTTPStatusError as exc:
+                    if not exc.response.status_code == status.HTTP_423_LOCKED:
+                        raise
 
+                    retry_after = exc.response.headers.get("Retry-After")
+                    if retry_after:
+                        retry_after = float(retry_after)
+                        await asyncio.sleep(retry_after)
                     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
+                        # 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

prefect/concurrency/v1/sync.py

@@ -1,31 +1,15 @@
+import asyncio
+from collections.abc import Generator
 from contextlib import contextmanager
-from typing import (
-    Generator,
-    List,
-    Optional,
-    TypeVar,
-    Union,
-    cast,
-)
+from typing import Optional, TypeVar, Union
 from uuid import UUID
 
 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 .asyncio import (
-    _acquire_concurrency_slots,
-    _release_concurrency_slots,
-)
-from .events import (
-    _emit_concurrency_acquisition_events,
-    _emit_concurrency_release_events,
+from ._asyncio import acquire_concurrency_slots, release_concurrency_slots
+from ._events import (
+    emit_concurrency_acquisition_events,
+    emit_concurrency_release_events,
 )
 
 T = TypeVar("T")
@@ -33,7 +17,7 @@ T = TypeVar("T")
 
 @contextmanager
 def concurrency(
-    names: Union[str, List[str]],
+    names: Union[str, list[str]],
     task_run_id: UUID,
     timeout_seconds: Optional[float] = None,
 ) -> Generator[None, None, None]:
@@ -69,23 +53,20 @@ def concurrency(
 
     names = names if isinstance(names, list) else [names]
 
-    limits: List[MinimalConcurrencyLimitResponse] = _acquire_concurrency_slots(
-        names,
-        timeout_seconds=timeout_seconds,
-        task_run_id=task_run_id,
-        _sync=True,
+    force = {"_sync": True}
+    result = acquire_concurrency_slots(
+        names, timeout_seconds=timeout_seconds, task_run_id=task_run_id, **force
     )
+    assert not asyncio.iscoroutine(result)
+    limits = result
     acquisition_time = pendulum.now("UTC")
-    emitted_events = _emit_concurrency_acquisition_events(limits, task_run_id)
+    emitted_events = emit_concurrency_acquisition_events(limits, task_run_id)
 
     try:
         yield
     finally:
-        occupancy_period = cast(Interval, pendulum.now("UTC") - acquisition_time)
-        _release_concurrency_slots(
-            names,
-            task_run_id,
-            occupancy_period.total_seconds(),
-            _sync=True,
+        occupancy_period = pendulum.now("UTC") - acquisition_time
+        release_concurrency_slots(
+            names, task_run_id, occupancy_period.total_seconds(), **force
         )
-        _emit_concurrency_release_events(limits, emitted_events, task_run_id)
+        emit_concurrency_release_events(limits, emitted_events, task_run_id)