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

prefect/results.py

@@ -8,6 +8,7 @@ from functools import partial
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
+    Annotated,
     Any,
     Callable,
     Dict,
@@ -25,8 +26,10 @@ from cachetools import LRUCache
 from pydantic import (
     BaseModel,
     ConfigDict,
+    Discriminator,
     Field,
     PrivateAttr,
+    Tag,
     ValidationError,
     model_serializer,
     model_validator,
@@ -47,6 +50,7 @@ from prefect.exceptions import (
 )
 from prefect.filesystems import (
     LocalFileSystem,
+    NullFileSystem,
     WritableFileSystem,
 )
 from prefect.locking.protocol import LockManager
@@ -218,6 +222,19 @@ def _format_user_supplied_storage_key(key: str) -> str:
 T = TypeVar("T")
 
 
+def result_storage_discriminator(x: Any) -> str:
+    if isinstance(x, dict):
+        if "block_type_slug" in x:
+            return "WritableFileSystem"
+        else:
+            return "NullFileSystem"
+    if isinstance(x, WritableFileSystem):
+        return "WritableFileSystem"
+    if isinstance(x, NullFileSystem):
+        return "NullFileSystem"
+    return "None"
+
+
 @deprecated_field(
     "persist_result",
     when=lambda x: x is not None,
@@ -245,7 +262,14 @@ class ResultStore(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     result_storage: Optional[WritableFileSystem] = Field(default=None)
-    metadata_storage: Optional[WritableFileSystem] = Field(default=None)
+    metadata_storage: Annotated[
+        Union[
+            Annotated[WritableFileSystem, Tag("WritableFileSystem")],
+            Annotated[NullFileSystem, Tag("NullFileSystem")],
+            Annotated[None, Tag("None")],
+        ],
+        Discriminator(result_storage_discriminator),
+    ] = Field(default=None)
     lock_manager: Optional[LockManager] = Field(default=None)
     cache_result_in_memory: bool = Field(default=True)
     serializer: Serializer = Field(default_factory=get_default_result_serializer)
@@ -281,6 +305,7 @@ class ResultStore(BaseModel):
             update["cache_result_in_memory"] = flow.cache_result_in_memory
         if self.result_storage is None and update.get("result_storage") is None:
             update["result_storage"] = await get_default_result_storage()
+        update["metadata_storage"] = NullFileSystem()
         return self.model_copy(update=update)
 
     @sync_compatible
@@ -294,6 +319,8 @@ class ResultStore(BaseModel):
         Returns:
             An updated result store.
         """
+        from prefect.transactions import get_transaction
+
         update = {}
         if task.result_storage is not None:
             update["result_storage"] = await resolve_result_storage(task.result_storage)
@@ -305,17 +332,30 @@ class ResultStore(BaseModel):
             update["storage_key_fn"] = partial(
                 _format_user_supplied_storage_key, task.result_storage_key
             )
+
+        # use the lock manager from a parent transaction if it exists
+        if (current_txn := get_transaction()) and isinstance(
+            current_txn.store, ResultStore
+        ):
+            update["lock_manager"] = current_txn.store.lock_manager
+
         if task.cache_policy is not None and task.cache_policy is not NotSet:
             if task.cache_policy.key_storage is not None:
                 storage = task.cache_policy.key_storage
                 if isinstance(storage, str) and not len(storage.split("/")) == 2:
                     storage = Path(storage)
                 update["metadata_storage"] = await resolve_result_storage(storage)
+            # if the cache policy has a lock manager, it takes precedence over the parent transaction
             if task.cache_policy.lock_manager is not None:
                 update["lock_manager"] = task.cache_policy.lock_manager
 
         if self.result_storage is None and update.get("result_storage") is None:
             update["result_storage"] = await get_default_result_storage()
+        if (
+            isinstance(self.metadata_storage, NullFileSystem)
+            and update.get("metadata_storage", NotSet) is NotSet
+        ):
+            update["metadata_storage"] = None
         return self.model_copy(update=update)
 
     @staticmethod
@@ -433,7 +473,9 @@ class ResultStore(BaseModel):
             )
         else:
             content = await self.result_storage.read_path(key)
-            result_record = ResultRecord.deserialize(content)
+            result_record = ResultRecord.deserialize(
+                content, backup_serializer=self.serializer
+            )
 
         if self.cache_result_in_memory:
             if self.result_storage_block_id is None and hasattr(
@@ -446,26 +488,36 @@ class ResultStore(BaseModel):
             self.cache[cache_key] = result_record
         return result_record
 
-    def read(self, key: str, holder: Optional[str] = None) -> "ResultRecord":
+    def read(
+        self,
+        key: str,
+        holder: Optional[str] = None,
+    ) -> "ResultRecord":
         """
         Read a result record from storage.
 
         Args:
             key: The key to read the result record from.
             holder: The holder of the lock if a lock was set on the record.
+
         Returns:
             A result record.
         """
         holder = holder or self.generate_default_holder()
         return self._read(key=key, holder=holder, _sync=True)
 
-    async def aread(self, key: str, holder: Optional[str] = None) -> "ResultRecord":
+    async def aread(
+        self,
+        key: str,
+        holder: Optional[str] = None,
+    ) -> "ResultRecord":
         """
         Read a result record from storage.
 
         Args:
             key: The key to read the result record from.
             holder: The holder of the lock if a lock was set on the record.
+
         Returns:
             A result record.
         """
@@ -1026,17 +1078,31 @@ class ResultRecord(BaseModel, Generic[R]):
         )
 
     @classmethod
-    def deserialize(cls, data: bytes) -> "ResultRecord[R]":
+    def deserialize(
+        cls, data: bytes, backup_serializer: Optional[Serializer] = None
+    ) -> "ResultRecord[R]":
         """
         Deserialize a record from bytes.
 
         Args:
             data: the serialized record
+            backup_serializer: The serializer to use to deserialize the result record. Only
+                necessary if the provided data does not specify a serializer.
 
         Returns:
             ResultRecord: the deserialized record
         """
-        instance = cls.model_validate_json(data)
+        try:
+            instance = cls.model_validate_json(data)
+        except ValidationError:
+            if backup_serializer is None:
+                raise
+            else:
+                result = backup_serializer.loads(data)
+                return cls(
+                    metadata=ResultRecordMetadata(serializer=backup_serializer),
+                    result=result,
+                )
         if isinstance(instance.result, bytes):
             instance.result = instance.serializer.loads(instance.result)
         elif isinstance(instance.result, str):
@@ -1089,7 +1155,7 @@ class BaseResult(BaseModel, abc.ABC, Generic[R]):
             try:
                 subcls = lookup_type(cls, dispatch_key=kwargs["type"])
             except KeyError as exc:
-                raise ValidationError(errors=[exc], model=cls)
+                raise ValueError(f"Invalid type: {kwargs['type']}") from exc
             return super().__new__(subcls)
         else:
             return super().__new__(cls)

prefect/runner/runner.py

@@ -74,7 +74,6 @@ from prefect.client.schemas.objects import Flow as APIFlow
 from prefect.concurrency.asyncio import (
     AcquireConcurrencySlotTimeoutError,
     ConcurrencySlotAcquisitionError,
-    concurrency,
 )
 from prefect.events import DeploymentTriggerTypes, TriggerTypes
 from prefect.events.related import tags_as_related_resources
@@ -92,7 +91,6 @@ from prefect.settings import (
     get_current_settings,
 )
 from prefect.states import (
-    AwaitingConcurrencySlot,
     Crashed,
     Pending,
     exception_to_failed_state,
@@ -1047,22 +1045,12 @@ 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.global_concurrency_limit:
-            limit_name = deployment.global_concurrency_limit.name
-            concurrency_ctx = concurrency
-        else:
-            limit_name = ""
-            concurrency_ctx = asyncnullcontext
-
         try:
-            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,
-                )
+            status_code = await self._run_process(
+                flow_run=flow_run,
+                task_status=task_status,
+                entrypoint=entrypoint,
+            )
         except (
             AcquireConcurrencySlotTimeoutError,
             ConcurrencySlotAcquisitionError,
@@ -1164,26 +1152,6 @@ 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: