prefect-client 3.0.0rc9__py3-none-any.whl → 3.0.0rc11__py3-none-any.whl

prefect/transactions.py

@@ -22,7 +22,7 @@ from prefect.records.result_store import ResultFactoryStore
 from prefect.results import (
     BaseResult,
     ResultFactory,
-    get_or_create_default_result_storage,
+    get_default_result_storage,
 )
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.collections import AutoEnum
@@ -187,8 +187,16 @@ class Transaction(ContextModel):
 
             for hook in self.on_commit_hooks:
                 hook_name = _get_hook_name(hook)
+                if self.logger:
+                    self.logger.info(f"Running commit hook {hook_name!r}")
+
                 hook(self)
 
+                if self.logger:
+                    self.logger.info(
+                        f"Commit hook {hook_name!r} finished running successfully"
+                    )
+
             if self.store and self.key:
                 self.store.write(key=self.key, value=self._staged_value)
             self.state = TransactionState.COMMITTED
@@ -235,8 +243,16 @@ class Transaction(ContextModel):
         try:
             for hook in reversed(self.on_rollback_hooks):
                 hook_name = _get_hook_name(hook)
+                if self.logger:
+                    self.logger.info(f"Running rollback hook {hook_name!r}")
+
                 hook(self)
 
+                if self.logger:
+                    self.logger.info(
+                        f"Rollback hook {hook_name!r} finished running successfully"
+                    )
+
             self.state = TransactionState.ROLLED_BACK
 
             for child in reversed(self.children):
@@ -297,7 +313,7 @@ def transaction(
                 }
             )
         else:
-            default_storage = get_or_create_default_result_storage(_sync=True)
+            default_storage = get_default_result_storage(_sync=True)
             if existing_factory:
                 new_factory = existing_factory.model_copy(
                     update={

prefect/utilities/asyncutils.py

@@ -184,7 +184,7 @@ def run_coro_as_sync(
     coroutine: Awaitable[R],
     force_new_thread: bool = False,
     wait_for_result: bool = True,
-) -> R:
+) -> Union[R, None]:
     """
     Runs a coroutine from a synchronous context, as if it were a synchronous
     function.
@@ -216,7 +216,7 @@ def run_coro_as_sync(
         else:
             raise TypeError("`coroutine` must be a coroutine object")
 
-    async def coroutine_wrapper():
+    async def coroutine_wrapper() -> Union[R, None]:
         """
         Set flags so that children (and grandchildren...) of this task know they are running in a new
         thread and do not try to run on the run_sync thread, which would cause a
@@ -245,7 +245,13 @@ def run_coro_as_sync(
         call = _cast_to_call(coroutine_wrapper)
         runner = get_run_sync_loop()
         runner.submit(call)
-        return call.result()
+        try:
+            return call.result()
+        except KeyboardInterrupt:
+            call.cancel()
+
+            logger.debug("Coroutine cancelled due to KeyboardInterrupt.")
+            raise
 
 
 async def run_sync_in_worker_thread(

prefect/utilities/engine.py

@@ -51,6 +51,7 @@ 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 (
@@ -558,7 +559,7 @@ def propose_state_sync(
         )
 
 
-def _dynamic_key_for_task_run(context: FlowRunContext, task: Task) -> int:
+def _dynamic_key_for_task_run(context: FlowRunContext, task: Task) -> Union[int, str]:
     if context.detached:  # this task is running on remote infrastructure
         return str(uuid4())
     elif context.flow_run is None:  # this is an autonomous task run
@@ -744,6 +745,12 @@ def emit_task_run_state_change_event(
                     "message": truncated_to(
                         state_message_truncation_length, initial_state.message
                     ),
+                    "state_details": initial_state.state_details.model_dump(
+                        mode="json",
+                        exclude_none=True,
+                        exclude_unset=True,
+                        exclude={"flow_run_id", "task_run_id"},
+                    ),
                 }
                 if initial_state
                 else None
@@ -754,7 +761,30 @@ def emit_task_run_state_change_event(
                 "message": truncated_to(
                     state_message_truncation_length, validated_state.message
                 ),
+                "state_details": validated_state.state_details.model_dump(
+                    mode="json",
+                    exclude_none=True,
+                    exclude_unset=True,
+                    exclude={"flow_run_id", "task_run_id"},
+                ),
+                "data": validated_state.data.model_dump(mode="json")
+                if isinstance(validated_state.data, BaseResult)
+                else None,
             },
+            "task_run": task_run.model_dump(
+                mode="json",
+                exclude_none=True,
+                exclude={
+                    "id",
+                    "created",
+                    "updated",
+                    "flow_run_id",
+                    "state_id",
+                    "state_type",
+                    "state_name",
+                    "state",
+                },
+            ),
         },
         resource={
             "prefect.resource.id": f"prefect.task-run.{task_run.id}",
@@ -769,6 +799,9 @@ 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",
         },
         follows=follows,
     )

prefect/utilities/importtools.py

@@ -422,7 +422,7 @@ def safe_load_namespace(source_code: str):
                 logger.debug("Failed to import from %s: %s", node.module, e)
 
     # Handle local definitions
-    for node in ast.walk(parsed_code):
+    for node in parsed_code.body:
         if isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.Assign)):
             try:
                 # Compile and execute each class and function definition and assignment

