prefect-client 2.20.2__py3-none-any.whl → 3.0.0__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 (288) hide show
  1. prefect/__init__.py +74 -110
  2. prefect/_internal/compatibility/deprecated.py +6 -115
  3. prefect/_internal/compatibility/experimental.py +4 -79
  4. prefect/_internal/compatibility/migration.py +166 -0
  5. prefect/_internal/concurrency/__init__.py +2 -2
  6. prefect/_internal/concurrency/api.py +1 -35
  7. prefect/_internal/concurrency/calls.py +0 -6
  8. prefect/_internal/concurrency/cancellation.py +0 -3
  9. prefect/_internal/concurrency/event_loop.py +0 -20
  10. prefect/_internal/concurrency/inspection.py +3 -3
  11. prefect/_internal/concurrency/primitives.py +1 -0
  12. prefect/_internal/concurrency/services.py +23 -0
  13. prefect/_internal/concurrency/threads.py +35 -0
  14. prefect/_internal/concurrency/waiters.py +0 -28
  15. prefect/_internal/integrations.py +7 -0
  16. prefect/_internal/pydantic/__init__.py +0 -45
  17. prefect/_internal/pydantic/annotations/pendulum.py +2 -2
  18. prefect/_internal/pydantic/v1_schema.py +21 -22
  19. prefect/_internal/pydantic/v2_schema.py +0 -2
  20. prefect/_internal/pydantic/v2_validated_func.py +18 -23
  21. prefect/_internal/pytz.py +1 -1
  22. prefect/_internal/retries.py +61 -0
  23. prefect/_internal/schemas/bases.py +45 -177
  24. prefect/_internal/schemas/fields.py +1 -43
  25. prefect/_internal/schemas/validators.py +47 -233
  26. prefect/agent.py +3 -695
  27. prefect/artifacts.py +173 -14
  28. prefect/automations.py +39 -4
  29. prefect/blocks/abstract.py +1 -1
  30. prefect/blocks/core.py +423 -164
  31. prefect/blocks/fields.py +2 -57
  32. prefect/blocks/notifications.py +43 -28
  33. prefect/blocks/redis.py +168 -0
  34. prefect/blocks/system.py +67 -20
  35. prefect/blocks/webhook.py +2 -9
  36. prefect/cache_policies.py +239 -0
  37. prefect/client/__init__.py +4 -0
  38. prefect/client/base.py +33 -27
  39. prefect/client/cloud.py +65 -20
  40. prefect/client/collections.py +1 -1
  41. prefect/client/orchestration.py +667 -440
  42. prefect/client/schemas/actions.py +115 -100
  43. prefect/client/schemas/filters.py +46 -52
  44. prefect/client/schemas/objects.py +228 -178
  45. prefect/client/schemas/responses.py +18 -36
  46. prefect/client/schemas/schedules.py +55 -36
  47. prefect/client/schemas/sorting.py +2 -0
  48. prefect/client/subscriptions.py +8 -7
  49. prefect/client/types/flexible_schedule_list.py +11 -0
  50. prefect/client/utilities.py +9 -6
  51. prefect/concurrency/asyncio.py +60 -11
  52. prefect/concurrency/context.py +24 -0
  53. prefect/concurrency/events.py +2 -2
  54. prefect/concurrency/services.py +46 -16
  55. prefect/concurrency/sync.py +51 -7
  56. prefect/concurrency/v1/asyncio.py +143 -0
  57. prefect/concurrency/v1/context.py +27 -0
  58. prefect/concurrency/v1/events.py +61 -0
  59. prefect/concurrency/v1/services.py +116 -0
  60. prefect/concurrency/v1/sync.py +92 -0
  61. prefect/context.py +246 -149
  62. prefect/deployments/__init__.py +33 -18
  63. prefect/deployments/base.py +10 -15
  64. prefect/deployments/deployments.py +2 -1048
  65. prefect/deployments/flow_runs.py +178 -0
  66. prefect/deployments/runner.py +72 -173
  67. prefect/deployments/schedules.py +31 -25
  68. prefect/deployments/steps/__init__.py +0 -1
  69. prefect/deployments/steps/core.py +7 -0
  70. prefect/deployments/steps/pull.py +15 -21
  71. prefect/deployments/steps/utility.py +2 -1
  72. prefect/docker/__init__.py +20 -0
  73. prefect/docker/docker_image.py +82 -0
  74. prefect/engine.py +15 -2466
  75. prefect/events/actions.py +17 -23
  76. prefect/events/cli/automations.py +20 -7
  77. prefect/events/clients.py +142 -80
  78. prefect/events/filters.py +14 -18
  79. prefect/events/related.py +74 -75
  80. prefect/events/schemas/__init__.py +0 -5
  81. prefect/events/schemas/automations.py +55 -46
  82. prefect/events/schemas/deployment_triggers.py +7 -197
  83. prefect/events/schemas/events.py +46 -65
  84. prefect/events/schemas/labelling.py +10 -14
  85. prefect/events/utilities.py +4 -5
  86. prefect/events/worker.py +23 -8
  87. prefect/exceptions.py +15 -0
  88. prefect/filesystems.py +30 -529
  89. prefect/flow_engine.py +827 -0
  90. prefect/flow_runs.py +379 -7
  91. prefect/flows.py +470 -360
  92. prefect/futures.py +382 -331
  93. prefect/infrastructure/__init__.py +5 -26
  94. prefect/infrastructure/base.py +3 -320
  95. prefect/infrastructure/provisioners/__init__.py +5 -3
  96. prefect/infrastructure/provisioners/cloud_run.py +13 -8
  97. prefect/infrastructure/provisioners/container_instance.py +14 -9
  98. prefect/infrastructure/provisioners/ecs.py +10 -8
  99. prefect/infrastructure/provisioners/modal.py +8 -5
  100. prefect/input/__init__.py +4 -0
  101. prefect/input/actions.py +2 -4
  102. prefect/input/run_input.py +9 -9
  103. prefect/logging/formatters.py +2 -4
  104. prefect/logging/handlers.py +9 -14
  105. prefect/logging/loggers.py +5 -5
  106. prefect/main.py +72 -0
  107. prefect/plugins.py +2 -64
  108. prefect/profiles.toml +16 -2
  109. prefect/records/__init__.py +1 -0
  110. prefect/records/base.py +223 -0
  111. prefect/records/filesystem.py +207 -0
  112. prefect/records/memory.py +178 -0
  113. prefect/records/result_store.py +64 -0
  114. prefect/results.py +577 -504
  115. prefect/runner/runner.py +124 -51
  116. prefect/runner/server.py +32 -34
  117. prefect/runner/storage.py +3 -12
  118. prefect/runner/submit.py +2 -10
  119. prefect/runner/utils.py +2 -2
  120. prefect/runtime/__init__.py +1 -0
  121. prefect/runtime/deployment.py +1 -0
  122. prefect/runtime/flow_run.py +40 -5
  123. prefect/runtime/task_run.py +1 -0
  124. prefect/serializers.py +28 -39
  125. prefect/server/api/collections_data/views/aggregate-worker-metadata.json +5 -14
  126. prefect/settings.py +209 -332
  127. prefect/states.py +160 -63
  128. prefect/task_engine.py +1478 -57
  129. prefect/task_runners.py +383 -287
  130. prefect/task_runs.py +240 -0
  131. prefect/task_worker.py +463 -0
  132. prefect/tasks.py +684 -374
  133. prefect/transactions.py +410 -0
  134. prefect/types/__init__.py +72 -86
  135. prefect/types/entrypoint.py +13 -0
  136. prefect/utilities/annotations.py +4 -3
  137. prefect/utilities/asyncutils.py +227 -148
  138. prefect/utilities/callables.py +138 -48
  139. prefect/utilities/collections.py +134 -86
  140. prefect/utilities/dispatch.py +27 -14
  141. prefect/utilities/dockerutils.py +11 -4
  142. prefect/utilities/engine.py +186 -32
  143. prefect/utilities/filesystem.py +4 -5
  144. prefect/utilities/importtools.py +26 -27
  145. prefect/utilities/pydantic.py +128 -38
  146. prefect/utilities/schema_tools/hydration.py +18 -1
  147. prefect/utilities/schema_tools/validation.py +30 -0
  148. prefect/utilities/services.py +35 -9
  149. prefect/utilities/templating.py +12 -2
  150. prefect/utilities/timeout.py +20 -5
  151. prefect/utilities/urls.py +195 -0
  152. prefect/utilities/visualization.py +1 -0
  153. prefect/variables.py +78 -59
  154. prefect/workers/__init__.py +0 -1
  155. prefect/workers/base.py +237 -244
  156. prefect/workers/block.py +5 -226
  157. prefect/workers/cloud.py +6 -0
  158. prefect/workers/process.py +265 -12
  159. prefect/workers/server.py +29 -11
  160. {prefect_client-2.20.2.dist-info → prefect_client-3.0.0.dist-info}/METADATA +30 -26
  161. prefect_client-3.0.0.dist-info/RECORD +201 -0
  162. {prefect_client-2.20.2.dist-info → prefect_client-3.0.0.dist-info}/WHEEL +1 -1
  163. prefect/_internal/pydantic/_base_model.py +0 -51
  164. prefect/_internal/pydantic/_compat.py +0 -82
  165. prefect/_internal/pydantic/_flags.py +0 -20
  166. prefect/_internal/pydantic/_types.py +0 -8
  167. prefect/_internal/pydantic/utilities/config_dict.py +0 -72
  168. prefect/_internal/pydantic/utilities/field_validator.py +0 -150
  169. prefect/_internal/pydantic/utilities/model_construct.py +0 -56
  170. prefect/_internal/pydantic/utilities/model_copy.py +0 -55
  171. prefect/_internal/pydantic/utilities/model_dump.py +0 -136
  172. prefect/_internal/pydantic/utilities/model_dump_json.py +0 -112
  173. prefect/_internal/pydantic/utilities/model_fields.py +0 -50
  174. prefect/_internal/pydantic/utilities/model_fields_set.py +0 -29
  175. prefect/_internal/pydantic/utilities/model_json_schema.py +0 -82
  176. prefect/_internal/pydantic/utilities/model_rebuild.py +0 -80
  177. prefect/_internal/pydantic/utilities/model_validate.py +0 -75
  178. prefect/_internal/pydantic/utilities/model_validate_json.py +0 -68
  179. prefect/_internal/pydantic/utilities/model_validator.py +0 -87
  180. prefect/_internal/pydantic/utilities/type_adapter.py +0 -71
  181. prefect/_vendor/fastapi/__init__.py +0 -25
  182. prefect/_vendor/fastapi/applications.py +0 -946
  183. prefect/_vendor/fastapi/background.py +0 -3
  184. prefect/_vendor/fastapi/concurrency.py +0 -44
  185. prefect/_vendor/fastapi/datastructures.py +0 -58
  186. prefect/_vendor/fastapi/dependencies/__init__.py +0 -0
  187. prefect/_vendor/fastapi/dependencies/models.py +0 -64
  188. prefect/_vendor/fastapi/dependencies/utils.py +0 -877
  189. prefect/_vendor/fastapi/encoders.py +0 -177
  190. prefect/_vendor/fastapi/exception_handlers.py +0 -40
  191. prefect/_vendor/fastapi/exceptions.py +0 -46
  192. prefect/_vendor/fastapi/logger.py +0 -3
  193. prefect/_vendor/fastapi/middleware/__init__.py +0 -1
  194. prefect/_vendor/fastapi/middleware/asyncexitstack.py +0 -25
  195. prefect/_vendor/fastapi/middleware/cors.py +0 -3
  196. prefect/_vendor/fastapi/middleware/gzip.py +0 -3
  197. prefect/_vendor/fastapi/middleware/httpsredirect.py +0 -3
  198. prefect/_vendor/fastapi/middleware/trustedhost.py +0 -3
  199. prefect/_vendor/fastapi/middleware/wsgi.py +0 -3
  200. prefect/_vendor/fastapi/openapi/__init__.py +0 -0
  201. prefect/_vendor/fastapi/openapi/constants.py +0 -2
  202. prefect/_vendor/fastapi/openapi/docs.py +0 -203
  203. prefect/_vendor/fastapi/openapi/models.py +0 -480
  204. prefect/_vendor/fastapi/openapi/utils.py +0 -485
  205. prefect/_vendor/fastapi/param_functions.py +0 -340
  206. prefect/_vendor/fastapi/params.py +0 -453
  207. prefect/_vendor/fastapi/py.typed +0 -0
  208. prefect/_vendor/fastapi/requests.py +0 -4
  209. prefect/_vendor/fastapi/responses.py +0 -40
  210. prefect/_vendor/fastapi/routing.py +0 -1331
  211. prefect/_vendor/fastapi/security/__init__.py +0 -15
  212. prefect/_vendor/fastapi/security/api_key.py +0 -98
  213. prefect/_vendor/fastapi/security/base.py +0 -6
  214. prefect/_vendor/fastapi/security/http.py +0 -172
  215. prefect/_vendor/fastapi/security/oauth2.py +0 -227
  216. prefect/_vendor/fastapi/security/open_id_connect_url.py +0 -34
  217. prefect/_vendor/fastapi/security/utils.py +0 -10
  218. prefect/_vendor/fastapi/staticfiles.py +0 -1
  219. prefect/_vendor/fastapi/templating.py +0 -3
  220. prefect/_vendor/fastapi/testclient.py +0 -1
  221. prefect/_vendor/fastapi/types.py +0 -3
  222. prefect/_vendor/fastapi/utils.py +0 -235
  223. prefect/_vendor/fastapi/websockets.py +0 -7
  224. prefect/_vendor/starlette/__init__.py +0 -1
  225. prefect/_vendor/starlette/_compat.py +0 -28
  226. prefect/_vendor/starlette/_exception_handler.py +0 -80
  227. prefect/_vendor/starlette/_utils.py +0 -88
  228. prefect/_vendor/starlette/applications.py +0 -261
  229. prefect/_vendor/starlette/authentication.py +0 -159
  230. prefect/_vendor/starlette/background.py +0 -43
  231. prefect/_vendor/starlette/concurrency.py +0 -59
  232. prefect/_vendor/starlette/config.py +0 -151
  233. prefect/_vendor/starlette/convertors.py +0 -87
  234. prefect/_vendor/starlette/datastructures.py +0 -707
  235. prefect/_vendor/starlette/endpoints.py +0 -130
  236. prefect/_vendor/starlette/exceptions.py +0 -60
  237. prefect/_vendor/starlette/formparsers.py +0 -276
  238. prefect/_vendor/starlette/middleware/__init__.py +0 -17
  239. prefect/_vendor/starlette/middleware/authentication.py +0 -52
  240. prefect/_vendor/starlette/middleware/base.py +0 -220
  241. prefect/_vendor/starlette/middleware/cors.py +0 -176
  242. prefect/_vendor/starlette/middleware/errors.py +0 -265
  243. prefect/_vendor/starlette/middleware/exceptions.py +0 -74
  244. prefect/_vendor/starlette/middleware/gzip.py +0 -113
  245. prefect/_vendor/starlette/middleware/httpsredirect.py +0 -19
  246. prefect/_vendor/starlette/middleware/sessions.py +0 -82
  247. prefect/_vendor/starlette/middleware/trustedhost.py +0 -64
  248. prefect/_vendor/starlette/middleware/wsgi.py +0 -147
  249. prefect/_vendor/starlette/py.typed +0 -0
  250. prefect/_vendor/starlette/requests.py +0 -328
  251. prefect/_vendor/starlette/responses.py +0 -347
  252. prefect/_vendor/starlette/routing.py +0 -933
  253. prefect/_vendor/starlette/schemas.py +0 -154
  254. prefect/_vendor/starlette/staticfiles.py +0 -248
  255. prefect/_vendor/starlette/status.py +0 -199
  256. prefect/_vendor/starlette/templating.py +0 -231
  257. prefect/_vendor/starlette/testclient.py +0 -804
  258. prefect/_vendor/starlette/types.py +0 -30
  259. prefect/_vendor/starlette/websockets.py +0 -193
  260. prefect/blocks/kubernetes.py +0 -119
  261. prefect/deprecated/__init__.py +0 -0
  262. prefect/deprecated/data_documents.py +0 -350
  263. prefect/deprecated/packaging/__init__.py +0 -12
  264. prefect/deprecated/packaging/base.py +0 -96
  265. prefect/deprecated/packaging/docker.py +0 -146
  266. prefect/deprecated/packaging/file.py +0 -92
  267. prefect/deprecated/packaging/orion.py +0 -80
  268. prefect/deprecated/packaging/serializers.py +0 -171
  269. prefect/events/instrument.py +0 -135
  270. prefect/infrastructure/container.py +0 -824
  271. prefect/infrastructure/kubernetes.py +0 -920
  272. prefect/infrastructure/process.py +0 -289
  273. prefect/manifests.py +0 -20
  274. prefect/new_flow_engine.py +0 -449
  275. prefect/new_task_engine.py +0 -423
  276. prefect/pydantic/__init__.py +0 -76
  277. prefect/pydantic/main.py +0 -39
  278. prefect/software/__init__.py +0 -2
  279. prefect/software/base.py +0 -50
  280. prefect/software/conda.py +0 -199
  281. prefect/software/pip.py +0 -122
  282. prefect/software/python.py +0 -52
  283. prefect/task_server.py +0 -322
  284. prefect_client-2.20.2.dist-info/RECORD +0 -294
  285. /prefect/{_internal/pydantic/utilities → client/types}/__init__.py +0 -0
  286. /prefect/{_vendor → concurrency/v1}/__init__.py +0 -0
  287. {prefect_client-2.20.2.dist-info → prefect_client-3.0.0.dist-info}/LICENSE +0 -0
  288. {prefect_client-2.20.2.dist-info → prefect_client-3.0.0.dist-info}/top_level.txt +0 -0
