prefect-client 3.1.6__py3-none-any.whl → 3.1.8__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 (55) hide show
  1. prefect/_experimental/__init__.py +0 -0
  2. prefect/_experimental/lineage.py +181 -0
  3. prefect/_internal/compatibility/async_dispatch.py +38 -9
  4. prefect/_internal/pydantic/v2_validated_func.py +15 -10
  5. prefect/_internal/retries.py +15 -6
  6. prefect/_internal/schemas/bases.py +2 -1
  7. prefect/_internal/schemas/validators.py +5 -4
  8. prefect/_version.py +3 -3
  9. prefect/blocks/core.py +144 -17
  10. prefect/blocks/system.py +2 -1
  11. prefect/client/orchestration.py +106 -0
  12. prefect/client/schemas/actions.py +5 -5
  13. prefect/client/schemas/filters.py +1 -1
  14. prefect/client/schemas/objects.py +5 -5
  15. prefect/client/schemas/responses.py +1 -2
  16. prefect/client/schemas/schedules.py +1 -1
  17. prefect/client/subscriptions.py +2 -1
  18. prefect/client/utilities.py +15 -1
  19. prefect/context.py +1 -1
  20. prefect/deployments/flow_runs.py +3 -3
  21. prefect/deployments/runner.py +14 -14
  22. prefect/deployments/steps/core.py +3 -1
  23. prefect/deployments/steps/pull.py +60 -12
  24. prefect/events/clients.py +55 -4
  25. prefect/events/filters.py +1 -1
  26. prefect/events/related.py +2 -1
  27. prefect/events/schemas/events.py +1 -1
  28. prefect/events/utilities.py +2 -0
  29. prefect/events/worker.py +8 -0
  30. prefect/flow_engine.py +41 -81
  31. prefect/flow_runs.py +4 -2
  32. prefect/flows.py +4 -6
  33. prefect/results.py +43 -22
  34. prefect/runner/runner.py +129 -18
  35. prefect/runner/storage.py +3 -3
  36. prefect/serializers.py +28 -24
  37. prefect/settings/__init__.py +1 -0
  38. prefect/settings/base.py +3 -2
  39. prefect/settings/models/api.py +4 -0
  40. prefect/settings/models/experiments.py +5 -0
  41. prefect/settings/models/runner.py +8 -0
  42. prefect/settings/models/server/api.py +7 -1
  43. prefect/task_engine.py +34 -26
  44. prefect/task_worker.py +43 -25
  45. prefect/tasks.py +118 -125
  46. prefect/telemetry/instrumentation.py +1 -1
  47. prefect/telemetry/processors.py +10 -7
  48. prefect/telemetry/run_telemetry.py +157 -33
  49. prefect/types/__init__.py +4 -1
  50. prefect/variables.py +127 -19
  51. {prefect_client-3.1.6.dist-info → prefect_client-3.1.8.dist-info}/METADATA +2 -1
  52. {prefect_client-3.1.6.dist-info → prefect_client-3.1.8.dist-info}/RECORD +55 -53
  53. {prefect_client-3.1.6.dist-info → prefect_client-3.1.8.dist-info}/LICENSE +0 -0
  54. {prefect_client-3.1.6.dist-info → prefect_client-3.1.8.dist-info}/WHEEL +0 -0
  55. {prefect_client-3.1.6.dist-info → prefect_client-3.1.8.dist-info}/top_level.txt +0 -0
prefect/tasks.py CHANGED
@@ -15,16 +15,11 @@ from typing import (
15
15
  Awaitable,
16
16
  Callable,
17
17
  Coroutine,
18
- Dict,
19
18
  Generic,
20
19
  Iterable,
21
- List,
22
20
  NoReturn,
23
21
  Optional,
24
22
  Protocol,
25
- Set,
26
- Tuple,
27
- Type,
28
23
  TypeVar,
29
24
  Union,
30
25
  cast,
@@ -32,7 +27,7 @@ from typing import (
32
27
  )
33
28
  from uuid import UUID, uuid4
34
29
 
35
- from typing_extensions import Literal, ParamSpec, Self, TypeIs
30
+ from typing_extensions import Literal, ParamSpec, Self, TypeAlias, TypeIs
36
31
 
37
32
  import prefect.states
38
33
  from prefect.cache_policies import DEFAULT, NONE, CachePolicy
@@ -64,10 +59,7 @@ from prefect.settings import (
64
59
  )
65
60
  from prefect.states import Pending, Scheduled, State
66
61
  from prefect.utilities.annotations import NotSet
67
- from prefect.utilities.asyncutils import (
68
- run_coro_as_sync,
69
- sync_compatible,
70
- )
62
+ from prefect.utilities.asyncutils import run_coro_as_sync, sync_compatible
71
63
  from prefect.utilities.callables import (
72
64
  expand_mapping_parameters,
73
65
  get_call_parameters,
@@ -92,7 +84,7 @@ logger = get_logger("tasks")
92
84
 
93
85
 
94
86
  def task_input_hash(
95
- context: "TaskRunContext", arguments: Dict[str, Any]
87
+ context: "TaskRunContext", arguments: dict[str, Any]
96
88
  ) -> Optional[str]:
97
89
  """
98
90
  A task cache key implementation which hashes all inputs to the task using a JSON or
@@ -117,7 +109,7 @@ def task_input_hash(
117
109
  )
118
110
 
119
111
 
120
- def exponential_backoff(backoff_factor: float) -> Callable[[int], List[float]]:
112
+ def exponential_backoff(backoff_factor: float) -> Callable[[int], list[float]]:
121
113
  """
122
114
  A task retry backoff utility that configures exponential backoff for task retries.
123
115
  The exponential backoff design matches the urllib3 implementation.
@@ -130,7 +122,7 @@ def exponential_backoff(backoff_factor: float) -> Callable[[int], List[float]]:
130
122
  a callable that can be passed to the task constructor
131
123
  """
132
124
 
133
- def retry_backoff_callable(retries: int) -> List[float]:
125
+ def retry_backoff_callable(retries: int) -> list[float]:
134
126
  # no more than 50 retry delays can be configured on a task
135
127
  retries = min(retries, 50)
136
128
 
@@ -142,8 +134,8 @@ def exponential_backoff(backoff_factor: float) -> Callable[[int], List[float]]:
142
134
  def _infer_parent_task_runs(
143
135
  flow_run_context: Optional[FlowRunContext],
144
136
  task_run_context: Optional[TaskRunContext],
145
- parameters: Dict[str, Any],
146
- ):
137
+ parameters: dict[str, Any],
138
+ ) -> list[TaskRunResult]:
147
139
  """
148
140
  Attempt to infer the parent task runs for this task run based on the
149
141
  provided flow run and task run contexts, as well as any parameters. It is
@@ -152,7 +144,7 @@ def _infer_parent_task_runs(
152
144
  a parent. This is expected to happen when task inputs are yielded from
153
145
  generator tasks.
154
146
  """
155
- parents = []
147
+ parents: list[TaskRunResult] = []
156
148
 
157
149
  # check if this task has a parent task run based on running in another
158
150
  # task run's existing context. A task run is only considered a parent if
@@ -234,6 +226,14 @@ class TaskRunNameCallbackWithParameters(Protocol):
234
226
  ...
235
227
 
236
228
 
229
+ StateHookCallable: TypeAlias = Callable[
230
+ ["Task[..., Any]", TaskRun, State], Union[Awaitable[None], None]
231
+ ]
232
+ TaskRunNameValueOrCallable: TypeAlias = Union[
233
+ Callable[[], str], TaskRunNameCallbackWithParameters, str
234
+ ]
235
+
236
+
237
237
  class Task(Generic[P, R]):
238
238
  """
239
239
  A Prefect task definition.
@@ -316,21 +316,19 @@ class Task(Generic[P, R]):
316
316
  description: Optional[str] = None,
317
317
  tags: Optional[Iterable[str]] = None,
318
318
  version: Optional[str] = None,
319
- cache_policy: Union[CachePolicy, Type[NotSet]] = NotSet,
319
+ cache_policy: Union[CachePolicy, type[NotSet]] = NotSet,
320
320
  cache_key_fn: Optional[
321
- Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
321
+ Callable[["TaskRunContext", dict[str, Any]], Optional[str]]
322
322
  ] = None,
323
323
  cache_expiration: Optional[datetime.timedelta] = None,
324
- task_run_name: Optional[
325
- Union[Callable[[], str], TaskRunNameCallbackWithParameters, str]
326
- ] = None,
324
+ task_run_name: Optional[TaskRunNameValueOrCallable] = None,
327
325
  retries: Optional[int] = None,
328
326
  retry_delay_seconds: Optional[
329
327
  Union[
330
328
  float,
331
329
  int,
332
- List[float],
333
- Callable[[int], List[float]],
330
+ list[float],
331
+ Callable[[int], list[float]],
334
332
  ]
335
333
  ] = None,
336
334
  retry_jitter_factor: Optional[float] = None,
@@ -342,11 +340,13 @@ class Task(Generic[P, R]):
342
340
  timeout_seconds: Union[int, float, None] = None,
343
341
  log_prints: Optional[bool] = False,
344
342
  refresh_cache: Optional[bool] = None,
345
- on_completion: Optional[List[Callable[["Task", TaskRun, State], None]]] = None,
346
- on_failure: Optional[List[Callable[["Task", TaskRun, State], None]]] = None,
347
- on_rollback: Optional[List[Callable[["Transaction"], None]]] = None,
348
- on_commit: Optional[List[Callable[["Transaction"], None]]] = None,
349
- retry_condition_fn: Optional[Callable[["Task", TaskRun, State], bool]] = None,
343
+ on_completion: Optional[list[StateHookCallable]] = None,
344
+ on_failure: Optional[list[StateHookCallable]] = None,
345
+ on_rollback: Optional[list[Callable[["Transaction"], None]]] = None,
346
+ on_commit: Optional[list[Callable[["Transaction"], None]]] = None,
347
+ retry_condition_fn: Optional[
348
+ Callable[["Task[..., Any]", TaskRun, State], bool]
349
+ ] = None,
350
350
  viz_return_value: Optional[Any] = None,
351
351
  ):
352
352
  # Validate if hook passed is list and contains callables
@@ -471,6 +471,10 @@ class Task(Generic[P, R]):
471
471
 
472
472
  if callable(retry_delay_seconds):
473
473
  self.retry_delay_seconds = retry_delay_seconds(retries)
474
+ elif not isinstance(retry_delay_seconds, (list, int, float, type(None))):
475
+ raise TypeError(
476
+ f"Invalid `retry_delay_seconds` provided; must be an int, float, list or callable. Received type {type(retry_delay_seconds)}"
477
+ )
474
478
  else:
475
479
  self.retry_delay_seconds = retry_delay_seconds
476
480
 
@@ -516,7 +520,7 @@ class Task(Generic[P, R]):
516
520
  def ismethod(self) -> bool:
517
521
  return hasattr(self.fn, "__prefect_self__")
518
522
 
519
- def __get__(self, instance, owner):
523
+ def __get__(self, instance: Any, owner: Any):
520
524
  """
521
525
  Implement the descriptor protocol so that the task can be used as an instance method.
522
526
  When an instance method is loaded, this method is called with the "self" instance as
@@ -531,7 +535,7 @@ class Task(Generic[P, R]):
531
535
  # of the task's function. This will allow it to be automatically added to the task's parameters
532
536
  else:
533
537
  bound_task = copy(self)
534
- bound_task.fn.__prefect_self__ = instance
538
+ bound_task.fn.__prefect_self__ = instance # type: ignore[attr-defined]
535
539
  return bound_task
536
540
 
537
541
  def with_options(
@@ -540,40 +544,36 @@ class Task(Generic[P, R]):
540
544
  name: Optional[str] = None,
541
545
  description: Optional[str] = None,
542
546
  tags: Optional[Iterable[str]] = None,
543
- cache_policy: Union[CachePolicy, Type[NotSet]] = NotSet,
547
+ cache_policy: Union[CachePolicy, type[NotSet]] = NotSet,
544
548
  cache_key_fn: Optional[
545
- Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
549
+ Callable[["TaskRunContext", dict[str, Any]], Optional[str]]
546
550
  ] = None,
547
551
  task_run_name: Optional[
548
- Union[
549
- Callable[[], str], TaskRunNameCallbackWithParameters, str, Type[NotSet]
550
- ]
552
+ Union[TaskRunNameValueOrCallable, type[NotSet]]
551
553
  ] = NotSet,
552
554
  cache_expiration: Optional[datetime.timedelta] = None,
553
- retries: Union[int, Type[NotSet]] = NotSet,
555
+ retries: Union[int, type[NotSet]] = NotSet,
554
556
  retry_delay_seconds: Union[
555
557
  float,
556
558
  int,
557
- List[float],
558
- Callable[[int], List[float]],
559
- Type[NotSet],
559
+ list[float],
560
+ Callable[[int], list[float]],
561
+ type[NotSet],
560
562
  ] = NotSet,
561
- retry_jitter_factor: Union[float, Type[NotSet]] = NotSet,
562
- persist_result: Union[bool, Type[NotSet]] = NotSet,
563
- result_storage: Union[ResultStorage, Type[NotSet]] = NotSet,
564
- result_serializer: Union[ResultSerializer, Type[NotSet]] = NotSet,
565
- result_storage_key: Union[str, Type[NotSet]] = NotSet,
563
+ retry_jitter_factor: Union[float, type[NotSet]] = NotSet,
564
+ persist_result: Union[bool, type[NotSet]] = NotSet,
565
+ result_storage: Union[ResultStorage, type[NotSet]] = NotSet,
566
+ result_serializer: Union[ResultSerializer, type[NotSet]] = NotSet,
567
+ result_storage_key: Union[str, type[NotSet]] = NotSet,
566
568
  cache_result_in_memory: Optional[bool] = None,
567
569
  timeout_seconds: Union[int, float, None] = None,
568
- log_prints: Union[bool, Type[NotSet]] = NotSet,
569
- refresh_cache: Union[bool, Type[NotSet]] = NotSet,
570
- on_completion: Optional[
571
- List[Callable[["Task", TaskRun, State], Union[Awaitable[None], None]]]
572
- ] = None,
573
- on_failure: Optional[
574
- List[Callable[["Task", TaskRun, State], Union[Awaitable[None], None]]]
570
+ log_prints: Union[bool, type[NotSet]] = NotSet,
571
+ refresh_cache: Union[bool, type[NotSet]] = NotSet,
572
+ on_completion: Optional[list[StateHookCallable]] = None,
573
+ on_failure: Optional[list[StateHookCallable]] = None,
574
+ retry_condition_fn: Optional[
575
+ Callable[["Task[..., Any]", TaskRun, State], bool]
575
576
  ] = None,
576
- retry_condition_fn: Optional[Callable[["Task", TaskRun, State], bool]] = None,
577
577
  viz_return_value: Optional[Any] = None,
578
578
  ):
579
579
  """
@@ -710,15 +710,11 @@ class Task(Generic[P, R]):
710
710
  viz_return_value=viz_return_value or self.viz_return_value,
711
711
  )
712
712
 
713
- def on_completion(
714
- self, fn: Callable[["Task", TaskRun, State], None]
715
- ) -> Callable[["Task", TaskRun, State], None]:
713
+ def on_completion(self, fn: StateHookCallable) -> StateHookCallable:
716
714
  self.on_completion_hooks.append(fn)
717
715
  return fn
718
716
 
719
- def on_failure(
720
- self, fn: Callable[["Task", TaskRun, State], None]
721
- ) -> Callable[["Task", TaskRun, State], None]:
717
+ def on_failure(self, fn: StateHookCallable) -> StateHookCallable:
722
718
  self.on_failure_hooks.append(fn)
723
719
  return fn
724
720
 
@@ -738,11 +734,11 @@ class Task(Generic[P, R]):
738
734
  self,
739
735
  client: Optional["PrefectClient"] = None,
740
736
  id: Optional[UUID] = None,
741
- parameters: Optional[Dict[str, Any]] = None,
737
+ parameters: Optional[dict[str, Any]] = None,
742
738
  flow_run_context: Optional[FlowRunContext] = None,
743
739
  parent_task_run_context: Optional[TaskRunContext] = None,
744
- wait_for: Optional[Iterable[PrefectFuture]] = None,
745
- extra_task_inputs: Optional[Dict[str, Set[TaskRunInput]]] = None,
740
+ wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
741
+ extra_task_inputs: Optional[dict[str, set[TaskRunInput]]] = None,
746
742
  deferred: bool = False,
747
743
  ) -> TaskRun:
