prefect-client 3.1.7__py3-none-any.whl → 3.1.9__py3-none-any.whl

prefect/results.py

@@ -233,6 +233,29 @@ def _format_user_supplied_storage_key(key: str) -> str:
     return key.format(**runtime_vars, parameters=prefect.runtime.task_run.parameters)
 
 
+async def _call_explicitly_async_block_method(
+    block: Union[WritableFileSystem, NullFileSystem],
+    method: str,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> Any:
+    """
+    TODO: remove this once we have explicit async methods on all storage blocks
+
+    see https://github.com/PrefectHQ/prefect/issues/15008
+    """
+    if hasattr(block, f"a{method}"):  # explicit async method
+        return await getattr(block.__class__.__name__, f"a{method}")(*args, **kwargs)
+    elif hasattr(getattr(block, method, None), "aio"):  # sync_compatible
+        return await getattr(block, method).aio(block, *args, **kwargs)
+    else:  # should not happen in prefect, but users can override impls
+        maybe_coro = getattr(block, method)(*args, **kwargs)
+        if inspect.isawaitable(maybe_coro):
+            return await maybe_coro
+        else:
+            return maybe_coro
+
+
 T = TypeVar("T")
 
 
@@ -405,7 +428,9 @@ class ResultStore(BaseModel):
             # TODO: Add an `exists` method to commonly used storage blocks
             # so the entire payload doesn't need to be read
             try:
-                metadata_content = await self.metadata_storage.read_path(key)
+                metadata_content = await _call_explicitly_async_block_method(
+                    self.metadata_storage, "read_path", (key,), {}
+                )
                 if metadata_content is None:
                     return False
                 metadata = ResultRecordMetadata.load_bytes(metadata_content)
@@ -414,7 +439,9 @@ class ResultStore(BaseModel):
                 return False
         else:
             try:
-                content = await self.result_storage.read_path(key)
+                content = await _call_explicitly_async_block_method(
+                    self.result_storage, "read_path", (key,), {}
+                )
                 if content is None:
                     return False
                 record = ResultRecord.deserialize(content)
@@ -491,12 +518,22 @@ class ResultStore(BaseModel):
             self.result_storage = await get_default_result_storage()
 
         if self.metadata_storage is not None:
-            metadata_content = await self.metadata_storage.read_path(key)
+            metadata_content = await _call_explicitly_async_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 = await self.result_storage.read_path(metadata.storage_key)
+            result_content = await _call_explicitly_async_block_method(
+                self.result_storage,
+                "read_path",
+                (metadata.storage_key,),
+                {},
+            )
             result_record: ResultRecord[
                 Any
             ] = ResultRecord.deserialize_from_result_and_metadata(
@@ -504,7 +541,12 @@ class ResultStore(BaseModel):
             )
             await emit_result_read_event(self, resolved_key_path)
         else:
-            content = await self.result_storage.read_path(key)
+            content = await _call_explicitly_async_block_method(
+                self.result_storage,
+                "read_path",
+                (key,),
+                {},
+            )
             result_record: ResultRecord[Any] = ResultRecord.deserialize(
                 content, backup_serializer=self.serializer
             )
@@ -555,7 +597,7 @@ class ResultStore(BaseModel):
         obj: Any,
         key: Optional[str] = None,
         expiration: Optional[DateTime] = None,
-    ) -> "ResultRecord":
+    ) -> "ResultRecord[Any]":
         """
         Create a result record.
 
@@ -671,19 +713,26 @@ class ResultStore(BaseModel):
 
         # If metadata storage is configured, write result and metadata separately
         if self.metadata_storage is not None:
-            await self.result_storage.write_path(
-                result_record.metadata.storage_key,
-                content=result_record.serialize_result(),
+            await _call_explicitly_async_block_method(
+                self.result_storage,
+                "write_path",
+                (result_record.metadata.storage_key,),
+                {"content": result_record.serialize_result()},
             )
-            await self.metadata_storage.write_path(
-                base_key,
-                content=result_record.serialize_metadata(),
+            await _call_explicitly_async_block_method(
+                self.metadata_storage,
+                "write_path",
+                (base_key,),
+                {"content": result_record.serialize_metadata()},
             )
             await emit_result_write_event(self, result_record.metadata.storage_key)
         # Otherwise, write the result metadata and result together
         else:
-            await self.result_storage.write_path(
-                result_record.metadata.storage_key, content=result_record.serialize()
+            await _call_explicitly_async_block_method(
+                self.result_storage,
+                "write_path",
+                (result_record.metadata.storage_key,),
+                {"content": result_record.serialize()},
             )
             await emit_result_write_event(self, result_record.metadata.storage_key)
         if self.cache_result_in_memory:
@@ -910,8 +959,11 @@ class ResultStore(BaseModel):
                 serializer=self.serializer, storage_key=str(identifier)
             ),
         )
-        await self.result_storage.write_path(
-            f"parameters/{identifier}", content=record.serialize()
+        await _call_explicitly_async_block_method(
+            self.result_storage,
+            "write_path",
+            (f"parameters/{identifier}",),
+            {"content": record.serialize()},
         )
 
     @sync_compatible
@@ -921,7 +973,12 @@ class ResultStore(BaseModel):
                 "Result store is not configured - must have a result storage block to read parameters"
             )
         record = ResultRecord.deserialize(
-            await self.result_storage.read_path(f"parameters/{identifier}")
+            await _call_explicitly_async_block_method(
+                self.result_storage,
+                "read_path",
+                (f"parameters/{identifier}",),
+                {},
+            )
         )
         return record.result
 
@@ -976,7 +1033,7 @@ class ResultRecordMetadata(BaseModel):
         """
         return cls.model_validate_json(data)
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         if not isinstance(other, ResultRecordMetadata):
             return False
         return (
@@ -1050,7 +1107,7 @@ class ResultRecord(BaseModel, Generic[R]):
 
     @model_validator(mode="before")
     @classmethod
-    def coerce_old_format(cls, value: Any):
+    def coerce_old_format(cls, value: Any) -> Any:
         if isinstance(value, dict):
             if "data" in value:
                 value["result"] = value.pop("data")

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:
@@ -1098,9 +1209,8 @@ class Runner:
                 )
                 # Mark the task as started to prevent agent crash
                 task_status.started(exc)