prefect/futures.py CHANGED
@@ -1,375 +1,447 @@
1
- """
2
- Futures represent the execution of a task and allow retrieval of the task run's state.
3
-
4
- This module contains the definition for futures as well as utilities for resolving
5
- futures in nested data structures.
6
- """
7
- import asyncio
8
- import warnings
1
+ import abc
2
+ import collections
3
+ import concurrent.futures
4
+ import inspect
5
+ import threading
6
+ import uuid
7
+ from collections.abc import Generator, Iterator
9
8
  from functools import partial
10
- from typing import (
11
- TYPE_CHECKING,
12
- Any,
13
- Awaitable,
14
- Callable,
15
- Generic,
16
- Optional,
17
- Set,
18
- TypeVar,
19
- Union,
20
- cast,
21
- overload,
22
- )
23
- from uuid import UUID
24
-
25
- import anyio
26
-
27
- from prefect._internal.concurrency.api import create_call, from_async, from_sync
28
- from prefect._internal.concurrency.event_loop import run_coroutine_in_loop_from_async
29
- from prefect.client.orchestration import PrefectClient
30
- from prefect.client.utilities import inject_client
31
- from prefect.states import State
32
- from prefect.utilities.annotations import quote
33
- from prefect.utilities.asyncutils import A, Async, Sync, sync
34
- from prefect.utilities.collections import StopVisiting, visit_collection
9
+ from typing import Any, Callable, Generic, List, Optional, Set, Union, cast
35
10
 