748
744
  from prefect.utilities._engine import dynamic_key_for_task_run
@@ -786,7 +782,7 @@ class Task(Generic[P, R]):
786
782
  result_storage=await get_or_create_default_task_scheduling_storage()
787
783
  ).update_for_task(self)
788
784
  context = serialize_context()
789
- data: Dict[str, Any] = {"context": context}
785
+ data: dict[str, Any] = {"context": context}
790
786
  if parameters:
791
787
  data["parameters"] = parameters
792
788
  if wait_for:
@@ -839,11 +835,11 @@ class Task(Generic[P, R]):
839
835
  self,
840
836
  client: Optional["PrefectClient"] = None,
841
837
  id: Optional[UUID] = None,
842
- parameters: Optional[Dict[str, Any]] = None,
838
+ parameters: Optional[dict[str, Any]] = None,
843
839
  flow_run_context: Optional[FlowRunContext] = None,
844
840
  parent_task_run_context: Optional[TaskRunContext] = None,
845
- wait_for: Optional[Iterable[PrefectFuture]] = None,
846
- extra_task_inputs: Optional[Dict[str, Set[TaskRunInput]]] = None,
841
+ wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
842
+ extra_task_inputs: Optional[dict[str, set[TaskRunInput]]] = None,
847
843
  deferred: bool = False,