prefect/utilities/timeout.py

@@ -1,6 +1,6 @@
 from asyncio import CancelledError
 from contextlib import contextmanager
-from typing import Optional
+from typing import Optional, Type
 
 from prefect._internal.concurrency.cancellation import (
     cancel_async_after,
@@ -8,8 +8,19 @@ from prefect._internal.concurrency.cancellation import (
 )
 
 
+def fail_if_not_timeout_error(timeout_exc_type: Type[Exception]) -> None:
+    if not issubclass(timeout_exc_type, TimeoutError):
+        raise ValueError(
+            "The `timeout_exc_type` argument must be a subclass of `TimeoutError`."
+        )
+
+
 @contextmanager
-def timeout_async(seconds: Optional[float] = None):
+def timeout_async(
+    seconds: Optional[float] = None, timeout_exc_type: Type[TimeoutError] = TimeoutError
+):
+    fail_if_not_timeout_error(timeout_exc_type)
+
     if seconds is None:
         yield
         return
@@ -18,11 +29,15 @@ def timeout_async(seconds: Optional[float] = None):
         with cancel_async_after(timeout=seconds):
             yield
     except CancelledError:
-        raise TimeoutError(f"Scope timed out after {seconds} second(s).")
+        raise timeout_exc_type(f"Scope timed out after {seconds} second(s).")
 
 
 @contextmanager
-def timeout(seconds: Optional[float] = None):
+def timeout(
+    seconds: Optional[float] = None, timeout_exc_type: Type[TimeoutError] = TimeoutError
+):
+    fail_if_not_timeout_error(timeout_exc_type)
+
     if seconds is None:
         yield
         return
@@ -31,4 +46,4 @@ def timeout(seconds: Optional[float] = None):
         with cancel_sync_after(timeout=seconds):
             yield
     except CancelledError:
-        raise TimeoutError(f"Scope timed out after {seconds} second(s).")
+        raise timeout_exc_type(f"Scope timed out after {seconds} second(s).")

prefect/workers/base.py

@@ -1,7 +1,9 @@
 import abc
 import inspect
-import warnings
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Type, Union
+import threading
+from contextlib import AsyncExitStack
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Type, Union
 from uuid import uuid4
 
 import anyio
@@ -12,49 +14,31 @@ from pydantic.json_schema import GenerateJsonSchema
 from typing_extensions import Literal
 
 import prefect
-from prefect._internal.compatibility.experimental import (
-    EXPERIMENTAL_WARNING,
-    ExperimentalFeature,
-    experiment_enabled,
-)
 from prefect._internal.schemas.validators import return_v_or_none
 from prefect.client.orchestration import PrefectClient, get_client
 from prefect.client.schemas.actions import WorkPoolCreate, WorkPoolUpdate
-from prefect.client.schemas.filters import (
-    FlowRunFilter,
-    FlowRunFilterId,
-    FlowRunFilterState,
-    FlowRunFilterStateName,
-    FlowRunFilterStateType,
-    WorkPoolFilter,
-    WorkPoolFilterName,
-    WorkQueueFilter,
-    WorkQueueFilterName,
-)
 from prefect.client.schemas.objects import StateType, WorkPool
 from prefect.client.utilities import inject_client
 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 (
     Abort,
-    InfrastructureNotAvailable,
-    InfrastructureNotFound,
     ObjectNotFound,
 )
 from prefect.logging.loggers import PrefectLogAdapter, flow_run_logger, get_logger
 from prefect.plugins import load_prefect_collections
 from prefect.settings import (
     PREFECT_API_URL,
-    PREFECT_EXPERIMENTAL_WARN,
-    PREFECT_EXPERIMENTAL_WARN_ENHANCED_CANCELLATION,
     PREFECT_TEST_MODE,
     PREFECT_WORKER_HEARTBEAT_SECONDS,
     PREFECT_WORKER_PREFETCH_SECONDS,
+    PREFECT_WORKER_QUERY_SECONDS,
     get_current_settings,
 )
 from prefect.states import Crashed, Pending, exception_to_failed_state
 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
 from prefect.utilities.slugify import slugify
 from prefect.utilities.templating import (
     apply_values,
@@ -237,22 +221,7 @@ class BaseJobConfiguration(BaseModel):
         """
         Generate a command for a flow run job.
         """
-        if experiment_enabled("enhanced_cancellation"):
-            if (
-                PREFECT_EXPERIMENTAL_WARN
-                and PREFECT_EXPERIMENTAL_WARN_ENHANCED_CANCELLATION
-            ):
-                warnings.warn(
-                    EXPERIMENTAL_WARNING.format(
-                        feature="Enhanced flow run cancellation",
-                        group="enhanced_cancellation",
-                        help="",
-                    ),
-                    ExperimentalFeature,
-                    stacklevel=3,
-                )
-            return "prefect flow-run execute"
-        return "python -m prefect.engine"
+        return "prefect flow-run execute"
 
     @staticmethod
     def _base_flow_run_labels(flow_run: "FlowRun") -> Dict[str, str]:
@@ -413,12 +382,14 @@ class BaseWorker(abc.ABC):
                 ensure that work pools are not created accidentally.
             limit: The maximum number of flow runs this worker should be running at
                 a given time.
+            heartbeat_interval_seconds: The number of seconds between worker heartbeats.
             base_job_template: If creating the work pool, provide the base job
                 template to use. Logs a warning if the pool already exists.
         """
         if name and ("/" in name or "%" in name):
             raise ValueError("Worker name cannot contain '/' or '%'")
         self.name = name or f"{self.__class__.__name__} {uuid4()}"
+        self._started_event: Optional[Event] = None
         self._logger = get_logger(f"worker.{self.__class__.type}.{self.name.lower()}")
 
         self.is_setup = False
@@ -435,6 +406,7 @@ class BaseWorker(abc.ABC):
         )
 
         self._work_pool: Optional[WorkPool] = None
+        self._exit_stack: AsyncExitStack = AsyncExitStack()
         self._runs_task_group: Optional[anyio.abc.TaskGroup] = None
         self._client: Optional[PrefectClient] = None
         self._last_polled_time: pendulum.DateTime = pendulum.now("utc")
@@ -511,6 +483,86 @@ class BaseWorker(abc.ABC):
             },
         )
 
+    async def start(
+        self,
+        run_once: bool = False,
+        with_healthcheck: bool = False,
+        printer: Callable[..., None] = print,
+    ):
+        """
+        Starts the worker and runs the main worker loops.
+
+        By default, the worker will run loops to poll for scheduled/cancelled flow
+        runs and sync with the Prefect API server.
+
+        If `run_once` is set, the worker will only run each loop once and then return.
+
+        If `with_healthcheck` is set, the worker will start a healthcheck server which
+        can be used to determine if the worker is still polling for flow runs and restart
+        the worker if necessary.
+
+        Args:
+            run_once: If set, the worker will only run each loop once then return.
+            with_healthcheck: If set, the worker will start a healthcheck server.
+            printer: A `print`-like function where logs will be reported.
+        """
+        healthcheck_server = None
+        healthcheck_thread = None
+        try:
+            async with self as worker:
+                # wait for an initial heartbeat to configure the worker
+                await worker.sync_with_backend()
+                # schedule the scheduled flow run polling loop
+                async with anyio.create_task_group() as loops_task_group:
+                    loops_task_group.start_soon(
+                        partial(
+                            critical_service_loop,
+                            workload=self.get_and_submit_flow_runs,
+                            interval=PREFECT_WORKER_QUERY_SECONDS.value(),
+                            run_once=run_once,
+                            jitter_range=0.3,
+                            backoff=4,  # Up to ~1 minute interval during backoff
+                        )
+                    )
+                    # schedule the sync loop
+                    loops_task_group.start_soon(
+                        partial(
+                            critical_service_loop,
+                            workload=self.sync_with_backend,
+                            interval=self.heartbeat_interval_seconds,
+                            run_once=run_once,
+                            jitter_range=0.3,
+                            backoff=4,
+                        )
+                    )
+
+                    self._started_event = await self._emit_worker_started_event()
+
+                    if with_healthcheck:
+                        from prefect.workers.server import build_healthcheck_server
+
+                        # we'll start the ASGI server in a separate thread so that
+                        # uvicorn does not block the main thread
+                        healthcheck_server = build_healthcheck_server(
+                            worker=worker,
+                            query_interval_seconds=PREFECT_WORKER_QUERY_SECONDS.value(),
+                        )
+                        healthcheck_thread = threading.Thread(
+                            name="healthcheck-server-thread",
+                            target=healthcheck_server.run,
+                            daemon=True,
+                        )
+                        healthcheck_thread.start()
+                    printer(f"Worker {worker.name!r} started!")
+        finally:
+            if healthcheck_server and healthcheck_thread:
+                self._logger.debug("Stopping healthcheck server...")
+                healthcheck_server.should_exit = True
+                healthcheck_thread.join()
+                self._logger.debug("Healthcheck server stopped.")
+
+        printer(f"Worker {worker.name!r} stopped!")
+
     @abc.abstractmethod
     async def run(
         self,
@@ -525,20 +577,6 @@ class BaseWorker(abc.ABC):
             "Workers must implement a method for running submitted flow runs"
         )
 
-    async def kill_infrastructure(
-        self,
-        infrastructure_pid: str,
-        configuration: BaseJobConfiguration,
-        grace_seconds: int = 30,
-    ):
-        """
-        Method for killing infrastructure created by a worker. Should be implemented by
-        individual workers if they support killing infrastructure.
-        """
-        raise NotImplementedError(
-            "This worker does not support killing infrastructure."
-        )
-
     @classmethod
     def __dispatch_key__(cls):
         if cls.__name__ == "BaseWorker":
@@ -557,8 +595,8 @@ class BaseWorker(abc.ABC):
             raise ValueError("`PREFECT_API_URL` must be set to start a Worker.")
 
         self._client = get_client()
-        await self._client.__aenter__()
-        await self._runs_task_group.__aenter__()
+        await self._exit_stack.enter_async_context(self._client)
+        await self._exit_stack.enter_async_context(self._runs_task_group)
 
         self.is_setup = True
 
@@ -568,14 +606,14 @@ class BaseWorker(abc.ABC):
         self.is_setup = False
         for scope in self._scheduled_task_scopes:
             scope.cancel()
-        if self._runs_task_group:
-            await self._runs_task_group.__aexit__(*exc_info)
-        if self._client:
-            await self._client.__aexit__(*exc_info)
+
+        await self._exit_stack.__aexit__(*exc_info)
+        if self._started_event:
+            await self._emit_worker_stopped_event(self._started_event)
         self._runs_task_group = None
         self._client = None
 
-    def is_worker_still_polling(self, query_interval_seconds: int) -> bool:
+    def is_worker_still_polling(self, query_interval_seconds: float) -> bool:
         """
         This method is invoked by a webserver healthcheck handler
         and returns a boolean indicating if the worker has recorded a
@@ -611,138 +649,6 @@ class BaseWorker(abc.ABC):
 
         return await self._submit_scheduled_flow_runs(flow_run_response=runs_response)
 
-    async def check_for_cancelled_flow_runs(self):
-        if not self.is_setup:
-            raise RuntimeError(
-                "Worker is not set up. Please make sure you are running this worker "
-                "as an async context manager."
-            )
-
-        self._logger.debug("Checking for cancelled flow runs...")
-
-        work_queue_filter = (
-            WorkQueueFilter(name=WorkQueueFilterName(any_=list(self._work_queues)))
-            if self._work_queues
-            else None
-        )
-
-        named_cancelling_flow_runs = await self._client.read_flow_runs(
-            flow_run_filter=FlowRunFilter(
-                state=FlowRunFilterState(
-                    type=FlowRunFilterStateType(any_=[StateType.CANCELLED]),
-                    name=FlowRunFilterStateName(any_=["Cancelling"]),
-                ),
-                # Avoid duplicate cancellation calls
-                id=FlowRunFilterId(not_any_=list(self._cancelling_flow_run_ids)),
-            ),
-            work_pool_filter=WorkPoolFilter(
-                name=WorkPoolFilterName(any_=[self._work_pool_name])
-            ),
-            work_queue_filter=work_queue_filter,
-        )
-
-        typed_cancelling_flow_runs = await self._client.read_flow_runs(
-            flow_run_filter=FlowRunFilter(
-                state=FlowRunFilterState(
-                    type=FlowRunFilterStateType(any_=[StateType.CANCELLING]),
-                ),
-                # Avoid duplicate cancellation calls
-                id=FlowRunFilterId(not_any_=list(self._cancelling_flow_run_ids)),
-            ),
-            work_pool_filter=WorkPoolFilter(
-                name=WorkPoolFilterName(any_=[self._work_pool_name])
-            ),
-            work_queue_filter=work_queue_filter,
-        )
-
-        cancelling_flow_runs = named_cancelling_flow_runs + typed_cancelling_flow_runs
-
-        if cancelling_flow_runs:
-            self._logger.info(
-                f"Found {len(cancelling_flow_runs)} flow runs awaiting cancellation."
-            )
-
-        for flow_run in cancelling_flow_runs:
-            self._cancelling_flow_run_ids.add(flow_run.id)
-            self._runs_task_group.start_soon(self.cancel_run, flow_run)
-
-        return cancelling_flow_runs
-
-    async def cancel_run(self, flow_run: "FlowRun"):
-        run_logger = self.get_flow_run_logger(flow_run)
-
-        try:
-            configuration = await self._get_configuration(flow_run)
-        except ObjectNotFound:
-            self._logger.warning(
-                f"Flow run {flow_run.id!r} cannot be cancelled by this worker:"
-                f" associated deployment {flow_run.deployment_id!r} does not exist."
-            )
-            await self._mark_flow_run_as_cancelled(
-                flow_run,
-                state_updates={
-                    "message": (
-                        "This flow run is missing infrastructure configuration information"
-                        " and cancellation cannot be guaranteed."
-                    )
-                },
-            )
-            return
-        else:
-            if configuration.is_using_a_runner:
-                self._logger.info(
-                    f"Skipping cancellation because flow run {str(flow_run.id)!r} is"
-                    " using enhanced cancellation. A dedicated runner will handle"
-                    " cancellation."
-                )
-                return
-
-        if not flow_run.infrastructure_pid:
-            run_logger.error(
-                f"Flow run '{flow_run.id}' does not have an infrastructure pid"
-                " attached. Cancellation cannot be guaranteed."
-            )
-            await self._mark_flow_run_as_cancelled(
-                flow_run,
-                state_updates={
-                    "message": (
-                        "This flow run is missing infrastructure tracking information"
-                        " and cancellation cannot be guaranteed."
-                    )
-                },
-            )
-            return
-
-        try:
-            await self.kill_infrastructure(
-                infrastructure_pid=flow_run.infrastructure_pid,
-                configuration=configuration,
-            )
-        except NotImplementedError:
-            self._logger.error(
-                f"Worker type {self.type!r} does not support killing created "
-                "infrastructure. Cancellation cannot be guaranteed."
-            )
-        except InfrastructureNotFound as exc:
-            self._logger.warning(f"{exc} Marking flow run as cancelled.")
-            await self._mark_flow_run_as_cancelled(flow_run)
-        except InfrastructureNotAvailable as exc:
-            self._logger.warning(f"{exc} Flow run cannot be cancelled by this worker.")
-        except Exception:
-            run_logger.exception(
-                "Encountered exception while killing infrastructure for flow run "
-                f"'{flow_run.id}'. Flow run may not be cancelled."
-            )
-            # We will try again on generic exceptions
-            self._cancelling_flow_run_ids.remove(flow_run.id)
-            return
-        else:
-            self._emit_flow_run_cancelled_event(
-                flow_run=flow_run, configuration=configuration
-            )
-            await self._mark_flow_run_as_cancelled(flow_run)
-            run_logger.info(f"Cancelled flow run '{flow_run.id}'!")
-
     async def _update_local_work_pool_info(self):
         try:
             work_pool = await self._client.read_work_pool(
@@ -1162,6 +1068,7 @@ class BaseWorker(abc.ABC):
     async def __aenter__(self):
         self._logger.debug("Entering worker context...")
         await self.setup()
+
         return self
 
     async def __aexit__(self, *exc_info):
@@ -1245,20 +1152,3 @@ class BaseWorker(abc.ABC):
             related=self._event_related_resources(),
             follows=started_event,
         )
-
-    def _emit_flow_run_cancelled_event(
-        self, flow_run: "FlowRun", configuration: BaseJobConfiguration
-    ):
-        related = self._event_related_resources(configuration=configuration)
-
-        for resource in related:
-            if resource.role == "flow-run":
-                resource["prefect.infrastructure.identifier"] = str(
-                    flow_run.infrastructure_pid
-                )
-
-        emit_event(
-            event="prefect.worker.cancelled-flow-run",
-            resource=self._event_resource(),
-            related=related,
-        )

prefect/workers/block.py

@@ -0,0 +1,6 @@
+"""
+2024-06-27: This surfaces an actionable error message for moved or removed objects in Prefect 3.0 upgrade.
+"""
+from prefect._internal.compatibility.migration import getattr_migration
+
+__getattr__ = getattr_migration(__name__)

prefect/workers/cloud.py

@@ -0,0 +1,6 @@
+"""
+2024-06-27: This surfaces an actionable error message for moved or removed objects in Prefect 3.0 upgrade.
+"""
+from prefect._internal.compatibility.migration import getattr_migration
+
+__getattr__ = getattr_migration(__name__)