36
- if TYPE_CHECKING:
37
- from prefect.task_runners import BaseTaskRunner
11
+ from typing_extensions import TypeVar
38
12
 
13
+ from prefect.client.orchestration import get_client
14
+ from prefect.client.schemas.objects import TaskRun
15
+ from prefect.exceptions import ObjectNotFound
16
+ from prefect.logging.loggers import get_logger, get_run_logger
17
+ from prefect.states import Pending, State
18
+ from prefect.task_runs import TaskRunWaiter
19
+ from prefect.utilities.annotations import quote
20
+ from prefect.utilities.asyncutils import run_coro_as_sync
21
+ from prefect.utilities.collections import StopVisiting, visit_collection
22
+ from prefect.utilities.timeout import timeout as timeout_context
39
23
 
24
+ F = TypeVar("F")
40
25
  R = TypeVar("R")
41
26
 
27
+ logger = get_logger(__name__)
28
+
42
29
 
43
- class PrefectFuture(Generic[R, A]):
30
+ class PrefectFuture(abc.ABC, Generic[R]):
31
+ """
32
+ Abstract base class for Prefect futures. A Prefect future is a handle to the
33
+ asynchronous execution of a task run. It provides methods to wait for the task
34
+ to complete and to retrieve the result of the task run.
44
35
  """
