prefect-client 3.0.1__py3-none-any.whl → 3.0.3__py3-none-any.whl

prefect/concurrency/asyncio.py

@@ -6,6 +6,8 @@ import anyio
 import httpx
 import pendulum
 
+from prefect._internal.compatibility.deprecated import deprecated_parameter
+
 try:
     from pendulum import Interval
 except ImportError:
@@ -14,6 +16,8 @@ except ImportError:
 
 from prefect.client.orchestration import get_client
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
+from prefect.logging.loggers import get_run_logger
+from prefect.utilities.asyncutils import sync_compatible
 
 from .context import ConcurrencyContext
 from .events import (
@@ -36,8 +40,9 @@ async 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,
+    create_if_missing: Optional[bool] = None,
+    strict: bool = False,
 ) -> AsyncGenerator[None, None]:
     """A context manager that acquires and releases concurrency slots from the
     given concurrency limits.
@@ -47,11 +52,13 @@ async 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.
+        strict: A boolean specifying whether to raise an error if the concurrency limit does not exist.
+            Defaults to `False`.
 
     Raises:
         TimeoutError: If the slots are not acquired within the given timeout.
+        ConcurrencySlotAcquisitionError: If the concurrency limit does not exist and `strict` is `True`.
 
     Example:
     A simple example of using the async `concurrency` context manager:
@@ -78,6 +85,7 @@ async def concurrency(
         timeout_seconds=timeout_seconds,
         create_if_missing=create_if_missing,
         max_retries=max_retries,
+        strict=strict,
     )
     acquisition_time = pendulum.now("UTC")
     emitted_events = _emit_concurrency_acquisition_events(limits, occupy)
@@ -106,7 +114,8 @@ async def rate_limit(
     names: Union[str, List[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
-    create_if_missing: Optional[bool] = True,
+    create_if_missing: Optional[bool] = None,
+    strict: bool = False,
 ) -> None:
     """Block execution until an `occupy` number of slots of the concurrency
     limits given in `names` are acquired. Requires that all given concurrency
