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

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/runner/submit.py

@@ -42,11 +42,8 @@ async def _submit_flow_to_runner(
     Returns:
         A `FlowRun` object representing the flow run that was submitted.
     """
-    from prefect.utilities.engine import (
-        _dynamic_key_for_task_run,
-        collect_task_run_inputs,
-        resolve_inputs,
-    )
+    from prefect.utilities._engine import dynamic_key_for_task_run
+    from prefect.utilities.engine import collect_task_run_inputs, resolve_inputs
 
     async with get_client() as client:
         if not retry_failed_submissions:
@@ -67,7 +64,7 @@ async def _submit_flow_to_runner(
                 parent_flow_run_context.flow_run.id if parent_flow_run_context else None
             ),
             dynamic_key=(
-                _dynamic_key_for_task_run(parent_flow_run_context, dummy_task)
+                dynamic_key_for_task_run(parent_flow_run_context, dummy_task)
                 if parent_flow_run_context
                 else str(uuid.uuid4())
             ),

prefect/runner/utils.py

@@ -1,5 +1,5 @@
 from copy import deepcopy
-from typing import Any, Dict
+from typing import Any
 
 from fastapi import FastAPI
 from fastapi.openapi.utils import get_openapi
@@ -8,8 +8,8 @@ from prefect import __version__ as PREFECT_VERSION
 
 
 def inject_schemas_into_openapi(
-    webserver: FastAPI, schemas_to_inject: Dict[str, Any]
-) -> Dict[str, Any]:
+    webserver: FastAPI, schemas_to_inject: dict[str, Any]
+) -> dict[str, Any]:
     """
     Augments the webserver's OpenAPI schema with additional schemas from deployments / flows / tasks.
 
@@ -29,8 +29,8 @@ def inject_schemas_into_openapi(
 
 
 def merge_definitions(
-    injected_schemas: Dict[str, Any], openapi_schema: Dict[str, Any]
-) -> Dict[str, Any]:
+    injected_schemas: dict[str, Any], openapi_schema: dict[str, Any]
+) -> dict[str, Any]:
     """
     Integrates definitions from injected schemas into the OpenAPI components.
 
@@ -69,7 +69,7 @@ def update_refs_in_schema(schema_item: Any, new_ref: str) -> None:
             update_refs_in_schema(item, new_ref)
 
 
-def update_refs_to_components(openapi_schema: Dict[str, Any]) -> Dict[str, Any]:
+def update_refs_to_components(openapi_schema: dict[str, Any]) -> dict[str, Any]:
     """
     Updates all `$ref` fields in the OpenAPI schema to reference the components section.
 

prefect/runtime/flow_run.py

@@ -44,6 +44,7 @@ __all__ = [
     "run_count",
     "api_url",
     "ui_url",
+    "job_variables",
 ]
 
 
@@ -317,6 +318,11 @@ def get_flow_run_ui_url() -> Optional[str]:
     return f"{PREFECT_UI_URL.value()}/flow-runs/flow-run/{flow_run_id}"
 
 
+def get_job_variables() -> Optional[Dict[str, Any]]:
+    flow_run_ctx = FlowRunContext.get()
+    return flow_run_ctx.flow_run.job_variables if flow_run_ctx else None
+
+
 FIELDS = {
     "id": get_id,
     "tags": get_tags,
@@ -331,4 +337,5 @@ FIELDS = {
     "api_url": get_flow_run_api_url,
     "ui_url": get_flow_run_ui_url,
     "flow_version": get_flow_version,
+    "job_variables": get_job_variables,
 }

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/constants.py

@@ -1,8 +1,8 @@
 from pathlib import Path
-from typing import Tuple, Type
+from typing import Any, Tuple, Type
 
 from pydantic import Secret, SecretStr
 
 DEFAULT_PREFECT_HOME = Path.home() / ".prefect"
 DEFAULT_PROFILES_PATH = Path(__file__).parent.joinpath("profiles.toml")
-_SECRET_TYPES: Tuple[Type, ...] = (Secret, SecretStr)
+_SECRET_TYPES: Tuple[Type[Any], ...] = (Secret, SecretStr)

prefect/settings/legacy.py

@@ -8,7 +8,7 @@ from pydantic_settings import BaseSettings
 from typing_extensions import Self
 
 from prefect.settings.base import PrefectBaseSettings
-from prefect.settings.constants import _SECRET_TYPES
+from prefect.settings.constants import _SECRET_TYPES  # type: ignore[reportPrivateUsage]
 from prefect.settings.context import get_current_settings
 from prefect.settings.models.root import Settings
 

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/server/events.py

@@ -135,3 +135,13 @@ class ServerEventsSettings(PrefectBaseSettings):
             "prefect_messaging_cache",
         ),
     )
+
+    maximum_event_name_length: int = Field(
+        default=1024,
+        gt=0,
+        description="The maximum length of an event name.",
+        validation_alias=AliasChoices(
+            AliasPath("maximum_event_name_length"),
+            "prefect_server_events_maximum_event_name_length",
+        ),
+    )

prefect/task_engine.py

@@ -4,7 +4,7 @@ import logging
 import threading
 import time
 from asyncio import CancelledError
-from contextlib import ExitStack, asynccontextmanager, contextmanager
+from contextlib import ExitStack, asynccontextmanager, contextmanager, nullcontext
 from dataclasses import dataclass, field
 from functools import partial
 from textwrap import dedent
@@ -29,6 +29,7 @@ from uuid import UUID
 
 import anyio
 import pendulum
+from opentelemetry import trace
 from typing_extensions import ParamSpec
 
 from prefect import Task
@@ -79,13 +80,14 @@ from prefect.states import (
     exception_to_failed_state,
     return_value_to_state,
 )
+from prefect.telemetry.run_telemetry import RunTelemetry
 from prefect.transactions import IsolationLevel, Transaction, transaction
+from prefect.utilities._engine import get_hook_name
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.callables import call_with_parameters, parameters_to_args_kwargs
 from prefect.utilities.collections import visit_collection
 from prefect.utilities.engine import (
-    _get_hook_name,
     emit_task_run_state_change_event,
     link_state_to_result,
     resolve_to_final_result,
@@ -120,6 +122,7 @@ class BaseTaskRunEngine(Generic[P, R]):
     _is_started: bool = False
     _task_name_set: bool = False
     _last_event: Optional[PrefectEvent] = None
+    _telemetry: RunTelemetry = field(default_factory=RunTelemetry)
 
     def __post_init__(self):
         if self.parameters is None:
@@ -193,11 +196,11 @@ class BaseTaskRunEngine(Generic[P, R]):
         self.parameters = resolved_parameters
 
     def _set_custom_task_run_name(self):
-        from prefect.utilities.engine import _resolve_custom_task_run_name
+        from prefect.utilities._engine import resolve_custom_task_run_name
 
         # update the task run name if necessary
         if not self._task_name_set and self.task.task_run_name:
-            task_run_name = _resolve_custom_task_run_name(
+            task_run_name = resolve_custom_task_run_name(
                 task=self.task, parameters=self.parameters or {}
             )
 
@@ -351,7 +354,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             hooks = None
 
         for hook in hooks or []:
-            hook_name = _get_hook_name(hook)
+            hook_name = get_hook_name(hook)
 
             try:
                 self.logger.info(
@@ -427,6 +430,11 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
         self.task_run.state = new_state = state
 
+        if last_state.timestamp == new_state.timestamp:
+            # Ensure that the state timestamp is unique, or at least not equal to the last state.
+            # This might occur especially on Windows where the timestamp resolution is limited.
+            new_state.timestamp += pendulum.duration(microseconds=1)
+
         # Ensure that the state_details are populated with the current run IDs
         new_state.state_details.task_run_id = self.task_run.id
         new_state.state_details.flow_run_id = self.task_run.flow_run_id
@@ -460,7 +468,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             validated_state=self.task_run.state,
             follows=self._last_event,
         )
-
+        self._telemetry.update_state(new_state)
         return new_state
 
     def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
@@ -514,6 +522,8 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self.record_terminal_state_timing(terminal_state)
         self.set_state(terminal_state)
         self._return_value = result
+
+        self._telemetry.end_span_on_success()
         return result
 
     def handle_retry(self, exc: Exception) -> bool:
@@ -562,6 +572,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
     def handle_exception(self, exc: Exception) -> None:
         # If the task fails, and we have retries left, set the task to retrying.
+        self._telemetry.record_exception(exc)
         if not self.handle_retry(exc):
             # If the task has no retries left, or the retry condition is not met, set the task to failed.
             state = run_coro_as_sync(
@@ -575,6 +586,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             self.record_terminal_state_timing(state)
             self.set_state(state)
             self._raised = exc
+            self._telemetry.end_span_on_failure(state.message if state else None)
 
     def handle_timeout(self, exc: TimeoutError) -> None:
         if not self.handle_retry(exc):
@@ -588,6 +600,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 message=message,
                 name="TimedOut",
             )
+            self.record_terminal_state_timing(state)
             self.set_state(state)
             self._raised = exc
 
@@ -598,6 +611,8 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self.record_terminal_state_timing(state)
         self.set_state(state, force=True)
         self._raised = exc
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(state.message if state else None)
 
     @contextmanager
     def setup_run_context(self, client: Optional[SyncPrefectClient] = None):
@@ -655,14 +670,17 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             with SyncClientContext.get_or_create() as client_ctx:
                 self._client = client_ctx.client
                 self._is_started = True
+                parent_flow_run_context = FlowRunContext.get()
+                parent_task_run_context = TaskRunContext.get()
+
                 try:
                     if not self.task_run:
                         self.task_run = run_coro_as_sync(
                             self.task.create_local_run(
                                 id=task_run_id,
                                 parameters=self.parameters,
-                                flow_run_context=FlowRunContext.get(),
-                                parent_task_run_context=TaskRunContext.get(),
+                                flow_run_context=parent_flow_run_context,
+                                parent_task_run_context=parent_task_run_context,
                                 wait_for=self.wait_for,
                                 extra_task_inputs=dependencies,
                             )
@@ -679,6 +697,14 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                         self.logger.debug(
                             f"Created task run {self.task_run.name!r} for task {self.task.name!r}"
                         )
+
+                        self._telemetry.start_span(
+                            run=self.task_run,
+                            name=self.task.name,
+                            client=self.client,
+                            parameters=self.parameters,
+                        )
+
                         yield self
 
                 except TerminationSignal as exc:
@@ -730,11 +756,14 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     ) -> Generator[None, None, None]:
         with self.initialize_run(task_run_id=task_run_id, dependencies=dependencies):
-            self.begin_run()
-            try:
-                yield
-            finally:
-                self.call_hooks()
+            with trace.use_span(
+                self._telemetry.span
+            ) if self._telemetry.span else nullcontext():
+                self.begin_run()
+                try:
+                    yield
+                finally:
+                    self.call_hooks()
 
     @contextmanager
     def transaction_context(self) -> Generator[Transaction, None, None]:
@@ -866,7 +895,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             hooks = None
 
         for hook in hooks or []:
-            hook_name = _get_hook_name(hook)
+            hook_name = get_hook_name(hook)
 
             try:
                 self.logger.info(
@@ -942,6 +971,11 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
         self.task_run.state = new_state = state
 
+        if last_state.timestamp == new_state.timestamp:
+            # Ensure that the state timestamp is unique, or at least not equal to the last state.
+            # This might occur especially on Windows where the timestamp resolution is limited.
+            new_state.timestamp += pendulum.duration(microseconds=1)
+
         # Ensure that the state_details are populated with the current run IDs
         new_state.state_details.task_run_id = self.task_run.id
         new_state.state_details.flow_run_id = self.task_run.flow_run_id
@@ -977,6 +1011,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             follows=self._last_event,
         )
 
+        self._telemetry.update_state(new_state)
         return new_state
 
     async def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
@@ -1025,6 +1060,9 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self.record_terminal_state_timing(terminal_state)
         await self.set_state(terminal_state)
         self._return_value = result
+
+        self._telemetry.end_span_on_success()
+
         return result
 
     async def handle_retry(self, exc: Exception) -> bool:
@@ -1073,6 +1111,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
     async def handle_exception(self, exc: Exception) -> None:
         # If the task fails, and we have retries left, set the task to retrying.
+        self._telemetry.record_exception(exc)
         if not await self.handle_retry(exc):
             # If the task has no retries left, or the retry condition is not met, set the task to failed.
             state = await exception_to_failed_state(
@@ -1084,7 +1123,10 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             await self.set_state(state)
             self._raised = exc
 
+            self._telemetry.end_span_on_failure(state.message)
+
     async def handle_timeout(self, exc: TimeoutError) -> None:
+        self._telemetry.record_exception(exc)
         if not await self.handle_retry(exc):
             if isinstance(exc, TaskRunTimeoutError):
                 message = f"Task run exceeded timeout of {self.task.timeout_seconds} second(s)"
@@ -1096,8 +1138,10 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 message=message,
                 name="TimedOut",
             )
+            self.record_terminal_state_timing(state)
             await self.set_state(state)
             self._raised = exc
+            self._telemetry.end_span_on_failure(state.message)
 
     async def handle_crash(self, exc: BaseException) -> None:
         state = await exception_to_crashed_state(exc)
@@ -1107,6 +1151,9 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         await self.set_state(state, force=True)
         self._raised = exc
 
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(state.message)
+
     @asynccontextmanager
     async def setup_run_context(self, client: Optional[PrefectClient] = None):
         from prefect.utilities.engine import (
@@ -1162,13 +1209,16 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             async with AsyncClientContext.get_or_create():
                 self._client = get_client()
                 self._is_started = True
+                parent_flow_run_context = FlowRunContext.get()
+                parent_task_run_context = TaskRunContext.get()
+
                 try:
                     if not self.task_run:
                         self.task_run = await self.task.create_local_run(
                             id=task_run_id,
                             parameters=self.parameters,
-                            flow_run_context=FlowRunContext.get(),
-                            parent_task_run_context=TaskRunContext.get(),
+                            flow_run_context=parent_flow_run_context,
+                            parent_task_run_context=parent_task_run_context,
                             wait_for=self.wait_for,
                             extra_task_inputs=dependencies,
                         )
@@ -1184,6 +1234,14 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                         self.logger.debug(
                             f"Created task run {self.task_run.name!r} for task {self.task.name!r}"
                         )
+
+                        await self._telemetry.async_start_span(
+                            run=self.task_run,
+                            name=self.task.name,
+                            client=self.client,
+                            parameters=self.parameters,
+                        )
+
                         yield self
 
                 except TerminationSignal as exc:
@@ -1237,11 +1295,14 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         async with self.initialize_run(
             task_run_id=task_run_id, dependencies=dependencies
         ):
-            await self.begin_run()
-            try:
-                yield
-            finally:
-                await self.call_hooks()
+            with trace.use_span(
+                self._telemetry.span
+            ) if self._telemetry.span else nullcontext():
+                await self.begin_run()
+                try:
+                    yield
+                finally:
+                    await self.call_hooks()
 
     @asynccontextmanager
     async def transaction_context(self) -> AsyncGenerator[Transaction, None]:
@@ -1317,7 +1378,7 @@ def run_task_sync(
     task_run_id: Optional[UUID] = None,
     task_run: Optional[TaskRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     return_type: Literal["state", "result"] = "result",
     dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     context: Optional[Dict[str, Any]] = None,
@@ -1344,7 +1405,7 @@ async def run_task_async(
     task_run_id: Optional[UUID] = None,
     task_run: Optional[TaskRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     return_type: Literal["state", "result"] = "result",
     dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     context: Optional[Dict[str, Any]] = None,
@@ -1371,7 +1432,7 @@ def run_generator_task_sync(
     task_run_id: Optional[UUID] = None,
     task_run: Optional[TaskRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     return_type: Literal["state", "result"] = "result",
     dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     context: Optional[Dict[str, Any]] = None,
@@ -1426,7 +1487,7 @@ async def run_generator_task_async(
     task_run_id: Optional[UUID] = None,
     task_run: Optional[TaskRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     return_type: Literal["state", "result"] = "result",
     dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     context: Optional[Dict[str, Any]] = None,
@@ -1482,7 +1543,7 @@ def run_task(
     task_run_id: Optional[UUID] = None,
     task_run: Optional[TaskRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     return_type: Literal["state", "result"] = "result",
     dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     context: Optional[Dict[str, Any]] = None,

prefect/task_runners.py

@@ -97,9 +97,9 @@ class TaskRunner(abc.ABC, Generic[F]):
 
     def map(
         self,
-        task: "Task",
+        task: "Task[P, R]",
         parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     ) -> PrefectFutureList[F]:
         """
         Submit multiple tasks to the task run engine.