prefect-client 2.20.4__py3-none-any.whl → 3.0.0__py3-none-any.whl

prefect/workers/block.py

@@ -1,227 +1,6 @@
-import shlex
-from typing import TYPE_CHECKING, Any, Dict, List, Optional
+"""
+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
 
-import anyio
-import anyio.abc
-
-from prefect._internal.compatibility.deprecated import deprecated_class
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-from prefect._internal.schemas.validators import validate_block_is_infrastructure
-from prefect.blocks.core import Block
-from prefect.client.schemas.objects import BlockDocument
-from prefect.utilities.collections import get_from_dict
-from prefect.workers.base import BaseWorker, BaseWorkerResult
-
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import BaseModel, Field, PrivateAttr, validator
-else:
-    from pydantic import BaseModel, Field, PrivateAttr, validator
-
-from prefect.client.orchestration import PrefectClient
-from prefect.client.utilities import inject_client
-from prefect.events import RelatedResource
-from prefect.events.related import object_as_related_resource, tags_as_related_resources
-from prefect.utilities.templating import apply_values
-
-if TYPE_CHECKING:
-    from prefect.client.schemas.objects import Flow, FlowRun
-    from prefect.client.schemas.responses import DeploymentResponse
-
-
-@deprecated_class(
-    start_date="Jun 2024",
-    help="Refer to the upgrade guide for more information: https://docs.prefect.io/latest/guides/upgrade-guide-agents-to-workers/",
-)
-class BlockWorkerJobConfiguration(BaseModel):
-    block: Block = Field(
-        default=..., description="The infrastructure block to use for job creation."
-    )
-
-    @validator("block")
-    def _validate_infrastructure_block(cls, v):
-        return validate_block_is_infrastructure(v)
-
-    _related_objects: Dict[str, Any] = PrivateAttr(default_factory=dict)
-
-    @property
-    def is_using_a_runner(self):
-        return (
-            self.block.command is not None
-            and "prefect flow-run execute" in shlex.join(self.block.command)
-        )
-
-    @staticmethod
-    def _get_base_config_defaults(variables: dict) -> dict:
-        """Get default values from base config for all variables that have them."""
-        defaults = dict()
-        for variable_name, attrs in variables.items():
-            if "default" in attrs:
-                defaults[variable_name] = attrs["default"]
-
-        return defaults
-
-    @classmethod
-    @inject_client
-    async def from_template_and_values(
-        cls, base_job_template: dict, values: dict, client: "PrefectClient" = None
-    ):
-        """Creates a valid worker configuration object from the provided base
-        configuration and overrides.
-
-        Important: this method expects that the base_job_template was already
-        validated server-side.
-        """
-        job_config: Dict[str, Any] = base_job_template["job_configuration"]
-        variables_schema = base_job_template["variables"]
-        variables = cls._get_base_config_defaults(
-            variables_schema.get("properties", {})
-        )
-        variables.update(values)
-
-        populated_configuration = apply_values(template=job_config, values=variables)
-
-        block_document_id = get_from_dict(
-            populated_configuration, "block.$ref.block_document_id"
-        )
-        if not block_document_id:
-            raise ValueError(
-                "Base job template is invalid for this worker type because it does not"
-                " contain a block_document_id after variable resolution."
-            )
-
-        block_document = await client.read_block_document(
-            block_document_id=block_document_id
-        )
-        infrastructure_block = Block._from_block_document(block_document)
-
-        populated_configuration["block"] = infrastructure_block
-
-        return cls(**populated_configuration)
-
-    @classmethod
-    def json_template(cls) -> dict:
-        """Returns a dict with job configuration as keys and the corresponding templates as values
-
-        Defaults to using the job configuration parameter name as the template variable name.
-
-        e.g.
-        {
-            key1: '{{ key1 }}',     # default variable template
-            key2: '{{ template2 }}', # `template2` specifically provide as template
-        }
-        """
-        configuration = {}
-        properties = cls.schema()["properties"]
-        for k, v in properties.items():
-            if v.get("template"):
-                template = v["template"]
-            else:
-                template = "{{ " + k + " }}"
-            configuration[k] = template
-
-        return configuration
-
-    def _related_resources(self) -> List[RelatedResource]:
-        tags = set()
-        related = []
-
-        for kind, obj in self._related_objects.items():
-            if obj is None:
-                continue
-            if hasattr(obj, "tags"):
-                tags.update(obj.tags)
-            related.append(object_as_related_resource(kind=kind, role=kind, object=obj))
-
-        return related + tags_as_related_resources(tags)
-
-    def prepare_for_flow_run(
-        self,
-        flow_run: "FlowRun",
-        deployment: Optional["DeploymentResponse"] = None,
-        flow: Optional["Flow"] = None,
-    ):
-        self.block = self.block.prepare_for_flow_run(
-            flow_run=flow_run, deployment=deployment, flow=flow
-        )
-
-
-class BlockWorkerResult(BaseWorkerResult):
-    """Result of a block worker job"""
-
-
-@deprecated_class(
-    start_date="Jun 2024",
-    help="Refer to the upgrade guide for more information: https://docs.prefect.io/latest/guides/upgrade-guide-agents-to-workers/",
-)
-class BlockWorker(BaseWorker):
-    type = "block"
-    job_configuration = BlockWorkerJobConfiguration
-
-    _description = "Execute flow runs using an infrastructure block as the job creator."
-    _display_name = "Block"
-
-    async def run(
-        self,
-        flow_run: "FlowRun",
-        configuration: BlockWorkerJobConfiguration,
-        task_status: Optional[anyio.abc.TaskStatus] = None,
-    ):
-        block = configuration.block
-
-        # logic for applying infra overrides taken from src/prefect/agent.py
-        deployment = await self._client.read_deployment(flow_run.deployment_id)
-        flow = await self._client.read_flow(deployment.flow_id)
-        infra_document = await self._client.read_block_document(
-            configuration.block._block_document_id
-        )
-
-        # this piece of logic applies any overrides that may have been set on the
-        # deployment; overrides are defined as dot.delimited paths on possibly nested
-        # attributes of the infrastructure block
-        doc_dict = infra_document.dict()
-        infra_dict = doc_dict.get("data", {})
-        for override, value in (deployment.job_variables or {}).items():
-            nested_fields = override.split(".")
-            if nested_fields == ["command"]:
-                value = shlex.split(value)
-            data = infra_dict
-            for field in nested_fields[:-1]:
-                data = data[field]
-
-            # once we reach the end, set the value
-            data[nested_fields[-1]] = value
-
-        # reconstruct the infra block
-        doc_dict["data"] = infra_dict
-        infra_document = BlockDocument(**doc_dict)
-        block = Block._from_block_document(infra_document)
-
-        block = block.prepare_for_flow_run(
-            flow_run=flow_run, deployment=deployment, flow=flow
-        )
-
-        result = await block.run(
-            task_status=task_status,
-        )
-        return BlockWorkerResult(
-            identifier=result.identifier, status_code=result.status_code
-        )
-
-    async def kill_infrastructure(
-        self,
-        infrastructure_pid: str,
-        configuration: BlockWorkerJobConfiguration,
-        grace_seconds: int = 30,
-    ):
-        block = configuration.block
-        if not hasattr(block, "kill"):
-            self._logger.error(
-                f"Flow run infrastructure block {block.type!r} "
-                "does not support killing created infrastructure. "
-                "Cancellation cannot be guaranteed."
-            )
-            return
-
-        await block.kill(
-            infrastructure_pid=infrastructure_pid, grace_seconds=grace_seconds
-        )
+__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__)

prefect/workers/process.py

@@ -21,23 +21,38 @@ import socket
 import subprocess
 import sys
 import tempfile
+import threading
+from functools import partial
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, Optional, Tuple
+from typing import TYPE_CHECKING, Callable, Dict, Optional, Tuple
 
 import anyio
 import anyio.abc
-
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import Field, validator
-else:
-    from pydantic import Field, validator
+from pydantic import Field, field_validator
 
 from prefect._internal.schemas.validators import validate_command
 from prefect.client.schemas import FlowRun
-from prefect.exceptions import InfrastructureNotAvailable, InfrastructureNotFound
+from prefect.client.schemas.filters import (
+    FlowRunFilter,
+    FlowRunFilterId,
+    FlowRunFilterState,
+    FlowRunFilterStateName,
+    FlowRunFilterStateType,
+    WorkPoolFilter,
+    WorkPoolFilterName,
+    WorkQueueFilter,
+    WorkQueueFilterName,
+)
+from prefect.client.schemas.objects import StateType
+from prefect.events.utilities import emit_event
+from prefect.exceptions import (
+    InfrastructureNotAvailable,
+    InfrastructureNotFound,
+    ObjectNotFound,
+)
+from prefect.settings import PREFECT_WORKER_QUERY_SECONDS
 from prefect.utilities.processutils import get_sys_executable, run_process
+from prefect.utilities.services import critical_service_loop
 from prefect.workers.base import (
     BaseJobConfiguration,
     BaseVariables,
@@ -68,7 +83,8 @@ class ProcessJobConfiguration(BaseJobConfiguration):
     stream_output: bool = Field(default=True)
     working_dir: Optional[Path] = Field(default=None)
 
-    @validator("working_dir")
+    @field_validator("working_dir")
+    @classmethod
     def validate_command(cls, v):
         return validate_command(v)
 
@@ -133,6 +149,96 @@ class ProcessWorker(BaseWorker):
     )
     _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/356e6766a91baf20e1d08bbe16e8b5aaef4d8643-48x48.png"
 
+    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,
+                        )
+                    )
+                    loops_task_group.start_soon(
+                        partial(
+                            critical_service_loop,
+                            workload=self.check_for_cancelled_flow_runs,
+                            interval=PREFECT_WORKER_QUERY_SECONDS.value() * 2,
+                            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!")
+
     async def run(
         self,
         flow_run: FlowRun,
@@ -214,10 +320,9 @@ class ProcessWorker(BaseWorker):
             status_code=process.returncode, identifier=str(process.pid)
         )
 
-    async def kill_infrastructure(
+    async def kill_process(
         self,
         infrastructure_pid: str,
-        configuration: ProcessJobConfiguration,
         grace_seconds: int = 30,
     ):
         hostname, pid = _parse_infrastructure_pid(infrastructure_pid)
@@ -268,3 +373,151 @@ class ProcessWorker(BaseWorker):
                 # We shouldn't ever end up here, but it's possible that the
                 # process ended right after the check above.
                 return
+
+    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_process(
+                infrastructure_pid=flow_run.infrastructure_pid,
+            )
+        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}'!")
+
+    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/server.py

@@ -1,8 +1,9 @@
 from typing import Union
 
 import uvicorn
-from prefect._vendor.fastapi import APIRouter, FastAPI, status
-from prefect._vendor.fastapi.responses import JSONResponse
+import uvicorn.server
+from fastapi import APIRouter, FastAPI, status
+from fastapi.responses import JSONResponse
 
 from prefect.settings import (
     PREFECT_WORKER_WEBSERVER_HOST,
@@ -12,19 +13,19 @@ from prefect.workers.base import BaseWorker
 from prefect.workers.process import ProcessWorker
 
 
-def start_healthcheck_server(
+def build_healthcheck_server(
     worker: Union[BaseWorker, ProcessWorker],
-    query_interval_seconds: int,
+    query_interval_seconds: float,
     log_level: str = "error",
-) -> None:
+):
     """