@@ -117,7 +126,12 @@ async def rate_limit(
         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.
+        strict: A boolean specifying whether to raise an error if the concurrency limit does not exist.
+            Defaults to `False`.
+
+    Raises:
+        TimeoutError: If the slots are not acquired within the given timeout.
+        ConcurrencySlotAcquisitionError: If the concurrency limit does not exist and `strict` is `True`.
     """
     if not names:
         return
@@ -130,17 +144,27 @@ async def rate_limit(
         mode="rate_limit",
         timeout_seconds=timeout_seconds,
         create_if_missing=create_if_missing,
+        strict=strict,
     )
     _emit_concurrency_acquisition_events(limits, occupy)
 
 
+@sync_compatible
+@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 _acquire_concurrency_slots(
     names: List[str],
     slots: int,
     mode: Union[Literal["concurrency"], Literal["rate_limit"]] = "concurrency",
     timeout_seconds: Optional[float] = None,
-    create_if_missing: Optional[bool] = True,
+    create_if_missing: Optional[bool] = None,
     max_retries: Optional[int] = None,
+    strict: bool = False,
 ) -> List[MinimalConcurrencyLimitResponse]:
     service = ConcurrencySlotAcquisitionService.instance(frozenset(names))
     future = service.send(
@@ -158,9 +182,24 @@ async def _acquire_concurrency_slots(
             f"Unable to acquire concurrency slots on {names!r}"
         ) from response_or_exception
 
-    return _response_to_minimal_concurrency_limit_response(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
 
 
+@sync_compatible
 async def _release_concurrency_slots(
     names: List[str], slots: int, occupancy_seconds: float
 ) -> List[MinimalConcurrencyLimitResponse]:

prefect/concurrency/services.py

@@ -63,7 +63,7 @@ class ConcurrencySlotAcquisitionService(QueueService):
         slots: int,
         mode: str,
         timeout_seconds: Optional[float] = None,
-        create_if_missing: Optional[bool] = False,
+        create_if_missing: Optional[bool] = None,
         max_retries: Optional[int] = None,
     ) -> httpx.Response:
         with timeout_async(seconds=timeout_seconds):

prefect/concurrency/sync.py

@@ -1,8 +1,5 @@
 from contextlib import contextmanager
 from typing import (
-    Any,
-    Awaitable,
-    Callable,
     Generator,
     List,
     Optional,
@@ -19,8 +16,6 @@ except ImportError:
     # pendulum < 3
     from pendulum.period import Period as Interval  # type: ignore
 
-from prefect._internal.concurrency.api import create_call, from_sync
-from prefect._internal.concurrency.event_loop import get_running_loop
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
 
 from .asyncio import (
@@ -40,8 +35,9 @@ 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,
+    strict: bool = False,
+    create_if_missing: Optional[bool] = None,
 ) -> Generator[None, None, None]:
     """A context manager that acquires and releases concurrency slots from the
     given concurrency limits.
@@ -51,11 +47,13 @@ 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.
+        strict: A boolean specifying whether to raise an error if the concurrency limit does not exist.
+            Defaults to `False`.
 
     Raises:
         TimeoutError: If the slots are not acquired within the given timeout.
+        ConcurrencySlotAcquisitionError: If the concurrency limit does not exist and `strict` is `True`.
 
     Example:
     A simple example of using the sync `concurrency` context manager:
@@ -76,13 +74,14 @@ def concurrency(
 
     names = names if isinstance(names, list) else [names]
 
-    limits: List[MinimalConcurrencyLimitResponse] = _call_async_function_from_sync(
-        _acquire_concurrency_slots,
+    limits: List[MinimalConcurrencyLimitResponse] = _acquire_concurrency_slots(
         names,
         occupy,
         timeout_seconds=timeout_seconds,
         create_if_missing=create_if_missing,
+        strict=strict,
         max_retries=max_retries,
+        _sync=True,
     )
     acquisition_time = pendulum.now("UTC")
     emitted_events = _emit_concurrency_acquisition_events(limits, occupy)
@@ -91,11 +90,11 @@ def concurrency(
         yield
     finally:
         occupancy_period = cast(Interval, pendulum.now("UTC") - acquisition_time)
-        _call_async_function_from_sync(
-            _release_concurrency_slots,
+        _release_concurrency_slots(
             names,
             occupy,
             occupancy_period.total_seconds(),
+            _sync=True,
         )
         _emit_concurrency_release_events(limits, occupy, emitted_events)
 
@@ -104,7 +103,8 @@ def rate_limit(
     names: Union[str, List[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
-    create_if_missing: Optional[bool] = True,
+    create_if_missing: Optional[bool] = None,
+    strict: bool = False,
 ) -> None:
     """Block execution until an `occupy` number of slots of the concurrency
     limits given in `names` are acquired. Requires that all given concurrency
@@ -115,31 +115,25 @@ def rate_limit(
         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.
+        strict: A boolean specifying whether to raise an error if the concurrency limit does not exist.
+            Defaults to `False`.
+
+    Raises:
+        TimeoutError: If the slots are not acquired within the given timeout.
+        ConcurrencySlotAcquisitionError: If the concurrency limit does not exist and `strict` is `True`.
     """
     if not names:
         return
 
     names = names if isinstance(names, list) else [names]
 
-    limits = _call_async_function_from_sync(
-        _acquire_concurrency_slots,
+    limits = _acquire_concurrency_slots(
         names,
         occupy,
         mode="rate_limit",
         timeout_seconds=timeout_seconds,
         create_if_missing=create_if_missing,
+        strict=strict,
+        _sync=True,
     )
     _emit_concurrency_acquisition_events(limits, occupy)
-
-
-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()  # type: ignore [return-value]

prefect/concurrency/v1/asyncio.py

@@ -16,6 +16,7 @@ except ImportError:
     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 (
@@ -98,6 +99,7 @@ async def concurrency(
         _emit_concurrency_release_events(limits, emitted_events, task_run_id)
 
 
+@sync_compatible
 async def _acquire_concurrency_slots(
     names: List[str],
     task_run_id: UUID,
@@ -120,6 +122,7 @@ async def _acquire_concurrency_slots(
     return _response_to_concurrency_limit_response(response_or_exception)
 
 
+@sync_compatible
 async def _release_concurrency_slots(
     names: List[str],
     task_run_id: UUID,

prefect/concurrency/v1/sync.py

@@ -12,7 +12,6 @@ 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
@@ -70,11 +69,11 @@ def concurrency(
 
     names = names if isinstance(names, list) else [names]
 
-    limits: List[MinimalConcurrencyLimitResponse] = _call_async_function_from_sync(
-        _acquire_concurrency_slots,
+    limits: List[MinimalConcurrencyLimitResponse] = _acquire_concurrency_slots(
         names,
         timeout_seconds=timeout_seconds,
         task_run_id=task_run_id,
+        _sync=True,
     )
     acquisition_time = pendulum.now("UTC")
     emitted_events = _emit_concurrency_acquisition_events(limits, task_run_id)
@@ -83,10 +82,10 @@ def concurrency(
         yield
     finally:
         occupancy_period = cast(Interval, pendulum.now("UTC") - acquisition_time)
-        _call_async_function_from_sync(
-            _release_concurrency_slots,
+        _release_concurrency_slots(
             names,
             task_run_id,
             occupancy_period.total_seconds(),
+            _sync=True,
         )
         _emit_concurrency_release_events(limits, emitted_events, task_run_id)

prefect/context.py

@@ -9,7 +9,6 @@ For more user-accessible information about the current run, see [`prefect.runtim
 import os
 import sys
 import warnings
-import weakref
 from contextlib import ExitStack, asynccontextmanager, contextmanager
 from contextvars import ContextVar, Token
 from pathlib import Path
@@ -40,7 +39,7 @@ from prefect.client.orchestration import PrefectClient, SyncPrefectClient, get_c
 from prefect.client.schemas import FlowRun, TaskRun
 from prefect.events.worker import EventsWorker
 from prefect.exceptions import MissingContextError
-from prefect.results import ResultStore
+from prefect.results import ResultStore, get_default_persist_setting
 from prefect.settings import PREFECT_HOME, Profile, Settings
 from prefect.states import State
 from prefect.task_runners import TaskRunner
@@ -343,6 +342,7 @@ class EngineContext(RunContext):
 
     # Result handling
     result_store: ResultStore
+    persist_result: bool = Field(default_factory=get_default_persist_setting)
 
     # Counter for task calls allowing unique
     task_run_dynamic_keys: Dict[str, int] = Field(default_factory=dict)
@@ -352,10 +352,7 @@ class EngineContext(RunContext):
 
     # Tracking for result from task runs in this flow run for dependency tracking
     # Holds the ID of the object returned by the task run and task run state
-    # This is a weakref dictionary to avoid undermining garbage collection
-    task_run_results: Mapping[int, State] = Field(
-        default_factory=weakref.WeakValueDictionary
-    )
+    task_run_results: Mapping[int, State] = Field(default_factory=dict)
 
     # Events worker to emit events
     events: Optional[EventsWorker] = None
@@ -372,6 +369,7 @@ class EngineContext(RunContext):
                 "start_time",
                 "input_keyset",
                 "result_store",
+                "persist_result",
             },
             exclude_unset=True,
         )
@@ -397,6 +395,7 @@ class TaskRunContext(RunContext):
 
     # Result handling
     result_store: ResultStore
+    persist_result: bool = Field(default_factory=get_default_persist_setting)
 
     __var__ = ContextVar("task_run")
 
@@ -410,6 +409,7 @@ class TaskRunContext(RunContext):
                 "start_time",
                 "input_keyset",
                 "result_store",
+                "persist_result",
             },
             exclude_unset=True,
         )

prefect/deployments/runner.py

@@ -54,6 +54,7 @@ from prefect._internal.schemas.validators import (
 )
 from prefect.client.orchestration import get_client
 from prefect.client.schemas.actions import DeploymentScheduleCreate
+from prefect.client.schemas.objects import ConcurrencyLimitConfig, ConcurrencyOptions
 from prefect.client.schemas.schedules import (
     SCHEDULE_TYPES,
     construct_schedule,
@@ -147,6 +148,10 @@ class RunnerDeployment(BaseModel):
         default=None,
         description="The maximum number of concurrent runs of this deployment.",
     )
+    concurrency_options: Optional[ConcurrencyOptions] = Field(
+        default=None,
+        description="The concurrency limit config for the deployment.",
+    )
     paused: Optional[bool] = Field(
         default=None, description="Whether or not the deployment is paused."
     )
@@ -279,6 +284,7 @@ class RunnerDeployment(BaseModel):
                 paused=self.paused,
                 schedules=self.schedules,
                 concurrency_limit=self.concurrency_limit,
+                concurrency_options=self.concurrency_options,
                 parameters=self.parameters,
                 description=self.description,
                 tags=self.tags,
@@ -437,7 +443,7 @@ class RunnerDeployment(BaseModel):
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
         schedules: Optional["FlexibleScheduleList"] = None,
-        concurrency_limit: Optional[int] = None,
+        concurrency_limit: Optional[Union[int, ConcurrencyLimitConfig, None]] = None,
         parameters: Optional[dict] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
@@ -462,6 +468,7 @@ class RunnerDeployment(BaseModel):
             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 concurrent runs this deployment will allow.
             triggers: A list of triggers that should kick of a run of this flow.
             parameters: A dictionary of default parameter values to pass to runs of this flow.
             description: A description for the created deployment. Defaults to the flow's
@@ -487,11 +494,20 @@ class RunnerDeployment(BaseModel):
 
         job_variables = job_variables or {}
 
+        if isinstance(concurrency_limit, ConcurrencyLimitConfig):
+            concurrency_options = {
+                "collision_strategy": concurrency_limit.collision_strategy
+            }
+            concurrency_limit = concurrency_limit.limit
+        else:
+            concurrency_options = None
+
         deployment = cls(
             name=Path(name).stem,
             flow_name=flow.name,
             schedules=constructed_schedules,
             concurrency_limit=concurrency_limit,
+            concurrency_options=concurrency_options,
             paused=paused,
             tags=tags or [],
             triggers=triggers or [],
@@ -558,6 +574,7 @@ class RunnerDeployment(BaseModel):
         cls,
         entrypoint: str,
         name: str,
+        flow_name: Optional[str] = None,
         interval: Optional[
             Union[Iterable[Union[int, float, timedelta]], int, float, timedelta]
         ] = None,
@@ -565,7 +582,7 @@ class RunnerDeployment(BaseModel):
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
         schedules: Optional["FlexibleScheduleList"] = None,
-        concurrency_limit: Optional[int] = None,
+        concurrency_limit: Optional[Union[int, ConcurrencyLimitConfig, None]] = None,
         parameters: Optional[dict] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
@@ -583,6 +600,7 @@ class RunnerDeployment(BaseModel):
             entrypoint:  The path to a file containing a flow and the name of the flow function in
                 the format `./path/to/file.py:flow_func_name`.
             name: A name for the deployment
+            flow_name: The name of the flow to deploy
             interval: An interval on which to execute the current flow. Accepts either a number
                 or a timedelta object. If a number is given, it will be interpreted as seconds.
             cron: A cron schedule of when to execute runs of this flow.
@@ -618,11 +636,20 @@ class RunnerDeployment(BaseModel):
             schedules=schedules,
         )
 
+        if isinstance(concurrency_limit, ConcurrencyLimitConfig):
+            concurrency_options = {
+                "collision_strategy": concurrency_limit.collision_strategy
+            }
+            concurrency_limit = concurrency_limit.limit
+        else:
+            concurrency_options = None
+
         deployment = cls(
             name=Path(name).stem,
-            flow_name=flow.name,
+            flow_name=flow_name or flow.name,
             schedules=constructed_schedules,
             concurrency_limit=concurrency_limit,
+            concurrency_options=concurrency_options,
             paused=paused,
             tags=tags or [],
             triggers=triggers or [],
@@ -648,6 +675,7 @@ class RunnerDeployment(BaseModel):
         storage: RunnerStorage,
         entrypoint: str,
         name: str,
+        flow_name: Optional[str] = None,
         interval: Optional[
             Union[Iterable[Union[int, float, timedelta]], int, float, timedelta]
         ] = None,
@@ -655,7 +683,7 @@ class RunnerDeployment(BaseModel):
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
         schedules: Optional["FlexibleScheduleList"] = None,
-        concurrency_limit: Optional[int] = None,
+        concurrency_limit: Optional[Union[int, ConcurrencyLimitConfig, None]] = None,
         parameters: Optional[dict] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
@@ -674,6 +702,7 @@ class RunnerDeployment(BaseModel):
             entrypoint:  The path to a file containing a flow and the name of the flow function in
                 the format `./path/to/file.py:flow_func_name`.
             name: A name for the deployment
+            flow_name: The name of the flow to deploy
             storage: A storage object to use for retrieving flow code. If not provided, a
                 URL must be provided.
             interval: An interval on which to execute the current flow. Accepts either a number
@@ -705,6 +734,14 @@ class RunnerDeployment(BaseModel):
             schedules=schedules,
         )
 
+        if isinstance(concurrency_limit, ConcurrencyLimitConfig):
+            concurrency_options = {
+                "collision_strategy": concurrency_limit.collision_strategy
+            }
+            concurrency_limit = concurrency_limit.limit
+        else:
+            concurrency_options = None
+
         job_variables = job_variables or {}
 
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -718,9 +755,10 @@ class RunnerDeployment(BaseModel):
 
         deployment = cls(
             name=Path(name).stem,
-            flow_name=flow.name,
+            flow_name=flow_name or flow.name,
             schedules=constructed_schedules,
             concurrency_limit=concurrency_limit,
+            concurrency_options=concurrency_options,
             paused=paused,
             tags=tags or [],
             triggers=triggers or [],

prefect/events/actions.py

@@ -113,6 +113,12 @@ class CancelFlowRun(Action):
     type: Literal["cancel-flow-run"] = "cancel-flow-run"
 
 
+class ResumeFlowRun(Action):
+    """Resumes a flow run associated with the trigger"""
+
+    type: Literal["resume-flow-run"] = "resume-flow-run"
+
+
 class SuspendFlowRun(Action):
     """Suspends a flow run associated with the trigger"""
 

prefect/flow_engine.py

@@ -47,7 +47,12 @@ from prefect.logging.loggers import (
     get_run_logger,
     patch_print,
 )
-from prefect.results import BaseResult, ResultStore, get_current_result_store
+from prefect.results import (
+    BaseResult,
+    ResultStore,
+    get_result_store,
+    should_persist_result,
+)
 from prefect.settings import PREFECT_DEBUG_MODE
 from prefect.states import (
     Failed,
@@ -202,7 +207,7 @@ class FlowRunEngine(Generic[P, R]):
                 self.handle_exception(
                     exc,
                     msg=message,
-                    result_store=get_current_result_store().update_for_flow(
+                    result_store=get_result_store().update_for_flow(
                         self.flow, _sync=True
                     ),
                 )
@@ -271,7 +276,7 @@ class FlowRunEngine(Generic[P, R]):
             return_value_to_state(
                 resolved_result,
                 result_store=result_store,
-                write_result=True,
+                write_result=should_persist_result(),
             )
         )
         self.set_state(terminal_state)
@@ -507,10 +512,13 @@ class FlowRunEngine(Generic[P, R]):
                     flow_run=self.flow_run,
                     parameters=self.parameters,
                     client=client,
-                    result_store=get_current_result_store().update_for_flow(
+                    result_store=get_result_store().update_for_flow(
                         self.flow, _sync=True
                     ),
                     task_runner=task_runner,
+                    persist_result=self.flow.persist_result
+                    if self.flow.persist_result is not None
+                    else should_persist_result(),
                 )
             )
             stack.enter_context(ConcurrencyContextV1())