45
- Represents the result of a computation happening in a task runner.
46
36
 
47
- When tasks are called, they are submitted to a task runner which creates a future
48
- for access to the state and result of the task.
37
+ def __init__(self, task_run_id: uuid.UUID):
38
+ self._task_run_id = task_run_id
39
+ self._final_state: Optional[State[R]] = None
49
40
 
50
- Examples:
51
- Define a task that returns a string
41
+ @property
42
+ def task_run_id(self) -> uuid.UUID:
43
+ """The ID of the task run associated with this future"""
44
+ return self._task_run_id
52
45
 
53
- >>> from prefect import flow, task
54
- >>> @task
55
- >>> def my_task() -> str:
56
- >>> return "hello"
46
+ @property
47
+ def state(self) -> State:
48
+ """The current state of the task run associated with this future"""
49
+ if self._final_state:
50
+ return self._final_state
51
+ client = get_client(sync_client=True)
52
+ try:
53
+ task_run = cast(TaskRun, client.read_task_run(task_run_id=self.task_run_id))
54
+ except ObjectNotFound:
55
+ # We'll be optimistic and assume this task will eventually start
56
+ # TODO: Consider using task run events to wait for the task to start
57
+ return Pending()
58
+ return task_run.state or Pending()
59
+
60
+ @abc.abstractmethod
61
+ def wait(self, timeout: Optional[float] = None) -> None:
62
+ ...
63
+ """
64
+ Wait for the task run to complete.
57
65
 
58
- Calls of this task in a flow will return a future
66
+ If the task run has already completed, this method will return immediately.
59
67
 
60
- >>> @flow
61
- >>> def my_flow():
62
- >>> future = my_task.submit() # PrefectFuture[str, Sync] includes result type
63
- >>> future.task_run.id # UUID for the task run
68
+ Args:
69
+ timeout: The maximum number of seconds to wait for the task run to complete.
70
+ If the task run has not completed after the timeout has elapsed, this method will return.
71
+ """
64
72
 
65
- Wait for the task to complete
73
+ @abc.abstractmethod
74
+ def result(
75
+ self,
76
+ timeout: Optional[float] = None,
77
+ raise_on_failure: bool = True,
78
+ ) -> R:
79
+ ...
80
+ """
81
+ Get the result of the task run associated with this future.
66
82
 
67
- >>> @flow
68
- >>> def my_flow():
69
- >>> future = my_task.submit()
70
- >>> final_state = future.wait()
83
+ If the task run has not completed, this method will wait for the task run to complete.
71
84
 
72
- Wait N seconds for the task to complete
85
+ Args:
86
+ timeout: The maximum number of seconds to wait for the task run to complete.
87
+ If the task run has not completed after the timeout has elapsed, this method will return.
88
+ raise_on_failure: If `True`, an exception will be raised if the task run fails.
73
89
 
74
- >>> @flow
75
- >>> def my_flow():
76
- >>> future = my_task.submit()
77
- >>> final_state = future.wait(0.1)
78
- >>> if final_state:
79
- >>> ... # Task done
80
- >>> else:
81
- >>> ... # Task not done yet
90
+ Returns:
91
+ The result of the task run.
92
+ """
82
93
 
