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

prefect/futures.py

@@ -1,17 +1,15 @@
 import abc
 import asyncio
-import collections
 import concurrent.futures
 import threading
 import uuid
 from collections.abc import Generator, Iterator
 from functools import partial
-from typing import Any, Callable, Generic, List, Optional, Set, Union, cast
+from typing import Any, Callable, Generic, Optional, Union
 
-from typing_extensions import TypeVar
+from typing_extensions import NamedTuple, Self, TypeVar
 
 from prefect.client.orchestration import get_client
-from prefect.client.schemas.objects import TaskRun
 from prefect.exceptions import ObjectNotFound
 from prefect.logging.loggers import get_logger, get_run_logger
 from prefect.states import Pending, State
@@ -50,7 +48,7 @@ class PrefectFuture(abc.ABC, Generic[R]):
             return self._final_state
         client = get_client(sync_client=True)
         try:
-            task_run = cast(TaskRun, client.read_task_run(task_run_id=self.task_run_id))
+            task_run = client.read_task_run(task_run_id=self.task_run_id)
         except ObjectNotFound:
             # We'll be optimistic and assume this task will eventually start
             # TODO: Consider using task run events to wait for the task to start
@@ -92,7 +90,7 @@ class PrefectFuture(abc.ABC, Generic[R]):
         """
 
     @abc.abstractmethod
-    def add_done_callback(self, fn):
+    def add_done_callback(self, fn: Callable[["PrefectFuture[R]"], None]):
         """
         Add a callback to be run when the future completes or is cancelled.
 
@@ -102,13 +100,17 @@ class PrefectFuture(abc.ABC, Generic[R]):
         ...
 
 
-class PrefectWrappedFuture(PrefectFuture, abc.ABC, Generic[R, F]):
+class PrefectWrappedFuture(PrefectFuture[R], abc.ABC, Generic[R, F]):
     """
     A Prefect future that wraps another future object.
+
+    Type Parameters:
+        R: The return type of the future
+        F: The type of the wrapped future
     """
 
     def __init__(self, task_run_id: uuid.UUID, wrapped_future: F):
-        self._wrapped_future = wrapped_future
+        self._wrapped_future: F = wrapped_future
         super().__init__(task_run_id)
 
     @property
@@ -116,10 +118,11 @@ class PrefectWrappedFuture(PrefectFuture, abc.ABC, Generic[R, F]):
         """The underlying future object wrapped by this Prefect future"""
         return self._wrapped_future
 
-    def add_done_callback(self, fn: Callable[[PrefectFuture[R]], None]):
+    def add_done_callback(self, fn: Callable[[PrefectFuture[R]], None]) -> None:
+        """Add a callback to be executed when the future completes."""
         if not self._final_state:
 
-            def call_with_self(future):
+            def call_with_self(future: F):
                 """Call the callback with self as the argument, this is necessary to ensure we remove the future from the pending set"""
                 fn(self)
 
@@ -128,7 +131,7 @@ class PrefectWrappedFuture(PrefectFuture, abc.ABC, Generic[R, F]):
         fn(self)
 
 
-class PrefectConcurrentFuture(PrefectWrappedFuture[R, concurrent.futures.Future]):
+class PrefectConcurrentFuture(PrefectWrappedFuture[R, concurrent.futures.Future[R]]):
     """
     A Prefect future that wraps a concurrent.futures.Future. This future is used
     when the task run is submitted to a ThreadPoolExecutor.
@@ -193,7 +196,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
     any task run scheduled in Prefect's API.
     """
 
-    done_callbacks: List[Callable[[PrefectFuture[R]], None]] = []
+    done_callbacks: list[Callable[[PrefectFuture[R]], None]] = []
     waiter = None
 
     def wait(self, timeout: Optional[float] = None) -> None:
@@ -270,7 +273,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
                 return
             TaskRunWaiter.add_done_callback(self._task_run_id, partial(fn, self))
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         if not isinstance(other, PrefectDistributedFuture):
             return False
         return self.task_run_id == other.task_run_id
@@ -279,7 +282,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
         return hash(self.task_run_id)
 
 
