prefect 3.6.12__py3-none-any.whl → 3.6.13__py3-none-any.whl

prefect/results.py

@@ -33,7 +33,10 @@ from typing_extensions import ParamSpec, Self
 import prefect
 import prefect.types._datetime
 from prefect._internal.compatibility.async_dispatch import async_dispatch
-from prefect._internal.compatibility.blocks import call_explicitly_async_block_method
+from prefect._internal.compatibility.blocks import (
+    call_explicitly_async_block_method,
+    call_explicitly_sync_block_method,
+)
 from prefect._internal.concurrency.event_loop import get_running_loop
 from prefect._result_records import R, ResultRecord, ResultRecordMetadata
 from prefect.blocks.core import Block
@@ -52,7 +55,6 @@ from prefect.serializers import Serializer
 from prefect.settings.context import get_current_settings
 from prefect.types import DateTime
 from prefect.utilities.annotations import NotSet
-from prefect.utilities.asyncutils import sync_compatible
 
 if TYPE_CHECKING:
     import logging
@@ -369,8 +371,7 @@ class ResultStore(BaseModel):
         result = await store.aread(metadata.storage_key)
         return result
 
-    @sync_compatible
-    async def update_for_flow(self, flow: "Flow[..., Any]") -> Self:
+    async def aupdate_for_flow(self, flow: "Flow[..., Any]") -> Self:
         """
         Create a new result store for a flow with updated settings.
 
@@ -393,8 +394,31 @@ class ResultStore(BaseModel):
         update["metadata_storage"] = NullFileSystem()
         return self.model_copy(update=update)
 
-    @sync_compatible
-    async def update_for_task(self: Self, task: "Task[P, R]") -> Self:
+    @async_dispatch(aupdate_for_flow)
+    def update_for_flow(self, flow: "Flow[..., Any]") -> Self:
+        """
+        Create a new result store for a flow with updated settings.
+
+        Args:
+            flow: The flow to update the result store for.
+
+        Returns:
+            An updated result store.
+        """
+        update: dict[str, Any] = {}
+        update["cache_result_in_memory"] = flow.cache_result_in_memory
+        if flow.result_storage is not None:
+            update["result_storage"] = resolve_result_storage(
+                flow.result_storage, _sync=True
+            )
+        if flow.result_serializer is not None:
+            update["serializer"] = resolve_serializer(flow.result_serializer)
+        if self.result_storage is None and update.get("result_storage") is None:
+            update["result_storage"] = get_default_result_storage(_sync=True)
+        update["metadata_storage"] = NullFileSystem()
+        return self.model_copy(update=update)
+
+    async def aupdate_for_task(self: Self, task: "Task[P, R]") -> Self:
         """
         Create a new result store for a task.
 
@@ -446,6 +470,59 @@ class ResultStore(BaseModel):
             update["metadata_storage"] = None
         return self.model_copy(update=update)
 
+    @async_dispatch(aupdate_for_task)
+    def update_for_task(self: Self, task: "Task[P, R]") -> Self:
+        """
+        Create a new result store for a task.
+
+        Args:
+            task: The task to update the result store for.
+
+        Returns:
+            An updated result store.
+        """
+        from prefect.transactions import get_transaction
+
+        update: dict[str, Any] = {}
+        update["cache_result_in_memory"] = task.cache_result_in_memory
+        if task.result_storage is not None:
+            update["result_storage"] = resolve_result_storage(
+                task.result_storage, _sync=True
+            )
+        if task.result_serializer is not None:
+            update["serializer"] = resolve_serializer(task.result_serializer)
+        if task.result_storage_key is not None:
+            update["storage_key_fn"] = partial(
+                _format_user_supplied_storage_key, task.result_storage_key
+            )
+
+        # use the lock manager from a parent transaction if it exists
+        if (current_txn := get_transaction()) and isinstance(
+            current_txn.store, ResultStore
+        ):
+            update["lock_manager"] = current_txn.store.lock_manager
+
+        from prefect.cache_policies import CachePolicy
+
+        if isinstance(task.cache_policy, CachePolicy):
+            if task.cache_policy.key_storage is not None:
+                storage = task.cache_policy.key_storage
+                if isinstance(storage, str) and not len(storage.split("/")) == 2:
+                    storage = Path(storage)
+                update["metadata_storage"] = resolve_result_storage(storage, _sync=True)
+            # if the cache policy has a lock manager, it takes precedence over the parent transaction
+            if task.cache_policy.lock_manager is not None:
+                update["lock_manager"] = task.cache_policy.lock_manager
+
+        if self.result_storage is None and update.get("result_storage") is None:
+            update["result_storage"] = get_default_result_storage(_sync=True)
+        if (
+            isinstance(self.metadata_storage, NullFileSystem)
+            and update.get("metadata_storage", NotSet) is NotSet
+        ):
+            update["metadata_storage"] = None
+        return self.model_copy(update=update)
+
     @staticmethod
     def generate_default_holder() -> str:
         """
@@ -467,8 +544,7 @@ class ResultStore(BaseModel):
                 return f"{hostname}:{pid}:{thread_id}:{thread_name}:{id(current_task)}"
         return f"{hostname}:{pid}:{thread_id}:{thread_name}"
 
-    @sync_compatible
-    async def _exists(self, key: str) -> bool:
+    async def _aexists(self, key: str) -> bool:
         """
         Check if a result record exists in storage.
 
@@ -511,6 +587,49 @@ class ResultStore(BaseModel):
             exists = True
         return exists
 
+    def _exists(self, key: str) -> bool:
+        """
+        Check if a result record exists in storage.
+
+        Args:
+            key: The key to check for the existence of a result record.
+
+        Returns:
+            bool: True if the result record exists, False otherwise.
+        """
+        if self.metadata_storage is not None:
+            # TODO: Add an `exists` method to commonly used storage blocks
+            # so the entire payload doesn't need to be read
+            try:
+                metadata_content = call_explicitly_sync_block_method(
+                    self.metadata_storage, "read_path", (key,), {}
+                )
+                if metadata_content is None:
+                    return False
+                metadata = ResultRecordMetadata.load_bytes(metadata_content)
+
+            except Exception:
+                return False
+        else:
+            try:
+                content = call_explicitly_sync_block_method(
+                    self.result_storage, "read_path", (key,), {}
+                )
+                if content is None:
+                    return False
+                record: ResultRecord[Any] = ResultRecord.deserialize(content)
+                metadata = record.metadata
+            except Exception:
+                return False
+
+        if metadata.expiration:
+            # if the result has an expiration,
+            # check if it is still in the future
+            exists = metadata.expiration > prefect.types._datetime.now("UTC")
+        else:
+            exists = True
+        return exists
+
     def exists(self, key: str) -> bool:
         """
         Check if a result record exists in storage.
@@ -521,7 +640,7 @@ class ResultStore(BaseModel):
         Returns:
             bool: True if the result record exists, False otherwise.
         """
-        return self._exists(key=key, _sync=True)
+        return self._exists(key=key)
 
     async def aexists(self, key: str) -> bool:
         """
@@ -533,7 +652,7 @@ class ResultStore(BaseModel):
         Returns:
             bool: True if the result record exists, False otherwise.
         """