83
- Wait for a task to complete and retrieve its result
94
+ @abc.abstractmethod
95
+ def add_done_callback(self, fn):
96
+ """
97
+ Add a callback to be run when the future completes or is cancelled.
84
98
 
85
- >>> @flow
86
- >>> def my_flow():
87
- >>> future = my_task.submit()
88
- >>> result = future.result()
89
- >>> assert result == "hello"
99
+ Args:
100
+ fn: A callable that will be called with this future as its only argument when the future completes or is cancelled.
101
+ """
102
+ ...
90
103
 
91
- Wait N seconds for a task to complete and retrieve its result
92
104
 
93
- >>> @flow
94
- >>> def my_flow():
95
- >>> future = my_task.submit()
96
- >>> result = future.result(timeout=5)
97
- >>> assert result == "hello"
105
+ class PrefectWrappedFuture(PrefectFuture, abc.ABC, Generic[R, F]):
106
+ """
107
+ A Prefect future that wraps another future object.
108
+ """
98
109
 
99
- Retrieve the state of a task without waiting for completion
110
+ def __init__(self, task_run_id: uuid.UUID, wrapped_future: F):
111
+ self._wrapped_future = wrapped_future
112
+ super().__init__(task_run_id)
100
113
 
101
- >>> @flow
102
- >>> def my_flow():
103
- >>> future = my_task.submit()
104
- >>> state = future.get_state()
105
- """
114
+ @property
115
+ def wrapped_future(self) -> F:
116
+ """The underlying future object wrapped by this Prefect future"""
117
+ return self._wrapped_future
106
118
 
107
- def __init__(
108
- self,
109
- name: str,
110
- key: UUID,
111
- task_runner: "BaseTaskRunner",
112
- asynchronous: A = True,
113
- _final_state: State[R] = None, # Exposed for testing
114
- ) -> None:
115
- self.key = key
116
- self.name = name
117
- self.asynchronous = asynchronous
118
- self.task_run = None
119
- self._final_state = _final_state
120
- self._exception: Optional[Exception] = None
121
- self._task_runner = task_runner
122
- self._submitted = anyio.Event()
123
-
124
- self._loop = asyncio.get_running_loop()
125
-
126
- @overload
127
- def wait(
128
- self: "PrefectFuture[R, Async]", timeout: None = None
129
- ) -> Awaitable[State[R]]:
130
- ...
119
+ def add_done_callback(self, fn: Callable[[PrefectFuture], None]):
120
+ if not self._final_state:
131
121
 
132
- @overload
133
- def wait(self: "PrefectFuture[R, Sync]", timeout: None = None) -> State[R]:
134
- ...
122
+ def call_with_self(future):
123
+ """Call the callback with self as the argument, this is necessary to ensure we remove the future from the pending set"""
124
+ fn(self)
135
125
 
136
- @overload
137
- def wait(
138
- self: "PrefectFuture[R, Async]", timeout: float
139
- ) -> Awaitable[Optional[State[R]]]:
140
- ...
126
+ self._wrapped_future.add_done_callback(call_with_self)
127
+ return
128
+ fn(self)
141
129
 
142
- @overload
143
- def wait(self: "PrefectFuture[R, Sync]", timeout: float) -> Optional[State[R]]:
144
- ...
145
130
 
146
- def wait(self, timeout=None):
147
- """
148
- Wait for the run to finish and return the final state
131
+ class PrefectConcurrentFuture(PrefectWrappedFuture[R, concurrent.futures.Future]):
132
+ """
133
+ A Prefect future that wraps a concurrent.futures.Future. This future is used
134
+ when the task run is submitted to a ThreadPoolExecutor.
135
+ """
149
136
 
150
- If the timeout is reached before the run reaches a final state,
151
- `None` is returned.
152
- """
153
- wait = create_call(self._wait, timeout=timeout)
154
- if self.asynchronous:
155
- return from_async.call_soon_in_loop_thread(wait).aresult()
156
- else:
157
- # type checking cannot handle the overloaded timeout passing
158
- return from_sync.call_soon_in_loop_thread(wait).result() # type: ignore
137
+ def wait(self, timeout: Optional[float] = None) -> None:
138
+ try:
139
+ result = self._wrapped_future.result(timeout=timeout)
140
+ except concurrent.futures.TimeoutError:
141
+ return
142
+ if isinstance(result, State):
143
+ self._final_state = result
159
144
 
160
- @overload
161
- async def _wait(self, timeout: None = None) -> State[R]:
162
- ...
145
+ def result(
146
+ self,
147
+ timeout: Optional[float] = None,
148
+ raise_on_failure: bool = True,
149
+ ) -> R:
150
+ if not self._final_state:
151
+ try:
152
+ future_result = self._wrapped_future.result(timeout=timeout)
153
+ except concurrent.futures.TimeoutError as exc:
154
+ raise TimeoutError(
155
+ f"Task run {self.task_run_id} did not complete within {timeout} seconds"
156
+ ) from exc
157
+
158
+ if isinstance(future_result, State):
159
+ self._final_state = future_result
160
+
161
+ else:
162
+ return future_result
163
+
164
+ _result = self._final_state.result(
165
+ raise_on_failure=raise_on_failure, fetch=True
166
+ )
167
+ # state.result is a `sync_compatible` function that may or may not return an awaitable
168
+ # depending on whether the parent frame is sync or not
169
+ if inspect.isawaitable(_result):
170
+ _result = run_coro_as_sync(_result)
171
+ return _result
172
+
173
+ def __del__(self):
174
+ if self._final_state or self._wrapped_future.done():
175
+ return
176
+ try:
177
+ local_logger = get_run_logger()
178
+ except Exception:
179
+ local_logger = logger
180
+ local_logger.warning(
181
+ "A future was garbage collected before it resolved."
182
+ " Please call `.wait()` or `.result()` on futures to ensure they resolve."
183
+ "\nSee https://docs-3.prefect.io/3.0rc/develop/task-runners for more details.",
184
+ )
163
185
 
164
- @overload
165
- async def _wait(self, timeout: float) -> Optional[State[R]]:
166
- ...
167
186
 
168
- async def _wait(self, timeout=None):
169
- """
170
- Async implementation for `wait`
171
- """
172
- await self._wait_for_submission()
187
+ class PrefectDistributedFuture(PrefectFuture[R]):
188
+ """
189
+ Represents the result of a computation happening anywhere.
173
190
 
