prefect-client 3.1.6__py3-none-any.whl → 3.1.8__py3-none-any.whl

prefect/runner/runner.py

@@ -47,11 +47,12 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
+    Coroutine,
     Dict,
     Iterable,
     List,
     Optional,
-    Set,
+    TypedDict,
     Union,
 )
 from uuid import UUID, uuid4
@@ -59,6 +60,7 @@ from uuid import UUID, uuid4
 import anyio
 import anyio.abc
 import pendulum
+from cachetools import LRUCache
 
 from prefect._internal.concurrency.api import (
     create_call,
@@ -94,8 +96,6 @@ from prefect.logging.loggers import PrefectLogAdapter, flow_run_logger, get_logg
 from prefect.runner.storage import RunnerStorage
 from prefect.settings import (
     PREFECT_API_URL,
-    PREFECT_RUNNER_POLL_FREQUENCY,
-    PREFECT_RUNNER_PROCESS_LIMIT,
     PREFECT_RUNNER_SERVER_ENABLE,
     get_current_settings,
 )
@@ -123,19 +123,25 @@ from prefect.utilities.services import (
 from prefect.utilities.slugify import slugify
 
 if TYPE_CHECKING:
-    from prefect.client.schemas.objects import Deployment
+    from prefect.client.schemas.responses import DeploymentResponse
     from prefect.client.types.flexible_schedule_list import FlexibleScheduleList
     from prefect.deployments.runner import RunnerDeployment
 
 __all__ = ["Runner"]
 
 
+class ProcessMapEntry(TypedDict):
+    flow_run: FlowRun
+    pid: int
+
+
 class Runner:
     def __init__(
         self,
         name: Optional[str] = None,
         query_seconds: Optional[float] = None,
         prefetch_seconds: float = 10,
+        heartbeat_seconds: Optional[float] = None,
         limit: Optional[int] = None,
         pause_on_shutdown: bool = True,
         webserver: bool = False,
@@ -149,6 +155,9 @@ class Runner:
             query_seconds: The number of seconds to wait between querying for
                 scheduled flow runs; defaults to `PREFECT_RUNNER_POLL_FREQUENCY`
             prefetch_seconds: The number of seconds to prefetch flow runs for.
+            heartbeat_seconds: The number of seconds to wait between emitting
+                flow run heartbeats. The runner will not emit heartbeats if the value is None.
+                Defaults to `PREFECT_RUNNER_HEARTBEAT_FREQUENCY`.
             limit: The maximum number of flow runs this runner should be running at
             pause_on_shutdown: A boolean for whether or not to automatically pause
                 deployment schedules on shutdown; defaults to `True`
@@ -180,6 +189,8 @@ class Runner:
                     asyncio.run(runner.start())
                 ```
         """
+        settings = get_current_settings()
+
         if name and ("/" in name or "%" in name):
             raise ValueError("Runner name cannot contain '/' or '%'")
         self.name = Path(name).stem if name is not None else f"runner-{uuid4()}"
@@ -188,19 +199,24 @@ class Runner:
         self.started = False
         self.stopping = False
         self.pause_on_shutdown = pause_on_shutdown
-        self.limit = limit or PREFECT_RUNNER_PROCESS_LIMIT.value()
+        self.limit = limit or settings.runner.process_limit
         self.webserver = webserver
 
-        self.query_seconds = query_seconds or PREFECT_RUNNER_POLL_FREQUENCY.value()
+        self.query_seconds = query_seconds or settings.runner.poll_frequency
         self._prefetch_seconds = prefetch_seconds
+        self.heartbeat_seconds = (
+            heartbeat_seconds or settings.runner.heartbeat_frequency
+        )
+        if self.heartbeat_seconds is not None and self.heartbeat_seconds < 30:
+            raise ValueError("Heartbeat must be 30 seconds or greater.")
 
         self._limiter: Optional[anyio.CapacityLimiter] = None
         self._client = get_client()
-        self._submitting_flow_run_ids = set()
-        self._cancelling_flow_run_ids = set()
-        self._scheduled_task_scopes = set()
-        self._deployment_ids: Set[UUID] = set()
-        self._flow_run_process_map: dict[UUID, dict[str, Any]] = dict()
+        self._submitting_flow_run_ids: set[UUID] = set()
+        self._cancelling_flow_run_ids: set[UUID] = set()
+        self._scheduled_task_scopes: set[UUID] = set()
+        self._deployment_ids: set[UUID] = set()
+        self._flow_run_process_map: dict[UUID, ProcessMapEntry] = dict()
 
         self._tmp_dir: Path = (
             Path(tempfile.gettempdir()) / "runner_storage" / str(uuid4())
@@ -210,6 +226,12 @@ class Runner:
 
         self._loop: Optional[asyncio.AbstractEventLoop] = None
 
+        # Caching
+        self._deployment_cache: LRUCache[UUID, "DeploymentResponse"] = LRUCache(
+            maxsize=100
+        )
+        self._flow_cache: LRUCache[UUID, "APIFlow"] = LRUCache(maxsize=100)
+
     @sync_compatible
     async def add_deployment(
         self,
@@ -234,7 +256,7 @@ class Runner:
     @sync_compatible
     async def add_flow(
         self,
-        flow: Flow,
+        flow: Flow[Any, Any],
         name: Optional[str] = None,
         interval: Optional[
             Union[
@@ -249,7 +271,7 @@ class Runner:
         paused: Optional[bool] = None,
         schedules: Optional["FlexibleScheduleList"] = None,
         concurrency_limit: Optional[Union[int, ConcurrencyLimitConfig, None]] = None,
-        parameters: Optional[dict] = None,
+        parameters: Optional[dict[str, Any]] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
         tags: Optional[List[str]] = None,
@@ -336,7 +358,7 @@ class Runner:
         else:
             return next(s for s in self._storage_objs if s == storage)
 
-    def handle_sigterm(self, signum, frame):
+    def handle_sigterm(self, **kwargs: Any) -> None:
         """
         Gracefully shuts down the runner when a SIGTERM is received.
         """
@@ -441,6 +463,16 @@ class Runner:
                         jitter_range=0.3,
                     )
                 )
+                if self.heartbeat_seconds is not None:
+                    loops_task_group.start_soon(
+                        partial(
+                            critical_service_loop,
+                            workload=runner._emit_flow_run_heartbeats,
+                            interval=self.heartbeat_seconds,
+                            run_once=run_once,
+                            jitter_range=0.3,
+                        )
+                    )
 
     def execute_in_background(
         self, func: Callable[..., Any], *args: Any, **kwargs: Any
@@ -538,6 +570,15 @@ class Runner:
                             jitter_range=0.3,
                         )
                     )
+                    if self.heartbeat_seconds is not None:
+                        tg.start_soon(
+                            partial(
+                                critical_service_loop,
+                                workload=self._emit_flow_run_heartbeats,
+                                interval=self.heartbeat_seconds,
+                                jitter_range=0.3,
+                            )
+                        )
 
     def _get_flow_run_logger(self, flow_run: "FlowRun") -> PrefectLogAdapter:
         return flow_run_logger(flow_run=flow_run).getChild(
@@ -850,18 +891,84 @@ class Runner:
                     "message": state_msg or "Flow run was cancelled successfully."
                 },
             )
+
+            flow, deployment = await self._get_flow_and_deployment(flow_run)
+            self._emit_flow_run_cancelled_event(
+                flow_run=flow_run, flow=flow, deployment=deployment
+            )
+            run_logger.info(f"Cancelled flow run '{flow_run.name}'!")
+
+    async def _get_flow_and_deployment(
+        self, flow_run: "FlowRun"
+    ) -> tuple[Optional["APIFlow"], Optional["DeploymentResponse"]]:
+        deployment: Optional["DeploymentResponse"] = (
+            self._deployment_cache.get(flow_run.deployment_id)
+            if flow_run.deployment_id
+            else None
+        )
+        flow: Optional["APIFlow"] = self._flow_cache.get(flow_run.flow_id)
+        if not deployment and flow_run.deployment_id is not None:
             try:
                 deployment = await self._client.read_deployment(flow_run.deployment_id)
+                self._deployment_cache[flow_run.deployment_id] = deployment
             except ObjectNotFound:
                 deployment = None
+        if not flow:
             try:
                 flow = await self._client.read_flow(flow_run.flow_id)
+                self._flow_cache[flow_run.flow_id] = flow
             except ObjectNotFound:
                 flow = None
-            self._emit_flow_run_cancelled_event(
-                flow_run=flow_run, flow=flow, deployment=deployment
+        return flow, deployment
+
+    async def _emit_flow_run_heartbeats(self):
+        coros: list[Coroutine[Any, Any, Any]] = []
+        for entry in self._flow_run_process_map.values():
+            coros.append(self._emit_flow_run_heartbeat(entry["flow_run"]))
+        await asyncio.gather(*coros)
+
+    async def _emit_flow_run_heartbeat(self, flow_run: "FlowRun"):
+        from prefect import __version__
+
+        related: list[RelatedResource] = []
+        tags: list[str] = []
+
+        flow, deployment = await self._get_flow_and_deployment(flow_run)
+        if deployment:
+            related.append(
+                RelatedResource(
+                    {
+                        "prefect.resource.id": f"prefect.deployment.{deployment.id}",
+                        "prefect.resource.role": "deployment",
+                        "prefect.resource.name": deployment.name,
+                    }
+                )
             )
-            run_logger.info(f"Cancelled flow run '{flow_run.name}'!")
+            tags.extend(deployment.tags)
+        if flow:
+            related.append(
+                RelatedResource(
+                    {
+                        "prefect.resource.id": f"prefect.flow.{flow.id}",
+                        "prefect.resource.role": "flow",
+                        "prefect.resource.name": flow.name,
+                    }
+                )
+            )
+        tags.extend(flow_run.tags)
+
+        related = [RelatedResource.model_validate(r) for r in related]
+        related += tags_as_related_resources(set(tags))
+
+        emit_event(
+            event="prefect.flow-run.heartbeat",
+            resource={
+                "prefect.resource.id": f"prefect.flow-run.{flow_run.id}",
+                "prefect.resource.name": flow_run.name,
+                "prefect.version": __version__,
+            },
+            related=related,
+        )
 
     def _event_resource(self):
         from prefect import __version__
@@ -876,7 +983,7 @@ class Runner:
         self,
         flow_run: "FlowRun",
         flow: "Optional[APIFlow]",
-        deployment: "Optional[Deployment]",
+        deployment: "Optional[DeploymentResponse]",
     ):
         related: list[RelatedResource] = []
         tags: list[str] = []
@@ -920,6 +1027,7 @@ class Runner:
             resource=self._event_resource(),
             related=related,
         )
+        self._logger.debug(f"Emitted flow run heartbeat event for {flow_run.id}")
 
     async def _get_scheduled_flow_runs(
         self,
@@ -1052,6 +1160,9 @@ class Runner:
                 self._flow_run_process_map[flow_run.id] = dict(
                     pid=readiness_result, flow_run=flow_run
                 )
+            # Heartbeats are opt-in and only emitted if a heartbeat frequency is set
+            if self.heartbeat_seconds is not None:
+                await self._emit_flow_run_heartbeat(flow_run)
 
             run_logger.info(f"Completed submission of flow run '{flow_run.id}'")
         else:

prefect/runner/storage.py

@@ -53,14 +53,14 @@ class RunnerStorage(Protocol):
         """
         ...
 
-    def to_pull_step(self) -> dict:
+    def to_pull_step(self) -> dict[str, Any]:
         """
         Returns a dictionary representation of the storage object that can be
         used as a deployment pull step.
         """
         ...
 
-    def __eq__(self, __value) -> bool:
+    def __eq__(self, __value: Any) -> bool:
         """
         Equality check for runner storage objects.
         """
@@ -69,7 +69,7 @@ class RunnerStorage(Protocol):
 
 class GitCredentials(TypedDict, total=False):
     username: str
-    access_token: Union[str, Secret]
+    access_token: Union[str, Secret[str]]
 
 
 class GitRepository:

prefect/serializers.py

@@ -13,7 +13,7 @@ bytes to an object respectively.
 
 import abc
 import base64
-from typing import Any, Dict, Generic, Optional, Type
+from typing import Any, Generic, Optional, Type, Union
 
 from pydantic import (
     BaseModel,
@@ -23,7 +23,7 @@ from pydantic import (
     ValidationError,
     field_validator,
 )
-from typing_extensions import Literal, Self, TypeVar
+from typing_extensions import Self, TypeVar
 
 from prefect._internal.schemas.validators import (
     cast_type_names_to_serializers,
@@ -54,7 +54,7 @@ def prefect_json_object_encoder(obj: Any) -> Any:
         }
 
 
-def prefect_json_object_decoder(result: dict):
+def prefect_json_object_decoder(result: dict[str, Any]):
     """
     `JSONDecoder.object_hook` for decoding objects from JSON when previously encoded
     with `prefect_json_object_encoder`
@@ -80,12 +80,16 @@ class Serializer(BaseModel, Generic[D], abc.ABC):
         data.setdefault("type", type_string)
         super().__init__(**data)
 
-    def __new__(cls: Type[Self], **kwargs) -> Self:
+    def __new__(cls: Type[Self], **kwargs: Any) -> Self:
         if "type" in kwargs:
             try:
                 subcls = lookup_type(cls, dispatch_key=kwargs["type"])
             except KeyError as exc:
-                raise ValidationError(errors=[exc], model=cls)
+                raise ValidationError.from_exception_data(
+                    title=cls.__name__,
+                    line_errors=[{"type": str(exc), "input": kwargs["type"]}],
+                    input_type="python",
+                )
 
             return super().__new__(subcls)
         else:
@@ -104,7 +108,7 @@ class Serializer(BaseModel, Generic[D], abc.ABC):
     model_config = ConfigDict(extra="forbid")
 
     @classmethod
-    def __dispatch_key__(cls) -> str:
+    def __dispatch_key__(cls) -> Optional[str]:
         type_str = cls.model_fields["type"].default
         return type_str if isinstance(type_str, str) else None
 
@@ -119,19 +123,15 @@ class PickleSerializer(Serializer):
     - Wraps pickles in base64 for safe transmission.
     """
 
-    type: Literal["pickle"] = "pickle"
+    type: str = Field(default="pickle", frozen=True)
 
     picklelib: str = "cloudpickle"
     picklelib_version: Optional[str] = None
 
     @field_validator("picklelib")
-    def check_picklelib(cls, value):
+    def check_picklelib(cls, value: str) -> str:
         return validate_picklelib(value)
 
-    # @model_validator(mode="before")
-    # def check_picklelib_version(cls, values):
-    #     return validate_picklelib_version(values)
-
     def dumps(self, obj: Any) -> bytes:
         pickler = from_qualified_name(self.picklelib)
         blob = pickler.dumps(obj)
@@ -151,7 +151,7 @@ class JSONSerializer(Serializer):
     Wraps the `json` library to serialize to UTF-8 bytes instead of string types.
     """
 
-    type: Literal["json"] = "json"
+    type: str = Field(default="json", frozen=True)
 
     jsonlib: str = "json"
     object_encoder: Optional[str] = Field(
@@ -171,23 +171,27 @@ class JSONSerializer(Serializer):
             "by our default `object_encoder`."
         ),
     )
-    dumps_kwargs: Dict[str, Any] = Field(default_factory=dict)
-    loads_kwargs: Dict[str, Any] = Field(default_factory=dict)
+    dumps_kwargs: dict[str, Any] = Field(default_factory=dict)
+    loads_kwargs: dict[str, Any] = Field(default_factory=dict)
 
     @field_validator("dumps_kwargs")
-    def dumps_kwargs_cannot_contain_default(cls, value):
+    def dumps_kwargs_cannot_contain_default(
+        cls, value: dict[str, Any]
+    ) -> dict[str, Any]:
         return validate_dump_kwargs(value)
 
     @field_validator("loads_kwargs")
-    def loads_kwargs_cannot_contain_object_hook(cls, value):
+    def loads_kwargs_cannot_contain_object_hook(
+        cls, value: dict[str, Any]
+    ) -> dict[str, Any]:
         return validate_load_kwargs(value)
 
-    def dumps(self, data: Any) -> bytes:
+    def dumps(self, obj: Any) -> bytes:
         json = from_qualified_name(self.jsonlib)
         kwargs = self.dumps_kwargs.copy()
         if self.object_encoder:
             kwargs["default"] = from_qualified_name(self.object_encoder)
-        result = json.dumps(data, **kwargs)
+        result = json.dumps(obj, **kwargs)
         if isinstance(result, str):
             # The standard library returns str but others may return bytes directly
             result = result.encode()
@@ -213,17 +217,17 @@ class CompressedSerializer(Serializer):
         level: If not null, the level of compression to pass to `compress`.
     """
 
-    type: Literal["compressed"] = "compressed"
+    type: str = Field(default="compressed", frozen=True)
 
     serializer: Serializer
     compressionlib: str = "lzma"
 
     @field_validator("serializer", mode="before")
-    def validate_serializer(cls, value):
+    def validate_serializer(cls, value: Union[str, Serializer]) -> Serializer:
         return cast_type_names_to_serializers(value)
 
     @field_validator("compressionlib")
-    def check_compressionlib(cls, value):
+    def check_compressionlib(cls, value: str) -> str:
         return validate_compressionlib(value)
 
     def dumps(self, obj: Any) -> bytes:
@@ -242,7 +246,7 @@ class CompressedPickleSerializer(CompressedSerializer):
     A compressed serializer preconfigured to use the pickle serializer.
     """
 
-    type: Literal["compressed/pickle"] = "compressed/pickle"
+    type: str = Field(default="compressed/pickle", frozen=True)
 
     serializer: Serializer = Field(default_factory=PickleSerializer)
 
@@ -252,6 +256,6 @@ class CompressedJSONSerializer(CompressedSerializer):
     A compressed serializer preconfigured to use the json serializer.
     """
 
-    type: Literal["compressed/json"] = "compressed/json"
+    type: str = Field(default="compressed/json", frozen=True)
 
     serializer: Serializer = Field(default_factory=JSONSerializer)

prefect/settings/__init__.py

@@ -53,6 +53,7 @@ __all__ = [  # noqa: F822
     "temporary_settings",
     "DEFAULT_PROFILES_PATH",
     # add public settings here for auto-completion
+    "PREFECT_API_AUTH_STRING",  # type: ignore
     "PREFECT_API_KEY",  # type: ignore
     "PREFECT_API_URL",  # type: ignore
     "PREFECT_UI_URL",  # type: ignore

prefect/settings/base.py

@@ -192,7 +192,7 @@ def _add_environment_variables(
 
 
 def _build_settings_config(
-    path: Tuple[str, ...] = tuple(),
+    path: Tuple[str, ...] = tuple(), frozen: bool = False
 ) -> PrefectSettingsConfigDict:
     env_prefix = f"PREFECT_{'_'.join(path).upper()}_" if path else "PREFECT_"
     return PrefectSettingsConfigDict(
@@ -202,7 +202,8 @@ def _build_settings_config(
         toml_file="prefect.toml",
         prefect_toml_table_header=path,
         pyproject_toml_table_header=("tool", "prefect", *path),
-        json_schema_extra=_add_environment_variables,
+        json_schema_extra=_add_environment_variables,  # type: ignore
+        frozen=frozen,
     )
 
 

prefect/settings/models/api.py

@@ -19,6 +19,10 @@ class APISettings(PrefectBaseSettings):
         default=None,
         description="The URL of the Prefect API. If not set, the client will attempt to infer it.",
     )
+    auth_string: Optional[SecretStr] = Field(
+        default=None,
+        description="The auth string used for basic authentication with a self-hosted Prefect API. Should be kept secret.",
+    )
     key: Optional[SecretStr] = Field(
         default=None,
         description="The API key used for authentication with the Prefect API. Should be kept secret.",

prefect/settings/models/experiments.py

@@ -22,3 +22,8 @@ class ExperimentsSettings(PrefectBaseSettings):
         default=False,
         description="Enables sending telemetry to Prefect Cloud.",
     )
+
+    lineage_events_enabled: bool = Field(
+        default=False,
+        description="If `True`, enables emitting lineage events. Set to `False` to disable lineage event emission.",
+    )

prefect/settings/models/runner.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from pydantic import Field
 
 from prefect.settings.base import PrefectBaseSettings, _build_settings_config
@@ -54,6 +56,12 @@ class RunnerSettings(PrefectBaseSettings):
         description="Number of seconds a runner should wait between queries for scheduled work.",
     )
 
+    heartbeat_frequency: Optional[int] = Field(
+        default=None,
+        description="Number of seconds a runner should wait between heartbeats for flow runs.",
+        ge=30,
+    )
+
     server: RunnerServerSettings = Field(
         default_factory=RunnerServerSettings,
         description="Settings for controlling runner server behavior",

prefect/settings/models/server/api.py

@@ -1,6 +1,7 @@
 from datetime import timedelta
+from typing import Optional
 
-from pydantic import AliasChoices, AliasPath, Field
+from pydantic import AliasChoices, AliasPath, Field, SecretStr
 
 from prefect.settings.base import PrefectBaseSettings, _build_settings_config
 
@@ -12,6 +13,11 @@ class ServerAPISettings(PrefectBaseSettings):
 
     model_config = _build_settings_config(("server", "api"))
 
+    auth_string: Optional[SecretStr] = Field(
+        default=None,
+        description="A string to use for basic authentication with the API; typically in the form 'user:password' but can be any string.",
+    )
+
     host: str = Field(
         default="127.0.0.1",
         description="The API's host address (defaults to `127.0.0.1`).",