-class PrefectFutureList(list, Iterator, Generic[F]):
+class PrefectFutureList(list[PrefectFuture[R]], Iterator[PrefectFuture[R]]):
     """
     A list of Prefect futures.
 
@@ -298,10 +301,10 @@ class PrefectFutureList(list, Iterator, Generic[F]):
         wait(self, timeout=timeout)
 
     def result(
-        self: "PrefectFutureList[R]",
+        self: Self,
         timeout: Optional[float] = None,
         raise_on_failure: bool = True,
-    ) -> List[R]:
+    ) -> list[R]:
         """
         Get the results of all task runs associated with the futures in the list.
 
@@ -331,21 +334,22 @@ class PrefectFutureList(list, Iterator, Generic[F]):
 
 
 def as_completed(
-    futures: List[PrefectFuture[R]], timeout: Optional[float] = None
+    futures: list[PrefectFuture[R]], timeout: Optional[float] = None
 ) -> Generator[PrefectFuture[R], None]:
-    unique_futures: Set[PrefectFuture[R]] = set(futures)
+    unique_futures: set[PrefectFuture[R]] = set(futures)
     total_futures = len(unique_futures)
+    pending = unique_futures
     try:
         with timeout_context(timeout):
-            done = {f for f in unique_futures if f._final_state}
+            done = {f for f in unique_futures if f._final_state}  # type: ignore[privateUsage]
             pending = unique_futures - done
             yield from done
 
             finished_event = threading.Event()
             finished_lock = threading.Lock()
-            finished_futures = []
+            finished_futures: list[PrefectFuture[R]] = []
 
-            def add_to_done(future):
+            def add_to_done(future: PrefectFuture[R]):
                 with finished_lock:
                     finished_futures.append(future)
                     finished_event.set()
@@ -370,10 +374,19 @@ def as_completed(
         )
 
 
-DoneAndNotDoneFutures = collections.namedtuple("DoneAndNotDoneFutures", "done not_done")
+class DoneAndNotDoneFutures(NamedTuple, Generic[R]):
+    """A named 2-tuple of sets.
+
+    multiple inheritance supported in 3.11+, use typing_extensions.NamedTuple
+    """
+
+    done: set[PrefectFuture[R]]
+    not_done: set[PrefectFuture[R]]
 
 
-def wait(futures: List[PrefectFuture[R]], timeout=None) -> DoneAndNotDoneFutures:
+def wait(
+    futures: list[PrefectFuture[R]], timeout: Optional[float] = None
+) -> DoneAndNotDoneFutures[R]:
     """
     Wait for the futures in the given sequence to complete.
 
@@ -431,9 +444,11 @@ def resolve_futures_to_states(
 
     Unsupported object types will be returned without modification.
     """
-    futures: Set[PrefectFuture[R]] = set()
+    futures: set[PrefectFuture[R]] = set()
 
-    def _collect_futures(futures, expr, context):
+    def _collect_futures(
+        futures: set[PrefectFuture[R]], expr: Any, context: Any
+    ) -> Union[PrefectFuture[R], Any]:
         # Expressions inside quotes should not be traversed
         if isinstance(context.get("annotation"), quote):
             raise StopVisiting()
@@ -455,14 +470,14 @@ def resolve_futures_to_states(
         return expr
 
     # Get final states for each future
-    states = []
+    states: list[State] = []
     for future in futures:
         future.wait()
         states.append(future.state)
 
     states_by_future = dict(zip(futures, states))
 
-    def replace_futures_with_states(expr, context):
+    def replace_futures_with_states(expr: Any, context: Any) -> Any:
         # Expressions inside quotes should not be modified
         if isinstance(context.get("annotation"), quote):
             raise StopVisiting()

prefect/infrastructure/__init__.py

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

prefect/infrastructure/base.py

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

prefect/locking/filesystem.py

@@ -1,6 +1,7 @@
 import time
+from logging import Logger
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Optional
 
 import anyio
 import pendulum
@@ -11,7 +12,7 @@ from prefect.logging.loggers import get_logger
 
 from .protocol import LockManager
 
-logger = get_logger(__name__)
+logger: Logger = get_logger(__name__)
 
 
 class _LockInfo(TypedDict):
@@ -37,11 +38,11 @@ class FileSystemLockManager(LockManager):
         lock_files_directory: the directory where lock files are stored
     """
 
-    def __init__(self, lock_files_directory: Path):
-        self.lock_files_directory = lock_files_directory.expanduser().resolve()
-        self._locks: Dict[str, _LockInfo] = {}
+    def __init__(self, lock_files_directory: Path) -> None:
+        self.lock_files_directory: Path = lock_files_directory.expanduser().resolve()
+        self._locks: dict[str, _LockInfo] = {}
 
-    def _ensure_lock_files_directory_exists(self):
+    def _ensure_lock_files_directory_exists(self) -> None:
         self.lock_files_directory.mkdir(parents=True, exist_ok=True)
 
     def _lock_path_for_key(self, key: str) -> Path:
@@ -49,7 +50,7 @@ class FileSystemLockManager(LockManager):
             return lock_info["path"]
         return self.lock_files_directory.joinpath(key).with_suffix(".lock")
 
-    def _get_lock_info(self, key: str, use_cache=True) -> Optional[_LockInfo]:
+    def _get_lock_info(self, key: str, use_cache: bool = True) -> Optional[_LockInfo]:
         if use_cache:
             if (lock_info := self._locks.get(key)) is not None:
                 return lock_info

prefect/locking/memory.py

@@ -1,6 +1,8 @@
 import asyncio
 import threading
-from typing import Dict, Optional, TypedDict
+from typing import Any, Optional, TypedDict
+
+from typing_extensions import Self
 
 from .protocol import LockManager
 
@@ -30,14 +32,14 @@ class MemoryLockManager(LockManager):
 
     _instance = None
 
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls, *args: Any, **kwargs: Any) -> Self:
         if cls._instance is None:
             cls._instance = super().__new__(cls)
         return cls._instance
 
     def __init__(self):
         self._locks_dict_lock = threading.Lock()
-        self._locks: Dict[str, _LockInfo] = {}
+        self._locks: dict[str, _LockInfo] = {}
 
     def _expire_lock(self, key: str):
         """

prefect/locking/protocol.py

@@ -57,7 +57,7 @@ class LockManager(Protocol):
         """
         ...
 
-    def release_lock(self, key: str, holder: str):
+    def release_lock(self, key: str, holder: str) -> None:
         """
         Releases the lock on the corresponding transaction record.
 

prefect/plugins.py

@@ -9,15 +9,15 @@ Currently supported entrypoints:
 """
 
 from types import ModuleType
-from typing import Any, Dict, Union
+from typing import Any, Union
 
 import prefect.settings
 from prefect.utilities.compat import EntryPoints, entry_points
 
-COLLECTIONS: Union[None, Dict[str, Union[ModuleType, Exception]]] = None
+_collections: Union[None, dict[str, Union[ModuleType, Exception]]] = None
 
 
-def safe_load_entrypoints(entrypoints: EntryPoints) -> Dict[str, Union[Exception, Any]]:
+def safe_load_entrypoints(entrypoints: EntryPoints) -> dict[str, Union[Exception, Any]]:
     """
     Load entry points for a group capturing any exceptions that occur.
     """
@@ -26,7 +26,7 @@ def safe_load_entrypoints(entrypoints: EntryPoints) -> Dict[str, Union[Exception
     #       also want to validate the type for the group for entrypoints that have
     #       a specific type we expect.
 
-    results = {}
+    results: dict[str, Union[Exception, Any]] = {}
 
     for entrypoint in entrypoints:
         result = None
@@ -40,18 +40,20 @@ def safe_load_entrypoints(entrypoints: EntryPoints) -> Dict[str, Union[Exception
     return results
 
 
-def load_prefect_collections() -> Dict[str, Union[ModuleType, Exception]]:
+def load_prefect_collections() -> dict[str, Union[ModuleType, Exception]]:
     """
     Load all Prefect collections that define an entrypoint in the group
     `prefect.collections`.
     """
-    global COLLECTIONS
+    global _collections
 
-    if COLLECTIONS is not None:
-        return COLLECTIONS
+    if _collections is not None:
+        return _collections
 
     collection_entrypoints: EntryPoints = entry_points(group="prefect.collections")
-    collections = safe_load_entrypoints(collection_entrypoints)
+    collections: dict[str, Union[Exception, Any]] = safe_load_entrypoints(
+        collection_entrypoints
+    )
 
     # TODO: Consider the utility of this once we've established this pattern.
     #       We cannot use a logger here because logging is not yet initialized.
@@ -68,5 +70,5 @@ def load_prefect_collections() -> Dict[str, Union[ModuleType, Exception]]:
             if prefect.settings.PREFECT_DEBUG_MODE:
                 print(f"Loaded collection {name!r}.")
 
-    COLLECTIONS = collections
+    _collections = collections
     return collections

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, 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

@@ -1209,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/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