174
- if self._final_state:
175
- return self._final_state
191
+ This class is typically used to interact with the result of a task run
192
+ scheduled to run in a Prefect task worker but can be used to interact with
193
+ any task run scheduled in Prefect's API.
194
+ """
176
195
 
177
- self._final_state = await self._task_runner.wait(self.key, timeout)
178
- return self._final_state
196
+ done_callbacks: List[Callable[[PrefectFuture], None]] = []
197
+ waiter = None
179
198
 
180
- @overload
181
- def result(
182
- self: "PrefectFuture[R, Sync]",
183
- timeout: float = None,
184
- raise_on_failure: bool = True,
185
- ) -> R:
186
- ...
199
+ def wait(self, timeout: Optional[float] = None) -> None:
200
+ return run_coro_as_sync(self.wait_async(timeout=timeout))
187
201
 
188
- @overload
189
- def result(
190
- self: "PrefectFuture[R, Sync]",
191
- timeout: float = None,
192
- raise_on_failure: bool = False,
193
- ) -> Union[R, Exception]:
194
- ...
202
+ async def wait_async(self, timeout: Optional[float] = None):
203
+ if self._final_state:
204
+ logger.debug(
205
+ "Final state already set for %s. Returning...", self.task_run_id
206
+ )
207
+ return
208
+
209
+ # Ask for the instance of TaskRunWaiter _now_ so that it's already running and
210
+ # can catch the completion event if it happens before we start listening for it.
211
+ TaskRunWaiter.instance()
212
+
213
+ # Read task run to see if it is still running
214
+ async with get_client() as client:
215
+ task_run = await client.read_task_run(task_run_id=self._task_run_id)
216
+ if task_run.state.is_final():
217
+ logger.debug(
218
+ "Task run %s already finished. Returning...",
219
+ self.task_run_id,
220
+ )
221
+ self._final_state = task_run.state
222
+ return
223
+
224
+ # If still running, wait for a completed event from the server
225
+ logger.debug(
226
+ "Waiting for completed event for task run %s...",
227
+ self.task_run_id,
228
+ )
229
+ await TaskRunWaiter.wait_for_task_run(self._task_run_id, timeout=timeout)
230
+ task_run = await client.read_task_run(task_run_id=self._task_run_id)
231
+ if task_run.state.is_final():
232
+ self._final_state = task_run.state
233
+ return
195
234
 
196
- @overload
197
235
  def result(
198
- self: "PrefectFuture[R, Async]",
199
- timeout: float = None,
236
+ self,
237
+ timeout: Optional[float] = None,
200
238
  raise_on_failure: bool = True,
201
- ) -> Awaitable[R]:
202
- ...
239
+ ) -> R:
240
+ return run_coro_as_sync(
241
+ self.result_async(timeout=timeout, raise_on_failure=raise_on_failure)
242
+ )
203
243
 
204
- @overload
205
- def result(
206
- self: "PrefectFuture[R, Async]",
207
- timeout: float = None,
208
- raise_on_failure: bool = False,
209
- ) -> Awaitable[Union[R, Exception]]:
210
- ...
244
+ async def result_async(
245
+ self,
246
+ timeout: Optional[float] = None,
247
+ raise_on_failure: bool = True,
248
+ ) -> R:
249
+ if not self._final_state:
250
+ await self.wait_async(timeout=timeout)
251
+ if not self._final_state:
252
+ raise TimeoutError(
253
+ f"Task run {self.task_run_id} did not complete within {timeout} seconds"
254
+ )
255
+
256
+ return await self._final_state.result(
257
+ raise_on_failure=raise_on_failure, fetch=True
258
+ )
211
259
 
212
- def result(self, timeout: float = None, raise_on_failure: bool = True):
213
- """
214
- Wait for the run to finish and return the final state.
260
+ def add_done_callback(self, fn: Callable[[PrefectFuture], None]):
261
+ if self._final_state:
262
+ fn(self)
263
+ return
264
+ TaskRunWaiter.instance()
265
+ with get_client(sync_client=True) as client:
266
+ task_run = client.read_task_run(task_run_id=self._task_run_id)
267
+ if task_run.state.is_final():
268
+ self._final_state = task_run.state
269
+ fn(self)
270
+ return
271
+ TaskRunWaiter.add_done_callback(self._task_run_id, partial(fn, self))
272
+
273
+ def __eq__(self, other):
274
+ if not isinstance(other, PrefectDistributedFuture):
275
+ return False
276
+ return self.task_run_id == other.task_run_id
277
+
278
+ def __hash__(self):
279
+ return hash(self.task_run_id)
280
+
281
+
282
+ class PrefectFutureList(list, Iterator, Generic[F]):
283
+ """
284
+ A list of Prefect futures.
215
285
 
216
- If the timeout is reached before the run reaches a final state, a `TimeoutError`
217
- will be raised.
286
+ This class provides methods to wait for all futures
287
+ in the list to complete and to retrieve the results of all task runs.
288
+ """
218
289
 
