prefect-client 3.0.0rc18__py3-none-any.whl → 3.0.0rc20__py3-none-any.whl

prefect/concurrency/services.py

@@ -36,13 +36,18 @@ class ConcurrencySlotAcquisitionService(QueueService):
     async def _handle(
         self,
         item: Tuple[
-            int, str, Optional[float], concurrent.futures.Future, Optional[bool]
+            int,
+            str,
+            Optional[float],
+            concurrent.futures.Future,
+            Optional[bool],
+            Optional[int],
         ],
     ) -> None:
-        occupy, mode, timeout_seconds, future, create_if_missing = item
+        occupy, mode, timeout_seconds, future, create_if_missing, max_retries = item
         try:
             response = await self.acquire_slots(
-                occupy, mode, timeout_seconds, create_if_missing
+                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
@@ -59,6 +64,7 @@ class ConcurrencySlotAcquisitionService(QueueService):
         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:
@@ -74,15 +80,19 @@ class ConcurrencySlotAcquisitionService(QueueService):
                         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], Optional[bool]]
+        self, item: Tuple[int, str, Optional[float], Optional[bool], Optional[int]]
     ) -> concurrent.futures.Future:
         with self._lock:
             if self._stopped:
@@ -91,9 +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, create_if_missing = item
+            occupy, mode, timeout_seconds, create_if_missing, max_retries = item
             self._queue.put_nowait(
-                (occupy, mode, timeout_seconds, future, create_if_missing)
+                (occupy, mode, timeout_seconds, future, create_if_missing, max_retries)
             )
 
         return future

prefect/concurrency/sync.py

@@ -40,7 +40,8 @@ def concurrency(
     names: Union[str, List[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
-    create_if_missing: Optional[bool] = True,
+    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.
@@ -51,6 +52,7 @@ def concurrency(
         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.
@@ -80,6 +82,7 @@ def concurrency(
         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)

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)

prefect/context.py

@@ -131,7 +131,7 @@ class ContextModel(BaseModel):
         extra="forbid",
     )
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         if self._token is not None:
             raise RuntimeError(
                 "Context already entered. Context enter calls cannot be nested."
@@ -290,7 +290,7 @@ class AsyncClientContext(ContextModel):
         if ctx:
             yield ctx
         else:
-            with cls() as ctx:
+            async with cls() as ctx:
                 yield ctx
 
 

prefect/deployments/flow_runs.py

@@ -5,7 +5,6 @@ from uuid import UUID
 import anyio
 import pendulum
 
-from prefect._internal.compatibility.deprecated import deprecated_parameter
 from prefect.client.schemas import FlowRun
 from prefect.client.utilities import inject_client
 from prefect.context import FlowRunContext, TaskRunContext
@@ -24,11 +23,6 @@ logger = get_logger(__name__)
 
 
 @sync_compatible
-@deprecated_parameter(
-    "infra_overrides",
-    start_date="Apr 2024",
-    help="Use `job_variables` instead.",
-)
 @inject_client
 async def run_deployment(
     name: Union[str, UUID],
@@ -42,7 +36,6 @@ async def run_deployment(
     idempotency_key: Optional[str] = None,
     work_queue_name: Optional[str] = None,
     as_subflow: Optional[bool] = True,
-    infra_overrides: Optional[dict] = None,
     job_variables: Optional[dict] = None,
 ) -> "FlowRun":
     """

prefect/deployments/runner.py

@@ -143,6 +143,10 @@ class RunnerDeployment(BaseModel):
         default=None,
         description="The schedules that should cause this deployment to run.",
     )
+    concurrency_limit: Optional[int] = Field(
+        default=None,
+        description="The maximum number of concurrent runs of this deployment.",
+    )
     paused: Optional[bool] = Field(
         default=None, description="Whether or not the deployment is paused."
     )
@@ -274,6 +278,7 @@ class RunnerDeployment(BaseModel):
                 version=self.version,
                 paused=self.paused,
                 schedules=self.schedules,
+                concurrency_limit=self.concurrency_limit,
                 parameters=self.parameters,
                 description=self.description,
                 tags=self.tags,
@@ -432,6 +437,7 @@ class RunnerDeployment(BaseModel):
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
         schedules: Optional["FlexibleScheduleList"] = None,
+        concurrency_limit: Optional[int] = None,
         parameters: Optional[dict] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
@@ -485,6 +491,7 @@ class RunnerDeployment(BaseModel):
             name=Path(name).stem,
             flow_name=flow.name,
             schedules=constructed_schedules,
+            concurrency_limit=concurrency_limit,
             paused=paused,
             tags=tags or [],
             triggers=triggers or [],
@@ -558,6 +565,7 @@ class RunnerDeployment(BaseModel):
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
         schedules: Optional["FlexibleScheduleList"] = None,
+        concurrency_limit: Optional[int] = None,
         parameters: Optional[dict] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
@@ -614,6 +622,7 @@ class RunnerDeployment(BaseModel):
             name=Path(name).stem,
             flow_name=flow.name,
             schedules=constructed_schedules,
+            concurrency_limit=concurrency_limit,
             paused=paused,
             tags=tags or [],
             triggers=triggers or [],
@@ -646,6 +655,7 @@ class RunnerDeployment(BaseModel):
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
         schedules: Optional["FlexibleScheduleList"] = None,
+        concurrency_limit: Optional[int] = None,
         parameters: Optional[dict] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
@@ -710,6 +720,7 @@ class RunnerDeployment(BaseModel):
             name=Path(name).stem,
             flow_name=flow.name,
             schedules=constructed_schedules,
+            concurrency_limit=concurrency_limit,
             paused=paused,
             tags=tags or [],
             triggers=triggers or [],

prefect/events/clients.py

@@ -346,6 +346,47 @@ class PrefectEventsClient(EventsClient):
                     await asyncio.sleep(1)
 
 
+class AssertingPassthroughEventsClient(PrefectEventsClient):
+    """A Prefect Events client that BOTH records all events sent to it for inspection
+    during tests AND sends them to a Prefect server."""
+
+    last: ClassVar["Optional[AssertingPassthroughEventsClient]"] = None
+    all: ClassVar[List["AssertingPassthroughEventsClient"]] = []
+
+    args: Tuple
+    kwargs: Dict[str, Any]
+    events: List[Event]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        AssertingPassthroughEventsClient.last = self
+        AssertingPassthroughEventsClient.all.append(self)
+        self.args = args
+        self.kwargs = kwargs
+
+    @classmethod
+    def reset(cls) -> None:
+        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:
+        # actually send the event to the server
+        await super()._emit(event)
+
+        # record the event for inspection
+        self.events.append(event)
+
+    async def __aenter__(self) -> Self:
+        await super().__aenter__()
+        self.events = []
+        return self
+
+
 class PrefectCloudEventsClient(PrefectEventsClient):
     """A Prefect Events client that streams events to a Prefect Cloud Workspace"""