prefect-client 3.1.11__py3-none-any.whl → 3.1.13__py3-none-any.whl

prefect/flows.py

@@ -2,6 +2,8 @@
 Module containing the base workflow class and decorator - for most use cases, using the [`@flow` decorator][prefect.flows.flow] is preferred.
 """
 
+from __future__ import annotations
+
 # This file requires type-checking with pyright because mypy does not yet support PEP612
 # See https://github.com/python/mypy/issues/8645
 import ast
@@ -27,6 +29,7 @@ from typing import (
     Iterable,
     NoReturn,
     Optional,
+    Protocol,
     Tuple,
     Type,
     TypeVar,
@@ -41,14 +44,14 @@ from pydantic.v1 import BaseModel as V1BaseModel
 from pydantic.v1.decorator import ValidatedFunction as V1ValidatedFunction
 from pydantic.v1.errors import ConfigError  # TODO
 from rich.console import Console
-from typing_extensions import Literal, ParamSpec, TypeAlias
+from typing_extensions import Literal, ParamSpec
 
+from prefect._experimental.sla.objects import SlaTypes
 from prefect._internal.concurrency.api import create_call, from_async
 from prefect.blocks.core import Block
 from prefect.client.schemas.actions import DeploymentScheduleCreate
-from prefect.client.schemas.filters import WorkerFilter
+from prefect.client.schemas.filters import WorkerFilter, WorkerFilterStatus
 from prefect.client.schemas.objects import ConcurrencyLimitConfig, FlowRun
-from prefect.client.schemas.objects import Flow as FlowSchema
 from prefect.client.utilities import client_injector
 from prefect.docker.docker_image import DockerImage
 from prefect.events import DeploymentTriggerTypes, TriggerTypes
@@ -83,6 +86,7 @@ from prefect.utilities.asyncutils import (
     sync_compatible,
 )
 from prefect.utilities.callables import (
+    ParameterSchema,
     get_call_parameters,
     parameter_schema,
     parameters_to_args_kwargs,
@@ -105,19 +109,31 @@ R = TypeVar("R")  # The return type of the user's function
 P = ParamSpec("P")  # The parameters of the flow
 F = TypeVar("F", bound="Flow[Any, Any]")  # The type of the flow
 
-StateHookCallable: TypeAlias = Callable[
-    [FlowSchema, FlowRun, State], Union[Awaitable[None], None]
-]
 
-logger = get_logger("flows")
+class FlowStateHook(Protocol, Generic[P, R]):
+    """
+    A callable that is invoked when a flow enters a given state.
+    """
+
+    __name__: str
+
+    def __call__(
+        self, flow: Flow[P, R], flow_run: FlowRun, state: State
+    ) -> Awaitable[None] | None:
+        ...
+
 
 if TYPE_CHECKING:
+    import logging
+
     from prefect.client.orchestration import PrefectClient
     from prefect.client.types.flexible_schedule_list import FlexibleScheduleList
     from prefect.deployments.runner import RunnerDeployment
     from prefect.flows import FlowRun
     from prefect.runner.storage import RunnerStorage
 
+logger: "logging.Logger" = get_logger("flows")
+
 
 class Flow(Generic[P, R]):
     """
@@ -187,7 +203,7 @@ class Flow(Generic[P, R]):
         retries: Optional[int] = None,
         retry_delay_seconds: Optional[Union[int, float]] = None,
         task_runner: Union[
-            Type[TaskRunner[PrefectFuture[R]]], TaskRunner[PrefectFuture[R]], None
+            Type[TaskRunner[PrefectFuture[Any]]], TaskRunner[PrefectFuture[Any]], None
         ] = None,
         description: Optional[str] = None,
         timeout_seconds: Union[int, float, None] = None,
@@ -197,13 +213,13 @@ class Flow(Generic[P, R]):
         result_serializer: Optional[ResultSerializer] = None,
         cache_result_in_memory: bool = True,
         log_prints: Optional[bool] = None,
-        on_completion: Optional[list[StateHookCallable]] = None,
-        on_failure: Optional[list[StateHookCallable]] = None,
-        on_cancellation: Optional[list[StateHookCallable]] = None,
-        on_crashed: Optional[list[StateHookCallable]] = None,
-        on_running: Optional[list[StateHookCallable]] = None,
+        on_completion: Optional[list[FlowStateHook[P, R]]] = None,
+        on_failure: Optional[list[FlowStateHook[P, R]]] = None,
+        on_cancellation: Optional[list[FlowStateHook[P, R]]] = None,
+        on_crashed: Optional[list[FlowStateHook[P, R]]] = None,
+        on_running: Optional[list[FlowStateHook[P, R]]] = None,
     ):
-        if name is not None and not isinstance(name, str):
+        if name is not None and not isinstance(name, str):  # pyright: ignore[reportUnnecessaryIsInstance]
             raise TypeError(
                 "Expected string for flow parameter 'name'; got {} instead. {}".format(
                     type(name).__name__,
@@ -257,7 +273,7 @@ class Flow(Generic[P, R]):
         if not callable(fn):
             raise TypeError("'fn' must be callable")
 
-        self.name = name or fn.__name__.replace("_", "-").replace(
+        self.name: str = name or fn.__name__.replace("_", "-").replace(
             "<lambda>",
             "unknown-lambda",  # prefect API will not accept "<" or ">" in flow names
         )
@@ -271,57 +287,64 @@ class Flow(Generic[P, R]):
                 )
         self.flow_run_name = flow_run_name
 
-        default_task_runner = ThreadPoolTaskRunner()
-        task_runner = task_runner or default_task_runner
-        self.task_runner = (
-            task_runner() if isinstance(task_runner, type) else task_runner
-        )
+        if task_runner is None:
+            self.task_runner: TaskRunner[PrefectFuture[Any]] = cast(
+                TaskRunner[PrefectFuture[Any]], ThreadPoolTaskRunner()
+            )
+        else:
+            self.task_runner: TaskRunner[PrefectFuture[Any]] = (
+                task_runner() if isinstance(task_runner, type) else task_runner
+            )
 
         self.log_prints = log_prints
 
-        self.description = description or inspect.getdoc(fn)
+        self.description: str | None = description or inspect.getdoc(fn)
         update_wrapper(self, fn)
         self.fn = fn
 
         # the flow is considered async if its function is async or an async
         # generator
-        self.isasync = asyncio.iscoroutinefunction(
+        self.isasync: bool = asyncio.iscoroutinefunction(
             self.fn
         ) or inspect.isasyncgenfunction(self.fn)
 
         # the flow is considered a generator if its function is a generator or
         # an async generator
-        self.isgenerator = inspect.isgeneratorfunction(
+        self.isgenerator: bool = inspect.isgeneratorfunction(
             self.fn
         ) or inspect.isasyncgenfunction(self.fn)
 
         raise_for_reserved_arguments(self.fn, ["return_state", "wait_for"])
 
         # Version defaults to a hash of the function's file
-        flow_file = inspect.getsourcefile(self.fn)
         if not version:
             try:
+                flow_file = inspect.getsourcefile(self.fn)
+                if flow_file is None:
+                    raise FileNotFoundError
                 version = file_hash(flow_file)
             except (FileNotFoundError, TypeError, OSError):
                 pass  # `getsourcefile` can return null values and "<stdin>" for objects in repls
         self.version = version
 
-        self.timeout_seconds = float(timeout_seconds) if timeout_seconds else None
+        self.timeout_seconds: float | None = (
+            float(timeout_seconds) if timeout_seconds else None
+        )
 
         # FlowRunPolicy settings
         # TODO: We can instantiate a `FlowRunPolicy` and add Pydantic bound checks to
         #       validate that the user passes positive numbers here
-        self.retries = (
+        self.retries: int = (
             retries if retries is not None else PREFECT_FLOW_DEFAULT_RETRIES.value()
         )
 
-        self.retry_delay_seconds = (
+        self.retry_delay_seconds: float | int = (
             retry_delay_seconds
             if retry_delay_seconds is not None
             else PREFECT_FLOW_DEFAULT_RETRY_DELAY_SECONDS.value()
         )
 
-        self.parameters = parameter_schema(self.fn)
+        self.parameters: ParameterSchema = parameter_schema(self.fn)
         self.should_validate_parameters = validate_parameters
 
         if self.should_validate_parameters:
@@ -352,11 +375,11 @@ class Flow(Generic[P, R]):
         self.result_storage = result_storage
         self.result_serializer = result_serializer
         self.cache_result_in_memory = cache_result_in_memory
-        self.on_completion_hooks = on_completion or []
-        self.on_failure_hooks = on_failure or []
-        self.on_cancellation_hooks = on_cancellation or []
-        self.on_crashed_hooks = on_crashed or []
-        self.on_running_hooks = on_running or []
+        self.on_completion_hooks: list[FlowStateHook[P, R]] = on_completion or []
+        self.on_failure_hooks: list[FlowStateHook[P, R]] = on_failure or []
+        self.on_cancellation_hooks: list[FlowStateHook[P, R]] = on_cancellation or []
+        self.on_crashed_hooks: list[FlowStateHook[P, R]] = on_crashed or []
+        self.on_running_hooks: list[FlowStateHook[P, R]] = on_running or []
 
         # Used for flows loaded from remote storage
         self._storage: Optional["RunnerStorage"] = None
@@ -373,7 +396,7 @@ class Flow(Generic[P, R]):
     def ismethod(self) -> bool:
         return hasattr(self.fn, "__prefect_self__")
 
-    def __get__(self, instance: Any, owner: Any):
+    def __get__(self, instance: Any, owner: Any) -> "Flow[P, R]":
         """
         Implement the descriptor protocol so that the flow can be used as an instance method.
         When an instance method is loaded, this method is called with the "self" instance as
@@ -388,7 +411,7 @@ class Flow(Generic[P, R]):
         # of the flow's function. This will allow it to be automatically added to the flow's parameters
         else:
             bound_flow = copy(self)
-            bound_flow.fn.__prefect_self__ = instance
+            setattr(bound_flow.fn, "__prefect_self__", instance)
             return bound_flow
 
     def with_options(
@@ -401,7 +424,7 @@ class Flow(Generic[P, R]):
         description: Optional[str] = None,
         flow_run_name: Optional[Union[Callable[[], str], str]] = None,
         task_runner: Union[
-            Type[TaskRunner[PrefectFuture[R]]], TaskRunner[PrefectFuture[R]], None
+            Type[TaskRunner[PrefectFuture[Any]]], TaskRunner[PrefectFuture[Any]], None
         ] = None,
         timeout_seconds: Union[int, float, None] = None,
         validate_parameters: Optional[bool] = None,
@@ -410,11 +433,11 @@ class Flow(Generic[P, R]):
         result_serializer: Optional[ResultSerializer] = NotSet,  # type: ignore
         cache_result_in_memory: Optional[bool] = None,
         log_prints: Optional[bool] = NotSet,  # type: ignore
-        on_completion: Optional[list[StateHookCallable]] = None,
-        on_failure: Optional[list[StateHookCallable]] = None,
-        on_cancellation: Optional[list[StateHookCallable]] = None,
-        on_crashed: Optional[list[StateHookCallable]] = None,
-        on_running: Optional[list[StateHookCallable]] = None,
+        on_completion: Optional[list[FlowStateHook[P, R]]] = None,
+        on_failure: Optional[list[FlowStateHook[P, R]]] = None,
+        on_cancellation: Optional[list[FlowStateHook[P, R]]] = None,
+        on_crashed: Optional[list[FlowStateHook[P, R]]] = None,
+        on_running: Optional[list[FlowStateHook[P, R]]] = None,
     ) -> "Flow[P, R]":
         """
         Create a new flow from the current object, updating provided options.
@@ -470,13 +493,18 @@ class Flow(Generic[P, R]):
             >>> state = my_flow.with_options(task_runner=ThreadPoolTaskRunner)(1, 3)
             >>> assert state.result() == 4
         """
+        new_task_runner = (
+            task_runner() if isinstance(task_runner, type) else task_runner
+        )
+        if new_task_runner is None:
+            new_task_runner = self.task_runner
         new_flow = Flow(
             fn=self.fn,
             name=name or self.name,
             description=description or self.description,
             flow_run_name=flow_run_name or self.flow_run_name,
             version=version or self.version,
-            task_runner=task_runner or self.task_runner,
+            task_runner=new_task_runner,
             retries=retries if retries is not None else self.retries,
             retry_delay_seconds=(
                 retry_delay_seconds
@@ -530,7 +558,7 @@ class Flow(Generic[P, R]):
             ParameterTypeError: if the provided parameters are not valid
         """
 
-        def resolve_block_reference(data: Any) -> Any:
+        def resolve_block_reference(data: Any | dict[str, Any]) -> Any:
             if isinstance(data, dict) and "$ref" in data:
                 return Block.load_from_ref(data["$ref"], _sync=True)
             return data
@@ -593,7 +621,9 @@ class Flow(Generic[P, R]):
         }
         return cast_parameters
 
-    def serialize_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]:
+    def serialize_parameters(
+        self, parameters: dict[str, Any | PrefectFuture[Any] | State]
+    ) -> dict[str, Any]:
         """
         Convert parameters to a serializable form.
 
@@ -601,10 +631,10 @@ class Flow(Generic[P, R]):
         converting everything directly to a string. This maintains basic types like
         integers during API roundtrips.
         """
-        serialized_parameters = {}
+        serialized_parameters: dict[str, Any] = {}
         for key, value in parameters.items():
             # do not serialize the bound self object
-            if self.ismethod and value is self.fn.__prefect_self__:
+            if self.ismethod and value is getattr(self.fn, "__prefect_self__", None):
                 continue
             if isinstance(value, (PrefectFuture, State)):
                 # Don't call jsonable_encoder() on a PrefectFuture or State to
@@ -651,6 +681,7 @@ class Flow(Generic[P, R]):
         work_queue_name: Optional[str] = None,
         job_variables: Optional[dict[str, Any]] = None,
         entrypoint_type: EntrypointType = EntrypointType.FILE_PATH,
+        _sla: Optional[Union[SlaTypes, list[SlaTypes]]] = None,  # experimental
     ) -> "RunnerDeployment":
         """
         Creates a runner deployment object for this flow.
@@ -681,6 +712,7 @@ class Flow(Generic[P, R]):
                 of the chosen work pool. Refer to the base job template of the chosen work pool for
             entrypoint_type: Type of entrypoint to use for the deployment. When using a module path
                 entrypoint, ensure that the module will be importable in the execution environment.
+            _sla: (Experimental) SLA configuration for the deployment. May be removed or modified at any time. Currently only supported on Prefect Cloud.
 
         Examples:
             Prepare two deployments and serve them:
@@ -728,6 +760,7 @@ class Flow(Generic[P, R]):
                 work_pool_name=work_pool_name,
                 work_queue_name=work_queue_name,
                 job_variables=job_variables,
+                _sla=_sla,
             )  # type: ignore # TODO: remove sync_compatible
         else:
             return RunnerDeployment.from_flow(
@@ -749,25 +782,26 @@ class Flow(Generic[P, R]):
                 work_queue_name=work_queue_name,
                 job_variables=job_variables,
                 entrypoint_type=entrypoint_type,
+                _sla=_sla,
             )
 
-    def on_completion(self, fn: StateHookCallable) -> StateHookCallable:
+    def on_completion(self, fn: FlowStateHook[P, R]) -> FlowStateHook[P, R]:
         self.on_completion_hooks.append(fn)
         return fn
 
-    def on_cancellation(self, fn: StateHookCallable) -> StateHookCallable:
+    def on_cancellation(self, fn: FlowStateHook[P, R]) -> FlowStateHook[P, R]:
         self.on_cancellation_hooks.append(fn)
         return fn
 
-    def on_crashed(self, fn: StateHookCallable) -> StateHookCallable:
+    def on_crashed(self, fn: FlowStateHook[P, R]) -> FlowStateHook[P, R]:
         self.on_crashed_hooks.append(fn)
         return fn
 
-    def on_running(self, fn: StateHookCallable) -> StateHookCallable:
+    def on_running(self, fn: FlowStateHook[P, R]) -> FlowStateHook[P, R]:
         self.on_running_hooks.append(fn)
         return fn
 
-    def on_failure(self, fn: StateHookCallable) -> StateHookCallable:
+    def on_failure(self, fn: FlowStateHook[P, R]) -> FlowStateHook[P, R]:
         self.on_failure_hooks.append(fn)
         return fn
 
@@ -798,7 +832,7 @@ class Flow(Generic[P, R]):
         limit: Optional[int] = None,
         webserver: bool = False,
         entrypoint_type: EntrypointType = EntrypointType.FILE_PATH,
-    ):
+    ) -> None:
         """
         Creates a deployment for this flow and starts a runner to monitor for scheduled work.
 
@@ -1061,6 +1095,7 @@ class Flow(Generic[P, R]):
         entrypoint_type: EntrypointType = EntrypointType.FILE_PATH,
         print_next_steps: bool = True,
         ignore_warnings: bool = False,
+        _sla: Optional[Union[SlaTypes, list[SlaTypes]]] = None,
     ) -> UUID:
         """
         Deploys a flow to run on dynamic infrastructure via a work pool.
@@ -1112,7 +1147,7 @@ class Flow(Generic[P, R]):
             print_next_steps_message: Whether or not to print a message with next steps
                 after deploying the deployments.
             ignore_warnings: Whether or not to ignore warnings about the work pool type.
-
+            _sla: (Experimental) SLA configuration for the deployment. May be removed or modified at any time. Currently only supported on Prefect Cloud.
         Returns:
             The ID of the created/updated deployment.
 
@@ -1165,7 +1200,9 @@ class Flow(Generic[P, R]):
                 work_pool = await client.read_work_pool(work_pool_name)
                 active_workers = await client.read_workers_for_work_pool(
                     work_pool_name,
-                    worker_filter=WorkerFilter(status={"any_": ["ONLINE"]}),
+                    worker_filter=WorkerFilter(
+                        status=WorkerFilterStatus(any_=["ONLINE"])
+                    ),
                 )
         except ObjectNotFound as exc:
             raise ValueError(
@@ -1173,7 +1210,7 @@ class Flow(Generic[P, R]):
                 " deploying this flow."
             ) from exc
 
-        deployment = await self.to_deployment(
+        to_deployment_coro = self.to_deployment(
             name=name,
             interval=interval,
             cron=cron,
@@ -1190,11 +1227,17 @@ class Flow(Generic[P, R]):
             work_queue_name=work_queue_name,
             job_variables=job_variables,
             entrypoint_type=entrypoint_type,
+            _sla=_sla,
         )
 
+        if TYPE_CHECKING:
+            assert inspect.isawaitable(to_deployment_coro)
+
+        deployment = await to_deployment_coro
+
         from prefect.deployments.runner import deploy
 
-        deployment_ids = await deploy(
+        deploy_coro = deploy(
             deployment,
             work_pool_name=work_pool_name,
             image=image,
@@ -1203,6 +1246,10 @@ class Flow(Generic[P, R]):
             print_next_steps_message=False,
             ignore_warnings=ignore_warnings,
         )
+        if TYPE_CHECKING:
+            assert inspect.isawaitable(deploy_coro)
+
+        deployment_ids = await deploy_coro
 
         if print_next_steps:
             console = Console()
@@ -1433,11 +1480,11 @@ class FlowDecorator:
         result_serializer: Optional[ResultSerializer] = None,
         cache_result_in_memory: bool = True,
         log_prints: Optional[bool] = None,
-        on_completion: Optional[list[StateHookCallable]] = None,
-        on_failure: Optional[list[StateHookCallable]] = None,
-        on_cancellation: Optional[list[StateHookCallable]] = None,
-        on_crashed: Optional[list[StateHookCallable]] = None,
-        on_running: Optional[list[StateHookCallable]] = None,
+        on_completion: Optional[list[FlowStateHook[..., Any]]] = None,
+        on_failure: Optional[list[FlowStateHook[..., Any]]] = None,
+        on_cancellation: Optional[list[FlowStateHook[..., Any]]] = None,
+        on_crashed: Optional[list[FlowStateHook[..., Any]]] = None,
+        on_running: Optional[list[FlowStateHook[..., Any]]] = None,
     ) -> Callable[[Callable[P, R]], Flow[P, R]]:
         ...
 
@@ -1460,11 +1507,11 @@ class FlowDecorator:
         result_serializer: Optional[ResultSerializer] = None,
         cache_result_in_memory: bool = True,
         log_prints: Optional[bool] = None,
-        on_completion: Optional[list[StateHookCallable]] = None,
-        on_failure: Optional[list[StateHookCallable]] = None,
-        on_cancellation: Optional[list[StateHookCallable]] = None,
-        on_crashed: Optional[list[StateHookCallable]] = None,
-        on_running: Optional[list[StateHookCallable]] = None,
+        on_completion: Optional[list[FlowStateHook[..., Any]]] = None,
+        on_failure: Optional[list[FlowStateHook[..., Any]]] = None,
+        on_cancellation: Optional[list[FlowStateHook[..., Any]]] = None,
+        on_crashed: Optional[list[FlowStateHook[..., Any]]] = None,
+        on_running: Optional[list[FlowStateHook[..., Any]]] = None,
     ) -> Callable[[Callable[P, R]], Flow[P, R]]:
         ...
 
@@ -1486,11 +1533,11 @@ class FlowDecorator:
         result_serializer: Optional[ResultSerializer] = None,
         cache_result_in_memory: bool = True,
         log_prints: Optional[bool] = None,
-        on_completion: Optional[list[StateHookCallable]] = None,
-        on_failure: Optional[list[StateHookCallable]] = None,
-        on_cancellation: Optional[list[StateHookCallable]] = None,
-        on_crashed: Optional[list[StateHookCallable]] = None,
-        on_running: Optional[list[StateHookCallable]] = None,
+        on_completion: Optional[list[FlowStateHook[..., Any]]] = None,
+        on_failure: Optional[list[FlowStateHook[..., Any]]] = None,
+        on_cancellation: Optional[list[FlowStateHook[..., Any]]] = None,
+        on_crashed: Optional[list[FlowStateHook[..., Any]]] = None,
+        on_running: Optional[list[FlowStateHook[..., Any]]] = None,
     ) -> Union[Flow[P, R], Callable[[Callable[P, R]], Flow[P, R]]]:
         """
         Decorator to designate a function as a Prefect workflow.
@@ -1664,7 +1711,7 @@ class FlowDecorator:
             ...
 
 
-flow = FlowDecorator()
+flow: FlowDecorator = FlowDecorator()
 
 
 def _raise_on_name_with_banned_characters(name: Optional[str]) -> Optional[str]:
@@ -1916,7 +1963,11 @@ async def aserve(
 
     runner = Runner(pause_on_shutdown=pause_on_shutdown, limit=limit, **kwargs)
     for deployment in args:
-        await runner.add_deployment(deployment)
+        add_deployment_coro = runner.add_deployment(deployment)
+        if TYPE_CHECKING:
+            assert inspect.isawaitable(add_deployment_coro)
+
+        await add_deployment_coro
 
     if print_starting_message:
         _display_serve_start_message(*args)
@@ -1962,13 +2013,16 @@ async def load_flow_from_flow_run(
     ignore_storage: bool = False,
     storage_base_path: Optional[str] = None,
     use_placeholder_flow: bool = True,
-) -> Flow[P, Any]:
+) -> Flow[..., Any]:
     """
     Load a flow from the location/script provided in a deployment's storage document.
 
     If `ignore_storage=True` is provided, no pull from remote storage occurs.  This flag
     is largely for testing, and assumes the flow is already available locally.
     """
+    if flow_run.deployment_id is None:
+        raise ValueError("Flow run does not have an associated deployment")
+
     deployment = await client.read_deployment(flow_run.deployment_id)
 
     if deployment.entrypoint is None:

prefect/futures.py

@@ -5,7 +5,7 @@ import threading
 import uuid
 from collections.abc import Generator, Iterator
 from functools import partial
-from typing import Any, Callable, Generic, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, Union
 
 from typing_extensions import NamedTuple, Self, TypeVar
 
@@ -22,7 +22,10 @@ from prefect.utilities.timeout import timeout as timeout_context
 F = TypeVar("F")
 R = TypeVar("R")
 
-logger = get_logger(__name__)
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger(__name__)
 
 
 class PrefectFuture(abc.ABC, Generic[R]):
@@ -90,7 +93,7 @@ class PrefectFuture(abc.ABC, Generic[R]):
         """
 
     @abc.abstractmethod
-    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 run when the future completes or is cancelled.
 
@@ -173,7 +176,7 @@ class PrefectConcurrentFuture(PrefectWrappedFuture[R, concurrent.futures.Future[
             _result = run_coro_as_sync(_result)
         return _result
 
-    def __del__(self):
+    def __del__(self) -> None:
         if self._final_state or self._wrapped_future.done():
             return
         try:
@@ -202,7 +205,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
     def wait(self, timeout: Optional[float] = None) -> None:
         return run_coro_as_sync(self.wait_async(timeout=timeout))
 
-    async def wait_async(self, timeout: Optional[float] = None):
+    async def wait_async(self, timeout: Optional[float] = None) -> None:
         if self._final_state:
             logger.debug(
                 "Final state already set for %s. Returning...", self.task_run_id
@@ -216,6 +219,10 @@ class PrefectDistributedFuture(PrefectFuture[R]):
         # Read task run to see if it is still running
         async with get_client() as client:
             task_run = await client.read_task_run(task_run_id=self._task_run_id)
+            if task_run.state is None:
+                raise RuntimeError(
+                    f"Task run {self.task_run_id} has no state which means it hasn't started yet."
+                )
             if task_run.state.is_final():
                 logger.debug(
                     "Task run %s already finished. Returning...",
@@ -260,7 +267,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
             raise_on_failure=raise_on_failure, fetch=True
         )
 
-    def add_done_callback(self, fn: Callable[[PrefectFuture[R]], None]):
+    def add_done_callback(self, fn: Callable[[PrefectFuture[R]], None]) -> None:
         if self._final_state:
             fn(self)
             return
@@ -278,7 +285,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
             return False
         return self.task_run_id == other.task_run_id
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         return hash(self.task_run_id)
 
 

prefect/infrastructure/provisioners/__init__.py

@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol, Type
 
+from prefect.infrastructure.provisioners.coiled import CoiledPushProvisioner
 from prefect.infrastructure.provisioners.modal import ModalPushProvisioner
 from .cloud_run import CloudRunPushProvisioner
 from .container_instance import ContainerInstancePushProvisioner
@@ -15,6 +16,7 @@ _provisioners = {
     "azure-container-instance:push": ContainerInstancePushProvisioner,
     "ecs:push": ElasticContainerServicePushProvisioner,
     "modal:push": ModalPushProvisioner,
+    "coiled:push": CoiledPushProvisioner,
 }
 
 

prefect/infrastructure/provisioners/cloud_run.py

@@ -33,7 +33,7 @@ if TYPE_CHECKING:
 
 class CloudRunPushProvisioner:
     def __init__(self):
-        self._console = Console()
+        self._console: Console = Console()
         self._project = None
         self._region = None
         self._service_account_name = "prefect-cloud-run"
@@ -41,14 +41,14 @@ class CloudRunPushProvisioner:
         self._image_repository_name = "prefect-images"
 
     @property
-    def console(self):
+    def console(self) -> Console:
         return self._console
 
     @console.setter
-    def console(self, value):
+    def console(self, value: Console) -> None:
         self._console = value
 
-    async def _run_command(self, command: str, *args, **kwargs):
+    async def _run_command(self, command: str, *args: Any, **kwargs: Any) -> str:
         result = await run_process(shlex.split(command), check=False, *args, **kwargs)
 
         if result.returncode != 0: