prefect-client 2.20.4__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 +405 -153
- 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 +650 -442
- 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 -2475
- 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 +117 -47
- 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 +137 -45
- 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.4.dist-info → prefect_client-3.0.0.dist-info}/METADATA +28 -24
- prefect_client-3.0.0.dist-info/RECORD +201 -0
- {prefect_client-2.20.4.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.4.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.4.dist-info → prefect_client-3.0.0.dist-info}/LICENSE +0 -0
- {prefect_client-2.20.4.dist-info → prefect_client-3.0.0.dist-info}/top_level.txt +0 -0
prefect/flow_engine.py
ADDED
@@ -0,0 +1,827 @@
|
|
1
|
+
import inspect
|
2
|
+
import logging
|
3
|
+
import os
|
4
|
+
import time
|
5
|
+
from contextlib import ExitStack, contextmanager
|
6
|
+
from dataclasses import dataclass, field
|
7
|
+
from typing import (
|
8
|
+
Any,
|
9
|
+
AsyncGenerator,
|
10
|
+
Coroutine,
|
11
|
+
Dict,
|
12
|
+
Generator,
|
13
|
+
Generic,
|
14
|
+
Iterable,
|
15
|
+
Literal,
|
16
|
+
Optional,
|
17
|
+
Tuple,
|
18
|
+
Type,
|
19
|
+
TypeVar,
|
20
|
+
Union,
|
21
|
+
cast,
|
22
|
+
)
|
23
|
+
from uuid import UUID
|
24
|
+
|
25
|
+
from typing_extensions import ParamSpec
|
26
|
+
|
27
|
+
from prefect import Task
|
28
|
+
from prefect.client.orchestration import SyncPrefectClient, get_client
|
29
|
+
from prefect.client.schemas import FlowRun, TaskRun
|
30
|
+
from prefect.client.schemas.filters import FlowRunFilter
|
31
|
+
from prefect.client.schemas.sorting import FlowRunSort
|
32
|
+
from prefect.concurrency.context import ConcurrencyContext
|
33
|
+
from prefect.concurrency.v1.context import ConcurrencyContext as ConcurrencyContextV1
|
34
|
+
from prefect.context import FlowRunContext, SyncClientContext, TagsContext
|
35
|
+
from prefect.exceptions import (
|
36
|
+
Abort,
|
37
|
+
Pause,
|
38
|
+
PrefectException,
|
39
|
+
TerminationSignal,
|
40
|
+
UpstreamTaskError,
|
41
|
+
)
|
42
|
+
from prefect.flows import Flow, load_flow_from_entrypoint, load_flow_from_flow_run
|
43
|
+
from prefect.futures import PrefectFuture, resolve_futures_to_states
|
44
|
+
from prefect.logging.loggers import (
|
45
|
+
flow_run_logger,
|
46
|
+
get_logger,
|
47
|
+
get_run_logger,
|
48
|
+
patch_print,
|
49
|
+
)
|
50
|
+
from prefect.results import BaseResult, ResultStore, get_current_result_store
|
51
|
+
from prefect.settings import PREFECT_DEBUG_MODE
|
52
|
+
from prefect.states import (
|
53
|
+
Failed,
|
54
|
+
Pending,
|
55
|
+
Running,
|
56
|
+
State,
|
57
|
+
exception_to_crashed_state,
|
58
|
+
exception_to_failed_state,
|
59
|
+
return_value_to_state,
|
60
|
+
)
|
61
|
+
from prefect.utilities.annotations import NotSet
|
62
|
+
from prefect.utilities.asyncutils import run_coro_as_sync
|
63
|
+
from prefect.utilities.callables import (
|
64
|
+
call_with_parameters,
|
65
|
+
get_call_parameters,
|
66
|
+
parameters_to_args_kwargs,
|
67
|
+
)
|
68
|
+
from prefect.utilities.collections import visit_collection
|
69
|
+
from prefect.utilities.engine import (
|
70
|
+
_get_hook_name,
|
71
|
+
_resolve_custom_flow_run_name,
|
72
|
+
capture_sigterm,
|
73
|
+
link_state_to_result,
|
74
|
+
propose_state_sync,
|
75
|
+
resolve_to_final_result,
|
76
|
+
)
|
77
|
+
from prefect.utilities.timeout import timeout, timeout_async
|
78
|
+
from prefect.utilities.urls import url_for
|
79
|
+
|
80
|
+
P = ParamSpec("P")
|
81
|
+
R = TypeVar("R")
|
82
|
+
|
83
|
+
|
84
|
+
class FlowRunTimeoutError(TimeoutError):
|
85
|
+
"""Raised when a flow run exceeds its defined timeout."""
|
86
|
+
|
87
|
+
|
88
|
+
def load_flow_and_flow_run(flow_run_id: UUID) -> Tuple[FlowRun, Flow]:
|
89
|
+
## TODO: add error handling to update state and log tracebacks
|
90
|
+
entrypoint = os.environ.get("PREFECT__FLOW_ENTRYPOINT")
|
91
|
+
|
92
|
+
client = cast(SyncPrefectClient, get_client(sync_client=True))
|
93
|
+
|
94
|
+
flow_run = client.read_flow_run(flow_run_id)
|
95
|
+
if entrypoint:
|
96
|
+
# we should not accept a placeholder flow at runtime
|
97
|
+
flow = load_flow_from_entrypoint(entrypoint, use_placeholder_flow=False)
|
98
|
+
else:
|
99
|
+
flow = run_coro_as_sync(
|
100
|
+
load_flow_from_flow_run(flow_run, use_placeholder_flow=False)
|
101
|
+
)
|
102
|
+
|
103
|
+
return flow_run, flow
|
104
|
+
|
105
|
+
|
106
|
+
@dataclass
|
107
|
+
class FlowRunEngine(Generic[P, R]):
|
108
|
+
flow: Union[Flow[P, R], Flow[P, Coroutine[Any, Any, R]]]
|
109
|
+
parameters: Optional[Dict[str, Any]] = None
|
110
|
+
flow_run: Optional[FlowRun] = None
|
111
|
+
flow_run_id: Optional[UUID] = None
|
112
|
+
logger: logging.Logger = field(default_factory=lambda: get_logger("engine"))
|
113
|
+
wait_for: Optional[Iterable[PrefectFuture]] = None
|
114
|
+
# holds the return value from the user code
|
115
|
+
_return_value: Union[R, Type[NotSet]] = NotSet
|
116
|
+
# holds the exception raised by the user code, if any
|
117
|
+
_raised: Union[Exception, Type[NotSet]] = NotSet
|
118
|
+
_is_started: bool = False
|
119
|
+
_client: Optional[SyncPrefectClient] = None
|
120
|
+
short_circuit: bool = False
|
121
|
+
_flow_run_name_set: bool = False
|
122
|
+
|
123
|
+
def __post_init__(self):
|
124
|
+
if self.flow is None and self.flow_run_id is None:
|
125
|
+
raise ValueError("Either a flow or a flow_run_id must be provided.")
|
126
|
+
|
127
|
+
if self.parameters is None:
|
128
|
+
self.parameters = {}
|
129
|
+
|
130
|
+
@property
|
131
|
+
def client(self) -> SyncPrefectClient:
|
132
|
+
if not self._is_started or self._client is None:
|
133
|
+
raise RuntimeError("Engine has not started.")
|
134
|
+
return self._client
|
135
|
+
|
136
|
+
@property
|
137
|
+
def state(self) -> State:
|
138
|
+
return self.flow_run.state # type: ignore
|
139
|
+
|
140
|
+
def _resolve_parameters(self):
|
141
|
+
if not self.parameters:
|
142
|
+
return {}
|
143
|
+
|
144
|
+
resolved_parameters = {}
|
145
|
+
for parameter, value in self.parameters.items():
|
146
|
+
try:
|
147
|
+
resolved_parameters[parameter] = visit_collection(
|
148
|
+
value,
|
149
|
+
visit_fn=resolve_to_final_result,
|
150
|
+
return_data=True,
|
151
|
+
max_depth=-1,
|
152
|
+
remove_annotations=True,
|
153
|
+
context={},
|
154
|
+
)
|
155
|
+
except UpstreamTaskError:
|
156
|
+
raise
|
157
|
+
except Exception as exc:
|
158
|
+
raise PrefectException(
|
159
|
+
f"Failed to resolve inputs in parameter {parameter!r}. If your"
|
160
|
+
" parameter type is not supported, consider using the `quote`"
|
161
|
+
" annotation to skip resolution of inputs."
|
162
|
+
) from exc
|
163
|
+
|
164
|
+
self.parameters = resolved_parameters
|
165
|
+
|
166
|
+
def _wait_for_dependencies(self):
|
167
|
+
if not self.wait_for:
|
168
|
+
return
|
169
|
+
|
170
|
+
visit_collection(
|
171
|
+
self.wait_for,
|
172
|
+
visit_fn=resolve_to_final_result,
|
173
|
+
return_data=False,
|
174
|
+
max_depth=-1,
|
175
|
+
remove_annotations=True,
|
176
|
+
context={},
|
177
|
+
)
|
178
|
+
|
179
|
+
def begin_run(self) -> State:
|
180
|
+
try:
|
181
|
+
self._resolve_parameters()
|
182
|
+
self._wait_for_dependencies()
|
183
|
+
except UpstreamTaskError as upstream_exc:
|
184
|
+
state = self.set_state(
|
185
|
+
Pending(
|
186
|
+
name="NotReady",
|
187
|
+
message=str(upstream_exc),
|
188
|
+
),
|
189
|
+
# if orchestrating a run already in a pending state, force orchestration to
|
190
|
+
# update the state name
|
191
|
+
force=self.state.is_pending(),
|
192
|
+
)
|
193
|
+
return state
|
194
|
+
|
195
|
+
# validate prior to context so that context receives validated params
|
196
|
+
if self.flow.should_validate_parameters:
|
197
|
+
try:
|
198
|
+
self.parameters = self.flow.validate_parameters(self.parameters or {})
|
199
|
+
except Exception as exc:
|
200
|
+
message = "Validation of flow parameters failed with error:"
|
201
|
+
self.logger.error("%s %s", message, exc)
|
202
|
+
self.handle_exception(
|
203
|
+
exc,
|
204
|
+
msg=message,
|
205
|
+
result_store=get_current_result_store().update_for_flow(
|
206
|
+
self.flow, _sync=True
|
207
|
+
),
|
208
|
+
)
|
209
|
+
self.short_circuit = True
|
210
|
+
self.call_hooks()
|
211
|
+
|
212
|
+
new_state = Running()
|
213
|
+
state = self.set_state(new_state)
|
214
|
+
while state.is_pending():
|
215
|
+
time.sleep(0.2)
|
216
|
+
state = self.set_state(new_state)
|
217
|
+
return state
|
218
|
+
|
219
|
+
def set_state(self, state: State, force: bool = False) -> State:
|
220
|
+
""" """
|
221
|
+
# prevents any state-setting activity
|
222
|
+
if self.short_circuit:
|
223
|
+
return self.state
|
224
|
+
|
225
|
+
state = propose_state_sync(
|
226
|
+
self.client, state, flow_run_id=self.flow_run.id, force=force
|
227
|
+
) # type: ignore
|
228
|
+
self.flow_run.state = state # type: ignore
|
229
|
+
self.flow_run.state_name = state.name # type: ignore
|
230
|
+
self.flow_run.state_type = state.type # type: ignore
|
231
|
+
return state
|
232
|
+
|
233
|
+
def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
|
234
|
+
if self._return_value is not NotSet and not isinstance(
|
235
|
+
self._return_value, State
|
236
|
+
):
|
237
|
+
if isinstance(self._return_value, BaseResult):
|
238
|
+
_result = self._return_value.get()
|
239
|
+
else:
|
240
|
+
_result = self._return_value
|
241
|
+
|
242
|
+
if inspect.isawaitable(_result):
|
243
|
+
# getting the value for a BaseResult may return an awaitable
|
244
|
+
# depending on whether the parent frame is sync or not
|
245
|
+
_result = run_coro_as_sync(_result)
|
246
|
+
return _result
|
247
|
+
|
248
|
+
if self._raised is not NotSet:
|
249
|
+
if raise_on_failure:
|
250
|
+
raise self._raised
|
251
|
+
return self._raised
|
252
|
+
|
253
|
+
# This is a fall through case which leans on the existing state result mechanics to get the
|
254
|
+
# return value. This is necessary because we currently will return a State object if the
|
255
|
+
# the State was Prefect-created.
|
256
|
+
# TODO: Remove the need to get the result from a State except in cases where the return value
|
257
|
+
# is a State object.
|
258
|
+
_result = self.state.result(raise_on_failure=raise_on_failure, fetch=True) # type: ignore
|
259
|
+
# state.result is a `sync_compatible` function that may or may not return an awaitable
|
260
|
+
# depending on whether the parent frame is sync or not
|
261
|
+
if inspect.isawaitable(_result):
|
262
|
+
_result = run_coro_as_sync(_result)
|
263
|
+
return _result
|
264
|
+
|
265
|
+
def handle_success(self, result: R) -> R:
|
266
|
+
result_store = getattr(FlowRunContext.get(), "result_store", None)
|
267
|
+
if result_store is None:
|
268
|
+
raise ValueError("Result store is not set")
|
269
|
+
resolved_result = resolve_futures_to_states(result)
|
270
|
+
terminal_state = run_coro_as_sync(
|
271
|
+
return_value_to_state(
|
272
|
+
resolved_result,
|
273
|
+
result_store=result_store,
|
274
|
+
write_result=True,
|
275
|
+
)
|
276
|
+
)
|
277
|
+
self.set_state(terminal_state)
|
278
|
+
self._return_value = resolved_result
|
279
|
+
return result
|
280
|
+
|
281
|
+
def handle_exception(
|
282
|
+
self,
|
283
|
+
exc: Exception,
|
284
|
+
msg: Optional[str] = None,
|
285
|
+
result_store: Optional[ResultStore] = None,
|
286
|
+
) -> State:
|
287
|
+
context = FlowRunContext.get()
|
288
|
+
terminal_state = run_coro_as_sync(
|
289
|
+
exception_to_failed_state(
|
290
|
+
exc,
|
291
|
+
message=msg or "Flow run encountered an exception:",
|
292
|
+
result_store=result_store or getattr(context, "result_store", None),
|
293
|
+
write_result=True,
|
294
|
+
)
|
295
|
+
)
|
296
|
+
state = self.set_state(terminal_state)
|
297
|
+
if self.state.is_scheduled():
|
298
|
+
self.logger.info(
|
299
|
+
(
|
300
|
+
f"Received non-final state {state.name!r} when proposing final"
|
301
|
+
f" state {terminal_state.name!r} and will attempt to run again..."
|
302
|
+
),
|
303
|
+
)
|
304
|
+
state = self.set_state(Running())
|
305
|
+
self._raised = exc
|
306
|
+
return state
|
307
|
+
|
308
|
+
def handle_timeout(self, exc: TimeoutError) -> None:
|
309
|
+
if isinstance(exc, FlowRunTimeoutError):
|
310
|
+
message = (
|
311
|
+
f"Flow run exceeded timeout of {self.flow.timeout_seconds} second(s)"
|
312
|
+
)
|
313
|
+
else:
|
314
|
+
message = f"Flow run failed due to timeout: {exc!r}"
|
315
|
+
self.logger.error(message)
|
316
|
+
state = Failed(
|
317
|
+
data=exc,
|
318
|
+
message=message,
|
319
|
+
name="TimedOut",
|
320
|
+
)
|
321
|
+
self.set_state(state)
|
322
|
+
self._raised = exc
|
323
|
+
|
324
|
+
def handle_crash(self, exc: BaseException) -> None:
|
325
|
+
state = run_coro_as_sync(exception_to_crashed_state(exc))
|
326
|
+
self.logger.error(f"Crash detected! {state.message}")
|
327
|
+
self.logger.debug("Crash details:", exc_info=exc)
|
328
|
+
self.set_state(state, force=True)
|
329
|
+
self._raised = exc
|
330
|
+
|
331
|
+
def load_subflow_run(
|
332
|
+
self,
|
333
|
+
parent_task_run: TaskRun,
|
334
|
+
client: SyncPrefectClient,
|
335
|
+
context: FlowRunContext,
|
336
|
+
) -> Union[FlowRun, None]:
|
337
|
+
"""
|
338
|
+
This method attempts to load an existing flow run for a subflow task
|
339
|
+
run, if appropriate.
|
340
|
+
|
341
|
+
If the parent task run is in a final but not COMPLETED state, and not
|
342
|
+
being rerun, then we attempt to load an existing flow run instead of
|
343
|
+
creating a new one. This will prevent the engine from running the
|
344
|
+
subflow again.
|
345
|
+
|
346
|
+
If no existing flow run is found, or if the subflow should be rerun,
|
347
|
+
then no flow run is returned.
|
348
|
+
"""
|
349
|
+
|
350
|
+
# check if the parent flow run is rerunning
|
351
|
+
rerunning = (
|
352
|
+
context.flow_run.run_count > 1
|
353
|
+
if getattr(context, "flow_run", None)
|
354
|
+
and isinstance(context.flow_run, FlowRun)
|
355
|
+
else False
|
356
|
+
)
|
357
|
+
|
358
|
+
# if the parent task run is in a final but not completed state, and
|
359
|
+
# not rerunning, then retrieve the most recent flow run instead of
|
360
|
+
# creating a new one. This effectively loads a cached flow run for
|
361
|
+
# situations where we are confident the flow should not be run
|
362
|
+
# again.
|
363
|
+
assert isinstance(parent_task_run.state, State)
|
364
|
+
if parent_task_run.state.is_final() and not (
|
365
|
+
rerunning and not parent_task_run.state.is_completed()
|
366
|
+
):
|
367
|
+
# return the most recent flow run, if it exists
|
368
|
+
flow_runs = client.read_flow_runs(
|
369
|
+
flow_run_filter=FlowRunFilter(
|
370
|
+
parent_task_run_id={"any_": [parent_task_run.id]}
|
371
|
+
),
|
372
|
+
sort=FlowRunSort.EXPECTED_START_TIME_ASC,
|
373
|
+
limit=1,
|
374
|
+
)
|
375
|
+
if flow_runs:
|
376
|
+
loaded_flow_run = flow_runs[-1]
|
377
|
+
self._return_value = loaded_flow_run.state
|
378
|
+
return loaded_flow_run
|
379
|
+
|
380
|
+
def create_flow_run(self, client: SyncPrefectClient) -> FlowRun:
|
381
|
+
flow_run_ctx = FlowRunContext.get()
|
382
|
+
parameters = self.parameters or {}
|
383
|
+
|
384
|
+
parent_task_run = None
|
385
|
+
|
386
|
+
# this is a subflow run
|
387
|
+
if flow_run_ctx:
|
388
|
+
# add a task to a parent flow run that represents the execution of a subflow run
|
389
|
+
parent_task = Task(
|
390
|
+
name=self.flow.name, fn=self.flow.fn, version=self.flow.version
|
391
|
+
)
|
392
|
+
|
393
|
+
parent_task_run = run_coro_as_sync(
|
394
|
+
parent_task.create_run(
|
395
|
+
flow_run_context=flow_run_ctx,
|
396
|
+
parameters=self.parameters,
|
397
|
+
wait_for=self.wait_for,
|
398
|
+
)
|
399
|
+
)
|
400
|
+
|
401
|
+
# check if there is already a flow run for this subflow
|
402
|
+
if subflow_run := self.load_subflow_run(
|
403
|
+
parent_task_run=parent_task_run, client=client, context=flow_run_ctx
|
404
|
+
):
|
405
|
+
return subflow_run
|
406
|
+
|
407
|
+
flow_run = client.create_flow_run(
|
408
|
+
flow=self.flow,
|
409
|
+
parameters=self.flow.serialize_parameters(parameters),
|
410
|
+
state=Pending(),
|
411
|
+
parent_task_run_id=getattr(parent_task_run, "id", None),
|
412
|
+
tags=TagsContext.get().current_tags,
|
413
|
+
)
|
414
|
+
if flow_run_ctx:
|
415
|
+
parent_logger = get_run_logger(flow_run_ctx)
|
416
|
+
parent_logger.info(
|
417
|
+
f"Created subflow run {flow_run.name!r} for flow {self.flow.name!r}"
|
418
|
+
)
|
419
|
+
else:
|
420
|
+
self.logger.info(
|
421
|
+
f"Created flow run {flow_run.name!r} for flow {self.flow.name!r}"
|
422
|
+
)
|
423
|
+
|
424
|
+
return flow_run
|
425
|
+
|
426
|
+
def call_hooks(self, state: Optional[State] = None):
|
427
|
+
if state is None:
|
428
|
+
state = self.state
|
429
|
+
flow = self.flow
|
430
|
+
flow_run = self.flow_run
|
431
|
+
|
432
|
+
if not flow_run:
|
433
|
+
raise ValueError("Flow run is not set")
|
434
|
+
|
435
|
+
enable_cancellation_and_crashed_hooks = (
|
436
|
+
os.environ.get(
|
437
|
+
"PREFECT__ENABLE_CANCELLATION_AND_CRASHED_HOOKS", "true"
|
438
|
+
).lower()
|
439
|
+
== "true"
|
440
|
+
)
|
441
|
+
|
442
|
+
if state.is_failed() and flow.on_failure_hooks:
|
443
|
+
hooks = flow.on_failure_hooks
|
444
|
+
elif state.is_completed() and flow.on_completion_hooks:
|
445
|
+
hooks = flow.on_completion_hooks
|
446
|
+
elif (
|
447
|
+
enable_cancellation_and_crashed_hooks
|
448
|
+
and state.is_cancelling()
|
449
|
+
and flow.on_cancellation_hooks
|
450
|
+
):
|
451
|
+
hooks = flow.on_cancellation_hooks
|
452
|
+
elif (
|
453
|
+
enable_cancellation_and_crashed_hooks
|
454
|
+
and state.is_crashed()
|
455
|
+
and flow.on_crashed_hooks
|
456
|
+
):
|
457
|
+
hooks = flow.on_crashed_hooks
|
458
|
+
elif state.is_running() and flow.on_running_hooks:
|
459
|
+
hooks = flow.on_running_hooks
|
460
|
+
else:
|
461
|
+
hooks = None
|
462
|
+
|
463
|
+
for hook in hooks or []:
|
464
|
+
hook_name = _get_hook_name(hook)
|
465
|
+
|
466
|
+
try:
|
467
|
+
self.logger.info(
|
468
|
+
f"Running hook {hook_name!r} in response to entering state"
|
469
|
+
f" {state.name!r}"
|
470
|
+
)
|
471
|
+
result = hook(flow, flow_run, state)
|
472
|
+
if inspect.isawaitable(result):
|
473
|
+
run_coro_as_sync(result)
|
474
|
+
except Exception:
|
475
|
+
self.logger.error(
|
476
|
+
f"An error was encountered while running hook {hook_name!r}",
|
477
|
+
exc_info=True,
|
478
|
+
)
|
479
|
+
else:
|
480
|
+
self.logger.info(f"Hook {hook_name!r} finished running successfully")
|
481
|
+
|
482
|
+
@contextmanager
|
483
|
+
def setup_run_context(self, client: Optional[SyncPrefectClient] = None):
|
484
|
+
from prefect.utilities.engine import (
|
485
|
+
should_log_prints,
|
486
|
+
)
|
487
|
+
|
488
|
+
if client is None:
|
489
|
+
client = self.client
|
490
|
+
if not self.flow_run:
|
491
|
+
raise ValueError("Flow run not set")
|
492
|
+
|
493
|
+
self.flow_run = client.read_flow_run(self.flow_run.id)
|
494
|
+
log_prints = should_log_prints(self.flow)
|
495
|
+
|
496
|
+
with ExitStack() as stack:
|
497
|
+
# TODO: Explore closing task runner before completing the flow to
|
498
|
+
# wait for futures to complete
|
499
|
+
stack.enter_context(capture_sigterm())
|
500
|
+
if log_prints:
|
501
|
+
stack.enter_context(patch_print())
|
502
|
+
task_runner = stack.enter_context(self.flow.task_runner.duplicate())
|
503
|
+
stack.enter_context(
|
504
|
+
FlowRunContext(
|
505
|
+
flow=self.flow,
|
506
|
+
log_prints=log_prints,
|
507
|
+
flow_run=self.flow_run,
|
508
|
+
parameters=self.parameters,
|
509
|
+
client=client,
|
510
|
+
result_store=get_current_result_store().update_for_flow(
|
511
|
+
self.flow, _sync=True
|
512
|
+
),
|
513
|
+
task_runner=task_runner,
|
514
|
+
)
|
515
|
+
)
|
516
|
+
stack.enter_context(ConcurrencyContextV1())
|
517
|
+
stack.enter_context(ConcurrencyContext())
|
518
|
+
|
519
|
+
# set the logger to the flow run logger
|
520
|
+
self.logger = flow_run_logger(flow_run=self.flow_run, flow=self.flow)
|
521
|
+
|
522
|
+
# update the flow run name if necessary
|
523
|
+
if not self._flow_run_name_set and self.flow.flow_run_name:
|
524
|
+
flow_run_name = _resolve_custom_flow_run_name(
|
525
|
+
flow=self.flow, parameters=self.parameters
|
526
|
+
)
|
527
|
+
self.client.set_flow_run_name(
|
528
|
+
flow_run_id=self.flow_run.id, name=flow_run_name
|
529
|
+
)
|
530
|
+
self.logger.extra["flow_run_name"] = flow_run_name
|
531
|
+
self.logger.debug(
|
532
|
+
f"Renamed flow run {self.flow_run.name!r} to {flow_run_name!r}"
|
533
|
+
)
|
534
|
+
self.flow_run.name = flow_run_name
|
535
|
+
self._flow_run_name_set = True
|
536
|
+
yield
|
537
|
+
|
538
|
+
@contextmanager
|
539
|
+
def initialize_run(self):
|
540
|
+
"""
|
541
|
+
Enters a client context and creates a flow run if needed.
|
542
|
+
"""
|
543
|
+
with SyncClientContext.get_or_create() as client_ctx:
|
544
|
+
self._client = client_ctx.client
|
545
|
+
self._is_started = True
|
546
|
+
|
547
|
+
if not self.flow_run:
|
548
|
+
self.flow_run = self.create_flow_run(self.client)
|
549
|
+
flow_run_url = url_for(self.flow_run)
|
550
|
+
|
551
|
+
if flow_run_url:
|
552
|
+
self.logger.info(
|
553
|
+
f"View at {flow_run_url}", extra={"send_to_api": False}
|
554
|
+
)
|
555
|
+
else:
|
556
|
+
# Update the empirical policy to match the flow if it is not set
|
557
|
+
if self.flow_run.empirical_policy.retry_delay is None:
|
558
|
+
self.flow_run.empirical_policy.retry_delay = (
|
559
|
+
self.flow.retry_delay_seconds
|
560
|
+
)
|
561
|
+
|
562
|
+
if self.flow_run.empirical_policy.retries is None:
|
563
|
+
self.flow_run.empirical_policy.retries = self.flow.retries
|
564
|
+
|
565
|
+
self.client.update_flow_run(
|
566
|
+
flow_run_id=self.flow_run.id,
|
567
|
+
flow_version=self.flow.version,
|
568
|
+
empirical_policy=self.flow_run.empirical_policy,
|
569
|
+
)
|
570
|
+
try:
|
571
|
+
yield self
|
572
|
+
|
573
|
+
except TerminationSignal as exc:
|
574
|
+
self.cancel_all_tasks()
|
575
|
+
self.handle_crash(exc)
|
576
|
+
raise
|
577
|
+
except Exception:
|
578
|
+
# regular exceptions are caught and re-raised to the user
|
579
|
+
raise
|
580
|
+
except (Abort, Pause):
|
581
|
+
raise
|
582
|
+
except GeneratorExit:
|
583
|
+
# Do not capture generator exits as crashes
|
584
|
+
raise
|
585
|
+
except BaseException as exc:
|
586
|
+
# BaseExceptions are caught and handled as crashes
|
587
|
+
self.handle_crash(exc)
|
588
|
+
raise
|
589
|
+
finally:
|
590
|
+
# If debugging, use the more complete `repr` than the usual `str` description
|
591
|
+
display_state = (
|
592
|
+
repr(self.state) if PREFECT_DEBUG_MODE else str(self.state)
|
593
|
+
)
|
594
|
+
self.logger.log(
|
595
|
+
level=logging.INFO if self.state.is_completed() else logging.ERROR,
|
596
|
+
msg=f"Finished in state {display_state}",
|
597
|
+
)
|
598
|
+
|
599
|
+
self._is_started = False
|
600
|
+
self._client = None
|
601
|
+
|
602
|
+
def is_running(self) -> bool:
|
603
|
+
if getattr(self, "flow_run", None) is None:
|
604
|
+
return False
|
605
|
+
return getattr(self, "flow_run").state.is_running()
|
606
|
+
|
607
|
+
def is_pending(self) -> bool:
|
608
|
+
if getattr(self, "flow_run", None) is None:
|
609
|
+
return False # TODO: handle this differently?
|
610
|
+
return getattr(self, "flow_run").state.is_pending()
|
611
|
+
|
612
|
+
def cancel_all_tasks(self):
|
613
|
+
if hasattr(self.flow.task_runner, "cancel_all"):
|
614
|
+
self.flow.task_runner.cancel_all() # type: ignore
|
615
|
+
|
616
|
+
# --------------------------
|
617
|
+
#
|
618
|
+
# The following methods compose the main task run loop
|
619
|
+
#
|
620
|
+
# --------------------------
|
621
|
+
|
622
|
+
@contextmanager
|
623
|
+
def start(self) -> Generator[None, None, None]:
|
624
|
+
with self.initialize_run():
|
625
|
+
self.begin_run()
|
626
|
+
|
627
|
+
if self.state.is_running():
|
628
|
+
self.call_hooks()
|
629
|
+
yield
|
630
|
+
|
631
|
+
@contextmanager
|
632
|
+
def run_context(self):
|
633
|
+
timeout_context = timeout_async if self.flow.isasync else timeout
|
634
|
+
# reenter the run context to ensure it is up to date for every run
|
635
|
+
with self.setup_run_context():
|
636
|
+
try:
|
637
|
+
with timeout_context(
|
638
|
+
seconds=self.flow.timeout_seconds,
|
639
|
+
timeout_exc_type=FlowRunTimeoutError,
|
640
|
+
):
|
641
|
+
self.logger.debug(
|
642
|
+
f"Executing flow {self.flow.name!r} for flow run {self.flow_run.name!r}..."
|
643
|
+
)
|
644
|
+
yield self
|
645
|
+
except TimeoutError as exc:
|
646
|
+
self.handle_timeout(exc)
|
647
|
+
except Exception as exc:
|
648
|
+
self.logger.exception("Encountered exception during execution: %r", exc)
|
649
|
+
self.handle_exception(exc)
|
650
|
+
finally:
|
651
|
+
if self.state.is_final() or self.state.is_cancelling():
|
652
|
+
self.call_hooks()
|
653
|
+
|
654
|
+
def call_flow_fn(self) -> Union[R, Coroutine[Any, Any, R]]:
|
655
|
+
"""
|
656
|
+
Convenience method to call the flow function. Returns a coroutine if the
|
657
|
+
flow is async.
|
658
|
+
"""
|
659
|
+
if self.flow.isasync:
|
660
|
+
|
661
|
+
async def _call_flow_fn():
|
662
|
+
result = await call_with_parameters(self.flow.fn, self.parameters)
|
663
|
+
self.handle_success(result)
|
664
|
+
|
665
|
+
return _call_flow_fn()
|
666
|
+
else:
|
667
|
+
result = call_with_parameters(self.flow.fn, self.parameters)
|
668
|
+
self.handle_success(result)
|
669
|
+
|
670
|
+
|
671
|
+
def run_flow_sync(
|
672
|
+
flow: Flow[P, R],
|
673
|
+
flow_run: Optional[FlowRun] = None,
|
674
|
+
parameters: Optional[Dict[str, Any]] = None,
|
675
|
+
wait_for: Optional[Iterable[PrefectFuture]] = None,
|
676
|
+
return_type: Literal["state", "result"] = "result",
|
677
|
+
) -> Union[R, State, None]:
|
678
|
+
engine = FlowRunEngine[P, R](
|
679
|
+
flow=flow,
|
680
|
+
parameters=parameters,
|
681
|
+
flow_run=flow_run,
|
682
|
+
wait_for=wait_for,
|
683
|
+
)
|
684
|
+
|
685
|
+
with engine.start():
|
686
|
+
while engine.is_running():
|
687
|
+
with engine.run_context():
|
688
|
+
engine.call_flow_fn()
|
689
|
+
|
690
|
+
return engine.state if return_type == "state" else engine.result()
|
691
|
+
|
692
|
+
|
693
|
+
async def run_flow_async(
|
694
|
+
flow: Flow[P, R],
|
695
|
+
flow_run: Optional[FlowRun] = None,
|
696
|
+
parameters: Optional[Dict[str, Any]] = None,
|
697
|
+
wait_for: Optional[Iterable[PrefectFuture]] = None,
|
698
|
+
return_type: Literal["state", "result"] = "result",
|
699
|
+
) -> Union[R, State, None]:
|
700
|
+
engine = FlowRunEngine[P, R](
|
701
|
+
flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
|
702
|
+
)
|
703
|
+
|
704
|
+
with engine.start():
|
705
|
+
while engine.is_running():
|
706
|
+
with engine.run_context():
|
707
|
+
await engine.call_flow_fn()
|
708
|
+
|
709
|
+
return engine.state if return_type == "state" else engine.result()
|
710
|
+
|
711
|
+
|
712
|
+
def run_generator_flow_sync(
|
713
|
+
flow: Flow[P, R],
|
714
|
+
flow_run: Optional[FlowRun] = None,
|
715
|
+
parameters: Optional[Dict[str, Any]] = None,
|
716
|
+
wait_for: Optional[Iterable[PrefectFuture]] = None,
|
717
|
+
return_type: Literal["state", "result"] = "result",
|
718
|
+
) -> Generator[R, None, None]:
|
719
|
+
if return_type != "result":
|
720
|
+
raise ValueError("The return_type for a generator flow must be 'result'")
|
721
|
+
|
722
|
+
engine = FlowRunEngine[P, R](
|
723
|
+
flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
|
724
|
+
)
|
725
|
+
|
726
|
+
with engine.start():
|
727
|
+
while engine.is_running():
|
728
|
+
with engine.run_context():
|
729
|
+
call_args, call_kwargs = parameters_to_args_kwargs(
|
730
|
+
flow.fn, engine.parameters or {}
|
731
|
+
)
|
732
|
+
gen = flow.fn(*call_args, **call_kwargs)
|
733
|
+
try:
|
734
|
+
while True:
|
735
|
+
gen_result = next(gen)
|
736
|
+
# link the current state to the result for dependency tracking
|
737
|
+
link_state_to_result(engine.state, gen_result)
|
738
|
+
yield gen_result
|
739
|
+
except StopIteration as exc:
|
740
|
+
engine.handle_success(exc.value)
|
741
|
+
except GeneratorExit as exc:
|
742
|
+
engine.handle_success(None)
|
743
|
+
gen.throw(exc)
|
744
|
+
|
745
|
+
return engine.result()
|
746
|
+
|
747
|
+
|
748
|
+
async def run_generator_flow_async(
|
749
|
+
flow: Flow[P, R],
|
750
|
+
flow_run: Optional[FlowRun] = None,
|
751
|
+
parameters: Optional[Dict[str, Any]] = None,
|
752
|
+
wait_for: Optional[Iterable[PrefectFuture]] = None,
|
753
|
+
return_type: Literal["state", "result"] = "result",
|
754
|
+
) -> AsyncGenerator[R, None]:
|
755
|
+
if return_type != "result":
|
756
|
+
raise ValueError("The return_type for a generator flow must be 'result'")
|
757
|
+
|
758
|
+
engine = FlowRunEngine[P, R](
|
759
|
+
flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
|
760
|
+
)
|
761
|
+
|
762
|
+
with engine.start():
|
763
|
+
while engine.is_running():
|
764
|
+
with engine.run_context():
|
765
|
+
call_args, call_kwargs = parameters_to_args_kwargs(
|
766
|
+
flow.fn, engine.parameters or {}
|
767
|
+
)
|
768
|
+
gen = flow.fn(*call_args, **call_kwargs)
|
769
|
+
try:
|
770
|
+
while True:
|
771
|
+
# can't use anext in Python < 3.10
|
772
|
+
gen_result = await gen.__anext__()
|
773
|
+
# link the current state to the result for dependency tracking
|
774
|
+
link_state_to_result(engine.state, gen_result)
|
775
|
+
yield gen_result
|
776
|
+
except (StopAsyncIteration, GeneratorExit) as exc:
|
777
|
+
engine.handle_success(None)
|
778
|
+
if isinstance(exc, GeneratorExit):
|
779
|
+
gen.throw(exc)
|
780
|
+
|
781
|
+
# async generators can't return, but we can raise failures here
|
782
|
+
if engine.state.is_failed():
|
783
|
+
engine.result()
|
784
|
+
|
785
|
+
|
786
|
+
def run_flow(
|
787
|
+
flow: Flow[P, R],
|
788
|
+
flow_run: Optional[FlowRun] = None,
|
789
|
+
parameters: Optional[Dict[str, Any]] = None,
|
790
|
+
wait_for: Optional[Iterable[PrefectFuture]] = None,
|
791
|
+
return_type: Literal["state", "result"] = "result",
|
792
|
+
) -> Union[R, State, None]:
|
793
|
+
kwargs = dict(
|
794
|
+
flow=flow,
|
795
|
+
flow_run=flow_run,
|
796
|
+
parameters=_flow_parameters(
|
797
|
+
flow=flow, flow_run=flow_run, parameters=parameters
|
798
|
+
),
|
799
|
+
wait_for=wait_for,
|
800
|
+
return_type=return_type,
|
801
|
+
)
|
802
|
+
|
803
|
+
if flow.isasync and flow.isgenerator:
|
804
|
+
return run_generator_flow_async(**kwargs)
|
805
|
+
elif flow.isgenerator:
|
806
|
+
return run_generator_flow_sync(**kwargs)
|
807
|
+
elif flow.isasync:
|
808
|
+
return run_flow_async(**kwargs)
|
809
|
+
else:
|
810
|
+
return run_flow_sync(**kwargs)
|
811
|
+
|
812
|
+
|
813
|
+
def _flow_parameters(
|
814
|
+
flow: Flow[P, R], flow_run: Optional[FlowRun], parameters: Optional[Dict[str, Any]]
|
815
|
+
) -> Dict[str, Any]:
|
816
|
+
if parameters:
|
817
|
+
# This path is taken when a flow is being called directly with
|
818
|
+
# parameters, in that case just return the parameters as-is.
|
819
|
+
return parameters
|
820
|
+
|
821
|
+
# Otherwise the flow is being executed indirectly and we may need to grab
|
822
|
+
# the parameters from the flow run. We also need to resolve any default
|
823
|
+
# parameters that are defined on the flow function itself.
|
824
|
+
|
825
|
+
parameters = flow_run.parameters if flow_run else {}
|
826
|
+
call_args, call_kwargs = parameters_to_args_kwargs(flow.fn, parameters)
|
827
|
+
return get_call_parameters(flow.fn, call_args, call_kwargs)
|