-    Run a healthcheck FastAPI server for a worker.
+    Build a healthcheck FastAPI server for a worker.
 
     Args:
         worker (BaseWorker | ProcessWorker): the worker whose health we will check
-        log_level (str): the log level to use for the server
+        log_level (str): the log
     """
-    webserver = FastAPI()
+    app = FastAPI()
     router = APIRouter()
 
     def perform_health_check():
@@ -41,11 +42,28 @@ def start_healthcheck_server(
 
     router.add_api_route("/health", perform_health_check, methods=["GET"])
 
-    webserver.include_router(router)
+    app.include_router(router)
 
-    uvicorn.run(
-        webserver,
+    config = uvicorn.Config(
+        app=app,
         host=PREFECT_WORKER_WEBSERVER_HOST.value(),
         port=PREFECT_WORKER_WEBSERVER_PORT.value(),
         log_level=log_level,
     )
+    return uvicorn.Server(config=config)
+
+
+def start_healthcheck_server(
+    worker: Union[BaseWorker, ProcessWorker],
+    query_interval_seconds: float,
+    log_level: str = "error",
+) -> None:
+    """
+    Run a healthcheck FastAPI server for a worker.
+
+    Args:
+        worker (BaseWorker | ProcessWorker): the worker whose health we will check
+        log_level (str): the log level to use for the server
+    """
+    server = build_healthcheck_server(worker, query_interval_seconds, log_level)
+    server.run()