219
- If `raise_on_failure` is `True` and the task run failed, the task run's
220
- exception will be raised.
290
+ def wait(self, timeout: Optional[float] = None) -> None:
221
291
  """
222
- result = create_call(
223
- self._result, timeout=timeout, raise_on_failure=raise_on_failure
224
- )
225
- if self.asynchronous:
226
- return from_async.call_soon_in_loop_thread(result).aresult()
227
- else:
228
- return from_sync.call_soon_in_loop_thread(result).result()
292
+ Wait for all futures in the list to complete.
229
293
 
230
- async def _result(self, timeout: float = None, raise_on_failure: bool = True):
231
- """
232
- Async implementation of `result`
294
+ Args:
295
+ timeout: The maximum number of seconds to wait for all futures to
296
+ complete. This method will not raise if the timeout is reached.
233
297
  """
234
- final_state = await self._wait(timeout=timeout)
235
- if not final_state:
236
- raise TimeoutError("Call timed out before task finished.")
237
- return await final_state.result(raise_on_failure=raise_on_failure, fetch=True)
238
-
239
- @overload
240
- def get_state(
241
- self: "PrefectFuture[R, Async]", client: PrefectClient = None
242
- ) -> Awaitable[State[R]]:
243
- ...
298
+ wait(self, timeout=timeout)
244
299
 
245
- @overload
246
- def get_state(
247
- self: "PrefectFuture[R, Sync]", client: PrefectClient = None
248
- ) -> State[R]:
249
- ...
250
-
251
- def get_state(self, client: PrefectClient = None):
252
- """
253
- Get the current state of the task run.
300
+ def result(
301
+ self,
302
+ timeout: Optional[float] = None,
303
+ raise_on_failure: bool = True,
304
+ ) -> List:
254
305
  """
255
- if self.asynchronous:
256
- return cast(Awaitable[State[R]], self._get_state(client=client))
257
- else:
258
- return cast(State[R], sync(self._get_state, client=client))
259
-
260
- @inject_client
261
- async def _get_state(self, client: PrefectClient = None) -> State[R]:
262
- assert client is not None # always injected
263
-
264
- # We must wait for the task run id to be populated
265
- await self._wait_for_submission()
306
+ Get the results of all task runs associated with the futures in the list.
266
307
 
267
- task_run = await client.read_task_run(self.task_run.id)
308
+ Args:
309
+ timeout: The maximum number of seconds to wait for all futures to
310
+ complete.
311
+ raise_on_failure: If `True`, an exception will be raised if any task run fails.
268
312
 
269
- if not task_run:
270
- raise RuntimeError("Future has no associated task run in the server.")
313
+ Returns:
314
+ A list of results of the task runs.
271
315
 
272
- # Update the task run reference
273
- self.task_run = task_run
274
- return task_run.state
275
-
276
- async def _wait_for_submission(self):
277
- await run_coroutine_in_loop_from_async(self._loop, self._submitted.wait())
278
-
279
- def __hash__(self) -> int:
280
- return hash(self.key)
281
-
282
- def __repr__(self) -> str:
283
- return f"PrefectFuture({self.name!r})"
284
-
285
- def __bool__(self) -> bool:
286
- warnings.warn(
287
- (
288
- "A 'PrefectFuture' from a task call was cast to a boolean; "
289
- "did you mean to check the result of the task instead? "
290
- "e.g. `if my_task().result(): ...`"
291
- ),
292
- stacklevel=2,
316
+ Raises:
317
+ TimeoutError: If the timeout is reached before all futures complete.
318
+ """
319
+ try:
320
+ with timeout_context(timeout):
321
+ return [
322
+ future.result(raise_on_failure=raise_on_failure) for future in self
323
+ ]
324
+ except TimeoutError as exc:
325
+ # timeout came from inside the task
326
+ if "Scope timed out after {timeout} second(s)." not in str(exc):
327
+ raise
328
+ raise TimeoutError(
329
+ f"Timed out waiting for all futures to complete within {timeout} seconds"
330
+ ) from exc
331
+
332
+
333
+ def as_completed(
334
+ futures: List[PrefectFuture], timeout: Optional[float] = None
335
+ ) -> Generator[PrefectFuture, None]:
336
+ unique_futures: Set[PrefectFuture] = set(futures)
337
+ total_futures = len(unique_futures)
338
+ try:
339
+ with timeout_context(timeout):
340
+ done = {f for f in unique_futures if f._final_state}
341
+ pending = unique_futures - done
342
+ yield from done
343
+
344
+ finished_event = threading.Event()
345
+ finished_lock = threading.Lock()
346
+ finished_futures = []
347
+
348
+ def add_to_done(future):
349
+ with finished_lock:
350
+ finished_futures.append(future)
351
+ finished_event.set()
352
+
353
+ for future in pending:
354
+ future.add_done_callback(add_to_done)
355
+
356
+ while pending:
357
+ finished_event.wait()
358
+ with finished_lock:
359
+ done = finished_futures
360
+ finished_futures = []
361
+ finished_event.clear()
362
+
363
+ for future in done:
364
+ pending.remove(future)
365
+ yield future
366
+
367
+ except TimeoutError:
368
+ raise TimeoutError(
369
+ "%d (of %d) futures unfinished" % (len(pending), total_futures)
293
370
  )
294
- return True
295
371
 
296
372
 
297
- def _collect_futures(futures, expr, context):
298
- # Expressions inside quotes should not be traversed
299
- if isinstance(context.get("annotation"), quote):
300
- raise StopVisiting()
373
+ DoneAndNotDoneFutures = collections.namedtuple("DoneAndNotDoneFutures", "done not_done")
374
+
301
375
 
302
- if isinstance(expr, PrefectFuture):
303
- futures.add(expr)
376
+ def wait(futures: List[PrefectFuture], timeout=None) -> DoneAndNotDoneFutures:
377
+ """
378
+ Wait for the futures in the given sequence to complete.
304
379
 
305
- return expr
380
+ Args:
381
+ futures: The sequence of Futures to wait upon.
382
+ timeout: The maximum number of seconds to wait. If None, then there
383
+ is no limit on the wait time.
306
384
 
385
+ Returns:
386
+ A named 2-tuple of sets. The first set, named 'done', contains the
387
+ futures that completed (is finished or cancelled) before the wait
388
+ completed. The second set, named 'not_done', contains uncompleted
389
+ futures. Duplicate futures given to *futures* are removed and will be
390
+ returned only once.
307
391
 
308
- async def resolve_futures_to_data(
309
- expr: Union[PrefectFuture[R, Any], Any],
310
- raise_on_failure: bool = True,
311
- ) -> Union[R, Any]:
392
+ Examples:
393
+ ```python
394
+ @task
395
+ def sleep_task(seconds):
396
+ sleep(seconds)
397
+ return 42
398
+
399
+ @flow
400
+ def flow():
401
+ futures = random_task.map(range(10))
402
+ done, not_done = wait(futures, timeout=5)
403
+ print(f"Done: {len(done)}")
404
+ print(f"Not Done: {len(not_done)}")
405
+ ```
406
+ """
407
+ futures = set(futures)
408
+ done = {f for f in futures if f._final_state}
409
+ not_done = futures - done
410
+ if len(done) == len(futures):
411
+ return DoneAndNotDoneFutures(done, not_done)
412
+ try:
413
+ with timeout_context(timeout):
414
+ for future in not_done.copy():
415
+ future.wait()
416
+ done.add(future)
417
+ not_done.remove(future)
418
+ return DoneAndNotDoneFutures(done, not_done)
419
+ except TimeoutError:
420
+ logger.debug("Timed out waiting for all futures to complete.")
421
+ return DoneAndNotDoneFutures(done, not_done)
422
+
423
+
424
+ def resolve_futures_to_states(
425
+ expr: Union[PrefectFuture, Any],
426
+ ) -> Union[State, Any]:
312
427
  """
313
428
  Given a Python built-in collection, recursively find `PrefectFutures` and build a
314
- new collection with the same structure with futures resolved to their results.
315
- Resolving futures to their results may wait for execution to complete and require
316
- communication with the API.
429
+ new collection with the same structure with futures resolved to their final states.
430
+ Resolving futures to their final states may wait for execution to complete.
317
431
 
318
432
  Unsupported object types will be returned without modification.
319
433
  """
320
434
  futures: Set[PrefectFuture] = set()
321
435
 
322
- maybe_expr = visit_collection(
323
- expr,
324
- visit_fn=partial(_collect_futures, futures),
325
- return_data=False,
326
- context={},
327
- )
328
- if maybe_expr is not None:
329
- expr = maybe_expr
330
-
331
- # Get results
332
- results = await asyncio.gather(
333
- *[
334
- # We must wait for the future in the thread it was created in
335
- from_async.call_soon_in_loop_thread(
336
- create_call(future._result, raise_on_failure=raise_on_failure)
337
- ).aresult()
338
- for future in futures
339
- ]
340
- )
341
-
342
- results_by_future = dict(zip(futures, results))
343
-
344
- def replace_futures_with_results(expr, context):
345
- # Expressions inside quotes should not be modified
436
+ def _collect_futures(futures, expr, context):
437
+ # Expressions inside quotes should not be traversed
346
438
  if isinstance(context.get("annotation"), quote):
347
439
  raise StopVisiting()
348
440
 
349
441
  if isinstance(expr, PrefectFuture):
350
- return results_by_future[expr]
351
- else:
352
- return expr
353
-
354
- return visit_collection(
355
- expr,
356
- visit_fn=replace_futures_with_results,
357
- return_data=True,
358
- context={},
359
- )
360
-
361
-
362
- async def resolve_futures_to_states(
363
- expr: Union[PrefectFuture[R, Any], Any],
364
- ) -> Union[State[R], Any]:
365
- """
366
- Given a Python built-in collection, recursively find `PrefectFutures` and build a
367
- new collection with the same structure with futures resolved to their final states.
368
- Resolving futures to their final states may wait for execution to complete.
442
+ futures.add(expr)
369
443
 
370
- Unsupported object types will be returned without modification.
371
- """
372
- futures: Set[PrefectFuture] = set()
444
+ return expr
373
445
 
374
446
  visit_collection(
375
447
  expr,
@@ -378,14 +450,15 @@ async def resolve_futures_to_states(
378
450
  context={},
379
451
  )
380
452
 
453
+ # if no futures were found, return the original expression
454
+ if not futures:
455
+ return expr
456
+
381
457
  # Get final states for each future
382
- states = await asyncio.gather(
383
- *[
384
- # We must wait for the future in the thread it was created in
385
- from_async.call_soon_in_loop_thread(create_call(future._wait)).aresult()
386
- for future in futures
387
- ]
388
- )
458
+ states = []
459
+ for future in futures:
460
+ future.wait()
461
+ states.append(future.state)
389
462
 
390
463
  states_by_future = dict(zip(futures, states))
391
464
 
@@ -405,25 +478,3 @@ async def resolve_futures_to_states(
405
478
  return_data=True,
406
479
  context={},
407
480
  )
408
-
409
-
410
- def call_repr(__fn: Callable, *args: Any, **kwargs: Any) -> str:
411
- """
412
- Generate a repr for a function call as "fn_name(arg_value, kwarg_name=kwarg_value)"
413
- """
414
-
415
- name = __fn.__name__
416
-
417
- # TODO: If this computation is concerningly expensive, we can iterate checking the
418
- # length at each arg or avoid calling `repr` on args with large amounts of
419
- # data
420
- call_args = ", ".join(
421
- [repr(arg) for arg in args]
422
- + [f"{key}={repr(val)}" for key, val in kwargs.items()]
423
- )
424
-
425
- # Enforce a maximum length
426
- if len(call_args) > 100:
427
- call_args = call_args[:100] + "..."
428
-
429
- return f"{name}({call_args})"