prefect-client 3.4.4.dev2__py3-none-any.whl → 3.4.5__py3-none-any.whl

prefect/task_engine.py

@@ -43,6 +43,7 @@ from prefect.concurrency.v1.asyncio import concurrency as aconcurrency
 from prefect.concurrency.v1.context import ConcurrencyContext as ConcurrencyContextV1
 from prefect.concurrency.v1.sync import concurrency
 from prefect.context import (
+    AssetContext,
     AsyncClientContext,
     FlowRunContext,
     SyncClientContext,
@@ -314,10 +315,13 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             raise RuntimeError("Engine has not started.")
         return self._client
 
-    def can_retry(self, exc: Exception) -> bool:
+    def can_retry(self, exc_or_state: Exception | State[R]) -> bool:
         retry_condition: Optional[
-            Callable[["Task[P, Coroutine[Any, Any, R]]", TaskRun, State], bool]
+            Callable[["Task[P, Coroutine[Any, Any, R]]", TaskRun, State[R]], bool]
         ] = self.task.retry_condition_fn
+
+        failure_type = "exception" if isinstance(exc_or_state, Exception) else "state"
+
         if not self.task_run:
             raise ValueError("Task run is not set")
         try:
@@ -326,8 +330,8 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 f" {self.task.name!r}"
             )
             state = Failed(
-                data=exc,
-                message=f"Task run encountered unexpected exception: {repr(exc)}",
+                data=exc_or_state,
+                message=f"Task run encountered unexpected {failure_type}: {repr(exc_or_state)}",
             )
             if asyncio.iscoroutinefunction(retry_condition):
                 should_retry = run_coro_as_sync(
@@ -449,7 +453,9 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             else:
                 result = state.data
 
-            link_state_to_result(state, result)
+            link_state_to_result(new_state, result)
+            if asset_context := AssetContext.get():
+                asset_context.emit_events(new_state)
 
         # emit a state change event
         self._last_event = emit_task_run_state_change_event(
@@ -476,7 +482,15 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             # otherwise, return the exception
             return self._raised
 
-    def handle_success(self, result: R, transaction: Transaction) -> R:
+    def handle_success(
+        self, result: R, transaction: Transaction
+    ) -> Union[ResultRecord[R], None, Coroutine[Any, Any, R], R]:
+        # Handle the case where the task explicitly returns a failed state, in
+        # which case we should retry the task if it has retries left.
+        if isinstance(result, State) and result.is_failed():
+            if self.handle_retry(result):
+                return None
+
         if self.task.cache_expiration is not None:
             expiration = prefect.types._datetime.now("UTC") + self.task.cache_expiration
         else:
@@ -508,16 +522,16 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self._return_value = result
 
         self._telemetry.end_span_on_success()
-        return result
 
-    def handle_retry(self, exc: Exception) -> bool:
+    def handle_retry(self, exc_or_state: Exception | State[R]) -> bool:
         """Handle any task run retries.
 
         - If the task has retries left, and the retry condition is met, set the task to retrying and return True.
         - If the task has a retry delay, place in AwaitingRetry state with a delayed scheduled time.
         - If the task has no retries left, or the retry condition is not met, return False.
         """
-        if self.retries < self.task.retries and self.can_retry(exc):
+        failure_type = "exception" if isinstance(exc_or_state, Exception) else "state"
+        if self.retries < self.task.retries and self.can_retry(exc_or_state):
             if self.task.retry_delay_seconds:
                 delay = (
                     self.task.retry_delay_seconds[
@@ -535,8 +549,9 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 new_state = Retrying()
 
             self.logger.info(
-                "Task run failed with exception: %r - Retry %s/%s will start %s",
-                exc,
+                "Task run failed with %s: %r - Retry %s/%s will start %s",
+                failure_type,
+                exc_or_state,
                 self.retries + 1,
                 self.task.retries,
                 str(delay) + " second(s) from now" if delay else "immediately",
@@ -552,7 +567,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 else "No retries configured for this task."
             )
             self.logger.error(
-                f"Task run failed with exception: {exc!r} - {retry_message_suffix}",
+                f"Task run failed with {failure_type}: {exc_or_state!r} - {retry_message_suffix}",
                 exc_info=True,
             )
             return False
@@ -625,6 +640,16 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 persist_result = settings.tasks.default_persist_result
             else:
                 persist_result = should_persist_result()
+
+            asset_context = AssetContext.get()
+            if not asset_context:
+                asset_context = AssetContext.from_task_and_inputs(
+                    task=self.task,
+                    task_run_id=self.task_run.id,
+                    task_inputs=self.task_run.task_inputs,
+                )
+                stack.enter_context(asset_context)
+
             stack.enter_context(
                 TaskRunContext(
                     task=self.task,
@@ -830,7 +855,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
     def call_task_fn(
         self, transaction: Transaction
-    ) -> Union[R, Coroutine[Any, Any, R]]:
+    ) -> Union[ResultRecord[Any], None, Coroutine[Any, Any, R], R]:
         """
         Convenience method to call the task function. Returns a coroutine if the
         task is async.
@@ -855,10 +880,13 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             raise RuntimeError("Engine has not started.")
         return self._client
 
-    async def can_retry(self, exc: Exception) -> bool:
+    async def can_retry(self, exc_or_state: Exception | State[R]) -> bool:
         retry_condition: Optional[
-            Callable[["Task[P, Coroutine[Any, Any, R]]", TaskRun, State], bool]
+            Callable[["Task[P, Coroutine[Any, Any, R]]", TaskRun, State[R]], bool]
         ] = self.task.retry_condition_fn
+
+        failure_type = "exception" if isinstance(exc_or_state, Exception) else "state"
+
         if not self.task_run:
             raise ValueError("Task run is not set")
         try:
@@ -867,8 +895,8 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 f" {self.task.name!r}"
             )
             state = Failed(
-                data=exc,
-                message=f"Task run encountered unexpected exception: {repr(exc)}",
+                data=exc_or_state,
+                message=f"Task run encountered unexpected {failure_type}: {repr(exc_or_state)}",
             )
             if asyncio.iscoroutinefunction(retry_condition):
                 should_retry = await retry_condition(self.task, self.task_run, state)
@@ -1004,6 +1032,8 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 result = new_state.data
 
             link_state_to_result(new_state, result)
+            if asset_context := AssetContext.get():
+                asset_context.emit_events(new_state)
 
         # emit a state change event
         self._last_event = emit_task_run_state_change_event(
@@ -1031,7 +1061,13 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             # otherwise, return the exception
             return self._raised
 
-    async def handle_success(self, result: R, transaction: AsyncTransaction) -> R:
+    async def handle_success(
+        self, result: R, transaction: AsyncTransaction
+    ) -> Union[ResultRecord[R], None, Coroutine[Any, Any, R], R]:
+        if isinstance(result, State) and result.is_failed():
+            if await self.handle_retry(result):
+                return None
+
         if self.task.cache_expiration is not None:
             expiration = prefect.types._datetime.now("UTC") + self.task.cache_expiration
         else:
@@ -1059,19 +1095,20 @@ 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:
+    async def handle_retry(self, exc_or_state: Exception | State[R]) -> bool:
         """Handle any task run retries.
 
         - If the task has retries left, and the retry condition is met, set the task to retrying and return True.
         - If the task has a retry delay, place in AwaitingRetry state with a delayed scheduled time.
         - If the task has no retries left, or the retry condition is not met, return False.
         """
-        if self.retries < self.task.retries and await self.can_retry(exc):
+        failure_type = "exception" if isinstance(exc_or_state, Exception) else "state"
+
+        if self.retries < self.task.retries and await self.can_retry(exc_or_state):
             if self.task.retry_delay_seconds:
                 delay = (
                     self.task.retry_delay_seconds[
@@ -1089,8 +1126,9 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 new_state = Retrying()
 
             self.logger.info(
-                "Task run failed with exception: %r - Retry %s/%s will start %s",
-                exc,
+                "Task run failed with %s: %r - Retry %s/%s will start %s",
+                failure_type,
+                exc_or_state,
                 self.retries + 1,
                 self.task.retries,
                 str(delay) + " second(s) from now" if delay else "immediately",
@@ -1106,7 +1144,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 else "No retries configured for this task."
             )
             self.logger.error(
-                f"Task run failed with exception: {exc!r} - {retry_message_suffix}",
+                f"Task run failed with {failure_type}: {exc_or_state!r} - {retry_message_suffix}",
                 exc_info=True,
             )
             return False
@@ -1180,6 +1218,16 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 persist_result = settings.tasks.default_persist_result
             else:
                 persist_result = should_persist_result()
+
+            asset_context = AssetContext.get()
+            if not asset_context:
+                asset_context = AssetContext.from_task_and_inputs(
+                    task=self.task,
+                    task_run_id=self.task_run.id,
+                    task_inputs=self.task_run.task_inputs,
+                )
+                stack.enter_context(asset_context)
+
             stack.enter_context(
                 TaskRunContext(
                     task=self.task,
@@ -1382,7 +1430,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
     async def call_task_fn(
         self, transaction: AsyncTransaction
-    ) -> Union[R, Coroutine[Any, Any, R]]:
+    ) -> Union[ResultRecord[Any], None, Coroutine[Any, Any, R], R]:
         """
         Convenience method to call the task function. Returns a coroutine if the
         task is async.

prefect/tasks.py

@@ -29,10 +29,20 @@ from typing import (
 )
 from uuid import UUID, uuid4
 
-from typing_extensions import Literal, ParamSpec, Self, TypeAlias, TypeIs
+from typing_extensions import (
+    Literal,
+    ParamSpec,
+    Self,
+    Sequence,
+    TypeAlias,
+    TypedDict,
+    TypeIs,
+    Unpack,
+)
 
 import prefect.states
 from prefect._internal.uuid7 import uuid7
+from prefect.assets import Asset
 from prefect.cache_policies import DEFAULT, NO_CACHE, CachePolicy
 from prefect.client.orchestration import get_client
 from prefect.client.schemas import TaskRun
@@ -90,6 +100,65 @@ OneOrManyFutureOrResult: TypeAlias = Union[
 ]
 
 
+class TaskRunNameCallbackWithParameters(Protocol):
+    @classmethod
+    def is_callback_with_parameters(cls, callable: Callable[..., str]) -> TypeIs[Self]:
+        sig = inspect.signature(callable)
+        return "parameters" in sig.parameters
+
+    def __call__(self, parameters: dict[str, Any]) -> str: ...
+
+
+StateHookCallable: TypeAlias = Callable[
+    ["Task[..., Any]", TaskRun, State], Union[Awaitable[None], None]
+]
+RetryConditionCallable: TypeAlias = Callable[
+    ["Task[..., Any]", TaskRun, State], Union[Awaitable[bool], bool]
+]
+TaskRunNameValueOrCallable: TypeAlias = Union[
+    Callable[[], str], TaskRunNameCallbackWithParameters, str
+]
+
+
+class TaskOptions(TypedDict, total=False):
+    """
+    A TypedDict representing all available task configuration options.
+
+    This can be used with `Unpack` to provide type hints for **kwargs.
+    """
+
+    name: Optional[str]
+    description: Optional[str]
+    tags: Optional[Iterable[str]]
+    version: Optional[str]
+    cache_policy: Union[CachePolicy, type[NotSet]]
+    cache_key_fn: Union[
+        Callable[["TaskRunContext", dict[str, Any]], Optional[str]], None
+    ]
+    cache_expiration: Optional[datetime.timedelta]
+    task_run_name: Optional[TaskRunNameValueOrCallable]
+    retries: Optional[int]
+    retry_delay_seconds: Union[
+        float, int, list[float], Callable[[int], list[float]], None
+    ]
+    retry_jitter_factor: Optional[float]
+    persist_result: Optional[bool]
+    result_storage: Optional[ResultStorage]
+    result_serializer: Optional[ResultSerializer]
+    result_storage_key: Optional[str]
+    cache_result_in_memory: bool
+    timeout_seconds: Union[int, float, None]
+    log_prints: Optional[bool]
+    refresh_cache: Optional[bool]
+    on_completion: Optional[list[StateHookCallable]]
+    on_failure: Optional[list[StateHookCallable]]
+    on_rollback: Optional[list[Callable[["Transaction"], None]]]
+    on_commit: Optional[list[Callable[["Transaction"], None]]]
+    retry_condition_fn: Optional[RetryConditionCallable]
+    viz_return_value: Any
+    asset_deps: Optional[list[Union[Asset, str]]]
+
+
 def task_input_hash(
     context: "TaskRunContext", arguments: dict[str, Any]
 ) -> Optional[str]:
@@ -223,23 +292,6 @@ def _generate_task_key(fn: Callable[..., Any]) -> str:
     return f"{qualname}-{code_hash}"
 
 
-class TaskRunNameCallbackWithParameters(Protocol):
-    @classmethod
-    def is_callback_with_parameters(cls, callable: Callable[..., str]) -> TypeIs[Self]:
-        sig = inspect.signature(callable)
-        return "parameters" in sig.parameters
-
-    def __call__(self, parameters: dict[str, Any]) -> str: ...
-
-
-StateHookCallable: TypeAlias = Callable[
-    ["Task[..., Any]", TaskRun, State], Union[Awaitable[None], None]
-]
-TaskRunNameValueOrCallable: TypeAlias = Union[
-    Callable[[], str], TaskRunNameCallbackWithParameters, str
-]
-
-
 class Task(Generic[P, R]):
     """
     A Prefect task definition.
@@ -311,6 +363,7 @@ class Task(Generic[P, R]):
             should end as failed. Defaults to `None`, indicating the task should always continue
             to its retry policy.
         viz_return_value: An optional value to return when the task dependency tree is visualized.
+        asset_deps: An optional list of upstream assets that this task depends on.
     """
 
     # NOTE: These parameters (types, defaults, and docstrings) should be duplicated
@@ -350,10 +403,9 @@ class Task(Generic[P, R]):
         on_failure: Optional[list[StateHookCallable]] = None,
         on_rollback: Optional[list[Callable[["Transaction"], None]]] = None,
         on_commit: Optional[list[Callable[["Transaction"], None]]] = None,
-        retry_condition_fn: Optional[
-            Callable[["Task[..., Any]", TaskRun, State], bool]
-        ] = None,
+        retry_condition_fn: Optional[RetryConditionCallable] = None,
         viz_return_value: Optional[Any] = None,
+        asset_deps: Optional[list[Union[str, Asset]]] = None,
     ):
         # Validate if hook passed is list and contains callables
         hook_categories = [on_completion, on_failure]
@@ -466,6 +518,10 @@ class Task(Generic[P, R]):
             ):
                 persist_result = True
 
+        # Check for global cache disable setting
+        if settings.tasks.disable_caching:
+            cache_policy = NO_CACHE
+
         if persist_result is False:
             self.cache_policy = None if cache_policy is None else NO_CACHE
             if cache_policy and cache_policy is not NotSet and cache_policy != NO_CACHE:
@@ -543,6 +599,14 @@ class Task(Generic[P, R]):
         self.retry_condition_fn = retry_condition_fn
         self.viz_return_value = viz_return_value
 
+        from prefect.assets import Asset
+
+        self.asset_deps: list[Asset] = (
+            [Asset(key=a) if isinstance(a, str) else a for a in asset_deps]
+            if asset_deps
+            else []
+        )
+
     @property
     def ismethod(self) -> bool:
         return hasattr(self.fn, "__prefect_self__")
@@ -609,10 +673,9 @@ class Task(Generic[P, R]):
         refresh_cache: Union[bool, type[NotSet]] = NotSet,
         on_completion: Optional[list[StateHookCallable]] = None,
         on_failure: Optional[list[StateHookCallable]] = None,
-        retry_condition_fn: Optional[
-            Callable[["Task[..., Any]", TaskRun, State], bool]
-        ] = None,
+        retry_condition_fn: Optional[RetryConditionCallable] = None,
         viz_return_value: Optional[Any] = None,
+        asset_deps: Optional[list[Union[str, Asset]]] = None,
     ) -> "Task[P, R]":
         """
         Create a new task from the current object, updating provided options.
@@ -746,6 +809,7 @@ class Task(Generic[P, R]):
             on_failure=on_failure or self.on_failure_hooks,
             retry_condition_fn=retry_condition_fn or self.retry_condition_fn,
             viz_return_value=viz_return_value or self.viz_return_value,
+            asset_deps=asset_deps or self.asset_deps,
         )
 
     def on_completion(self, fn: StateHookCallable) -> StateHookCallable:
@@ -883,7 +947,9 @@ class Task(Generic[P, R]):
         deferred: bool = False,
     ) -> TaskRun:
         from prefect.utilities._engine import dynamic_key_for_task_run
-        from prefect.utilities.engine import collect_task_run_inputs_sync
+        from prefect.utilities.engine import (
+            collect_task_run_inputs_sync,
+        )
 
         if flow_run_context is None:
             flow_run_context = FlowRunContext.get()
@@ -923,7 +989,7 @@ class Task(Generic[P, R]):
 
                 store = await ResultStore(
                     result_storage=await get_or_create_default_task_scheduling_storage()
-                ).update_for_task(task)
+                ).update_for_task(self)
                 context = serialize_context()
                 data: dict[str, Any] = {"context": context}
                 if parameters:
@@ -959,6 +1025,7 @@ class Task(Generic[P, R]):
                 else None
             )
             task_run_id = id or uuid7()
+
             state = prefect.states.Pending(
                 state_details=StateDetails(
                     task_run_id=task_run_id,
@@ -1660,8 +1727,9 @@ def task(
     refresh_cache: Optional[bool] = None,
     on_completion: Optional[list[StateHookCallable]] = None,
     on_failure: Optional[list[StateHookCallable]] = None,
-    retry_condition_fn: Literal[None] = None,
+    retry_condition_fn: Optional[RetryConditionCallable] = None,
     viz_return_value: Any = None,
+    asset_deps: Optional[list[Union[str, Asset]]] = None,
 ) -> Callable[[Callable[P, R]], Task[P, R]]: ...
 
 
@@ -1695,8 +1763,9 @@ def task(
     refresh_cache: Optional[bool] = None,
     on_completion: Optional[list[StateHookCallable]] = None,
     on_failure: Optional[list[StateHookCallable]] = None,
-    retry_condition_fn: Optional[Callable[[Task[P, R], TaskRun, State], bool]] = None,
+    retry_condition_fn: Optional[RetryConditionCallable] = None,
     viz_return_value: Any = None,
+    asset_deps: Optional[list[Union[str, Asset]]] = None,
 ) -> Callable[[Callable[P, R]], Task[P, R]]: ...
 
 
@@ -1731,8 +1800,9 @@ def task(
     refresh_cache: Optional[bool] = None,
     on_completion: Optional[list[StateHookCallable]] = None,
     on_failure: Optional[list[StateHookCallable]] = None,
-    retry_condition_fn: Optional[Callable[[Task[P, Any], TaskRun, State], bool]] = None,
+    retry_condition_fn: Optional[RetryConditionCallable] = None,
     viz_return_value: Any = None,
+    asset_deps: Optional[list[Union[str, Asset]]] = None,
 ) -> Callable[[Callable[P, R]], Task[P, R]]: ...
 
 
@@ -1764,8 +1834,9 @@ def task(
     refresh_cache: Optional[bool] = None,
     on_completion: Optional[list[StateHookCallable]] = None,
     on_failure: Optional[list[StateHookCallable]] = None,
-    retry_condition_fn: Optional[Callable[[Task[P, Any], TaskRun, State], bool]] = None,
+    retry_condition_fn: Optional[RetryConditionCallable] = None,
     viz_return_value: Any = None,
+    asset_deps: Optional[list[Union[str, Asset]]] = None,
 ):
     """
     Decorator to designate a function as a task in a Prefect workflow.
@@ -1826,6 +1897,7 @@ def task(
             should end as failed. Defaults to `None`, indicating the task should always continue
             to its retry policy.
         viz_return_value: An optional value to return when the task dependency tree is visualized.
+        asset_deps: An optional list of upstream assets that this task depends on.
 
     Returns:
         A callable `Task` object which, when called, will submit the task for execution.
@@ -1902,6 +1974,7 @@ def task(
             on_failure=on_failure,
             retry_condition_fn=retry_condition_fn,
             viz_return_value=viz_return_value,
+            asset_deps=asset_deps,
         )
     else:
         return cast(
@@ -1931,5 +2004,32 @@ def task(
                 on_failure=on_failure,
                 retry_condition_fn=retry_condition_fn,
                 viz_return_value=viz_return_value,
+                asset_deps=asset_deps,
             ),
         )
+
+
+class MaterializingTask(Task[P, R]):
+    """
+    A task that materializes Assets.
+
+    Args:
+        assets: List of Assets that this task materializes (can be str or Asset)
+        materialized_by: An optional tool that materialized the asset e.g. "dbt" or "spark"
+        **task_kwargs: All other Task arguments
+    """
+
+    def __init__(
+        self,
+        fn: Callable[P, R],
+        *,
+        assets: Sequence[Union[str, Asset]],
+        materialized_by: str | None = None,
+        **task_kwargs: Unpack[TaskOptions],
+    ):
+        super().__init__(fn=fn, **task_kwargs)
+
+        self.assets: list[Asset] = [
+            Asset(key=a) if isinstance(a, str) else a for a in assets
+        ]
+        self.materialized_by = materialized_by

prefect/types/__init__.py

@@ -14,6 +14,7 @@ from .names import (
     BANNED_CHARACTERS,
     WITHOUT_BANNED_CHARACTERS,
     MAX_VARIABLE_NAME_LENGTH,
+    URILike,
 )
 from pydantic import (
     BeforeValidator,
@@ -219,4 +220,5 @@ __all__ = [
     "StatusCode",
     "StrictVariableValue",
     "TaskRetryDelaySeconds",
+    "URILike",
 ]

prefect/types/_datetime.py

@@ -223,10 +223,10 @@ def travel_to(dt: Any):
 
 def in_local_tz(dt: datetime.datetime) -> datetime.datetime:
     if sys.version_info >= (3, 13):
-        from whenever import LocalDateTime, ZonedDateTime
+        from whenever import PlainDateTime, ZonedDateTime
 
         if dt.tzinfo is None:
-            wdt = LocalDateTime.from_py_datetime(dt)
+            wdt = PlainDateTime.from_py_datetime(dt)
         else:
             if not isinstance(dt.tzinfo, ZoneInfo):
                 if key := getattr(dt.tzinfo, "key", None):

prefect/types/names.py

@@ -137,3 +137,26 @@ VariableName = Annotated[
         examples=["my_variable"],
     ),
 ]
+
+
+# URI validation
+URI_REGEX = re.compile(r"^[a-z0-9]+://")
+
+
+def validate_uri(value: str) -> str:
+    """Validate that a string is a valid URI with lowercase protocol."""
+    if not URI_REGEX.match(value):
+        raise ValueError(
+            "Key must be a valid URI, e.g. storage://bucket/folder/asset.csv"
+        )
+    return value
+
+
+URILike = Annotated[
+    str,
+    AfterValidator(validate_uri),
+    Field(
+        description="A URI-like string with a lowercase protocol",
+        examples=["s3://bucket/folder/data.csv", "postgres://dbtable"],
+    ),
+]

prefect/utilities/callables.py

@@ -286,8 +286,10 @@ def process_v1_params(
     docstrings: dict[str, str],
     aliases: dict[str, str],
 ) -> tuple[str, Any, Any]:
+    import pydantic.v1 as pydantic_v1
+
     # Pydantic model creation will fail if names collide with the BaseModel type
-    if hasattr(pydantic.BaseModel, param.name):
+    if hasattr(pydantic_v1.BaseModel, param.name):
         name = param.name + "__"
         aliases[name] = param.name
     else:
@@ -296,10 +298,9 @@ def process_v1_params(
     type_ = Any if param.annotation is inspect.Parameter.empty else param.annotation
 
     with warnings.catch_warnings():
-        warnings.filterwarnings(
-            "ignore", category=pydantic.warnings.PydanticDeprecatedSince20
-        )
-        field: Any = pydantic.Field(  # type: ignore   # this uses the v1 signature, not v2
+        # Note: pydantic.v1 doesn't have the warnings module, so we can't suppress them
+        warnings.filterwarnings("ignore", category=DeprecationWarning)
+        field: Any = pydantic_v1.Field(
             default=... if param.default is param.empty else param.default,
             title=param.name,
             description=docstrings.get(param.name, None),
@@ -312,18 +313,19 @@ def process_v1_params(
 def create_v1_schema(
     name_: str, model_cfg: type[Any], model_fields: Optional[dict[str, Any]] = None
 ) -> dict[str, Any]:
+    import pydantic.v1 as pydantic_v1
+
     with warnings.catch_warnings():
-        warnings.filterwarnings(
-            "ignore", category=pydantic.warnings.PydanticDeprecatedSince20
-        )
+        # Note: pydantic.v1 doesn't have the warnings module, so we can't suppress them
+        warnings.filterwarnings("ignore", category=DeprecationWarning)
 
         model_fields = model_fields or {}
-        model: type[pydantic.BaseModel] = pydantic.create_model(  # type: ignore   # this uses the v1 signature, not v2
+        model: type[pydantic_v1.BaseModel] = pydantic_v1.create_model(
             name_,
-            __config__=model_cfg,  # type: ignore   # this uses the v1 signature, not v2
+            __config__=model_cfg,
             **model_fields,
         )
-        return model.schema(by_alias=True)  # type: ignore   # this uses the v1 signature, not v2
+        return model.schema(by_alias=True)
 
 
 def parameter_schema(fn: Callable[..., Any]) -> ParameterSchema:

prefect/utilities/engine.py

@@ -80,7 +80,11 @@ async def collect_task_run_inputs(expr: Any, max_depth: int = -1) -> set[TaskRun
             inputs.add(TaskRunResult(id=obj.task_run_id))
         elif isinstance(obj, State):
             if obj.state_details.task_run_id:
-                inputs.add(TaskRunResult(id=obj.state_details.task_run_id))
+                inputs.add(
+                    TaskRunResult(
+                        id=obj.state_details.task_run_id,
+                    )
+                )
         # Expressions inside quotes should not be traversed
         elif isinstance(obj, quote):
             raise StopVisiting
@@ -118,10 +122,18 @@ def collect_task_run_inputs_sync(
 
     def add_futures_and_states_to_inputs(obj: Any) -> None:
         if isinstance(obj, future_cls) and hasattr(obj, "task_run_id"):
-            inputs.add(TaskRunResult(id=obj.task_run_id))
+            inputs.add(
+                TaskRunResult(
+                    id=obj.task_run_id,
+                )
+            )
         elif isinstance(obj, State):
             if obj.state_details.task_run_id:
-                inputs.add(TaskRunResult(id=obj.state_details.task_run_id))
+                inputs.add(
+                    TaskRunResult(
+                        id=obj.state_details.task_run_id,
+                    )
+                )
         # Expressions inside quotes should not be traversed
         elif isinstance(obj, quote):
             raise StopVisiting