prefect-client 3.0.0rc19__py3-none-any.whl → 3.0.1__py3-none-any.whl

prefect/transactions.py

@@ -17,17 +17,13 @@ from typing import (
 from pydantic import Field, PrivateAttr
 from typing_extensions import Self
 
-from prefect.context import ContextModel, FlowRunContext, TaskRunContext
-from prefect.exceptions import MissingContextError
+from prefect.context import ContextModel
+from prefect.exceptions import MissingContextError, SerializationError
 from prefect.logging.loggers import get_logger, get_run_logger
 from prefect.records import RecordStore
-from prefect.results import (
-    BaseResult,
-    ResultFactory,
-    get_default_result_storage,
-)
+from prefect.records.base import TransactionRecord
+from prefect.results import BaseResult, ResultRecord, ResultStore
 from prefect.utilities.annotations import NotSet
-from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.collections import AutoEnum
 from prefect.utilities.engine import _get_hook_name
 
@@ -56,7 +52,7 @@ class Transaction(ContextModel):
     A base model for transaction state.
     """
 
-    store: Optional[RecordStore] = None
+    store: Union[RecordStore, ResultStore, None] = None
     key: Optional[str] = None
     children: List["Transaction"] = Field(default_factory=list)
     commit_mode: Optional[CommitMode] = None
@@ -177,10 +173,14 @@ class Transaction(ContextModel):
         ):
             self.state = TransactionState.COMMITTED
 
-    def read(self) -> Optional[BaseResult]:
+    def read(self) -> Union["BaseResult", ResultRecord, None]:
         if self.store and self.key:
             record = self.store.read(key=self.key)
-            if record is not None:
+            if isinstance(record, ResultRecord):
+                return record
+            # for backwards compatibility, if we encounter a transaction record, return the result
+            # This happens when the transaction is using a `ResultStore`
+            if isinstance(record, TransactionRecord):
                 return record.result
         return None
 
@@ -230,7 +230,13 @@ class Transaction(ContextModel):
                 self.run_hook(hook, "commit")
 
             if self.store and self.key:
-                self.store.write(key=self.key, result=self._staged_value)
+                if isinstance(self.store, ResultStore):
+                    if isinstance(self._staged_value, BaseResult):
+                        self.store.write(self.key, self._staged_value.get(_sync=True))
+                    else:
+                        self.store.write(self.key, self._staged_value)
+                else:
+                    self.store.write(self.key, self._staged_value)
             self.state = TransactionState.COMMITTED
             if (
                 self.store
@@ -240,6 +246,14 @@ class Transaction(ContextModel):
                 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(
@@ -251,23 +265,29 @@ 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,
-        value: BaseResult,
+        value: Union["BaseResult", Any],
         on_rollback_hooks: Optional[List] = None,
         on_commit_hooks: Optional[List] = None,
     ) -> None:
@@ -325,7 +345,7 @@ def get_transaction() -> Optional[Transaction]:
 @contextmanager
 def transaction(
     key: Optional[str] = None,
-    store: Optional[RecordStore] = None,
+    store: Union[RecordStore, ResultStore, None] = None,
     commit_mode: Optional[CommitMode] = None,
     isolation_level: Optional[IsolationLevel] = None,
     overwrite: bool = False,
@@ -347,23 +367,26 @@ def transaction(
     """
     # if there is no key, we won't persist a record
     if key and not store:
+        from prefect.context import FlowRunContext, TaskRunContext
+        from prefect.results import ResultStore, get_default_result_storage
+
         flow_run_context = FlowRunContext.get()
         task_run_context = TaskRunContext.get()
-        existing_factory = getattr(task_run_context, "result_factory", None) or getattr(
-            flow_run_context, "result_factory", None
+        existing_store = getattr(task_run_context, "result_store", None) or getattr(
+            flow_run_context, "result_store", None
         )
 
-        new_factory: ResultFactory
-        if existing_factory and existing_factory.storage_block_id:
-            new_factory = existing_factory.model_copy(
+        new_store: ResultStore
+        if existing_store and existing_store.result_storage_block_id:
+            new_store = existing_store.model_copy(
                 update={
                     "persist_result": True,
                 }
             )
         else:
             default_storage = get_default_result_storage(_sync=True)
-            if existing_factory:
-                new_factory = existing_factory.model_copy(
+            if existing_store:
+                new_store = existing_store.model_copy(
                     update={
                         "persist_result": True,
                         "storage_block": default_storage,
@@ -371,16 +394,14 @@ def transaction(
                     }
                 )
             else:
-                new_factory = run_coro_as_sync(
-                    ResultFactory.default_factory(
-                        persist_result=True,
-                        result_storage=default_storage,
-                    )
+                new_store = ResultStore(
+                    persist_result=True,
+                    result_storage=default_storage,
                 )
-        from prefect.records.result_store import ResultFactoryStore
+        from prefect.records.result_store import ResultRecordStore
 
-        store = ResultFactoryStore(
-            result_factory=new_factory,
+        store = ResultRecordStore(
+            result_store=new_store,
         )
 
     try:

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

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

@@ -263,9 +263,7 @@ def parameter_docstrings(docstring: Optional[str]) -> Dict[str, str]:
     if not docstring:
         return param_docstrings
 
-    with disable_logger("griffe.docstrings.google"), disable_logger(
-        "griffe.agents.nodes"
-    ):
+    with disable_logger("griffe"):
         parsed = parse(Docstring(docstring), Parser.google)
         for section in parsed:
             if section.kind != DocstringSectionKind.parameters:

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/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/variables.py

@@ -1,19 +1,18 @@
-from typing import List, Optional, Union
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
 
 from prefect._internal.compatibility.migration import getattr_migration
-from prefect.client.schemas.actions import VariableCreate as VariableRequest
-from prefect.client.schemas.actions import VariableUpdate as VariableUpdateRequest
-from prefect.client.schemas.objects import Variable as VariableResponse
+from prefect.client.schemas.actions import VariableCreate, VariableUpdate
 from prefect.client.utilities import get_or_create_client
 from prefect.exceptions import ObjectNotFound
-from prefect.types import StrictVariableValue
+from prefect.types import MAX_VARIABLE_NAME_LENGTH, StrictVariableValue
 from prefect.utilities.asyncutils import sync_compatible
 
 
-class Variable(VariableRequest):
+class Variable(BaseModel):
     """
-    Variables are named, mutable string values, much like environment variables. Variables are scoped to a Prefect server instance or a single workspace in Prefect Cloud.
-    https://docs.prefect.io/latest/concepts/variables/
+    Variables are named, mutable JSON values that can be shared across tasks and flows.
 
     Arguments:
         name: A string identifying the variable.
@@ -21,6 +20,19 @@ class Variable(VariableRequest):
         tags: An optional list of strings to associate with the variable.
     """
 
+    name: str = Field(
+        default=...,
+        description="The name of the variable",
+        examples=["my_variable"],
+        max_length=MAX_VARIABLE_NAME_LENGTH,
+    )
+    value: StrictVariableValue = Field(
+        default=...,
+        description="The value of the variable",
+        examples=["my-value"],
+    )
+    tags: Optional[List[str]] = Field(default=None)
+
     @classmethod
     @sync_compatible
     async def set(
@@ -29,22 +41,21 @@ class Variable(VariableRequest):
         value: StrictVariableValue,
         tags: Optional[List[str]] = None,
         overwrite: bool = False,
-        as_object: bool = False,
-    ):
+    ) -> "Variable":
         """
         Sets a new variable. If one exists with the same name, must pass `overwrite=True`
 
-        Returns the newly set value. If `as_object=True`, return the full Variable object
+        Returns the newly set variable object.
 
         Args:
             - name: The name of the variable to set.
             - value: The value of the variable to set.
             - tags: An optional list of strings to associate with the variable.
             - overwrite: Whether to overwrite the variable if it already exists.
-            - as_object: Whether to return the full Variable object.
 
         Example:
             Set a new variable and overwrite it if it already exists.
+
             ```
             from prefect.variables import Variable
 
@@ -62,14 +73,17 @@ class Variable(VariableRequest):
                 raise ValueError(
                     f"Variable {name!r} already exists. Use `overwrite=True` to update it."
                 )
-            await client.update_variable(variable=VariableUpdateRequest(**var_dict))
+            await client.update_variable(variable=VariableUpdate(**var_dict))
             variable = await client.read_variable_by_name(name)
+            var_dict = {
+                "name": variable.name,
+                "value": variable.value,
+                "tags": variable.tags or [],
+            }
         else:
-            variable = await client.create_variable(
-                variable=VariableRequest(**var_dict)
-            )
+            await client.create_variable(variable=VariableCreate(**var_dict))
 
-        return variable if as_object else variable.value
+        return cls(**var_dict)
 
     @classmethod
     @sync_compatible
@@ -77,19 +91,15 @@ class Variable(VariableRequest):
         cls,
         name: str,
         default: StrictVariableValue = None,
-        as_object: bool = False,
-    ) -> Union[StrictVariableValue, VariableResponse]:
+    ) -> StrictVariableValue:
         """
         Get a variable's value by name.
 
         If the variable does not exist, return the default value.
 
-        If `as_object=True`, return the full variable object. `default` is ignored in this case.
-
         Args:
-            - name: The name of the variable to get.
+            - name: The name of the variable value to get.
             - default: The default value to return if the variable does not exist.
-            - as_object: Whether to return the full variable object.
 
         Example:
             Get a variable's value by name.
@@ -105,7 +115,7 @@ class Variable(VariableRequest):
         client, _ = get_or_create_client()
         variable = await client.read_variable_by_name(name)
 
-        return variable if as_object else (variable.value if variable else default)
+        return variable.value if variable else default
 
     @classmethod
     @sync_compatible

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/workers/process.py

@@ -144,9 +144,7 @@ class ProcessWorker(BaseWorker):
         " when first getting started."
     )
     _display_name = "Process"
-    _documentation_url = (
-        "https://docs.prefect.io/latest/api-ref/prefect/workers/process/"
-    )
+    _documentation_url = "https://docs.prefect.io/latest/get-started/quickstart"
     _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/356e6766a91baf20e1d08bbe16e8b5aaef4d8643-48x48.png"
 
     async def start(

{prefect_client-3.0.0rc19.dist-info → prefect_client-3.0.1.dist-info}/METADATA

@@ -1,12 +1,12 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.0.0rc19
+Version: 3.0.1
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.
 Author-email: help@prefect.io
 License: UNKNOWN
-Project-URL: Changelog, https://github.com/PrefectHQ/prefect/blob/main/RELEASE-NOTES.md
+Project-URL: Changelog, https://github.com/PrefectHQ/prefect/releases
 Project-URL: Documentation, https://docs.prefect.io
 Project-URL: Source, https://github.com/PrefectHQ/prefect
 Project-URL: Tracker, https://github.com/PrefectHQ/prefect/issues