-                await self._propose_crashed_state(
-                    flow_run, "Flow run process could not be started"
-                )
+                message = f"Flow run process could not be started:\n{exc!r}"
+                await self._propose_crashed_state(flow_run, message)
             else:
                 run_logger.exception(
                     f"An error occurred while monitoring flow run '{flow_run.id}'. "

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/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`).",

prefect/states.py

@@ -11,6 +11,7 @@ from typing import Any, Dict, Iterable, Optional, Type
 import anyio
 import httpx
 import pendulum
+from opentelemetry import propagate
 from typing_extensions import TypeGuard
 
 from prefect._internal.compatibility import deprecated
@@ -588,6 +589,16 @@ class StateGroup:
         return f"StateGroup<{self.counts_message()}>"
 
 
+def _traced(cls: Type[State[R]], **kwargs: Any) -> State[R]:
+    state_details = StateDetails.model_validate(kwargs.pop("state_details", {}))
+
+    carrier = {}
+    propagate.inject(carrier)
+    state_details.traceparent = carrier.get("traceparent")
+
+    return cls(**kwargs, state_details=state_details)
+
+
 def Scheduled(
     cls: Type[State[R]] = State,
     scheduled_time: Optional[datetime.datetime] = None,
@@ -605,7 +616,7 @@ def Scheduled(
         raise ValueError("An extra scheduled_time was provided in state_details")
     state_details.scheduled_time = scheduled_time
 
-    return cls(type=StateType.SCHEDULED, state_details=state_details, **kwargs)
+    return _traced(cls, type=StateType.SCHEDULED, state_details=state_details, **kwargs)
 
 
 def Completed(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
@@ -614,7 +625,8 @@ def Completed(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     Returns:
         State: a Completed state
     """
-    return cls(type=StateType.COMPLETED, **kwargs)
+
+    return _traced(cls, type=StateType.COMPLETED, **kwargs)
 
 
 def Running(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
@@ -623,7 +635,7 @@ def Running(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     Returns:
         State: a Running state
     """
-    return cls(type=StateType.RUNNING, **kwargs)
+    return _traced(cls, type=StateType.RUNNING, **kwargs)
 
 
 def Failed(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
@@ -632,7 +644,7 @@ def Failed(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     Returns:
         State: a Failed state
     """
-    return cls(type=StateType.FAILED, **kwargs)
+    return _traced(cls, type=StateType.FAILED, **kwargs)
 
 
 def Crashed(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
@@ -641,7 +653,7 @@ def Crashed(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     Returns:
         State: a Crashed state
     """
-    return cls(type=StateType.CRASHED, **kwargs)
+    return _traced(cls, type=StateType.CRASHED, **kwargs)
 
 
 def Cancelling(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
@@ -650,7 +662,7 @@ def Cancelling(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     Returns:
         State: a Cancelling state
     """
-    return cls(type=StateType.CANCELLING, **kwargs)
+    return _traced(cls, type=StateType.CANCELLING, **kwargs)
 
 
 def Cancelled(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
@@ -659,7 +671,7 @@ def Cancelled(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     Returns:
         State: a Cancelled state
     """
-    return cls(type=StateType.CANCELLED, **kwargs)
+    return _traced(cls, type=StateType.CANCELLED, **kwargs)
 
 
 def Pending(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
@@ -668,7 +680,7 @@ def Pending(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     Returns:
         State: a Pending state
     """
-    return cls(type=StateType.PENDING, **kwargs)
+    return _traced(cls, type=StateType.PENDING, **kwargs)
 
 
 def Paused(
@@ -704,7 +716,7 @@ def Paused(
     state_details.pause_reschedule = reschedule
     state_details.pause_key = pause_key
 
-    return cls(type=StateType.PAUSED, state_details=state_details, **kwargs)
+    return _traced(cls, type=StateType.PAUSED, state_details=state_details, **kwargs)
 
 
 def Suspended(
@@ -766,7 +778,7 @@ def Retrying(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     Returns:
         State: a Retrying state
     """
-    return cls(type=StateType.RUNNING, name="Retrying", **kwargs)
+    return _traced(cls, type=StateType.RUNNING, name="Retrying", **kwargs)
 
 
 def Late(

prefect/task_engine.py

@@ -182,7 +182,7 @@ class BaseTaskRunEngine(Generic[P, R]):
                     return_data=True,
                     max_depth=-1,
                     remove_annotations=True,
-                    context={},
+                    context={"parameter_name": parameter},
                 )
             except UpstreamTaskError:
                 raise