prefect-client 3.0.0rc17__py3-none-any.whl → 3.0.0rc19__py3-none-any.whl

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"""
 

prefect/events/related.py

@@ -21,6 +21,7 @@ from .schemas.events import RelatedResource
 
 if TYPE_CHECKING:
     from prefect._internal.schemas.bases import ObjectBaseModel
+    from prefect.client.orchestration import PrefectClient
 
 ResourceCacheEntry = Dict[str, Union[str, "ObjectBaseModel", None]]
 RelatedResourceCache = Dict[str, Tuple[ResourceCacheEntry, DateTime]]
@@ -54,9 +55,9 @@ def object_as_related_resource(kind: str, role: str, object: Any) -> RelatedReso
 
 
 async def related_resources_from_run_context(
+    client: "PrefectClient",
     exclude: Optional[Set[str]] = None,
 ) -> List[RelatedResource]:
-    from prefect.client.orchestration import get_client
     from prefect.client.schemas.objects import FlowRun
     from prefect.context import FlowRunContext, TaskRunContext
 
@@ -77,86 +78,84 @@ async def related_resources_from_run_context(
 
     related_objects: List[ResourceCacheEntry] = []
 
-    async with get_client() as client:
+    async def dummy_read():
+        return {}
 
-        async def dummy_read():
-            return {}
-
-        if flow_run_context:
-            related_objects.append(
-                {
-                    "kind": "flow-run",
-                    "role": "flow-run",
-                    "object": flow_run_context.flow_run,
-                },
-            )
-        else:
-            related_objects.append(
-                await _get_and_cache_related_object(
-                    kind="flow-run",
-                    role="flow-run",
-                    client_method=client.read_flow_run,
-                    obj_id=flow_run_id,
-                    cache=RESOURCE_CACHE,
-                )
+    if flow_run_context:
+        related_objects.append(
+            {
+                "kind": "flow-run",
+                "role": "flow-run",
+                "object": flow_run_context.flow_run,
+            },
+        )
+    else:
+        related_objects.append(
+            await _get_and_cache_related_object(
+                kind="flow-run",
+                role="flow-run",
+                client_method=client.read_flow_run,
+                obj_id=flow_run_id,
+                cache=RESOURCE_CACHE,
             )
+        )
 
-        if task_run_context:
-            related_objects.append(
-                {
-                    "kind": "task-run",
-                    "role": "task-run",
-                    "object": task_run_context.task_run,
-                },
-            )
+    if task_run_context:
+        related_objects.append(
+            {
+                "kind": "task-run",
+                "role": "task-run",
+                "object": task_run_context.task_run,
+            },
+        )
 
-        flow_run = related_objects[0]["object"]
+    flow_run = related_objects[0]["object"]
 
-        if isinstance(flow_run, FlowRun):
-            related_objects += list(
-                await asyncio.gather(
+    if isinstance(flow_run, FlowRun):
+        related_objects += list(
+            await asyncio.gather(
+                _get_and_cache_related_object(
+                    kind="flow",
+                    role="flow",
+                    client_method=client.read_flow,
+                    obj_id=flow_run.flow_id,
+                    cache=RESOURCE_CACHE,
+                ),
+                (
+                    _get_and_cache_related_object(
+                        kind="deployment",
+                        role="deployment",
+                        client_method=client.read_deployment,
+                        obj_id=flow_run.deployment_id,
+                        cache=RESOURCE_CACHE,
+                    )
+                    if flow_run.deployment_id
+                    else dummy_read()
+                ),
+                (
+                    _get_and_cache_related_object(
+                        kind="work-queue",
+                        role="work-queue",
+                        client_method=client.read_work_queue,
+                        obj_id=flow_run.work_queue_id,
+                        cache=RESOURCE_CACHE,
+                    )
+                    if flow_run.work_queue_id
+                    else dummy_read()
+                ),
+                (
                     _get_and_cache_related_object(
-                        kind="flow",
-                        role="flow",
-                        client_method=client.read_flow,
-                        obj_id=flow_run.flow_id,
+                        kind="work-pool",
+                        role="work-pool",
+                        client_method=client.read_work_pool,
+                        obj_id=flow_run.work_pool_name,
                         cache=RESOURCE_CACHE,
-                    ),
-                    (
-                        _get_and_cache_related_object(
-                            kind="deployment",
-                            role="deployment",
-                            client_method=client.read_deployment,
-                            obj_id=flow_run.deployment_id,
-                            cache=RESOURCE_CACHE,
-                        )
-                        if flow_run.deployment_id
-                        else dummy_read()
-                    ),
-                    (
-                        _get_and_cache_related_object(
-                            kind="work-queue",
-                            role="work-queue",
-                            client_method=client.read_work_queue,
-                            obj_id=flow_run.work_queue_id,
-                            cache=RESOURCE_CACHE,
-                        )
-                        if flow_run.work_queue_id
-                        else dummy_read()
-                    ),
-                    (
-                        _get_and_cache_related_object(
-                            kind="work-pool",
-                            role="work-pool",
-                            client_method=client.read_work_pool,
-                            obj_id=flow_run.work_pool_name,
-                            cache=RESOURCE_CACHE,
-                        )
-                        if flow_run.work_pool_name
-                        else dummy_read()
-                    ),
-                )
+                    )
+                    if flow_run.work_pool_name
+                    else dummy_read()
+                ),
             )
+        )
 
     related = []
     tags = set()

prefect/events/utilities.py

@@ -7,6 +7,7 @@ from pydantic_extra_types.pendulum_dt import DateTime
 
 from .clients import (
     AssertingEventsClient,
+    AssertingPassthroughEventsClient,
     PrefectCloudEventsClient,
     PrefectEventsClient,
 )
@@ -49,6 +50,7 @@ def emit_event(
         return None
 
     operational_clients = [
+        AssertingPassthroughEventsClient,
         AssertingEventsClient,
         PrefectCloudEventsClient,
         PrefectEventsClient,

prefect/events/worker.py

@@ -1,6 +1,6 @@
 from contextlib import asynccontextmanager
 from contextvars import Context, copy_context
-from typing import Any, Dict, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Type
 from uuid import UUID
 
 from typing_extensions import Self
@@ -22,6 +22,9 @@ from .clients import (
 from .related import related_resources_from_run_context
 from .schemas.events import Event
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
 
 def should_emit_events() -> bool:
     return (
@@ -55,14 +58,18 @@ class EventsWorker(QueueService[Event]):
         self.client_type = client_type
         self.client_options = client_options
         self._client: EventsClient
+        self._orchestration_client: "PrefectClient"
         self._context_cache: Dict[UUID, Context] = {}
 
     @asynccontextmanager
     async def _lifespan(self):
         self._client = self.client_type(**{k: v for k, v in self.client_options})
+        from prefect.client.orchestration import get_client
 
+        self._orchestration_client = get_client()
         async with self._client:
-            yield
+            async with self._orchestration_client:
+                yield
 
     def _prepare_item(self, event: Event) -> Event:
         self._context_cache[event.id] = copy_context()
@@ -77,7 +84,9 @@ class EventsWorker(QueueService[Event]):
 
     async def attach_related_resources_from_context(self, event: Event):
         exclude = {resource.id for resource in event.involved_resources}
-        event.related += await related_resources_from_run_context(exclude=exclude)
+        event.related += await related_resources_from_run_context(
+            client=self._orchestration_client, exclude=exclude
+        )
 
     @classmethod
     def instance(

prefect/flow_engine.py

@@ -30,6 +30,7 @@ from prefect.client.schemas import FlowRun, TaskRun
 from prefect.client.schemas.filters import FlowRunFilter
 from prefect.client.schemas.sorting import FlowRunSort
 from prefect.concurrency.context import ConcurrencyContext
+from prefect.concurrency.v1.context import ConcurrencyContext as ConcurrencyContextV1
 from prefect.context import FlowRunContext, SyncClientContext, TagsContext
 from prefect.exceptions import (
     Abort,
@@ -506,6 +507,7 @@ class FlowRunEngine(Generic[P, R]):
                     task_runner=task_runner,
                 )
             )
+            stack.enter_context(ConcurrencyContextV1())
             stack.enter_context(ConcurrencyContext())
 
             # set the logger to the flow run logger

prefect/flows.py

@@ -643,6 +643,7 @@ class Flow(Generic[P, R]):
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
         schedules: Optional[List["FlexibleScheduleList"]] = None,
+        concurrency_limit: Optional[int] = None,
         parameters: Optional[dict] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
@@ -666,6 +667,7 @@ class Flow(Generic[P, R]):
             paused: Whether or not to set this deployment as paused.
             schedules: A list of schedule objects defining when to execute runs of this deployment.
                 Used to define multiple schedules or additional scheduling options such as `timezone`.
+            concurrency_limit: The maximum number of runs of this deployment that can run at the same time.
             parameters: A dictionary of default parameter values to pass to runs of this deployment.
             triggers: A list of triggers that will kick off runs of this deployment.
             description: A description for the created deployment. Defaults to the flow's
@@ -718,6 +720,7 @@ class Flow(Generic[P, R]):
                 rrule=rrule,
                 paused=paused,
                 schedules=schedules,
+                concurrency_limit=concurrency_limit,
                 tags=tags,
                 triggers=triggers,
                 parameters=parameters or {},
@@ -737,6 +740,7 @@ class Flow(Generic[P, R]):
                 rrule=rrule,
                 paused=paused,
                 schedules=schedules,
+                concurrency_limit=concurrency_limit,
                 tags=tags,
                 triggers=triggers,
                 parameters=parameters or {},
@@ -1055,6 +1059,7 @@ class Flow(Generic[P, R]):
         rrule: Optional[str] = None,
         paused: Optional[bool] = None,
         schedules: Optional[List[DeploymentScheduleCreate]] = None,
+        concurrency_limit: Optional[int] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         parameters: Optional[dict] = None,
         description: Optional[str] = None,
@@ -1101,6 +1106,7 @@ class Flow(Generic[P, R]):
             paused: Whether or not to set this deployment as paused.
             schedules: A list of schedule objects defining when to execute runs of this deployment.
                 Used to define multiple schedules or additional scheduling options like `timezone`.
+            concurrency_limit: The maximum number of runs that can be executed concurrently.
             parameters: A dictionary of default parameter values to pass to runs of this deployment.
             description: A description for the created deployment. Defaults to the flow's
                 description if not provided.
@@ -1175,6 +1181,7 @@ class Flow(Generic[P, R]):
             cron=cron,
             rrule=rrule,
             schedules=schedules,
+            concurrency_limit=concurrency_limit,
             paused=paused,
             triggers=triggers,
             parameters=parameters,

prefect/records/__init__.py

@@ -1 +1 @@
-from .store import RecordStore
+from .base import RecordStore