848
844
  ) -> TaskRun:
849
845
  from prefect.utilities._engine import dynamic_key_for_task_run
@@ -887,7 +883,7 @@ class Task(Generic[P, R]):
887
883
  result_storage=await get_or_create_default_task_scheduling_storage()
888
884
  ).update_for_task(task)
889
885
  context = serialize_context()
890
- data: Dict[str, Any] = {"context": context}
886
+ data: dict[str, Any] = {"context": context}
891
887
  if parameters:
892
888
  data["parameters"] = parameters
893
889
  if wait_for:
@@ -983,7 +979,7 @@ class Task(Generic[P, R]):
983
979
  self,
984
980
  *args: P.args,
985
981
  return_state: bool = False,
986
- wait_for: Optional[Iterable[PrefectFuture]] = None,
982
+ wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
987
983
  **kwargs: P.kwargs,
988
984
  ):
989
985
  """
@@ -1063,7 +1059,7 @@ class Task(Generic[P, R]):
1063
1059
  self,
1064
1060
  *args: Any,
1065
1061
  return_state: bool = False,
1066
- wait_for: Optional[Iterable[PrefectFuture]] = None,
1062
+ wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
1067
1063
  **kwargs: Any,
1068
1064
  ):
1069
1065
  """
@@ -1191,7 +1187,7 @@ class Task(Generic[P, R]):
1191
1187
  wait_for: Optional[Iterable[Union[PrefectFuture[R], R]]] = ...,