-        return await self._exists(key=key, _sync=False)
+        return await self._aexists(key=key)
 
     def _resolved_key_path(self, key: str) -> str:
         if self.result_storage_block_id is None and (
@@ -546,8 +665,7 @@ class ResultStore(BaseModel):
                 return key
         return key
 
-    @sync_compatible
-    async def _read(self, key: str, holder: str) -> "ResultRecord[Any]":
+    async def _aread(self, key: str, holder: str) -> "ResultRecord[Any]":
         """
         Read a result record from storage.
 
@@ -610,6 +728,69 @@ class ResultStore(BaseModel):
             self.cache[resolved_key_path] = result_record
         return result_record
 
+    def _read(self, key: str, holder: str) -> "ResultRecord[Any]":
+        """
+        Read a result record from storage.
+
+        This is the internal implementation. Use `read` or `aread` for synchronous and
+        asynchronous result reading respectively.
+
+        Args:
+            key: The key to read the result record from.
+            holder: The holder of the lock if a lock was set on the record.
+
+        Returns:
+            A result record.
+        """
+
+        if self.lock_manager is not None and not self.is_lock_holder(key, holder):
+            self.wait_for_lock(key)
+
+        resolved_key_path = self._resolved_key_path(key)
+
+        if resolved_key_path in self.cache:
+            return self.cache[resolved_key_path]
+
+        if self.result_storage is None:
+            self.result_storage = get_default_result_storage(_sync=True)
+
+        if self.metadata_storage is not None:
+            metadata_content = call_explicitly_sync_block_method(
+                self.metadata_storage,
+                "read_path",
+                (key,),
+                {},
+            )
+            metadata = ResultRecordMetadata.load_bytes(metadata_content)
+            assert metadata.storage_key is not None, (
+                "Did not find storage key in metadata"
+            )
+            result_content = call_explicitly_sync_block_method(
+                self.result_storage,
+                "read_path",
+                (metadata.storage_key,),
+                {},
+            )
+            result_record: ResultRecord[Any] = (
+                ResultRecord.deserialize_from_result_and_metadata(
+                    result=result_content, metadata=metadata_content
+                )
+            )
+        else:
+            content = call_explicitly_sync_block_method(
+                self.result_storage,
+                "read_path",
+                (key,),
+                {},
+            )
+            result_record: ResultRecord[Any] = ResultRecord.deserialize(
+                content, backup_serializer=self.serializer
+            )
+
+        if self.cache_result_in_memory:
+            self.cache[resolved_key_path] = result_record
+        return result_record
+
     def read(
         self,
         key: str,
@@ -626,7 +807,7 @@ class ResultStore(BaseModel):
             A result record.
         """
         holder = holder or self.generate_default_holder()
-        return self._read(key=key, holder=holder, _sync=True)
+        return self._read(key=key, holder=holder)
 
     async def aread(
         self,
@@ -644,7 +825,7 @@ class ResultStore(BaseModel):
             A result record.
         """
         holder = holder or self.generate_default_holder()
-        return await self._read(key=key, holder=holder, _sync=False)
+        return await self._aread(key=key, holder=holder)
 
     def create_result_record(
         self,
@@ -734,8 +915,7 @@ class ResultStore(BaseModel):
             holder=holder,
         )
 
-    @sync_compatible
-    async def _persist_result_record(
+    async def _apersist_result_record(
         self, result_record: "ResultRecord[Any]", holder: str
     ) -> None:
         """
@@ -798,6 +978,69 @@ class ResultStore(BaseModel):
         if self.cache_result_in_memory:
             self.cache[key] = result_record
 
+    def _persist_result_record(
+        self, result_record: "ResultRecord[Any]", holder: str
+    ) -> None:
+        """
+        Persist a result record to storage.
+
+        Args:
+            result_record: The result record to persist.
+            holder: The holder of the lock if a lock was set on the record.
+        """
+        assert result_record.metadata.storage_key is not None, (
+            "Storage key is required on result record"
+        )
+
+        key = result_record.metadata.storage_key
+        if result_record.metadata.storage_block_id is None:
+            basepath = (
+                _resolve_path("")
+                if (
+                    _resolve_path := getattr(self.result_storage, "_resolve_path", None)
+                )
+                else Path(".").resolve()
+            )
+            base_key = key if basepath is None else str(Path(key).relative_to(basepath))
+        else:
+            base_key = key
+        if (
+            self.lock_manager is not None
+            and self.is_locked(base_key)
+            and not self.is_lock_holder(base_key, holder)
+        ):
+            raise RuntimeError(
+                f"Cannot write to result record with key {base_key} because it is locked by "
+                f"another holder."
+            )
+        if self.result_storage is None:
+            self.result_storage = get_default_result_storage(_sync=True)
+
+        # If metadata storage is configured, write result and metadata separately
+        if self.metadata_storage is not None:
+            call_explicitly_sync_block_method(
+                self.result_storage,
+                "write_path",
+                (result_record.metadata.storage_key,),
+                {"content": result_record.serialize_result()},
+            )
+            call_explicitly_sync_block_method(
+                self.metadata_storage,
+                "write_path",
+                (base_key,),
+                {"content": result_record.serialize_metadata()},
+            )
+        # Otherwise, write the result metadata and result together
+        else:
+            call_explicitly_sync_block_method(
+                self.result_storage,
+                "write_path",
+                (result_record.metadata.storage_key,),
+                {"content": result_record.serialize()},
+            )
+        if self.cache_result_in_memory:
+            self.cache[key] = result_record
+
     def persist_result_record(
         self, result_record: "ResultRecord[Any]", holder: str | None = None
     ) -> None:
@@ -808,9 +1051,7 @@ class ResultStore(BaseModel):
             result_record: The result record to persist.
         """
         holder = holder or self.generate_default_holder()
-        return self._persist_result_record(
-            result_record=result_record, holder=holder, _sync=True
-        )
+        return self._persist_result_record(result_record=result_record, holder=holder)
 
     async def apersist_result_record(
         self, result_record: "ResultRecord[Any]", holder: str | None = None
@@ -822,8 +1063,8 @@ class ResultStore(BaseModel):
             result_record: The result record to persist.
         """
         holder = holder or self.generate_default_holder()
-        return await self._persist_result_record(
-            result_record=result_record, holder=holder, _sync=False
+        return await self._apersist_result_record(
+            result_record=result_record, holder=holder
         )
 
     def supports_isolation_level(self, level: "IsolationLevel") -> bool:

prefect/runner/runner.py

@@ -56,7 +56,6 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
-    Coroutine,
     Dict,
     Iterable,
     List,
@@ -167,9 +166,10 @@ 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`.
+            heartbeat_seconds: The number of seconds between heartbeat events emitted
+                by flow runs managed by this runner. If not provided, the value of
+                `PREFECT_FLOWS_HEARTBEAT_FREQUENCY` will be used. Heartbeats are used
+                to detect crashed flow runs.
             limit: The maximum number of flow runs this runner should be running at. Provide `None` for no limit.
                 If not provided, the runner will use the value of `PREFECT_RUNNER_PROCESS_LIMIT`.
             pause_on_shutdown: A boolean for whether or not to automatically pause
@@ -202,6 +202,8 @@ class Runner:
                     asyncio.run(runner.start())
                 ```
         """
+        self._heartbeat_seconds = heartbeat_seconds
+
         settings = get_current_settings()
 
         if name and ("/" in name or "%" in name):
@@ -221,12 +223,6 @@ class Runner:
 
         self.query_seconds: float = query_seconds or settings.runner.poll_frequency
         self._prefetch_seconds: float = prefetch_seconds
-        self.heartbeat_seconds: float | None = (
-            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._heartbeat_task: asyncio.Task[None] | None = None
         self._events_client: EventsClient = get_events_client(checkpoint_every=1)
 
         self._exit_stack = AsyncExitStack()
@@ -644,9 +640,6 @@ class Runner:
 
             task_status.started(process.pid)
 
-            if self.heartbeat_seconds is not None:
-                await self._emit_flow_run_heartbeat(flow_run)
-
             # Only add the process to the map if it is still running
             # The process may be a multiprocessing.context.SpawnProcess, in which case it will have an `exitcode`` attribute
             # but no `returncode` attribute
@@ -682,6 +675,11 @@ class Runner:
 
         flow_run = FlowRun.model_validate(bundle["flow_run"])
 
+        # Add heartbeat_seconds to env if configured
+        if self._heartbeat_seconds is not None:
+            env = env or {}
+            env["PREFECT_FLOWS_HEARTBEAT_FREQUENCY"] = str(int(self._heartbeat_seconds))
+
         async with context:
             if not self._acquire_limit_slot(flow_run.id):
                 return
@@ -695,9 +693,6 @@ class Runner:
                 await self._propose_crashed_state(flow_run, msg)
                 raise RuntimeError(msg)
 
-            if self.heartbeat_seconds is not None:
-                await self._emit_flow_run_heartbeat(flow_run)
-
             await self._add_flow_run_process_map_entry(
                 flow_run.id, ProcessMapEntry(pid=process.pid, flow_run=flow_run)
             )
@@ -796,7 +791,14 @@ class Runner:
         if flow_run.deployment_id is not None:
             flow = self._deployment_flow_map.get(flow_run.deployment_id)
             if flow:
-                process = run_flow_in_subprocess(flow, flow_run=flow_run)
+                subprocess_env: dict[str, str] = {}
+                if self._heartbeat_seconds is not None:
+                    subprocess_env["PREFECT_FLOWS_HEARTBEAT_FREQUENCY"] = str(
+                        int(self._heartbeat_seconds)
+                    )
+                process = run_flow_in_subprocess(
+                    flow, flow_run=flow_run, env=subprocess_env or None
+                )
                 task_status.started(process)
                 await anyio.to_thread.run_sync(process.join)
                 return process.exitcode
@@ -829,6 +831,15 @@ class Runner:
                     "PREFECT__ENABLE_CANCELLATION_AND_CRASHED_HOOKS": "false",
                 },
                 **({"PREFECT__FLOW_ENTRYPOINT": entrypoint} if entrypoint else {}),
+                **(
+                    {
+                        "PREFECT_FLOWS_HEARTBEAT_FREQUENCY": str(
+                            int(self._heartbeat_seconds)
+                        )
+                    }
+                    if self._heartbeat_seconds is not None
+                    else {}
+                ),
             }
         )
         env.update(**os.environ)  # is this really necessary??
@@ -1049,51 +1060,6 @@ class Runner:
                 flow = None
         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(deployment.as_related_resource())
-            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))
-
-        await self._events_client.emit(
-            Event(
-                event="prefect.flow-run.heartbeat",
-                resource=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__
 
@@ -1145,7 +1111,7 @@ class Runner:
                 related=related,
             )
         )
-        self._logger.debug(f"Emitted flow run heartbeat event for {flow_run.id}")
+        self._logger.debug(f"Emitted cancelled-flow-run event for {flow_run.id}")
 
     async def _get_scheduled_flow_runs(
         self,
@@ -1289,9 +1255,6 @@ class Runner:
                     flow_run.id,
                     ProcessMapEntry(pid=readiness_result.pid, 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:
@@ -1605,15 +1568,6 @@ class Runner:
         if not hasattr(self, "_loops_task_group") or not self._loops_task_group:
             self._loops_task_group: anyio.abc.TaskGroup = anyio.create_task_group()
 
-        if self.heartbeat_seconds is not None:
-            self._heartbeat_task = asyncio.create_task(
-                critical_service_loop(
-                    workload=self._emit_flow_run_heartbeats,
-                    interval=self.heartbeat_seconds,
-                    jitter_range=0.3,
-                )
-            )
-
         self.started = True
         return self
 
@@ -1632,13 +1586,6 @@ class Runner:
         shutil.rmtree(str(self._tmp_dir), ignore_errors=True)
         del self._runs_task_group, self._loops_task_group
 
-        if self._heartbeat_task:
-            self._heartbeat_task.cancel()
-            try:
-                await self._heartbeat_task
-            except asyncio.CancelledError:
-                pass
-
     def __repr__(self) -> str:
         return f"Runner(name={self.name!r})"
 

prefect/runner/storage.py

@@ -20,6 +20,7 @@ from anyio import run_process
 from pydantic import SecretStr
 
 from prefect._internal.concurrency.api import create_call, from_async
+from prefect._internal.urls import strip_auth_from_url
 from prefect.blocks.core import Block, BlockNotSavedError
 from prefect.blocks.system import Secret
 from prefect.filesystems import ReadableDeploymentStorage, WritableDeploymentStorage
@@ -344,12 +345,13 @@ class GitRepository:
                 cwd=str(self.destination),
             )
             existing_repo_url = None
-            existing_repo_url = _strip_auth_from_url(result.stdout.decode().strip())
+            existing_repo_url = strip_auth_from_url(result.stdout.decode().strip())
+            configured_repo_url = strip_auth_from_url(self._url)
 
-            if existing_repo_url != self._url:
+            if existing_repo_url != configured_repo_url:
                 raise ValueError(
                     f"The existing repository at {str(self.destination)} "
-                    f"does not match the configured repository {self._url}"
+                    f"does not match the configured repository {configured_repo_url}"
                 )
 
             # Sparsely checkout the repository if directories are specified and the repo is not in sparse-checkout mode already
@@ -458,7 +460,7 @@ class GitRepository:
                 else exc
             )
             raise RuntimeError(
-                f"Failed to clone repository {_strip_auth_from_url(self._url)!r} with exit code"
+                f"Failed to clone repository {strip_auth_from_url(self._url)!r} with exit code"
                 f" {exc.returncode}."
             ) from exc_chain
 
@@ -966,24 +968,3 @@ def _format_token_from_credentials(
 
     # GitHub and other providers: plain token
     return user_provided_token
-
-
-def _strip_auth_from_url(url: str) -> str:
-    parsed = urlparse(url)
-
-    # Construct a new netloc without the auth info
-    netloc = parsed.hostname
-    if parsed.port and netloc:
-        netloc += f":{parsed.port}"
-
-    # Build the sanitized URL
-    return urlunparse(
-        (
-            parsed.scheme,
-            netloc,
-            parsed.path,
-            parsed.params,
-            parsed.query,
-            parsed.fragment,
-        )
-    )