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

prefect/results.py

@@ -5,6 +5,7 @@ import socket
 import threading
 import uuid
 from functools import partial
+from pathlib import Path
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -19,6 +20,8 @@ from typing import (
 )
 from uuid import UUID
 
+import pendulum
+from cachetools import LRUCache
 from pydantic import (
     BaseModel,
     ConfigDict,
@@ -33,6 +36,8 @@ from pydantic_extra_types.pendulum_dt import DateTime
 from typing_extensions import ParamSpec, Self
 
 import prefect
+from prefect._internal.compatibility import deprecated
+from prefect._internal.compatibility.deprecated import deprecated_field
 from prefect.blocks.core import Block
 from prefect.client.utilities import inject_client
 from prefect.exceptions import (
@@ -97,7 +102,7 @@ async def get_default_result_storage() -> WritableFileSystem:
 
 @sync_compatible
 async def resolve_result_storage(
-    result_storage: ResultStorage,
+    result_storage: Union[ResultStorage, UUID],
 ) -> WritableFileSystem:
     """
     Resolve one of the valid `ResultStorage` input types into a saved block
@@ -118,6 +123,9 @@ async def resolve_result_storage(
         storage_block = await Block.load(result_storage, client=client)
         storage_block_id = storage_block._block_document_id
         assert storage_block_id is not None, "Loaded storage blocks must have ids"
+    elif isinstance(result_storage, UUID):
+        block_document = await client.read_block_document(result_storage)
+        storage_block = Block._from_block_document(block_document)
     else:
         raise TypeError(
             "Result storage must be one of the following types: 'UUID', 'Block', "
@@ -171,6 +179,25 @@ def get_default_persist_setting() -> bool:
     return PREFECT_RESULTS_PERSIST_BY_DEFAULT.value()
 
 
+def should_persist_result() -> bool:
+    """
+    Return the default option for result persistence determined by the current run context.
+
+    If there is no current run context, the default value set by
+    `PREFECT_RESULTS_PERSIST_BY_DEFAULT` will be returned.
+    """
+    from prefect.context import FlowRunContext, TaskRunContext
+
+    task_run_context = TaskRunContext.get()
+    if task_run_context is not None:
+        return task_run_context.persist_result
+    flow_run_context = FlowRunContext.get()
+    if flow_run_context is not None:
+        return flow_run_context.persist_result
+
+    return PREFECT_RESULTS_PERSIST_BY_DEFAULT.value()
+
+
 def _format_user_supplied_storage_key(key: str) -> str:
     # Note here we are pinning to task runs since flow runs do not support storage keys
     # yet; we'll need to split logic in the future or have two separate functions
@@ -178,6 +205,16 @@ def _format_user_supplied_storage_key(key: str) -> str:
     return key.format(**runtime_vars, parameters=prefect.runtime.task_run.parameters)
 
 
+T = TypeVar("T")
+
+
+@deprecated_field(
+    "persist_result",
+    when=lambda x: x is not None,
+    when_message="use the `should_persist_result` utility function instead",
+    start_date="Sep 2024",
+    end_date="Nov 2024",
+)
 class ResultStore(BaseModel):
     """
     Manages the storage and retrieval of results.
@@ -200,10 +237,13 @@ class ResultStore(BaseModel):
     result_storage: Optional[WritableFileSystem] = Field(default=None)
     metadata_storage: Optional[WritableFileSystem] = Field(default=None)
     lock_manager: Optional[LockManager] = Field(default=None)
-    persist_result: bool = Field(default_factory=get_default_persist_setting)
     cache_result_in_memory: bool = Field(default=True)
     serializer: Serializer = Field(default_factory=get_default_result_serializer)
     storage_key_fn: Callable[[], str] = Field(default=DEFAULT_STORAGE_KEY_FN)
+    cache: LRUCache = Field(default_factory=lambda: LRUCache(maxsize=1000))
+
+    # Deprecated fields
+    persist_result: Optional[bool] = Field(default=None)
 
     @property
     def result_storage_block_id(self) -> Optional[UUID]:
@@ -227,8 +267,6 @@ class ResultStore(BaseModel):
             update["result_storage"] = await resolve_result_storage(flow.result_storage)
         if flow.result_serializer is not None:
             update["serializer"] = resolve_serializer(flow.result_serializer)
-        if flow.persist_result is not None:
-            update["persist_result"] = flow.persist_result
         if flow.cache_result_in_memory is not None:
             update["cache_result_in_memory"] = flow.cache_result_in_memory
         if self.result_storage is None and update.get("result_storage") is None:
@@ -251,8 +289,6 @@ class ResultStore(BaseModel):
             update["result_storage"] = await resolve_result_storage(task.result_storage)
         if task.result_serializer is not None:
             update["serializer"] = resolve_serializer(task.result_serializer)
-        if task.persist_result is not None:
-            update["persist_result"] = task.persist_result
         if task.cache_result_in_memory is not None:
             update["cache_result_in_memory"] = task.cache_result_in_memory
         if task.result_storage_key is not None:
@@ -293,16 +329,30 @@ class ResultStore(BaseModel):
             # so the entire payload doesn't need to be read
             try:
                 metadata_content = await self.metadata_storage.read_path(key)
-                return metadata_content is not None
+                if metadata_content is None:
+                    return False
+                metadata = ResultRecordMetadata.load_bytes(metadata_content)
+
             except Exception:
                 return False
         else:
             try:
                 content = await self.result_storage.read_path(key)
-                return content is not None
+                if content is None:
+                    return False
+                record = ResultRecord.deserialize(content)
+                metadata = record.metadata
             except Exception:
                 return False
 
+        if metadata.expiration:
+            # if the result has an expiration,
+            # check if it is still in the future
+            exists = metadata.expiration > pendulum.now("utc")
+        else:
+            exists = True
+        return exists
+
     def exists(self, key: str) -> bool:
         """
         Check if a result record exists in storage.
@@ -342,9 +392,13 @@ class ResultStore(BaseModel):
         Returns:
             A result record.
         """
+
         if self.lock_manager is not None and not self.is_lock_holder(key, holder):
             await self.await_for_lock(key)
 
+        if key in self.cache:
+            return self.cache[key]
+
         if self.result_storage is None:
             self.result_storage = await get_default_result_storage()
 
@@ -355,12 +409,23 @@ class ResultStore(BaseModel):
                 metadata.storage_key is not None
             ), "Did not find storage key in metadata"
             result_content = await self.result_storage.read_path(metadata.storage_key)
-            return ResultRecord.deserialize_from_result_and_metadata(
+            result_record = ResultRecord.deserialize_from_result_and_metadata(
                 result=result_content, metadata=metadata_content
             )
         else:
             content = await self.result_storage.read_path(key)
-            return ResultRecord.deserialize(content)
+            result_record = ResultRecord.deserialize(content)
+
+        if self.cache_result_in_memory:
+            if self.result_storage_block_id is None and hasattr(
+                self.result_storage, "_resolve_path"
+            ):
+                cache_key = str(self.result_storage._resolve_path(key))
+            else:
+                cache_key = key
+
+            self.cache[cache_key] = result_record
+        return result_record
 
     def read(self, key: str, holder: Optional[str] = None) -> "ResultRecord":
         """
@@ -390,10 +455,10 @@ class ResultStore(BaseModel):
 
     def create_result_record(
         self,
-        key: str,
         obj: Any,
+        key: Optional[str] = None,
         expiration: Optional[DateTime] = None,
-    ):
+    ) -> "ResultRecord":
         """
         Create a result record.
 
@@ -404,6 +469,13 @@ class ResultStore(BaseModel):
         """
         key = key or self.storage_key_fn()
 
+        if self.result_storage is None:
+            self.result_storage = get_default_result_storage(_sync=True)
+
+        if self.result_storage_block_id is None:
+            if hasattr(self.result_storage, "_resolve_path"):
+                key = str(self.result_storage._resolve_path(key))
+
         return ResultRecord(
             result=obj,
             metadata=ResultRecordMetadata(
@@ -416,8 +488,8 @@ class ResultStore(BaseModel):
 
     def write(
         self,
-        key: str,
         obj: Any,
+        key: Optional[str] = None,
         expiration: Optional[DateTime] = None,
         holder: Optional[str] = None,
     ):
@@ -433,17 +505,18 @@ class ResultStore(BaseModel):
             holder: The holder of the lock if a lock was set on the record.
         """
         holder = holder or self.generate_default_holder()
+        result_record = self.create_result_record(
+            key=key, obj=obj, expiration=expiration
+        )
         return self.persist_result_record(
-            result_record=self.create_result_record(
-                key=key, obj=obj, expiration=expiration
-            ),
+            result_record=result_record,
             holder=holder,
         )
 
     async def awrite(
         self,
-        key: str,
         obj: Any,
+        key: Optional[str] = None,
         expiration: Optional[DateTime] = None,
         holder: Optional[str] = None,
     ):
@@ -478,13 +551,22 @@ class ResultStore(BaseModel):
         ), "Storage key is required on result record"
 
         key = result_record.metadata.storage_key
+        if result_record.metadata.storage_block_id is None:
+            basepath = (
+                self.result_storage._resolve_path("")
+                if hasattr(self.result_storage, "_resolve_path")
+                else Path(".").resolve()
+            )
+            base_key = str(Path(key).relative_to(basepath))
+        else:
+            base_key = key
         if (
             self.lock_manager is not None
-            and self.is_locked(key)
-            and not self.is_lock_holder(key, holder)
+            and self.is_locked(base_key)
+            and not self.is_lock_holder(base_key, holder)
         ):
             raise RuntimeError(
-                f"Cannot write to result record with key {key} because it is locked by "
+                f"Cannot write to result record with key {base_key} because it is locked by "
                 f"another holder."
             )
         if self.result_storage is None:
@@ -497,7 +579,7 @@ class ResultStore(BaseModel):
                 content=result_record.serialize_result(),
             )
             await self.metadata_storage.write_path(
-                result_record.metadata.storage_key,
+                base_key,
                 content=result_record.serialize_metadata(),
             )
         # Otherwise, write the result metadata and result together
@@ -506,6 +588,9 @@ class ResultStore(BaseModel):
                 result_record.metadata.storage_key, content=result_record.serialize()
             )
 
+        if self.cache_result_in_memory:
+            self.cache[key] = result_record
+
     def persist_result_record(
         self, result_record: "ResultRecord", holder: Optional[str] = None
     ):
@@ -671,6 +756,11 @@ class ResultStore(BaseModel):
             )
         return await self.lock_manager.await_for_lock(key, timeout)
 
+    @deprecated.deprecated_callable(
+        start_date="Sep 2024",
+        end_date="Nov 2024",
+        help="Use `create_result_record` instead.",
+    )
     @sync_compatible
     async def create_result(
         self,
@@ -683,6 +773,11 @@ class ResultStore(BaseModel):
         """
         # Null objects are "cached" in memory at no cost
         should_cache_object = self.cache_result_in_memory or obj is None
+        should_persist_result = (
+            self.persist_result
+            if self.persist_result is not None
+            else not should_cache_object
+        )
 
         if key:
 
@@ -704,7 +799,7 @@ class ResultStore(BaseModel):
             serializer=self.serializer,
             cache_object=should_cache_object,
             expiration=expiration,
-            serialize_to_none=not self.persist_result,
+            serialize_to_none=not should_persist_result,
         )
 
     # TODO: These two methods need to find a new home
@@ -729,7 +824,7 @@ class ResultStore(BaseModel):
         return record.result
 
 
-def get_current_result_store() -> ResultStore:
+def get_result_store() -> ResultStore:
     """
     Get the current result store.
     """
@@ -779,6 +874,17 @@ class ResultRecordMetadata(BaseModel):
         """
         return cls.model_validate_json(data)
 
+    def __eq__(self, other):
+        if not isinstance(other, ResultRecordMetadata):
+            return False
+        return (
+            self.storage_key == other.storage_key
+            and self.expiration == other.expiration
+            and self.serializer == other.serializer
+            and self.prefect_version == other.prefect_version
+            and self.storage_block_id == other.storage_block_id
+        )
+
 
 class ResultRecord(BaseModel, Generic[R]):
     """
@@ -856,6 +962,31 @@ class ResultRecord(BaseModel, Generic[R]):
                 value["metadata"]["prefect_version"] = value.pop("prefect_version")
         return value
 
+    @classmethod
+    async def _from_metadata(cls, metadata: ResultRecordMetadata) -> "ResultRecord[R]":
+        """
+        Create a result record from metadata.
+
+        Will use the result record metadata to fetch data via a result store.
+
+        Args:
+            metadata: The metadata to create the result record from.
+
+        Returns:
+            ResultRecord: The result record.
+        """
+        if metadata.storage_block_id is None:
+            storage_block = None
+        else:
+            storage_block = await resolve_result_storage(
+                metadata.storage_block_id, _sync=False
+            )
+        store = ResultStore(
+            result_storage=storage_block, serializer=metadata.serializer
+        )
+        result = await store.aread(metadata.storage_key)
+        return result
+
     def serialize_metadata(self) -> bytes:
         return self.metadata.dump_bytes()
 
@@ -913,7 +1044,15 @@ class ResultRecord(BaseModel, Generic[R]):
             result=result_record_metadata.serializer.loads(result),
         )
 
+    def __eq__(self, other):
+        if not isinstance(other, ResultRecord):
+            return False
+        return self.metadata == other.metadata and self.result == other.result
 
+
+@deprecated.deprecated_class(
+    start_date="Sep 2024", end_date="Nov 2024", help="Use `ResultRecord` instead."
+)
 @register_base_type
 class BaseResult(BaseModel, abc.ABC, Generic[R]):
     model_config = ConfigDict(extra="forbid")
@@ -964,6 +1103,9 @@ class BaseResult(BaseModel, abc.ABC, Generic[R]):
         return cls.__name__ if isinstance(default, PydanticUndefinedType) else default
 
 
+@deprecated.deprecated_class(
+    start_date="Sep 2024", end_date="Nov 2024", help="Use `ResultRecord` instead."
+)
 class PersistedResult(BaseResult):
     """
     Result type which stores a reference to a persisted result.
@@ -1057,7 +1199,6 @@ class PersistedResult(BaseResult):
         """
         Write the result to the storage block.
         """
-
         if self._persisted or self.serialize_to_none:
             # don't double write or overwrite
             return
@@ -1078,7 +1219,10 @@ class PersistedResult(BaseResult):
             # this could error if the serializer requires kwargs
             serializer = Serializer(type=self.serializer_type)
 
-        result_store = ResultStore(result_storage=storage_block, serializer=serializer)
+        result_store = ResultStore(
+            result_storage=storage_block,
+            serializer=serializer,
+        )
         await result_store.awrite(
             obj=obj, key=self.storage_key, expiration=self.expiration
         )

prefect/runner/runner.py

@@ -66,6 +66,11 @@ from prefect.client.schemas.filters import (
 )
 from prefect.client.schemas.objects import Flow as APIFlow
 from prefect.client.schemas.objects import FlowRun, State, StateType
+from prefect.concurrency.asyncio import (
+    AcquireConcurrencySlotTimeoutError,
+    ConcurrencySlotAcquisitionError,
+    concurrency,
+)
 from prefect.events import DeploymentTriggerTypes, TriggerTypes
 from prefect.events.related import tags_as_related_resources
 from prefect.events.schemas.events import RelatedResource
@@ -81,7 +86,12 @@ from prefect.settings import (
     PREFECT_RUNNER_SERVER_ENABLE,
     get_current_settings,
 )
-from prefect.states import Crashed, Pending, exception_to_failed_state
+from prefect.states import (
+    AwaitingConcurrencySlot,
+    Crashed,
+    Pending,
+    exception_to_failed_state,
+)
 from prefect.types.entrypoint import EntrypointType
 from prefect.utilities.asyncutils import (
     asyncnullcontext,
@@ -226,6 +236,7 @@ class Runner:
         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,
@@ -248,6 +259,10 @@ class Runner:
                 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.
             rrule: An rrule schedule of when to execute runs of this flow.
+            paused: Whether or not to set the created deployment as paused.
+            schedules: A list of schedule objects defining when to execute runs of this flow.
+                Used to define multiple schedules or additional scheduling options like `timezone`.
+            concurrency_limit: The maximum number of concurrent runs of this flow to 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
@@ -280,6 +295,7 @@ class Runner:
             version=version,
             enforce_parameter_schema=enforce_parameter_schema,
             entrypoint_type=entrypoint_type,
+            concurrency_limit=concurrency_limit,
         )
         return await self.add_deployment(deployment)
 
@@ -959,6 +975,7 @@ class Runner:
         """
         submittable_flow_runs = flow_run_response
         submittable_flow_runs.sort(key=lambda run: run.next_scheduled_start_time)
+
         for i, flow_run in enumerate(submittable_flow_runs):
             if flow_run.id in self._submitting_flow_run_ids:
                 continue
@@ -1025,12 +1042,44 @@ class Runner:
     ) -> Union[Optional[int], Exception]:
         run_logger = self._get_flow_run_logger(flow_run)
 
+        if flow_run.deployment_id:
+            deployment = await self._client.read_deployment(flow_run.deployment_id)
+        if deployment and deployment.concurrency_limit:
+            limit_name = f"deployment:{deployment.id}"
+            concurrency_ctx = concurrency
+
+            # ensure that the global concurrency limit is available
+            # and up-to-date before attempting to acquire a slot
+            await self._client.upsert_global_concurrency_limit_by_name(
+                limit_name, deployment.concurrency_limit
+            )
+        else:
+            limit_name = ""
+            concurrency_ctx = asyncnullcontext
+
         try:
-            status_code = await self._run_process(
-                flow_run=flow_run,
-                task_status=task_status,
-                entrypoint=entrypoint,
+            async with concurrency_ctx(limit_name, max_retries=0, strict=True):
+                status_code = await self._run_process(
+                    flow_run=flow_run,
+                    task_status=task_status,
+                    entrypoint=entrypoint,
+                )
+        except (
+            AcquireConcurrencySlotTimeoutError,
+            ConcurrencySlotAcquisitionError,
+        ) as exc:
+            self._logger.info(
+                (
+                    "Deployment %s reached its concurrency limit when attempting to execute flow run %s. Will attempt to execute later."
+                ),
+                flow_run.deployment_id,
+                flow_run.name,
             )
+            await self._propose_scheduled_state(flow_run)
+
+            if not task_status._future.done():
+                task_status.started(exc)
+            return exc
         except Exception as exc:
             if not task_status._future.done():
                 # This flow run was being submitted and did not start successfully
@@ -1116,6 +1165,26 @@ class Runner:
                 exc_info=True,
             )
 
+    async def _propose_scheduled_state(self, flow_run: "FlowRun") -> None:
+        run_logger = self._get_flow_run_logger(flow_run)
+        try:
+            state = await propose_state(
+                self._client,
+                AwaitingConcurrencySlot(),
+                flow_run_id=flow_run.id,
+            )
+            self._logger.info(f"Flow run {flow_run.id} now has state {state.name}")
+        except Abort as exc:
+            run_logger.info(
+                (
+                    f"Aborted rescheduling of flow run '{flow_run.id}'. "
+                    f"Server sent an abort signal: {exc}"
+                ),
+            )
+            pass
+        except Exception:
+            run_logger.exception(f"Failed to update state of flow run '{flow_run.id}'")
+
     async def _propose_crashed_state(self, flow_run: "FlowRun", message: str) -> None:
         run_logger = self._get_flow_run_logger(flow_run)
         try:

prefect/settings.py

@@ -637,7 +637,7 @@ PREFECT_API_KEY = Setting(
 )
 """API key used to authenticate with a the Prefect API. Defaults to `None`."""
 
-PREFECT_API_ENABLE_HTTP2 = Setting(bool, default=True)
+PREFECT_API_ENABLE_HTTP2 = Setting(bool, default=False)
 """
 If true, enable support for HTTP/2 for communicating with an API.