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

prefect/transactions.py

@@ -2,6 +2,7 @@ import copy
 import logging
 from contextlib import contextmanager
 from contextvars import ContextVar, Token
+from functools import partial
 from typing import (
     Any,
     Callable,
@@ -17,10 +18,9 @@ from pydantic import Field, PrivateAttr
 from typing_extensions import Self
 
 from prefect.context import ContextModel, FlowRunContext, TaskRunContext
-from prefect.exceptions import MissingContextError
-from prefect.logging.loggers import PrefectLogAdapter, get_logger, get_run_logger
+from prefect.exceptions import MissingContextError, SerializationError
+from prefect.logging.loggers import get_logger, get_run_logger
 from prefect.records import RecordStore
-from prefect.records.result_store import ResultFactoryStore
 from prefect.results import (
     BaseResult,
     ResultFactory,
@@ -60,13 +60,16 @@ class Transaction(ContextModel):
     key: Optional[str] = None
     children: List["Transaction"] = Field(default_factory=list)
     commit_mode: Optional[CommitMode] = None
+    isolation_level: Optional[IsolationLevel] = IsolationLevel.READ_COMMITTED
     state: TransactionState = TransactionState.PENDING
     on_commit_hooks: List[Callable[["Transaction"], None]] = Field(default_factory=list)
     on_rollback_hooks: List[Callable[["Transaction"], None]] = Field(
         default_factory=list
     )
     overwrite: bool = False
-    logger: Union[logging.Logger, logging.LoggerAdapter, None] = None
+    logger: Union[logging.Logger, logging.LoggerAdapter] = Field(
+        default_factory=partial(get_logger, "transactions")
+    )
     _stored_values: Dict[str, Any] = PrivateAttr(default_factory=dict)
     _staged_value: Any = None
     __var__: ContextVar = ContextVar("transaction")
@@ -101,16 +104,27 @@ class Transaction(ContextModel):
             raise RuntimeError(
                 "Context already entered. Context enter calls cannot be nested."
             )
-        # set default commit behavior
+        parent = get_transaction()
+        if parent:
+            self._stored_values = copy.deepcopy(parent._stored_values)
+        # set default commit behavior; either inherit from parent or set a default of eager
         if self.commit_mode is None:
-            parent = get_transaction()
+            self.commit_mode = parent.commit_mode if parent else CommitMode.LAZY
+        # set default isolation level; either inherit from parent or set a default of read committed
+        if self.isolation_level is None:
+            self.isolation_level = (
+                parent.isolation_level if parent else IsolationLevel.READ_COMMITTED
+            )
 
-            # either inherit from parent or set a default of eager
-            if parent:
-                self.commit_mode = parent.commit_mode
-                self._stored_values = copy.deepcopy(parent._stored_values)
-            else:
-                self.commit_mode = CommitMode.LAZY
+        assert self.isolation_level is not None, "Isolation level was not set correctly"
+        if (
+            self.store
+            and self.key
+            and not self.store.supports_isolation_level(self.isolation_level)
+        ):
+            raise ValueError(
+                f"Isolation level {self.isolation_level.name} is not supported by record store type {self.store.__class__.__name__}"
+            )
 
         # this needs to go before begin, which could set the state to committed
         self.state = TransactionState.ACTIVE
@@ -148,8 +162,13 @@ class Transaction(ContextModel):
         self.reset()
 
     def begin(self):
-        # currently we only support READ_COMMITTED isolation
-        # i.e., no locking behavior
+        if (
+            self.store
+            and self.key
+            and self.isolation_level == IsolationLevel.SERIALIZABLE
+        ):
+            self.logger.debug(f"Acquiring lock for transaction {self.key!r}")
+            self.store.acquire_lock(self.key)
         if (
             not self.overwrite
             and self.store
@@ -158,11 +177,12 @@ class Transaction(ContextModel):
         ):
             self.state = TransactionState.COMMITTED
 
-    def read(self) -> BaseResult:
+    def read(self) -> Optional[BaseResult]:
         if self.store and self.key:
-            return self.store.read(key=self.key)
-        else:
-            return {}  # TODO: Determine what this should be
+            record = self.store.read(key=self.key)
+            if record is not None:
+                return record.result
+        return None
 
     def reset(self) -> None:
         parent = self.get_parent()
@@ -192,6 +212,14 @@ class Transaction(ContextModel):
 
     def commit(self) -> bool:
         if self.state in [TransactionState.ROLLED_BACK, TransactionState.COMMITTED]:
+            if (
+                self.store
+                and self.key
+                and self.isolation_level == IsolationLevel.SERIALIZABLE
+            ):
+                self.logger.debug(f"Releasing lock for transaction {self.key!r}")
+                self.store.release_lock(self.key)
+
             return False
 
         try:
@@ -202,9 +230,24 @@ class Transaction(ContextModel):
                 self.run_hook(hook, "commit")
 
             if self.store and self.key:
-                self.store.write(key=self.key, value=self._staged_value)
+                self.store.write(key=self.key, result=self._staged_value)
             self.state = TransactionState.COMMITTED
+            if (
+                self.store
+                and self.key
+                and self.isolation_level == IsolationLevel.SERIALIZABLE
+            ):
+                self.logger.debug(f"Releasing lock for transaction {self.key!r}")
+                self.store.release_lock(self.key)
             return True
+        except SerializationError as exc:
+            if self.logger:
+                self.logger.warning(
+                    f"Encountered an error while serializing result for transaction {self.key!r}: {exc}"
+                    " Code execution will continue, but the transaction will not be committed.",
+                )
+            self.rollback()
+            return False
         except Exception:
             if self.logger:
                 self.logger.exception(
@@ -216,19 +259,25 @@ class Transaction(ContextModel):
 
     def run_hook(self, hook, hook_type: str) -> None:
         hook_name = _get_hook_name(hook)
-        self.logger.info(f"Running {hook_type} hook {hook_name!r}")
+        # Undocumented way to disable logging for a hook. Subject to change.
+        should_log = getattr(hook, "log_on_run", True)
+
+        if should_log:
+            self.logger.info(f"Running {hook_type} hook {hook_name!r}")
 
         try:
             hook(self)
         except Exception as exc:
-            self.logger.error(
-                f"An error was encountered while running {hook_type} hook {hook_name!r}",
-            )
+            if should_log:
+                self.logger.error(
+                    f"An error was encountered while running {hook_type} hook {hook_name!r}",
+                )
             raise exc
         else:
-            self.logger.info(
-                f"{hook_type.capitalize()} hook {hook_name!r} finished running successfully"
-            )
+            if should_log:
+                self.logger.info(
+                    f"{hook_type.capitalize()} hook {hook_name!r} finished running successfully"
+                )
 
     def stage(
         self,
@@ -269,6 +318,14 @@ class Transaction(ContextModel):
                     exc_info=True,
                 )
             return False
+        finally:
+            if (
+                self.store
+                and self.key
+                and self.isolation_level == IsolationLevel.SERIALIZABLE
+            ):
+                self.logger.debug(f"Releasing lock for transaction {self.key!r}")
+                self.store.release_lock(self.key)
 
     @classmethod
     def get_active(cls: Type[Self]) -> Optional[Self]:
@@ -284,8 +341,9 @@ def transaction(
     key: Optional[str] = None,
     store: Optional[RecordStore] = None,
     commit_mode: Optional[CommitMode] = None,
+    isolation_level: Optional[IsolationLevel] = None,
     overwrite: bool = False,
-    logger: Optional[PrefectLogAdapter] = None,
+    logger: Union[logging.Logger, logging.LoggerAdapter, None] = None,
 ) -> Generator[Transaction, None, None]:
     """
     A context manager for opening and managing a transaction.
@@ -309,6 +367,7 @@ def transaction(
             flow_run_context, "result_factory", None
         )
 
+        new_factory: ResultFactory
         if existing_factory and existing_factory.storage_block_id:
             new_factory = existing_factory.model_copy(
                 update={
@@ -332,6 +391,8 @@ def transaction(
                         result_storage=default_storage,
                     )
                 )
+        from prefect.records.result_store import ResultFactoryStore
+
         store = ResultFactoryStore(
             result_factory=new_factory,
         )
@@ -345,6 +406,7 @@ def transaction(
         key=key,
         store=store,
         commit_mode=commit_mode,
+        isolation_level=isolation_level,
         overwrite=overwrite,
         logger=logger,
     ) as txn:

prefect/utilities/annotations.py

@@ -21,8 +21,8 @@ class BaseAnnotation(
     def rewrap(self, value: T) -> "BaseAnnotation[T]":
         return type(self)(value)
 
-    def __eq__(self, other: object) -> bool:
-        if not type(self) == type(other):
+    def __eq__(self, other: "BaseAnnotation[T]") -> bool:
+        if type(self) is not type(other):
             return False
         return self.unwrap() == other.unwrap()
 
@@ -90,10 +90,11 @@ class quote(BaseAnnotation[T]):
 class Quote(quote):
     def __init__(self, expr):
         warnings.warn(
-            DeprecationWarning,
             "Use of `Quote` is deprecated. Use `quote` instead.",
+            DeprecationWarning,
             stacklevel=2,
         )
+        super().__init__(expr)
 
 
 class NotSet:

prefect/utilities/asyncutils.py

@@ -314,20 +314,20 @@ def in_async_main_thread() -> bool:
 
 @overload
 def sync_compatible(
-    async_fn: Callable[..., Coroutine[Any, Any, R]], force_sync: bool = False
+    async_fn: Callable[..., Coroutine[Any, Any, R]],
 ) -> Callable[..., R]:
     ...
 
 
 @overload
 def sync_compatible(
-    async_fn: Callable[..., Coroutine[Any, Any, R]], force_sync: bool = False
+    async_fn: Callable[..., Coroutine[Any, Any, R]],
 ) -> Callable[..., Coroutine[Any, Any, R]]:
     ...
 
 
 def sync_compatible(
-    async_fn: Callable[..., Coroutine[Any, Any, R]], force_sync: bool = False
+    async_fn: Callable[..., Coroutine[Any, Any, R]],
 ) -> Callable[..., Union[R, Coroutine[Any, Any, R]]]:
     """
     Converts an async function into a dual async and sync function.
@@ -403,7 +403,7 @@ def sync_compatible(
 
 
 @asynccontextmanager
-async def asyncnullcontext(value=None):
+async def asyncnullcontext(value=None, *args, **kwargs):
     yield value
 
 

prefect/utilities/callables.py

@@ -12,9 +12,7 @@ from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
 
 import cloudpickle
 import pydantic
-from griffe.dataclasses import Docstring
-from griffe.docstrings.dataclasses import DocstringSectionKind
-from griffe.docstrings.parsers import Parser, parse
+from griffe import Docstring, DocstringSectionKind, Parser, parse
 from typing_extensions import Literal
 
 from prefect._internal.pydantic.v1_schema import has_v1_type_as_param

prefect/utilities/dispatch.py

@@ -162,17 +162,22 @@ def register_type(cls: T) -> T:
     key = get_dispatch_key(cls)
     existing_value = registry.get(key)
     if existing_value is not None and id(existing_value) != id(cls):
-        # Get line numbers for debugging
-        file = inspect.getsourcefile(cls)
-        line_number = inspect.getsourcelines(cls)[1]
-        existing_file = inspect.getsourcefile(existing_value)
-        existing_line_number = inspect.getsourcelines(existing_value)[1]
-        warnings.warn(
-            f"Type {cls.__name__!r} at {file}:{line_number} has key {key!r} that "
-            f"matches existing registered type {existing_value.__name__!r} from "
-            f"{existing_file}:{existing_line_number}. The existing type will be "
-            "overridden."
-        )
+        try:
+            # Get line numbers for debugging
+            file = inspect.getsourcefile(cls)
+            line_number = inspect.getsourcelines(cls)[1]
+            existing_file = inspect.getsourcefile(existing_value)
+            existing_line_number = inspect.getsourcelines(existing_value)[1]
+            warnings.warn(
+                f"Type {cls.__name__!r} at {file}:{line_number} has key {key!r} that "
+                f"matches existing registered type {existing_value.__name__!r} from "
+                f"{existing_file}:{existing_line_number}. The existing type will be "
+                "overridden."
+            )
+        except OSError:
+            # If we can't get the source, another actor is loading this class via eval
+            # and we shouldn't update the registry
+            return cls
 
     # Add to the registry
     registry[key] = cls

prefect/utilities/engine.py

@@ -51,7 +51,6 @@ from prefect.logging.loggers import (
 )
 from prefect.results import BaseResult
 from prefect.settings import (
-    PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION,
     PREFECT_LOGGING_LOG_PRINTS,
 )
 from prefect.states import (
@@ -806,9 +805,7 @@ def emit_task_run_state_change_event(
                 else ""
             ),
             "prefect.state-type": str(validated_state.type.value),
-            "prefect.orchestration": "client"
-            if PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION
-            else "server",
+            "prefect.orchestration": "client",
         },
         follows=follows,
     )

prefect/utilities/schema_tools/hydration.py

@@ -202,6 +202,11 @@ def json_handler(obj: dict, ctx: HydrationContext):
             dehydrated_json = _hydrate(obj["value"], ctx)
         else:
             dehydrated_json = obj["value"]
+
+        # If the result is a Placeholder, we should return it as is
+        if isinstance(dehydrated_json, Placeholder):
+            return dehydrated_json
+
         try:
             return json.loads(dehydrated_json)
         except (json.decoder.JSONDecodeError, TypeError) as e:
@@ -224,6 +229,10 @@ def jinja_handler(obj: dict, ctx: HydrationContext):
         else:
             dehydrated_jinja = obj["template"]
 
+        # If the result is a Placeholder, we should return it as is
+        if isinstance(dehydrated_jinja, Placeholder):
+            return dehydrated_jinja
+
         try:
             validate_user_template(dehydrated_jinja)
         except (jinja2.exceptions.TemplateSyntaxError, TemplateSecurityError) as exc:
@@ -245,6 +254,10 @@ def workspace_variable_handler(obj: dict, ctx: HydrationContext):
         else:
             dehydrated_variable = obj["variable_name"]
 
+        # If the result is a Placeholder, we should return it as is
+        if isinstance(dehydrated_variable, Placeholder):
+            return dehydrated_variable
+
         if not ctx.render_workspace_variables:
             return WorkspaceVariable(variable_name=obj["variable_name"])
 

prefect/workers/base.py

@@ -19,6 +19,11 @@ from prefect.client.orchestration import PrefectClient, get_client
 from prefect.client.schemas.actions import WorkPoolCreate, WorkPoolUpdate
 from prefect.client.schemas.objects import StateType, WorkPool
 from prefect.client.utilities import inject_client
+from prefect.concurrency.asyncio import (
+    AcquireConcurrencySlotTimeoutError,
+    ConcurrencySlotAcquisitionError,
+    concurrency,
+)
 from prefect.events import Event, RelatedResource, emit_event
 from prefect.events.related import object_as_related_resource, tags_as_related_resources
 from prefect.exceptions import (
@@ -35,7 +40,13 @@ from prefect.settings import (
     PREFECT_WORKER_QUERY_SECONDS,
     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.utilities.asyncutils import asyncnullcontext
 from prefect.utilities.dispatch import get_registry_for_type, register_base_type
 from prefect.utilities.engine import propose_state
 from prefect.utilities.services import critical_service_loop
@@ -654,6 +665,7 @@ class BaseWorker(abc.ABC):
             work_pool = await self._client.read_work_pool(
                 work_pool_name=self._work_pool_name
             )
+
         except ObjectNotFound:
             if self._create_pool_if_not_found:
                 wp = WorkPoolCreate(
@@ -747,11 +759,10 @@ class BaseWorker(abc.ABC):
         for execution by the worker.
         """
         submittable_flow_runs = [entry.flow_run for entry in flow_run_response]
-        submittable_flow_runs.sort(key=lambda run: run.next_scheduled_start_time)
+
         for flow_run in submittable_flow_runs:
             if flow_run.id in self._submitting_flow_run_ids:
                 continue
-
             try:
                 if self._limiter:
                     self._limiter.acquire_on_behalf_of_nowait(flow_run.id)
@@ -796,8 +807,6 @@ class BaseWorker(abc.ABC):
                     " Please use an agent to execute this flow run."
                 )
 
-            #
-
     async def _submit_run(self, flow_run: "FlowRun") -> None:
         """
         Submits a given flow run for execution by the worker.
@@ -837,28 +846,59 @@ class BaseWorker(abc.ABC):
                         "not be cancellable."
                     )
 
-            run_logger.info(f"Completed submission of flow run '{flow_run.id}'")
+                run_logger.info(f"Completed submission of flow run '{flow_run.id}'")
 
-        else:
-            # If the run is not ready to submit, release the concurrency slot
-            if self._limiter:
-                self._limiter.release_on_behalf_of(flow_run.id)
+            else:
+                # If the run is not ready to submit, release the concurrency slot
+                if self._limiter:
+                    self._limiter.release_on_behalf_of(flow_run.id)
 
-        self._submitting_flow_run_ids.remove(flow_run.id)
+            self._submitting_flow_run_ids.remove(flow_run.id)
 
     async def _submit_run_and_capture_errors(
         self, flow_run: "FlowRun", task_status: Optional[anyio.abc.TaskStatus] = None
     ) -> Union[BaseWorkerResult, Exception]:
         run_logger = self.get_flow_run_logger(flow_run)
+        deployment = None
+
+        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_limit = deployment.concurrency_limit
+            concurrency_ctx = concurrency
+        else:
+            limit_name = None
+            concurrency_limit = None
+            concurrency_ctx = asyncnullcontext
 
         try:
-            configuration = await self._get_configuration(flow_run)
-            submitted_event = self._emit_flow_run_submitted_event(configuration)
-            result = await self.run(
-                flow_run=flow_run,
-                task_status=task_status,
-                configuration=configuration,
+            async with concurrency_ctx(
+                limit_name, occupy=concurrency_limit, max_retries=0
+            ):
+                configuration = await self._get_configuration(flow_run, deployment)
+                submitted_event = self._emit_flow_run_submitted_event(configuration)
+                result = await self.run(
+                    flow_run=flow_run,
+                    task_status=task_status,
+                    configuration=configuration,
+                )
+        except (
+            AcquireConcurrencySlotTimeoutError,
+            ConcurrencySlotAcquisitionError,
+        ) as exc:
+            self._logger.info(
+                (
+                    "Deployment %s has reached its concurrency limit when submitting flow run %s"
+                ),
+                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
@@ -924,8 +964,13 @@ class BaseWorker(abc.ABC):
     async def _get_configuration(
         self,
         flow_run: "FlowRun",
+        deployment: Optional["DeploymentResponse"] = None,
     ) -> BaseJobConfiguration:
-        deployment = await self._client.read_deployment(flow_run.deployment_id)
+        deployment = (
+            deployment
+            if deployment
+            else await self._client.read_deployment(flow_run.deployment_id)
+        )
         flow = await self._client.read_flow(flow_run.flow_id)
 
         deployment_vars = deployment.job_variables or {}
@@ -979,6 +1024,21 @@ class BaseWorker(abc.ABC):
 
         return 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:
+            # Flow run already marked as failed
+            pass
+        except Exception:
+            run_logger.exception(f"Failed to update state of flow run '{flow_run.id}'")
+
     async def _propose_failed_state(self, flow_run: "FlowRun", exc: Exception) -> None:
         run_logger = self.get_flow_run_logger(flow_run)
         try:

{prefect_client-3.0.0rc18.dist-info → prefect_client-3.0.0rc20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.0.0rc18
+Version: 3.0.0rc20
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.
@@ -34,10 +34,9 @@ Requires-Dist: exceptiongroup>=1.0.0
 Requires-Dist: fastapi<1.0.0,>=0.111.0
 Requires-Dist: fsspec>=2022.5.0
 Requires-Dist: graphviz>=0.20.1
-Requires-Dist: griffe<0.48.0,>=0.20.0
+Requires-Dist: griffe<2.0.0,>=0.49.0
 Requires-Dist: httpcore<2.0.0,>=1.0.5
 Requires-Dist: httpx[http2]!=0.23.2,>=0.23
-Requires-Dist: importlib-resources<6.2.0,>=6.1.3
 Requires-Dist: jsonpatch<2.0,>=1.32
 Requires-Dist: jsonschema<5.0.0,>=4.0.0
 Requires-Dist: orjson<4.0,>=3.7
@@ -60,7 +59,7 @@ Requires-Dist: toml>=0.10.0
 Requires-Dist: typing-extensions<5.0.0,>=4.5.0
 Requires-Dist: ujson<6.0.0,>=5.8.0
 Requires-Dist: uvicorn!=0.29.0,>=0.14.0
-Requires-Dist: websockets<13.0,>=10.4
+Requires-Dist: websockets<14.0,>=10.4
 Requires-Dist: importlib-metadata>=4.4; python_version < "3.10"
 Provides-Extra: notifications
 Requires-Dist: apprise<2.0.0,>=1.1.0; extra == "notifications"