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.
- prefect/__init__.py +74 -110
- prefect/_internal/compatibility/deprecated.py +6 -115
- prefect/_internal/compatibility/experimental.py +4 -79
- prefect/_internal/compatibility/migration.py +166 -0
- prefect/_internal/concurrency/__init__.py +2 -2
- prefect/_internal/concurrency/api.py +1 -35
- prefect/_internal/concurrency/calls.py +0 -6
- prefect/_internal/concurrency/cancellation.py +0 -3
- prefect/_internal/concurrency/event_loop.py +0 -20
- prefect/_internal/concurrency/inspection.py +3 -3
- prefect/_internal/concurrency/primitives.py +1 -0
- prefect/_internal/concurrency/services.py +23 -0
- prefect/_internal/concurrency/threads.py +35 -0
- prefect/_internal/concurrency/waiters.py +0 -28
- prefect/_internal/integrations.py +7 -0
- prefect/_internal/pydantic/__init__.py +0 -45
- prefect/_internal/pydantic/annotations/pendulum.py +2 -2
- prefect/_internal/pydantic/v1_schema.py +21 -22
- prefect/_internal/pydantic/v2_schema.py +0 -2
- prefect/_internal/pydantic/v2_validated_func.py +18 -23
- prefect/_internal/pytz.py +1 -1
- prefect/_internal/retries.py +61 -0
- prefect/_internal/schemas/bases.py +45 -177
- prefect/_internal/schemas/fields.py +1 -43
- prefect/_internal/schemas/validators.py +47 -233
- prefect/agent.py +3 -695
- prefect/artifacts.py +173 -14
- prefect/automations.py +39 -4
- prefect/blocks/abstract.py +1 -1
- prefect/blocks/core.py +423 -164
- prefect/blocks/fields.py +2 -57
- prefect/blocks/notifications.py +43 -28
- prefect/blocks/redis.py +168 -0
- prefect/blocks/system.py +67 -20
- prefect/blocks/webhook.py +2 -9
- prefect/cache_policies.py +239 -0
- prefect/client/__init__.py +4 -0
- prefect/client/base.py +33 -27
- prefect/client/cloud.py +65 -20
- prefect/client/collections.py +1 -1
- prefect/client/orchestration.py +667 -440
- prefect/client/schemas/actions.py +115 -100
- prefect/client/schemas/filters.py +46 -52
- prefect/client/schemas/objects.py +228 -178
- prefect/client/schemas/responses.py +18 -36
- prefect/client/schemas/schedules.py +55 -36
- prefect/client/schemas/sorting.py +2 -0
- prefect/client/subscriptions.py +8 -7
- prefect/client/types/flexible_schedule_list.py +11 -0
- prefect/client/utilities.py +9 -6
- prefect/concurrency/asyncio.py +60 -11
- prefect/concurrency/context.py +24 -0
- prefect/concurrency/events.py +2 -2
- prefect/concurrency/services.py +46 -16
- prefect/concurrency/sync.py +51 -7
- prefect/concurrency/v1/asyncio.py +143 -0
- prefect/concurrency/v1/context.py +27 -0
- prefect/concurrency/v1/events.py +61 -0
- prefect/concurrency/v1/services.py +116 -0
- prefect/concurrency/v1/sync.py +92 -0
- prefect/context.py +246 -149
- prefect/deployments/__init__.py +33 -18
- prefect/deployments/base.py +10 -15
- prefect/deployments/deployments.py +2 -1048
- prefect/deployments/flow_runs.py +178 -0
- prefect/deployments/runner.py +72 -173
- prefect/deployments/schedules.py +31 -25
- prefect/deployments/steps/__init__.py +0 -1
- prefect/deployments/steps/core.py +7 -0
- prefect/deployments/steps/pull.py +15 -21
- prefect/deployments/steps/utility.py +2 -1
- prefect/docker/__init__.py +20 -0
- prefect/docker/docker_image.py +82 -0
- prefect/engine.py +15 -2466
- prefect/events/actions.py +17 -23
- prefect/events/cli/automations.py +20 -7
- prefect/events/clients.py +142 -80
- prefect/events/filters.py +14 -18
- prefect/events/related.py +74 -75
- prefect/events/schemas/__init__.py +0 -5
- prefect/events/schemas/automations.py +55 -46
- prefect/events/schemas/deployment_triggers.py +7 -197
- prefect/events/schemas/events.py +46 -65
- prefect/events/schemas/labelling.py +10 -14
- prefect/events/utilities.py +4 -5
- prefect/events/worker.py +23 -8
- prefect/exceptions.py +15 -0
- prefect/filesystems.py +30 -529
- prefect/flow_engine.py +827 -0
- prefect/flow_runs.py +379 -7
- prefect/flows.py +470 -360
- prefect/futures.py +382 -331
- prefect/infrastructure/__init__.py +5 -26
- prefect/infrastructure/base.py +3 -320
- prefect/infrastructure/provisioners/__init__.py +5 -3
- prefect/infrastructure/provisioners/cloud_run.py +13 -8
- prefect/infrastructure/provisioners/container_instance.py +14 -9
- prefect/infrastructure/provisioners/ecs.py +10 -8
- prefect/infrastructure/provisioners/modal.py +8 -5
- prefect/input/__init__.py +4 -0
- prefect/input/actions.py +2 -4
- prefect/input/run_input.py +9 -9
- prefect/logging/formatters.py +2 -4
- prefect/logging/handlers.py +9 -14
- prefect/logging/loggers.py +5 -5
- prefect/main.py +72 -0
- prefect/plugins.py +2 -64
- prefect/profiles.toml +16 -2
- prefect/records/__init__.py +1 -0
- prefect/records/base.py +223 -0
- prefect/records/filesystem.py +207 -0
- prefect/records/memory.py +178 -0
- prefect/records/result_store.py +64 -0
- prefect/results.py +577 -504
- prefect/runner/runner.py +124 -51
- prefect/runner/server.py +32 -34
- prefect/runner/storage.py +3 -12
- prefect/runner/submit.py +2 -10
- prefect/runner/utils.py +2 -2
- prefect/runtime/__init__.py +1 -0
- prefect/runtime/deployment.py +1 -0
- prefect/runtime/flow_run.py +40 -5
- prefect/runtime/task_run.py +1 -0
- prefect/serializers.py +28 -39
- prefect/server/api/collections_data/views/aggregate-worker-metadata.json +5 -14
- prefect/settings.py +209 -332
- prefect/states.py +160 -63
- prefect/task_engine.py +1478 -57
- prefect/task_runners.py +383 -287
- prefect/task_runs.py +240 -0
- prefect/task_worker.py +463 -0
- prefect/tasks.py +684 -374
- prefect/transactions.py +410 -0
- prefect/types/__init__.py +72 -86
- prefect/types/entrypoint.py +13 -0
- prefect/utilities/annotations.py +4 -3
- prefect/utilities/asyncutils.py +227 -148
- prefect/utilities/callables.py +138 -48
- prefect/utilities/collections.py +134 -86
- prefect/utilities/dispatch.py +27 -14
- prefect/utilities/dockerutils.py +11 -4
- prefect/utilities/engine.py +186 -32
- prefect/utilities/filesystem.py +4 -5
- prefect/utilities/importtools.py +26 -27
- prefect/utilities/pydantic.py +128 -38
- prefect/utilities/schema_tools/hydration.py +18 -1
- prefect/utilities/schema_tools/validation.py +30 -0
- prefect/utilities/services.py +35 -9
- prefect/utilities/templating.py +12 -2
- prefect/utilities/timeout.py +20 -5
- prefect/utilities/urls.py +195 -0
- prefect/utilities/visualization.py +1 -0
- prefect/variables.py +78 -59
- prefect/workers/__init__.py +0 -1
- prefect/workers/base.py +237 -244
- prefect/workers/block.py +5 -226
- prefect/workers/cloud.py +6 -0
- prefect/workers/process.py +265 -12
- prefect/workers/server.py +29 -11
- {prefect_client-2.20.2.dist-info → prefect_client-3.0.0.dist-info}/METADATA +30 -26
- prefect_client-3.0.0.dist-info/RECORD +201 -0
- {prefect_client-2.20.2.dist-info → prefect_client-3.0.0.dist-info}/WHEEL +1 -1
- prefect/_internal/pydantic/_base_model.py +0 -51
- prefect/_internal/pydantic/_compat.py +0 -82
- prefect/_internal/pydantic/_flags.py +0 -20
- prefect/_internal/pydantic/_types.py +0 -8
- prefect/_internal/pydantic/utilities/config_dict.py +0 -72
- prefect/_internal/pydantic/utilities/field_validator.py +0 -150
- prefect/_internal/pydantic/utilities/model_construct.py +0 -56
- prefect/_internal/pydantic/utilities/model_copy.py +0 -55
- prefect/_internal/pydantic/utilities/model_dump.py +0 -136
- prefect/_internal/pydantic/utilities/model_dump_json.py +0 -112
- prefect/_internal/pydantic/utilities/model_fields.py +0 -50
- prefect/_internal/pydantic/utilities/model_fields_set.py +0 -29
- prefect/_internal/pydantic/utilities/model_json_schema.py +0 -82
- prefect/_internal/pydantic/utilities/model_rebuild.py +0 -80
- prefect/_internal/pydantic/utilities/model_validate.py +0 -75
- prefect/_internal/pydantic/utilities/model_validate_json.py +0 -68
- prefect/_internal/pydantic/utilities/model_validator.py +0 -87
- prefect/_internal/pydantic/utilities/type_adapter.py +0 -71
- prefect/_vendor/fastapi/__init__.py +0 -25
- prefect/_vendor/fastapi/applications.py +0 -946
- prefect/_vendor/fastapi/background.py +0 -3
- prefect/_vendor/fastapi/concurrency.py +0 -44
- prefect/_vendor/fastapi/datastructures.py +0 -58
- prefect/_vendor/fastapi/dependencies/__init__.py +0 -0
- prefect/_vendor/fastapi/dependencies/models.py +0 -64
- prefect/_vendor/fastapi/dependencies/utils.py +0 -877
- prefect/_vendor/fastapi/encoders.py +0 -177
- prefect/_vendor/fastapi/exception_handlers.py +0 -40
- prefect/_vendor/fastapi/exceptions.py +0 -46
- prefect/_vendor/fastapi/logger.py +0 -3
- prefect/_vendor/fastapi/middleware/__init__.py +0 -1
- prefect/_vendor/fastapi/middleware/asyncexitstack.py +0 -25
- prefect/_vendor/fastapi/middleware/cors.py +0 -3
- prefect/_vendor/fastapi/middleware/gzip.py +0 -3
- prefect/_vendor/fastapi/middleware/httpsredirect.py +0 -3
- prefect/_vendor/fastapi/middleware/trustedhost.py +0 -3
- prefect/_vendor/fastapi/middleware/wsgi.py +0 -3
- prefect/_vendor/fastapi/openapi/__init__.py +0 -0
- prefect/_vendor/fastapi/openapi/constants.py +0 -2
- prefect/_vendor/fastapi/openapi/docs.py +0 -203
- prefect/_vendor/fastapi/openapi/models.py +0 -480
- prefect/_vendor/fastapi/openapi/utils.py +0 -485
- prefect/_vendor/fastapi/param_functions.py +0 -340
- prefect/_vendor/fastapi/params.py +0 -453
- prefect/_vendor/fastapi/py.typed +0 -0
- prefect/_vendor/fastapi/requests.py +0 -4
- prefect/_vendor/fastapi/responses.py +0 -40
- prefect/_vendor/fastapi/routing.py +0 -1331
- prefect/_vendor/fastapi/security/__init__.py +0 -15
- prefect/_vendor/fastapi/security/api_key.py +0 -98
- prefect/_vendor/fastapi/security/base.py +0 -6
- prefect/_vendor/fastapi/security/http.py +0 -172
- prefect/_vendor/fastapi/security/oauth2.py +0 -227
- prefect/_vendor/fastapi/security/open_id_connect_url.py +0 -34
- prefect/_vendor/fastapi/security/utils.py +0 -10
- prefect/_vendor/fastapi/staticfiles.py +0 -1
- prefect/_vendor/fastapi/templating.py +0 -3
- prefect/_vendor/fastapi/testclient.py +0 -1
- prefect/_vendor/fastapi/types.py +0 -3
- prefect/_vendor/fastapi/utils.py +0 -235
- prefect/_vendor/fastapi/websockets.py +0 -7
- prefect/_vendor/starlette/__init__.py +0 -1
- prefect/_vendor/starlette/_compat.py +0 -28
- prefect/_vendor/starlette/_exception_handler.py +0 -80
- prefect/_vendor/starlette/_utils.py +0 -88
- prefect/_vendor/starlette/applications.py +0 -261
- prefect/_vendor/starlette/authentication.py +0 -159
- prefect/_vendor/starlette/background.py +0 -43
- prefect/_vendor/starlette/concurrency.py +0 -59
- prefect/_vendor/starlette/config.py +0 -151
- prefect/_vendor/starlette/convertors.py +0 -87
- prefect/_vendor/starlette/datastructures.py +0 -707
- prefect/_vendor/starlette/endpoints.py +0 -130
- prefect/_vendor/starlette/exceptions.py +0 -60
- prefect/_vendor/starlette/formparsers.py +0 -276
- prefect/_vendor/starlette/middleware/__init__.py +0 -17
- prefect/_vendor/starlette/middleware/authentication.py +0 -52
- prefect/_vendor/starlette/middleware/base.py +0 -220
- prefect/_vendor/starlette/middleware/cors.py +0 -176
- prefect/_vendor/starlette/middleware/errors.py +0 -265
- prefect/_vendor/starlette/middleware/exceptions.py +0 -74
- prefect/_vendor/starlette/middleware/gzip.py +0 -113
- prefect/_vendor/starlette/middleware/httpsredirect.py +0 -19
- prefect/_vendor/starlette/middleware/sessions.py +0 -82
- prefect/_vendor/starlette/middleware/trustedhost.py +0 -64
- prefect/_vendor/starlette/middleware/wsgi.py +0 -147
- prefect/_vendor/starlette/py.typed +0 -0
- prefect/_vendor/starlette/requests.py +0 -328
- prefect/_vendor/starlette/responses.py +0 -347
- prefect/_vendor/starlette/routing.py +0 -933
- prefect/_vendor/starlette/schemas.py +0 -154
- prefect/_vendor/starlette/staticfiles.py +0 -248
- prefect/_vendor/starlette/status.py +0 -199
- prefect/_vendor/starlette/templating.py +0 -231
- prefect/_vendor/starlette/testclient.py +0 -804
- prefect/_vendor/starlette/types.py +0 -30
- prefect/_vendor/starlette/websockets.py +0 -193
- prefect/blocks/kubernetes.py +0 -119
- prefect/deprecated/__init__.py +0 -0
- prefect/deprecated/data_documents.py +0 -350
- prefect/deprecated/packaging/__init__.py +0 -12
- prefect/deprecated/packaging/base.py +0 -96
- prefect/deprecated/packaging/docker.py +0 -146
- prefect/deprecated/packaging/file.py +0 -92
- prefect/deprecated/packaging/orion.py +0 -80
- prefect/deprecated/packaging/serializers.py +0 -171
- prefect/events/instrument.py +0 -135
- prefect/infrastructure/container.py +0 -824
- prefect/infrastructure/kubernetes.py +0 -920
- prefect/infrastructure/process.py +0 -289
- prefect/manifests.py +0 -20
- prefect/new_flow_engine.py +0 -449
- prefect/new_task_engine.py +0 -423
- prefect/pydantic/__init__.py +0 -76
- prefect/pydantic/main.py +0 -39
- prefect/software/__init__.py +0 -2
- prefect/software/base.py +0 -50
- prefect/software/conda.py +0 -199
- prefect/software/pip.py +0 -122
- prefect/software/python.py +0 -52
- prefect/task_server.py +0 -322
- prefect_client-2.20.2.dist-info/RECORD +0 -294
- /prefect/{_internal/pydantic/utilities → client/types}/__init__.py +0 -0
- /prefect/{_vendor → concurrency/v1}/__init__.py +0 -0
- {prefect_client-2.20.2.dist-info → prefect_client-3.0.0.dist-info}/LICENSE +0 -0
- {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
|
-
|
3
|
-
|
4
|
-
|
5
|
-
|
6
|
-
|
7
|
-
import
|
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
|
-
|
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
|
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
|
-
|
48
|
-
|
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
|
-
|
51
|
-
|
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
|
-
|
54
|
-
|
55
|
-
|
56
|
-
|
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
|
-
|
66
|
+
If the task run has already completed, this method will return immediately.
|
59
67
|
|
60
|
-
|
61
|
-
|
62
|
-
|
63
|
-
|
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
|
-
|
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
|
-
|
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
|
-
|
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
|
-
|
75
|
-
|
76
|
-
|
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
|
-
|
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
|
-
|
86
|
-
|
87
|
-
|
88
|
-
|
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
|
-
|
94
|
-
|
95
|
-
|
96
|
-
|
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
|
-
|
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
|
-
|
102
|
-
|
103
|
-
|
104
|
-
|
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
|
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
|
-
|
133
|
-
|
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
|
-
|
137
|
-
|
138
|
-
self
|
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
|
-
|
147
|
-
|
148
|
-
|
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
|
-
|
151
|
-
|
152
|
-
|
153
|
-
|
154
|
-
|
155
|
-
|
156
|
-
|
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
|
-
|
161
|
-
|
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
|
-
|
169
|
-
|
170
|
-
|
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
|
-
|
175
|
-
|
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
|
-
|
178
|
-
|
196
|
+
done_callbacks: List[Callable[[PrefectFuture], None]] = []
|
197
|
+
waiter = None
|
179
198
|
|
180
|
-
|
181
|
-
|
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
|
-
|
189
|
-
|
190
|
-
|
191
|
-
|
192
|
-
|
193
|
-
|
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
|
199
|
-
timeout: float = None,
|
236
|
+
self,
|
237
|
+
timeout: Optional[float] = None,
|
200
238
|
raise_on_failure: bool = True,
|
201
|
-
) ->
|
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
|
-
|
205
|
-
|
206
|
-
|
207
|
-
|
208
|
-
|
209
|
-
|
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
|
213
|
-
|
214
|
-
|
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
|
-
|
217
|
-
|
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
|
-
|
220
|
-
exception will be raised.
|
290
|
+
def wait(self, timeout: Optional[float] = None) -> None:
|
221
291
|
"""
|
222
|
-
|
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
|
-
|
231
|
-
|
232
|
-
|
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
|
-
|
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
|
-
|
246
|
-
|
247
|
-
|
248
|
-
|
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
|
-
|
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
|
-
|
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
|
-
|
270
|
-
|
313
|
+
Returns:
|
314
|
+
A list of results of the task runs.
|
271
315
|
|
272
|
-
|
273
|
-
|
274
|
-
|
275
|
-
|
276
|
-
|
277
|
-
|
278
|
-
|
279
|
-
|
280
|
-
|
281
|
-
|
282
|
-
|
283
|
-
|
284
|
-
|
285
|
-
|
286
|
-
|
287
|
-
|
288
|
-
|
289
|
-
|
290
|
-
|
291
|
-
|
292
|
-
|
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
|
-
|
298
|
-
|
299
|
-
if isinstance(context.get("annotation"), quote):
|
300
|
-
raise StopVisiting()
|
373
|
+
DoneAndNotDoneFutures = collections.namedtuple("DoneAndNotDoneFutures", "done not_done")
|
374
|
+
|
301
375
|
|
302
|
-
|
303
|
-
|
376
|
+
def wait(futures: List[PrefectFuture], timeout=None) -> DoneAndNotDoneFutures:
|
377
|
+
"""
|
378
|
+
Wait for the futures in the given sequence to complete.
|
304
379
|
|
305
|
-
|
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
|
-
|
309
|
-
|
310
|
-
|
311
|
-
)
|
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
|
315
|
-
Resolving futures to their
|
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
|
-
|
323
|
-
|
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
|
-
|
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
|
-
|
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 =
|
383
|
-
|
384
|
-
|
385
|
-
|
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})"
|