1192
1188
  deferred: bool = ...,
1193
1189
  **kwargs: Any,
1194
- ) -> List[State[R]]:
1190
+ ) -> list[State[R]]:
1195
1191
  ...
1196
1192
 
1197
1193
  @overload
@@ -1212,7 +1208,7 @@ class Task(Generic[P, R]):
1212
1208
  wait_for: Optional[Iterable[Union[PrefectFuture[R], R]]] = ...,
1213
1209
  deferred: bool = ...,
1214
1210
  **kwargs: Any,
1215
- ) -> List[State[R]]:
1211
+ ) -> list[State[R]]:
1216
1212
  ...
1217
1213
 
1218
1214
  @overload
@@ -1233,7 +1229,7 @@ class Task(Generic[P, R]):
1233
1229
  wait_for: Optional[Iterable[Union[PrefectFuture[R], R]]] = ...,
1234
1230
  deferred: bool = ...,
1235
1231
  **kwargs: Any,
1236
- ) -> List[State[R]]:
1232
+ ) -> list[State[R]]:
1237
1233
  ...
1238
1234
 
1239
1235
  @overload
@@ -1254,7 +1250,7 @@ class Task(Generic[P, R]):
1254
1250
  wait_for: Optional[Iterable[Union[PrefectFuture[R], R]]] = None,
1255
1251
  deferred: bool = False,
1256
1252
  **kwargs: Any,
1257
- ) -> Union[List[State[R]], PrefectFutureList[R]]:
1253
+ ) -> Union[list[State[R]], PrefectFutureList[R]]:
1258
1254
  """
1259
1255
  Submit a mapped run of the task to a worker.
1260
1256
 
@@ -1413,10 +1409,10 @@ class Task(Generic[P, R]):
1413
1409
 
1414
1410
  def apply_async(
1415
1411
  self,
1416
- args: Optional[Tuple[Any, ...]] = None,
1417
- kwargs: Optional[Dict[str, Any]] = None,
1412
+ args: Optional[tuple[Any, ...]] = None,
1413
+ kwargs: Optional[dict[str, Any]] = None,
1418
1414
  wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
1419
- dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
1415
+ dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
1420
1416
  ) -> PrefectDistributedFuture[R]:
1421
1417
  """
1422
1418
  Create a pending task run for a task worker to execute.
@@ -1516,7 +1512,7 @@ class Task(Generic[P, R]):
1516
1512
 
1517
1513
  return PrefectDistributedFuture(task_run_id=task_run.id)
1518
1514
 
1519
- def delay(self, *args: P.args, **kwargs: P.kwargs) -> PrefectDistributedFuture:
1515
+ def delay(self, *args: P.args, **kwargs: P.kwargs) -> PrefectDistributedFuture[R]:
1520
1516
  """
1521
1517
  An alias for `apply_async` with simpler calling semantics.
1522
1518
 
@@ -1584,6 +1580,14 @@ def task(__fn: Callable[P, R]) -> Task[P, R]:
1584
1580
  ...
1585
1581
 
1586
1582
 
1583
+ # see https://github.com/PrefectHQ/prefect/issues/16380
1584
+ @overload
1585
+ def task(
1586
+ __fn: Literal[None] = None,
1587
+ ) -> Callable[[Callable[P, R]], Task[P, R]]:
1588
+ ...
1589
+
1590
+
1587
1591
  @overload
1588
1592
  def task(
1589
1593
  *,
@@ -1591,20 +1595,18 @@ def task(
1591
1595
  description: Optional[str] = None,
1592
1596
  tags: Optional[Iterable[str]] = None,
1593
1597
  version: Optional[str] = None,
1594
- cache_policy: Union[CachePolicy, Type[NotSet]] = NotSet,
1598
+ cache_policy: Union[CachePolicy, type[NotSet]] = NotSet,
1595
1599
  cache_key_fn: Optional[
1596
- Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
1600
+ Callable[["TaskRunContext", dict[str, Any]], Optional[str]]
1597
1601
  ] = None,
1598
1602
  cache_expiration: Optional[datetime.timedelta] = None,
1599
- task_run_name: Optional[
1600
- Union[Callable[[], str], TaskRunNameCallbackWithParameters, str]
1601
- ] = None,
1603
+ task_run_name: Optional[TaskRunNameValueOrCallable] = None,
1602
1604
  retries: int = 0,
1603
1605
  retry_delay_seconds: Union[
1604
1606
  float,
1605
1607
  int,
1606
- List[float],
1607
- Callable[[int], List[float]],
1608
+ list[float],
1609
+ Callable[[int], list[float]],
1608
1610
  ] = 0,
1609
1611
  retry_jitter_factor: Optional[float] = None,
1610
1612
  persist_result: Optional[bool] = None,
@@ -1615,10 +1617,8 @@ def task(
1615
1617
  timeout_seconds: Union[int, float, None] = None,
1616
1618
  log_prints: Optional[bool] = None,
1617
1619
  refresh_cache: Optional[bool] = None,
1618
- on_completion: Optional[
1619
- List[Callable[["Task[P, R]", TaskRun, State], None]]
1620
- ] = None,
1621
- on_failure: Optional[List[Callable[["Task[P, R]", TaskRun, State], None]]] = None,
1620
+ on_completion: Optional[list[StateHookCallable]] = None,
1621
+ on_failure: Optional[list[StateHookCallable]] = None,
1622
1622
  retry_condition_fn: Optional[Callable[["Task[P, R]", TaskRun, State], bool]] = None,
1623
1623
  viz_return_value: Any = None,
1624
1624
  ) -> Callable[[Callable[P, R]], Task[P, R]]:
@@ -1632,17 +1632,15 @@ def task(
1632
1632
  description: Optional[str] = None,
1633
1633
  tags: Optional[Iterable[str]] = None,
1634
1634
  version: Optional[str] = None,
1635
- cache_policy: Union[CachePolicy, Type[NotSet]] = NotSet,
1635
+ cache_policy: Union[CachePolicy, type[NotSet]] = NotSet,
1636
1636
  cache_key_fn: Union[
1637
- Callable[["TaskRunContext", Dict[str, Any]], Optional[str]], None
1637
+ Callable[["TaskRunContext", dict[str, Any]], Optional[str]], None
1638
1638
  ] = None,
1639
1639
  cache_expiration: Optional[datetime.timedelta] = None,
1640
- task_run_name: Optional[
1641
- Union[Callable[[], str], TaskRunNameCallbackWithParameters, str]
1642
- ] = None,
1640
+ task_run_name: Optional[TaskRunNameValueOrCallable] = None,
1643
1641
  retries: Optional[int] = None,
1644
1642
  retry_delay_seconds: Union[
1645
- float, int, List[float], Callable[[int], List[float]], None
1643
+ float, int, list[float], Callable[[int], list[float]], None
1646
1644
  ] = None,
1647
1645
  retry_jitter_factor: Optional[float] = None,
1648
1646
  persist_result: Optional[bool] = None,
@@ -1653,10 +1651,8 @@ def task(
1653
1651
  timeout_seconds: Union[int, float, None] = None,
1654
1652
  log_prints: Optional[bool] = None,
1655
1653
  refresh_cache: Optional[bool] = None,
1656
- on_completion: Optional[
1657
- List[Callable[["Task[P, R]", TaskRun, State], None]]
1658
- ] = None,
1659
- on_failure: Optional[List[Callable[["Task[P, R]", TaskRun, State], None]]] = None,
1654
+ on_completion: Optional[list[StateHookCallable]] = None,
1655
+ on_failure: Optional[list[StateHookCallable]] = None,
1660
1656
  retry_condition_fn: Optional[Callable[["Task[P, R]", TaskRun, State], bool]] = None,
1661
1657
  viz_return_value: Any = None,
1662
1658
  ):
@@ -1773,34 +1769,31 @@ def task(
1773
1769
  if isinstance(__fn, (classmethod, staticmethod)):
1774
1770
  method_decorator = type(__fn).__name__
1775
1771
  raise TypeError(f"@{method_decorator} should be applied on top of @task")
1776
- return cast(
1777
- Task[P, R],
1778
- Task(
1779
- fn=__fn,
1780
- name=name,
1781
- description=description,
1782
- tags=tags,
1783
- version=version,
1784
- cache_policy=cache_policy,
1785
- cache_key_fn=cache_key_fn,
1786
- cache_expiration=cache_expiration,
1787
- task_run_name=task_run_name,
1788
- retries=retries,
1789
- retry_delay_seconds=retry_delay_seconds,
1790
- retry_jitter_factor=retry_jitter_factor,
1791
- persist_result=persist_result,
1792
- result_storage=result_storage,
1793
- result_storage_key=result_storage_key,
1794
- result_serializer=result_serializer,
1795
- cache_result_in_memory=cache_result_in_memory,
1796
- timeout_seconds=timeout_seconds,
1797
- log_prints=log_prints,
1798
- refresh_cache=refresh_cache,
1799
- on_completion=on_completion,
1800
- on_failure=on_failure,
1801
- retry_condition_fn=retry_condition_fn,
1802
- viz_return_value=viz_return_value,
1803
- ),
1772
+ return Task(
1773
+ fn=__fn,
1774
+ name=name,
1775
+ description=description,
1776
+ tags=tags,
1777
+ version=version,
1778
+ cache_policy=cache_policy,
1779
+ cache_key_fn=cache_key_fn,
1780
+ cache_expiration=cache_expiration,
1781
+ task_run_name=task_run_name,
1782
+ retries=retries,
1783
+ retry_delay_seconds=retry_delay_seconds,
1784
+ retry_jitter_factor=retry_jitter_factor,
1785
+ persist_result=persist_result,
1786
+ result_storage=result_storage,
1787
+ result_storage_key=result_storage_key,
1788
+ result_serializer=result_serializer,
1789
+ cache_result_in_memory=cache_result_in_memory,
1790
+ timeout_seconds=timeout_seconds,
1791
+ log_prints=log_prints,
1792
+ refresh_cache=refresh_cache,
1793
+ on_completion=on_completion,
1794
+ on_failure=on_failure,
1795
+ retry_condition_fn=retry_condition_fn,
1796
+ viz_return_value=viz_return_value,
1804
1797
  )
1805
1798
  else:
1806
1799
  return cast(
@@ -55,7 +55,7 @@ def _url_join(base_url: str, path: str) -> str:
55
55
 
56
56
  def setup_exporters(
57
57
  api_url: str, api_key: str
58
- ) -> tuple[TracerProvider, MeterProvider, "LoggerProvider"]:
58
+ ) -> "tuple[TracerProvider, MeterProvider, LoggerProvider]":
59
59
  account_id, workspace_id = extract_account_and_workspace_id(api_url)
60
60
  telemetry_url = _url_join(api_url, "telemetry/")
61
61
 
@@ -1,14 +1,17 @@
1
1
  import time
2
2
  from threading import Event, Lock, Thread
3
- from typing import Dict, Optional
3
+ from typing import TYPE_CHECKING, Dict, Optional
4
4
 
5
5
  from opentelemetry.context import Context
6
- from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
7
- from opentelemetry.sdk.trace.export import SpanExporter
6
+ from opentelemetry.sdk.trace import Span, SpanProcessor
7
+
8
+ if TYPE_CHECKING:
9
+ from opentelemetry.sdk.trace import ReadableSpan, Span
10
+ from opentelemetry.sdk.trace.export import SpanExporter
8
11
 
9
12
 
10
13
  class InFlightSpanProcessor(SpanProcessor):
11
- def __init__(self, span_exporter: SpanExporter):
14
+ def __init__(self, span_exporter: "SpanExporter"):
12
15
  self.span_exporter = span_exporter
13
16
  self._in_flight: Dict[int, Span] = {}
14
17
  self._lock = Lock()
@@ -26,7 +29,7 @@ class InFlightSpanProcessor(SpanProcessor):
26
29
  if to_export:
27
30
  self.span_exporter.export(to_export)
28
31
 
29
- def _readable_span(self, span: Span) -> ReadableSpan:
32
+ def _readable_span(self, span: "Span") -> "ReadableSpan":
30
33
  readable = span._readable_span()
31
34
  readable._end_time = time.time_ns()
32
35
  readable._attributes = {
@@ -35,13 +38,13 @@ class InFlightSpanProcessor(SpanProcessor):
35
38
  }
36
39
  return readable
37
40
 
38
- def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
41
+ def on_start(self, span: "Span", parent_context: Optional[Context] = None) -> None:
39
42
  if not span.context or not span.context.trace_flags.sampled:
40
43
  return
41
44
  with self._lock:
42
45
  self._in_flight[span.context.span_id] = span
43
46
 
44
- def on_end(self, span: ReadableSpan) -> None:
47
+ def on_end(self, span: "ReadableSpan") -> None:
45
48
  if not span.context or not span.context.trace_flags.sampled:
46
49
  return
47
50
  with self._lock: