prefect-client 3.1.12__py3-none-any.whl → 3.1.14__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (111) hide show
  1. prefect/_experimental/lineage.py +63 -0
  2. prefect/_experimental/sla/client.py +53 -27
  3. prefect/_experimental/sla/objects.py +10 -2
  4. prefect/_internal/concurrency/services.py +2 -2
  5. prefect/_internal/concurrency/threads.py +6 -0
  6. prefect/_internal/retries.py +6 -3
  7. prefect/_internal/schemas/validators.py +6 -4
  8. prefect/_version.py +3 -3
  9. prefect/artifacts.py +4 -1
  10. prefect/automations.py +1 -1
  11. prefect/blocks/abstract.py +5 -2
  12. prefect/blocks/notifications.py +1 -0
  13. prefect/cache_policies.py +70 -22
  14. prefect/client/orchestration/_automations/client.py +4 -0
  15. prefect/client/orchestration/_deployments/client.py +3 -3
  16. prefect/client/utilities.py +3 -3
  17. prefect/context.py +16 -6
  18. prefect/deployments/base.py +7 -4
  19. prefect/deployments/flow_runs.py +5 -1
  20. prefect/deployments/runner.py +6 -11
  21. prefect/deployments/steps/core.py +1 -1
  22. prefect/deployments/steps/pull.py +8 -3
  23. prefect/deployments/steps/utility.py +2 -2
  24. prefect/docker/docker_image.py +13 -9
  25. prefect/engine.py +19 -10
  26. prefect/events/cli/automations.py +4 -4
  27. prefect/events/clients.py +17 -14
  28. prefect/events/filters.py +34 -34
  29. prefect/events/schemas/automations.py +12 -8
  30. prefect/events/schemas/events.py +5 -1
  31. prefect/events/worker.py +1 -1
  32. prefect/filesystems.py +1 -1
  33. prefect/flow_engine.py +172 -123
  34. prefect/flows.py +119 -74
  35. prefect/futures.py +14 -7
  36. prefect/infrastructure/provisioners/__init__.py +2 -0
  37. prefect/infrastructure/provisioners/cloud_run.py +4 -4
  38. prefect/infrastructure/provisioners/coiled.py +249 -0
  39. prefect/infrastructure/provisioners/container_instance.py +4 -3
  40. prefect/infrastructure/provisioners/ecs.py +55 -43
  41. prefect/infrastructure/provisioners/modal.py +5 -4
  42. prefect/input/actions.py +5 -1
  43. prefect/input/run_input.py +157 -43
  44. prefect/logging/configuration.py +5 -8
  45. prefect/logging/filters.py +2 -2
  46. prefect/logging/formatters.py +15 -11
  47. prefect/logging/handlers.py +24 -14
  48. prefect/logging/highlighters.py +5 -5
  49. prefect/logging/loggers.py +29 -20
  50. prefect/main.py +3 -1
  51. prefect/results.py +166 -86
  52. prefect/runner/runner.py +112 -84
  53. prefect/runner/server.py +3 -1
  54. prefect/runner/storage.py +18 -18
  55. prefect/runner/submit.py +19 -12
  56. prefect/runtime/deployment.py +15 -8
  57. prefect/runtime/flow_run.py +19 -6
  58. prefect/runtime/task_run.py +7 -3
  59. prefect/settings/base.py +17 -7
  60. prefect/settings/legacy.py +4 -4
  61. prefect/settings/models/api.py +4 -3
  62. prefect/settings/models/cli.py +4 -3
  63. prefect/settings/models/client.py +7 -4
  64. prefect/settings/models/cloud.py +4 -3
  65. prefect/settings/models/deployments.py +4 -3
  66. prefect/settings/models/experiments.py +4 -3
  67. prefect/settings/models/flows.py +4 -3
  68. prefect/settings/models/internal.py +4 -3
  69. prefect/settings/models/logging.py +8 -6
  70. prefect/settings/models/results.py +4 -3
  71. prefect/settings/models/root.py +11 -16
  72. prefect/settings/models/runner.py +8 -5
  73. prefect/settings/models/server/api.py +6 -3
  74. prefect/settings/models/server/database.py +120 -25
  75. prefect/settings/models/server/deployments.py +4 -3
  76. prefect/settings/models/server/ephemeral.py +7 -4
  77. prefect/settings/models/server/events.py +6 -3
  78. prefect/settings/models/server/flow_run_graph.py +4 -3
  79. prefect/settings/models/server/root.py +4 -3
  80. prefect/settings/models/server/services.py +15 -12
  81. prefect/settings/models/server/tasks.py +7 -4
  82. prefect/settings/models/server/ui.py +4 -3
  83. prefect/settings/models/tasks.py +10 -5
  84. prefect/settings/models/testing.py +4 -3
  85. prefect/settings/models/worker.py +7 -4
  86. prefect/settings/profiles.py +13 -12
  87. prefect/settings/sources.py +20 -19
  88. prefect/states.py +17 -13
  89. prefect/task_engine.py +43 -33
  90. prefect/task_runners.py +35 -23
  91. prefect/task_runs.py +20 -11
  92. prefect/task_worker.py +12 -7
  93. prefect/tasks.py +67 -25
  94. prefect/telemetry/bootstrap.py +4 -1
  95. prefect/telemetry/run_telemetry.py +15 -13
  96. prefect/transactions.py +3 -3
  97. prefect/types/__init__.py +9 -6
  98. prefect/types/_datetime.py +19 -0
  99. prefect/utilities/_deprecated.py +38 -0
  100. prefect/utilities/engine.py +11 -4
  101. prefect/utilities/filesystem.py +2 -2
  102. prefect/utilities/generics.py +1 -1
  103. prefect/utilities/pydantic.py +21 -36
  104. prefect/workers/base.py +52 -30
  105. prefect/workers/process.py +20 -15
  106. prefect/workers/server.py +4 -5
  107. {prefect_client-3.1.12.dist-info → prefect_client-3.1.14.dist-info}/METADATA +2 -2
  108. {prefect_client-3.1.12.dist-info → prefect_client-3.1.14.dist-info}/RECORD +111 -108
  109. {prefect_client-3.1.12.dist-info → prefect_client-3.1.14.dist-info}/LICENSE +0 -0
  110. {prefect_client-3.1.12.dist-info → prefect_client-3.1.14.dist-info}/WHEEL +0 -0
  111. {prefect_client-3.1.12.dist-info → prefect_client-3.1.14.dist-info}/top_level.txt +0 -0
prefect/flows.py CHANGED
@@ -29,6 +29,7 @@ from typing import (
29
29
  Iterable,
30
30
  NoReturn,
31
31
  Optional,
32
+ Protocol,
32
33
  Tuple,
33
34
  Type,
34
35
  TypeVar,
@@ -43,15 +44,14 @@ from pydantic.v1 import BaseModel as V1BaseModel
43
44
  from pydantic.v1.decorator import ValidatedFunction as V1ValidatedFunction
44
45
  from pydantic.v1.errors import ConfigError # TODO
45
46
  from rich.console import Console
46
- from typing_extensions import Literal, ParamSpec, TypeAlias
47
+ from typing_extensions import Literal, ParamSpec
47
48
 
48
49
  from prefect._experimental.sla.objects import SlaTypes
49
50
  from prefect._internal.concurrency.api import create_call, from_async
50
51
  from prefect.blocks.core import Block
51
52
  from prefect.client.schemas.actions import DeploymentScheduleCreate
52
- from prefect.client.schemas.filters import WorkerFilter
53
+ from prefect.client.schemas.filters import WorkerFilter, WorkerFilterStatus
53
54
  from prefect.client.schemas.objects import ConcurrencyLimitConfig, FlowRun
54
- from prefect.client.schemas.objects import Flow as FlowSchema
55
55
  from prefect.client.utilities import client_injector
56
56
  from prefect.docker.docker_image import DockerImage
57
57
  from prefect.events import DeploymentTriggerTypes, TriggerTypes
@@ -86,6 +86,7 @@ from prefect.utilities.asyncutils import (
86
86
  sync_compatible,
87
87
  )
88
88
  from prefect.utilities.callables import (
89
+ ParameterSchema,
89
90
  get_call_parameters,
90
91
  parameter_schema,
91
92
  parameters_to_args_kwargs,
@@ -108,19 +109,31 @@ R = TypeVar("R") # The return type of the user's function
108
109
  P = ParamSpec("P") # The parameters of the flow
109
110
  F = TypeVar("F", bound="Flow[Any, Any]") # The type of the flow
110
111
 
111
- StateHookCallable: TypeAlias = Callable[
112
- [FlowSchema, FlowRun, State], Union[Awaitable[None], None]
113
- ]
114
112
 
115
- logger = get_logger("flows")
113
+ class FlowStateHook(Protocol, Generic[P, R]):
114
+ """
115
+ A callable that is invoked when a flow enters a given state.
116
+ """
117
+
118
+ __name__: str
119
+
120
+ def __call__(
121
+ self, flow: Flow[P, R], flow_run: FlowRun, state: State
122
+ ) -> Awaitable[None] | None:
123
+ ...
124
+
116
125
 
117
126
  if TYPE_CHECKING:
127
+ import logging
128
+
118
129
  from prefect.client.orchestration import PrefectClient
130
+ from prefect.client.schemas.objects import FlowRun
119
131
  from prefect.client.types.flexible_schedule_list import FlexibleScheduleList
120
132
  from prefect.deployments.runner import RunnerDeployment
121
- from prefect.flows import FlowRun
122
133
  from prefect.runner.storage import RunnerStorage
123
134
 
135
+ logger: "logging.Logger" = get_logger("flows")
136
+
124
137
 
125
138
  class Flow(Generic[P, R]):
126
139
  """
@@ -190,7 +203,7 @@ class Flow(Generic[P, R]):
190
203
  retries: Optional[int] = None,
191
204
  retry_delay_seconds: Optional[Union[int, float]] = None,
192
205
  task_runner: Union[
193
- Type[TaskRunner[PrefectFuture[R]]], TaskRunner[PrefectFuture[R]], None
206
+ Type[TaskRunner[PrefectFuture[Any]]], TaskRunner[PrefectFuture[Any]], None
194
207
  ] = None,
195
208
  description: Optional[str] = None,
196
209
  timeout_seconds: Union[int, float, None] = None,
@@ -200,13 +213,13 @@ class Flow(Generic[P, R]):
200
213
  result_serializer: Optional[ResultSerializer] = None,
201
214
  cache_result_in_memory: bool = True,
202
215
  log_prints: Optional[bool] = None,
203
- on_completion: Optional[list[StateHookCallable]] = None,
204
- on_failure: Optional[list[StateHookCallable]] = None,
205
- on_cancellation: Optional[list[StateHookCallable]] = None,
206
- on_crashed: Optional[list[StateHookCallable]] = None,
207
- on_running: Optional[list[StateHookCallable]] = None,
216
+ on_completion: Optional[list[FlowStateHook[P, R]]] = None,
217
+ on_failure: Optional[list[FlowStateHook[P, R]]] = None,
218
+ on_cancellation: Optional[list[FlowStateHook[P, R]]] = None,
219
+ on_crashed: Optional[list[FlowStateHook[P, R]]] = None,
220
+ on_running: Optional[list[FlowStateHook[P, R]]] = None,
208
221
  ):
209
- if name is not None and not isinstance(name, str):
222
+ if name is not None and not isinstance(name, str): # pyright: ignore[reportUnnecessaryIsInstance]
210
223
  raise TypeError(
211
224
  "Expected string for flow parameter 'name'; got {} instead. {}".format(
212
225
  type(name).__name__,
@@ -260,7 +273,7 @@ class Flow(Generic[P, R]):
260
273
  if not callable(fn):
261
274
  raise TypeError("'fn' must be callable")
262
275
 
263
- self.name = name or fn.__name__.replace("_", "-").replace(
276
+ self.name: str = name or fn.__name__.replace("_", "-").replace(
264
277
  "<lambda>",
265
278
  "unknown-lambda", # prefect API will not accept "<" or ">" in flow names
266
279
  )
@@ -274,57 +287,64 @@ class Flow(Generic[P, R]):
274
287
  )
275
288
  self.flow_run_name = flow_run_name
276
289
 
277
- default_task_runner = ThreadPoolTaskRunner()
278
- task_runner = task_runner or default_task_runner
279
- self.task_runner = (
280
- task_runner() if isinstance(task_runner, type) else task_runner
281
- )
290
+ if task_runner is None:
291
+ self.task_runner: TaskRunner[PrefectFuture[Any]] = cast(
292
+ TaskRunner[PrefectFuture[Any]], ThreadPoolTaskRunner()
293
+ )
294
+ else:
295
+ self.task_runner: TaskRunner[PrefectFuture[Any]] = (
296
+ task_runner() if isinstance(task_runner, type) else task_runner
297
+ )
282
298
 
283
299
  self.log_prints = log_prints
284
300
 
285
- self.description = description or inspect.getdoc(fn)
301
+ self.description: str | None = description or inspect.getdoc(fn)
286
302
  update_wrapper(self, fn)
287
303
  self.fn = fn
288
304
 
289
305
  # the flow is considered async if its function is async or an async
290
306
  # generator
291
- self.isasync = asyncio.iscoroutinefunction(
307
+ self.isasync: bool = asyncio.iscoroutinefunction(
292
308
  self.fn
293
309
  ) or inspect.isasyncgenfunction(self.fn)
294
310
 
295
311
  # the flow is considered a generator if its function is a generator or
296
312
  # an async generator
297
- self.isgenerator = inspect.isgeneratorfunction(
313
+ self.isgenerator: bool = inspect.isgeneratorfunction(
298
314
  self.fn
299
315
  ) or inspect.isasyncgenfunction(self.fn)
300
316
 
301
317
  raise_for_reserved_arguments(self.fn, ["return_state", "wait_for"])
302
318
 
303
319
  # Version defaults to a hash of the function's file
304
- flow_file = inspect.getsourcefile(self.fn)
305
320
  if not version:
306
321
  try:
322
+ flow_file = inspect.getsourcefile(self.fn)
323
+ if flow_file is None:
324
+ raise FileNotFoundError
307
325
  version = file_hash(flow_file)
308
326
  except (FileNotFoundError, TypeError, OSError):
309
327
  pass # `getsourcefile` can return null values and "<stdin>" for objects in repls
310
328
  self.version = version
311
329
 
312
- self.timeout_seconds = float(timeout_seconds) if timeout_seconds else None
330
+ self.timeout_seconds: float | None = (
331
+ float(timeout_seconds) if timeout_seconds else None
332
+ )
313
333
 
314
334
  # FlowRunPolicy settings
315
335
  # TODO: We can instantiate a `FlowRunPolicy` and add Pydantic bound checks to
316
336
  # validate that the user passes positive numbers here
317
- self.retries = (
337
+ self.retries: int = (
318
338
  retries if retries is not None else PREFECT_FLOW_DEFAULT_RETRIES.value()
319
339
  )
320
340
 
321
- self.retry_delay_seconds = (
341
+ self.retry_delay_seconds: float | int = (
322
342
  retry_delay_seconds
323
343
  if retry_delay_seconds is not None
324
344
  else PREFECT_FLOW_DEFAULT_RETRY_DELAY_SECONDS.value()
325
345
  )
326
346
 
327
- self.parameters = parameter_schema(self.fn)
347
+ self.parameters: ParameterSchema = parameter_schema(self.fn)
328
348
  self.should_validate_parameters = validate_parameters
329
349
 
330
350
  if self.should_validate_parameters:
@@ -355,11 +375,11 @@ class Flow(Generic[P, R]):
355
375
  self.result_storage = result_storage
356
376
  self.result_serializer = result_serializer
357
377
  self.cache_result_in_memory = cache_result_in_memory
358
- self.on_completion_hooks = on_completion or []
359
- self.on_failure_hooks = on_failure or []
360
- self.on_cancellation_hooks = on_cancellation or []
361
- self.on_crashed_hooks = on_crashed or []
362
- self.on_running_hooks = on_running or []
378
+ self.on_completion_hooks: list[FlowStateHook[P, R]] = on_completion or []
379
+ self.on_failure_hooks: list[FlowStateHook[P, R]] = on_failure or []
380
+ self.on_cancellation_hooks: list[FlowStateHook[P, R]] = on_cancellation or []
381
+ self.on_crashed_hooks: list[FlowStateHook[P, R]] = on_crashed or []
382
+ self.on_running_hooks: list[FlowStateHook[P, R]] = on_running or []
363
383
 
364
384
  # Used for flows loaded from remote storage
365
385
  self._storage: Optional["RunnerStorage"] = None
@@ -376,7 +396,7 @@ class Flow(Generic[P, R]):
376
396
  def ismethod(self) -> bool:
377
397
  return hasattr(self.fn, "__prefect_self__")
378
398
 
379
- def __get__(self, instance: Any, owner: Any):
399
+ def __get__(self, instance: Any, owner: Any) -> "Flow[P, R]":
380
400
  """
381
401
  Implement the descriptor protocol so that the flow can be used as an instance method.
382
402
  When an instance method is loaded, this method is called with the "self" instance as
@@ -391,7 +411,7 @@ class Flow(Generic[P, R]):
391
411
  # of the flow's function. This will allow it to be automatically added to the flow's parameters
392
412
  else:
393
413
  bound_flow = copy(self)
394
- bound_flow.fn.__prefect_self__ = instance
414
+ setattr(bound_flow.fn, "__prefect_self__", instance)
395
415
  return bound_flow
396
416
 
397
417
  def with_options(
@@ -404,7 +424,7 @@ class Flow(Generic[P, R]):
404
424
  description: Optional[str] = None,
405
425
  flow_run_name: Optional[Union[Callable[[], str], str]] = None,
406
426
  task_runner: Union[
407
- Type[TaskRunner[PrefectFuture[R]]], TaskRunner[PrefectFuture[R]], None
427
+ Type[TaskRunner[PrefectFuture[Any]]], TaskRunner[PrefectFuture[Any]], None
408
428
  ] = None,
409
429
  timeout_seconds: Union[int, float, None] = None,
410
430
  validate_parameters: Optional[bool] = None,
@@ -413,11 +433,11 @@ class Flow(Generic[P, R]):
413
433
  result_serializer: Optional[ResultSerializer] = NotSet, # type: ignore
414
434
  cache_result_in_memory: Optional[bool] = None,
415
435
  log_prints: Optional[bool] = NotSet, # type: ignore
416
- on_completion: Optional[list[StateHookCallable]] = None,
417
- on_failure: Optional[list[StateHookCallable]] = None,
418
- on_cancellation: Optional[list[StateHookCallable]] = None,
419
- on_crashed: Optional[list[StateHookCallable]] = None,
420
- on_running: Optional[list[StateHookCallable]] = None,
436
+ on_completion: Optional[list[FlowStateHook[P, R]]] = None,
437
+ on_failure: Optional[list[FlowStateHook[P, R]]] = None,
438
+ on_cancellation: Optional[list[FlowStateHook[P, R]]] = None,
439
+ on_crashed: Optional[list[FlowStateHook[P, R]]] = None,
440
+ on_running: Optional[list[FlowStateHook[P, R]]] = None,
421
441
  ) -> "Flow[P, R]":
422
442
  """
423
443
  Create a new flow from the current object, updating provided options.
@@ -473,13 +493,18 @@ class Flow(Generic[P, R]):
473
493
  >>> state = my_flow.with_options(task_runner=ThreadPoolTaskRunner)(1, 3)
474
494
  >>> assert state.result() == 4
475
495
  """
496
+ new_task_runner = (
497
+ task_runner() if isinstance(task_runner, type) else task_runner
498
+ )
499
+ if new_task_runner is None:
500
+ new_task_runner = self.task_runner
476
501
  new_flow = Flow(
477
502
  fn=self.fn,
478
503
  name=name or self.name,
479
504
  description=description or self.description,
480
505
  flow_run_name=flow_run_name or self.flow_run_name,
481
506
  version=version or self.version,
482
- task_runner=task_runner or self.task_runner,
507
+ task_runner=new_task_runner,
483
508
  retries=retries if retries is not None else self.retries,
484
509
  retry_delay_seconds=(
485
510
  retry_delay_seconds
@@ -533,7 +558,7 @@ class Flow(Generic[P, R]):
533
558
  ParameterTypeError: if the provided parameters are not valid
534
559
  """
535
560
 
536
- def resolve_block_reference(data: Any) -> Any:
561
+ def resolve_block_reference(data: Any | dict[str, Any]) -> Any:
537
562
  if isinstance(data, dict) and "$ref" in data:
538
563
  return Block.load_from_ref(data["$ref"], _sync=True)
539
564
  return data
@@ -596,7 +621,9 @@ class Flow(Generic[P, R]):
596
621
  }
597
622
  return cast_parameters
598
623
 
599
- def serialize_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]:
624
+ def serialize_parameters(
625
+ self, parameters: dict[str, Any | PrefectFuture[Any] | State]
626
+ ) -> dict[str, Any]:
600
627
  """
601
628
  Convert parameters to a serializable form.
602
629
 
@@ -604,10 +631,10 @@ class Flow(Generic[P, R]):
604
631
  converting everything directly to a string. This maintains basic types like
605
632
  integers during API roundtrips.
606
633
  """
607
- serialized_parameters = {}
634
+ serialized_parameters: dict[str, Any] = {}
608
635
  for key, value in parameters.items():
609
636
  # do not serialize the bound self object
610
- if self.ismethod and value is self.fn.__prefect_self__:
637
+ if self.ismethod and value is getattr(self.fn, "__prefect_self__", None):
611
638
  continue
612
639
  if isinstance(value, (PrefectFuture, State)):
613
640
  # Don't call jsonable_encoder() on a PrefectFuture or State to
@@ -758,23 +785,23 @@ class Flow(Generic[P, R]):
758
785
  _sla=_sla,
759
786
  )
760
787
 
761
- def on_completion(self, fn: StateHookCallable) -> StateHookCallable:
788
+ def on_completion(self, fn: FlowStateHook[P, R]) -> FlowStateHook[P, R]:
762
789
  self.on_completion_hooks.append(fn)
763
790
  return fn
764
791
 
765
- def on_cancellation(self, fn: StateHookCallable) -> StateHookCallable:
792
+ def on_cancellation(self, fn: FlowStateHook[P, R]) -> FlowStateHook[P, R]:
766
793
  self.on_cancellation_hooks.append(fn)
767
794
  return fn
768
795
 
769
- def on_crashed(self, fn: StateHookCallable) -> StateHookCallable:
796
+ def on_crashed(self, fn: FlowStateHook[P, R]) -> FlowStateHook[P, R]:
770
797
  self.on_crashed_hooks.append(fn)
771
798
  return fn
772
799
 
773
- def on_running(self, fn: StateHookCallable) -> StateHookCallable:
800
+ def on_running(self, fn: FlowStateHook[P, R]) -> FlowStateHook[P, R]:
774
801
  self.on_running_hooks.append(fn)
775
802
  return fn
776
803
 
777
- def on_failure(self, fn: StateHookCallable) -> StateHookCallable:
804
+ def on_failure(self, fn: FlowStateHook[P, R]) -> FlowStateHook[P, R]:
778
805
  self.on_failure_hooks.append(fn)
779
806
  return fn
780
807
 
@@ -805,7 +832,7 @@ class Flow(Generic[P, R]):
805
832
  limit: Optional[int] = None,
806
833
  webserver: bool = False,
807
834
  entrypoint_type: EntrypointType = EntrypointType.FILE_PATH,
808
- ):
835
+ ) -> None:
809
836
  """
810
837
  Creates a deployment for this flow and starts a runner to monitor for scheduled work.
811
838
 
@@ -1173,7 +1200,9 @@ class Flow(Generic[P, R]):
1173
1200
  work_pool = await client.read_work_pool(work_pool_name)
1174
1201
  active_workers = await client.read_workers_for_work_pool(
1175
1202
  work_pool_name,
1176
- worker_filter=WorkerFilter(status={"any_": ["ONLINE"]}),
1203
+ worker_filter=WorkerFilter(
1204
+ status=WorkerFilterStatus(any_=["ONLINE"])
1205
+ ),
1177
1206
  )
1178
1207
  except ObjectNotFound as exc:
1179
1208
  raise ValueError(
@@ -1181,7 +1210,7 @@ class Flow(Generic[P, R]):
1181
1210
  " deploying this flow."
1182
1211
  ) from exc
1183
1212
 
1184
- deployment = await self.to_deployment(
1213
+ to_deployment_coro = self.to_deployment(
1185
1214
  name=name,
1186
1215
  interval=interval,
1187
1216
  cron=cron,
@@ -1201,9 +1230,14 @@ class Flow(Generic[P, R]):
1201
1230
  _sla=_sla,
1202
1231
  )
1203
1232
 
1233
+ if TYPE_CHECKING:
1234
+ assert inspect.isawaitable(to_deployment_coro)
1235
+
1236
+ deployment = await to_deployment_coro
1237
+
1204
1238
  from prefect.deployments.runner import deploy
1205
1239
 
1206
- deployment_ids = await deploy(
1240
+ deploy_coro = deploy(
1207
1241
  deployment,
1208
1242
  work_pool_name=work_pool_name,
1209
1243
  image=image,
@@ -1212,6 +1246,10 @@ class Flow(Generic[P, R]):
1212
1246
  print_next_steps_message=False,
1213
1247
  ignore_warnings=ignore_warnings,
1214
1248
  )
1249
+ if TYPE_CHECKING:
1250
+ assert inspect.isawaitable(deploy_coro)
1251
+
1252
+ deployment_ids = await deploy_coro
1215
1253
 
1216
1254
  if print_next_steps:
1217
1255
  console = Console()
@@ -1442,11 +1480,11 @@ class FlowDecorator:
1442
1480
  result_serializer: Optional[ResultSerializer] = None,
1443
1481
  cache_result_in_memory: bool = True,
1444
1482
  log_prints: Optional[bool] = None,
1445
- on_completion: Optional[list[StateHookCallable]] = None,
1446
- on_failure: Optional[list[StateHookCallable]] = None,
1447
- on_cancellation: Optional[list[StateHookCallable]] = None,
1448
- on_crashed: Optional[list[StateHookCallable]] = None,
1449
- on_running: Optional[list[StateHookCallable]] = None,
1483
+ on_completion: Optional[list[FlowStateHook[..., Any]]] = None,
1484
+ on_failure: Optional[list[FlowStateHook[..., Any]]] = None,
1485
+ on_cancellation: Optional[list[FlowStateHook[..., Any]]] = None,
1486
+ on_crashed: Optional[list[FlowStateHook[..., Any]]] = None,
1487
+ on_running: Optional[list[FlowStateHook[..., Any]]] = None,
1450
1488
  ) -> Callable[[Callable[P, R]], Flow[P, R]]:
1451
1489
  ...
1452
1490
 
@@ -1469,11 +1507,11 @@ class FlowDecorator:
1469
1507
  result_serializer: Optional[ResultSerializer] = None,
1470
1508
  cache_result_in_memory: bool = True,
1471
1509
  log_prints: Optional[bool] = None,
1472
- on_completion: Optional[list[StateHookCallable]] = None,
1473
- on_failure: Optional[list[StateHookCallable]] = None,
1474
- on_cancellation: Optional[list[StateHookCallable]] = None,
1475
- on_crashed: Optional[list[StateHookCallable]] = None,
1476
- on_running: Optional[list[StateHookCallable]] = None,
1510
+ on_completion: Optional[list[FlowStateHook[..., Any]]] = None,
1511
+ on_failure: Optional[list[FlowStateHook[..., Any]]] = None,
1512
+ on_cancellation: Optional[list[FlowStateHook[..., Any]]] = None,
1513
+ on_crashed: Optional[list[FlowStateHook[..., Any]]] = None,
1514
+ on_running: Optional[list[FlowStateHook[..., Any]]] = None,
1477
1515
  ) -> Callable[[Callable[P, R]], Flow[P, R]]:
1478
1516
  ...
1479
1517
 
@@ -1495,11 +1533,11 @@ class FlowDecorator:
1495
1533
  result_serializer: Optional[ResultSerializer] = None,
1496
1534
  cache_result_in_memory: bool = True,
1497
1535
  log_prints: Optional[bool] = None,
1498
- on_completion: Optional[list[StateHookCallable]] = None,
1499
- on_failure: Optional[list[StateHookCallable]] = None,
1500
- on_cancellation: Optional[list[StateHookCallable]] = None,
1501
- on_crashed: Optional[list[StateHookCallable]] = None,
1502
- on_running: Optional[list[StateHookCallable]] = None,
1536
+ on_completion: Optional[list[FlowStateHook[..., Any]]] = None,
1537
+ on_failure: Optional[list[FlowStateHook[..., Any]]] = None,
1538
+ on_cancellation: Optional[list[FlowStateHook[..., Any]]] = None,
1539
+ on_crashed: Optional[list[FlowStateHook[..., Any]]] = None,
1540
+ on_running: Optional[list[FlowStateHook[..., Any]]] = None,
1503
1541
  ) -> Union[Flow[P, R], Callable[[Callable[P, R]], Flow[P, R]]]:
1504
1542
  """
1505
1543
  Decorator to designate a function as a Prefect workflow.
@@ -1673,7 +1711,7 @@ class FlowDecorator:
1673
1711
  ...
1674
1712
 
1675
1713
 
1676
- flow = FlowDecorator()
1714
+ flow: FlowDecorator = FlowDecorator()
1677
1715
 
1678
1716
 
1679
1717
  def _raise_on_name_with_banned_characters(name: Optional[str]) -> Optional[str]:
@@ -1925,7 +1963,11 @@ async def aserve(
1925
1963
 
1926
1964
  runner = Runner(pause_on_shutdown=pause_on_shutdown, limit=limit, **kwargs)
1927
1965
  for deployment in args:
1928
- await runner.add_deployment(deployment)
1966
+ add_deployment_coro = runner.add_deployment(deployment)
1967
+ if TYPE_CHECKING:
1968
+ assert inspect.isawaitable(add_deployment_coro)
1969
+
1970
+ await add_deployment_coro
1929
1971
 
1930
1972
  if print_starting_message:
1931
1973
  _display_serve_start_message(*args)
@@ -1971,13 +2013,16 @@ async def load_flow_from_flow_run(
1971
2013
  ignore_storage: bool = False,
1972
2014
  storage_base_path: Optional[str] = None,
1973
2015
  use_placeholder_flow: bool = True,
1974
- ) -> Flow[P, Any]:
2016
+ ) -> Flow[..., Any]:
1975
2017
  """
1976
2018
  Load a flow from the location/script provided in a deployment's storage document.
1977
2019
 
1978
2020
  If `ignore_storage=True` is provided, no pull from remote storage occurs. This flag
1979
2021
  is largely for testing, and assumes the flow is already available locally.
1980
2022
  """
2023
+ if flow_run.deployment_id is None:
2024
+ raise ValueError("Flow run does not have an associated deployment")
2025
+
1981
2026
  deployment = await client.read_deployment(flow_run.deployment_id)
1982
2027
 
1983
2028
  if deployment.entrypoint is None:
prefect/futures.py CHANGED
@@ -5,7 +5,7 @@ import threading
5
5
  import uuid
6
6
  from collections.abc import Generator, Iterator
7
7
  from functools import partial
8
- from typing import Any, Callable, Generic, Optional, Union
8
+ from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, Union
9
9
 
10
10
  from typing_extensions import NamedTuple, Self, TypeVar
11
11
 
@@ -22,7 +22,10 @@ from prefect.utilities.timeout import timeout as timeout_context
22
22
  F = TypeVar("F")
23
23
  R = TypeVar("R")
24
24
 
25
- logger = get_logger(__name__)
25
+ if TYPE_CHECKING:
26
+ import logging
27
+
28
+ logger: "logging.Logger" = get_logger(__name__)
26
29
 
27
30
 
28
31
  class PrefectFuture(abc.ABC, Generic[R]):
@@ -90,7 +93,7 @@ class PrefectFuture(abc.ABC, Generic[R]):
90
93
  """
91
94
 
92
95
  @abc.abstractmethod
93
- def add_done_callback(self, fn: Callable[["PrefectFuture[R]"], None]):
96
+ def add_done_callback(self, fn: Callable[["PrefectFuture[R]"], None]) -> None:
94
97
  """
95
98
  Add a callback to be run when the future completes or is cancelled.
96
99
 
@@ -173,7 +176,7 @@ class PrefectConcurrentFuture(PrefectWrappedFuture[R, concurrent.futures.Future[
173
176
  _result = run_coro_as_sync(_result)
174
177
  return _result
175
178
 
176
- def __del__(self):
179
+ def __del__(self) -> None:
177
180
  if self._final_state or self._wrapped_future.done():
178
181
  return
179
182
  try:
@@ -202,7 +205,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
202
205
  def wait(self, timeout: Optional[float] = None) -> None:
203
206
  return run_coro_as_sync(self.wait_async(timeout=timeout))
204
207
 
205
- async def wait_async(self, timeout: Optional[float] = None):
208
+ async def wait_async(self, timeout: Optional[float] = None) -> None:
206
209
  if self._final_state:
207
210
  logger.debug(
208
211
  "Final state already set for %s. Returning...", self.task_run_id
@@ -216,6 +219,10 @@ class PrefectDistributedFuture(PrefectFuture[R]):
216
219
  # Read task run to see if it is still running
217
220
  async with get_client() as client:
218
221
  task_run = await client.read_task_run(task_run_id=self._task_run_id)
222
+ if task_run.state is None:
223
+ raise RuntimeError(
224
+ f"Task run {self.task_run_id} has no state which means it hasn't started yet."
225
+ )
219
226
  if task_run.state.is_final():
220
227
  logger.debug(
221
228
  "Task run %s already finished. Returning...",
@@ -260,7 +267,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
260
267
  raise_on_failure=raise_on_failure, fetch=True
261
268
  )
262
269
 
263
- def add_done_callback(self, fn: Callable[[PrefectFuture[R]], None]):
270
+ def add_done_callback(self, fn: Callable[[PrefectFuture[R]], None]) -> None:
264
271
  if self._final_state:
265
272
  fn(self)
266
273
  return
@@ -278,7 +285,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
278
285
  return False
279
286
  return self.task_run_id == other.task_run_id
280
287
 
281
- def __hash__(self):
288
+ def __hash__(self) -> int:
282
289
  return hash(self.task_run_id)
283
290
 
284
291
 
@@ -1,5 +1,6 @@
1
1
  from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol, Type
2
2
 
3
+ from prefect.infrastructure.provisioners.coiled import CoiledPushProvisioner
3
4
  from prefect.infrastructure.provisioners.modal import ModalPushProvisioner
4
5
  from .cloud_run import CloudRunPushProvisioner
5
6
  from .container_instance import ContainerInstancePushProvisioner
@@ -15,6 +16,7 @@ _provisioners = {
15
16
  "azure-container-instance:push": ContainerInstancePushProvisioner,
16
17
  "ecs:push": ElasticContainerServicePushProvisioner,
17
18
  "modal:push": ModalPushProvisioner,
19
+ "coiled:push": CoiledPushProvisioner,
18
20
  }
19
21
 
20
22
 
@@ -33,7 +33,7 @@ if TYPE_CHECKING:
33
33
 
34
34
  class CloudRunPushProvisioner:
35
35
  def __init__(self):
36
- self._console = Console()
36
+ self._console: Console = Console()
37
37
  self._project = None
38
38
  self._region = None
39
39
  self._service_account_name = "prefect-cloud-run"
@@ -41,14 +41,14 @@ class CloudRunPushProvisioner:
41
41
  self._image_repository_name = "prefect-images"
42
42
 
43
43
  @property
44
- def console(self):
44
+ def console(self) -> Console:
45
45
  return self._console
46
46
 
47
47
  @console.setter
48
- def console(self, value):
48
+ def console(self, value: Console) -> None:
49
49
  self._console = value
50
50
 
51
- async def _run_command(self, command: str, *args, **kwargs):
51
+ async def _run_command(self, command: str, *args: Any, **kwargs: Any) -> str:
52
52
  result = await run_process(shlex.split(command), check=False, *args, **kwargs)
53
53
 
54
54
  if result.returncode != 0: