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

prefect/states.py

@@ -18,12 +18,13 @@ from prefect.exceptions import (
     CancelledRun,
     CrashedRun,
     FailedRun,
+    MissingContextError,
     MissingResult,
     PausedRun,
     TerminationSignal,
     UnfinishedRun,
 )
-from prefect.logging.loggers import get_logger
+from prefect.logging.loggers import get_logger, get_run_logger
 from prefect.results import BaseResult, R, ResultFactory
 from prefect.settings import PREFECT_ASYNC_FETCH_STATE_RESULT
 from prefect.utilities.annotations import BaseAnnotation
@@ -218,11 +219,17 @@ async def exception_to_crashed_state(
 async def exception_to_failed_state(
     exc: Optional[BaseException] = None,
     result_factory: Optional[ResultFactory] = None,
+    write_result: bool = False,
     **kwargs,
 ) -> State:
     """
     Convenience function for creating `Failed` states from exceptions
     """
+    try:
+        local_logger = get_run_logger()
+    except MissingContextError:
+        local_logger = logger
+
     if not exc:
         _, exc, _ = sys.exc_info()
         if exc is None:
@@ -234,6 +241,14 @@ async def exception_to_failed_state(
 
     if result_factory:
         data = await result_factory.create_result(exc)
+        if write_result:
+            try:
+                await data.write()
+            except Exception as exc:
+                local_logger.warning(
+                    "Failed to write result: %s Execution will continue, but the result has not been written",
+                    exc,
+                )
     else:
         # Attach the exception for local usage, will not be available when retrieved
         # from the API
@@ -258,7 +273,7 @@ async def return_value_to_state(
     result_factory: ResultFactory,
     key: Optional[str] = None,
     expiration: Optional[datetime.datetime] = None,
-    defer_persistence: bool = False,
+    write_result: bool = False,
 ) -> State[R]:
     """
     Given a return value from a user's function, create a `State` the run should
@@ -280,6 +295,10 @@ async def return_value_to_state(
     Callers should resolve all futures into states before passing return values to this
     function.
     """
+    try:
+        local_logger = get_run_logger()
+    except MissingContextError:
+        local_logger = logger
 
     if (
         isinstance(retval, State)
@@ -291,13 +310,20 @@ async def return_value_to_state(
         # Unless the user has already constructed a result explicitly, use the factory
         # to update the data to the correct type
         if not isinstance(state.data, BaseResult):
-            state.data = await result_factory.create_result(
+            result = await result_factory.create_result(
                 state.data,
                 key=key,
                 expiration=expiration,
-                defer_persistence=defer_persistence,
             )
-
+            if write_result:
+                try:
+                    await result.write()
+                except Exception as exc:
+                    local_logger.warning(
+                        "Encountered an error while persisting result: %s Execution will continue, but the result has not been persisted",
+                        exc,
+                    )
+            state.data = result
         return state
 
     # Determine a new state from the aggregate of contained states
@@ -333,15 +359,23 @@ async def return_value_to_state(
         # TODO: We may actually want to set the data to a `StateGroup` object and just
         #       allow it to be unpacked into a tuple and such so users can interact with
         #       it
+        result = await result_factory.create_result(
+            retval,
+            key=key,
+            expiration=expiration,
+        )
+        if write_result:
+            try:
+                await result.write()
+            except Exception as exc:
+                local_logger.warning(
+                    "Encountered an error while persisting result: %s Execution will continue, but the result has not been persisted",
+                    exc,
+                )
         return State(
             type=new_state_type,
             message=message,
-            data=await result_factory.create_result(
-                retval,
-                key=key,
-                expiration=expiration,
-                defer_persistence=defer_persistence,
-            ),
+            data=result,
         )
 
     # Generators aren't portable, implicitly convert them to a list.
@@ -354,14 +388,20 @@ async def return_value_to_state(
     if isinstance(data, BaseResult):
         return Completed(data=data)
     else:
-        return Completed(
-            data=await result_factory.create_result(
-                data,
-                key=key,
-                expiration=expiration,
-                defer_persistence=defer_persistence,
-            )
+        result = await result_factory.create_result(
+            data,
+            key=key,
+            expiration=expiration,
         )
+        if write_result:
+            try:
+                await result.write()
+            except Exception as exc:
+                local_logger.warning(
+                    "Encountered an error while persisting result: %s Execution will continue, but the result has not been persisted",
+                    exc,
+                )
+        return Completed(data=result)
 
 
 @sync_compatible
@@ -684,6 +724,21 @@ def AwaitingRetry(
     )
 
 
+def AwaitingConcurrencySlot(
+    cls: Type[State[R]] = State,
+    scheduled_time: Optional[datetime.datetime] = None,
+    **kwargs: Any,
+) -> State[R]:
+    """Convenience function for creating `AwaitingConcurrencySlot` states.
+
+    Returns:
+        State: a AwaitingConcurrencySlot state
+    """
+    return Scheduled(
+        cls=cls, scheduled_time=scheduled_time, name="AwaitingConcurrencySlot", **kwargs
+    )
+
+
 def Retrying(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     """Convenience function for creating `Retrying` states.
 

prefect/task_engine.py

@@ -5,6 +5,7 @@ import time
 from asyncio import CancelledError
 from contextlib import ExitStack, asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
+from functools import partial
 from textwrap import dedent
 from typing import (
     Any,
@@ -464,13 +465,16 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 result_factory=result_factory,
                 key=transaction.key,
                 expiration=expiration,
-                # defer persistence to transaction commit
-                defer_persistence=True,
             )
         )
+
+        # Avoid logging when running this rollback hook since it is not user-defined
+        handle_rollback = partial(self.handle_rollback)
+        handle_rollback.log_on_run = False
+
         transaction.stage(
             terminal_state.data,
-            on_rollback_hooks=[self.handle_rollback] + self.task.on_rollback_hooks,
+            on_rollback_hooks=[handle_rollback] + self.task.on_rollback_hooks,
             on_commit_hooks=self.task.on_commit_hooks,
         )
         if transaction.is_committed():
@@ -536,6 +540,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                     exc,
                     message="Task run encountered an exception",
                     result_factory=getattr(context, "result_factory", None),
+                    write_result=True,
                 )
             )
             self.record_terminal_state_timing(state)
@@ -705,17 +710,22 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
     @contextmanager
     def transaction_context(self) -> Generator[Transaction, None, None]:
-        result_factory = getattr(TaskRunContext.get(), "result_factory", None)
-
         # refresh cache setting is now repurposes as overwrite transaction record
         overwrite = (
             self.task.refresh_cache
             if self.task.refresh_cache is not None
             else PREFECT_TASKS_REFRESH_CACHE.value()
         )
+
+        result_factory = getattr(TaskRunContext.get(), "result_factory", None)
+        if result_factory and result_factory.persist_result:
+            store = ResultFactoryStore(result_factory=result_factory)
+        else:
+            store = None
+
         with transaction(
             key=self.compute_transaction_key(),
-            store=ResultFactoryStore(result_factory=result_factory),
+            store=store,
             overwrite=overwrite,
             logger=self.logger,
         ) as txn:
@@ -964,12 +974,15 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             result_factory=result_factory,
             key=transaction.key,
             expiration=expiration,
-            # defer persistence to transaction commit
-            defer_persistence=True,
         )
+
+        # Avoid logging when running this rollback hook since it is not user-defined
+        handle_rollback = partial(self.handle_rollback)
+        handle_rollback.log_on_run = False
+
         transaction.stage(
             terminal_state.data,
-            on_rollback_hooks=[self.handle_rollback] + self.task.on_rollback_hooks,
+            on_rollback_hooks=[handle_rollback] + self.task.on_rollback_hooks,
             on_commit_hooks=self.task.on_commit_hooks,
         )
         if transaction.is_committed():
@@ -1199,17 +1212,21 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
     @asynccontextmanager
     async def transaction_context(self) -> AsyncGenerator[Transaction, None]:
-        result_factory = getattr(TaskRunContext.get(), "result_factory", None)
-
         # refresh cache setting is now repurposes as overwrite transaction record
         overwrite = (
             self.task.refresh_cache
             if self.task.refresh_cache is not None
             else PREFECT_TASKS_REFRESH_CACHE.value()
         )
+        result_factory = getattr(TaskRunContext.get(), "result_factory", None)
+        if result_factory and result_factory.persist_result:
+            store = ResultFactoryStore(result_factory=result_factory)
+        else:
+            store = None
+
         with transaction(
             key=self.compute_transaction_key(),
-            store=ResultFactoryStore(result_factory=result_factory),
+            store=store,
             overwrite=overwrite,
             logger=self.logger,
         ) as txn:

prefect/transactions.py

@@ -18,7 +18,7 @@ from pydantic import Field, PrivateAttr
 from typing_extensions import Self
 
 from prefect.context import ContextModel, FlowRunContext, TaskRunContext
-from prefect.exceptions import MissingContextError
+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 (
@@ -240,6 +240,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,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,

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/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/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.0rc19.dist-info → prefect_client-3.0.0rc20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.0.0rc19
+Version: 3.0.0rc20
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.