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/workers/base.py
CHANGED
@@ -1,62 +1,55 @@
|
|
1
1
|
import abc
|
2
2
|
import inspect
|
3
|
-
import
|
4
|
-
from
|
3
|
+
import threading
|
4
|
+
from contextlib import AsyncExitStack
|
5
|
+
from functools import partial
|
6
|
+
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Type, Union
|
5
7
|
from uuid import uuid4
|
6
8
|
|
7
9
|
import anyio
|
8
10
|
import anyio.abc
|
9
11
|
import pendulum
|
10
|
-
|
11
|
-
from
|
12
|
-
from
|
13
|
-
|
14
|
-
if HAS_PYDANTIC_V2:
|
15
|
-
from pydantic.v1 import BaseModel, Field, PrivateAttr, validator
|
16
|
-
else:
|
17
|
-
from pydantic import BaseModel, Field, PrivateAttr, validator
|
12
|
+
from pydantic import BaseModel, Field, PrivateAttr, field_validator
|
13
|
+
from pydantic.json_schema import GenerateJsonSchema
|
14
|
+
from typing_extensions import Literal
|
18
15
|
|
19
16
|
import prefect
|
20
|
-
from prefect._internal.
|
21
|
-
EXPERIMENTAL_WARNING,
|
22
|
-
ExperimentalFeature,
|
23
|
-
experiment_enabled,
|
24
|
-
)
|
17
|
+
from prefect._internal.schemas.validators import return_v_or_none
|
25
18
|
from prefect.client.orchestration import PrefectClient, get_client
|
26
19
|
from prefect.client.schemas.actions import WorkPoolCreate, WorkPoolUpdate
|
27
|
-
from prefect.client.schemas.filters import (
|
28
|
-
FlowRunFilter,
|
29
|
-
FlowRunFilterId,
|
30
|
-
FlowRunFilterState,
|
31
|
-
FlowRunFilterStateName,
|
32
|
-
FlowRunFilterStateType,
|
33
|
-
WorkPoolFilter,
|
34
|
-
WorkPoolFilterName,
|
35
|
-
WorkQueueFilter,
|
36
|
-
WorkQueueFilterName,
|
37
|
-
)
|
38
20
|
from prefect.client.schemas.objects import StateType, WorkPool
|
39
21
|
from prefect.client.utilities import inject_client
|
40
|
-
from prefect.
|
22
|
+
from prefect.concurrency.asyncio import (
|
23
|
+
AcquireConcurrencySlotTimeoutError,
|
24
|
+
ConcurrencySlotAcquisitionError,
|
25
|
+
concurrency,
|
26
|
+
)
|
41
27
|
from prefect.events import Event, RelatedResource, emit_event
|
42
28
|
from prefect.events.related import object_as_related_resource, tags_as_related_resources
|
43
29
|
from prefect.exceptions import (
|
44
30
|
Abort,
|
45
|
-
InfrastructureNotAvailable,
|
46
|
-
InfrastructureNotFound,
|
47
31
|
ObjectNotFound,
|
48
32
|
)
|
49
33
|
from prefect.logging.loggers import PrefectLogAdapter, flow_run_logger, get_logger
|
50
34
|
from prefect.plugins import load_prefect_collections
|
51
35
|
from prefect.settings import (
|
52
|
-
|
53
|
-
|
36
|
+
PREFECT_API_URL,
|
37
|
+
PREFECT_TEST_MODE,
|
54
38
|
PREFECT_WORKER_HEARTBEAT_SECONDS,
|
55
39
|
PREFECT_WORKER_PREFETCH_SECONDS,
|
40
|
+
PREFECT_WORKER_QUERY_SECONDS,
|
56
41
|
get_current_settings,
|
57
42
|
)
|
58
|
-
from prefect.states import
|
43
|
+
from prefect.states import (
|
44
|
+
AwaitingConcurrencySlot,
|
45
|
+
Crashed,
|
46
|
+
Pending,
|
47
|
+
exception_to_failed_state,
|
48
|
+
)
|
49
|
+
from prefect.utilities.asyncutils import asyncnullcontext
|
59
50
|
from prefect.utilities.dispatch import get_registry_for_type, register_base_type
|
51
|
+
from prefect.utilities.engine import propose_state
|
52
|
+
from prefect.utilities.services import critical_service_loop
|
60
53
|
from prefect.utilities.slugify import slugify
|
61
54
|
from prefect.utilities.templating import (
|
62
55
|
apply_values,
|
@@ -107,16 +100,27 @@ class BaseJobConfiguration(BaseModel):
|
|
107
100
|
def is_using_a_runner(self):
|
108
101
|
return self.command is not None and "prefect flow-run execute" in self.command
|
109
102
|
|
110
|
-
@
|
103
|
+
@field_validator("command")
|
104
|
+
@classmethod
|
111
105
|
def _coerce_command(cls, v):
|
112
106
|
return return_v_or_none(v)
|
113
107
|
|
108
|
+
@field_validator("env", mode="before")
|
109
|
+
@classmethod
|
110
|
+
def _coerce_env(cls, v):
|
111
|
+
return {k: str(v) if v is not None else None for k, v in v.items()}
|
112
|
+
|
114
113
|
@staticmethod
|
115
114
|
def _get_base_config_defaults(variables: dict) -> dict:
|
116
115
|
"""Get default values from base config for all variables that have them."""
|
117
116
|
defaults = dict()
|
118
117
|
for variable_name, attrs in variables.items():
|
119
|
-
|
118
|
+
# We remote `None` values because we don't want to use them in templating.
|
119
|
+
# The currently logic depends on keys not existing to populate the correct value
|
120
|
+
# in some cases.
|
121
|
+
# Pydantic will provide default values if the keys are missing when creating
|
122
|
+
# a configuration class.
|
123
|
+
if "default" in attrs and attrs.get("default") is not None:
|
120
124
|
defaults[variable_name] = attrs["default"]
|
121
125
|
|
122
126
|
return defaults
|
@@ -124,7 +128,10 @@ class BaseJobConfiguration(BaseModel):
|
|
124
128
|
@classmethod
|
125
129
|
@inject_client
|
126
130
|
async def from_template_and_values(
|
127
|
-
cls,
|
131
|
+
cls,
|
132
|
+
base_job_template: dict,
|
133
|
+
values: dict,
|
134
|
+
client: Optional["PrefectClient"] = None,
|
128
135
|
):
|
129
136
|
"""Creates a valid worker configuration object from the provided base
|
130
137
|
configuration and overrides.
|
@@ -161,7 +168,7 @@ class BaseJobConfiguration(BaseModel):
|
|
161
168
|
}
|
162
169
|
"""
|
163
170
|
configuration = {}
|
164
|
-
properties = cls.
|
171
|
+
properties = cls.model_json_schema()["properties"]
|
165
172
|
for k, v in properties.items():
|
166
173
|
if v.get("template"):
|
167
174
|
template = v["template"]
|
@@ -225,22 +232,7 @@ class BaseJobConfiguration(BaseModel):
|
|
225
232
|
"""
|
226
233
|
Generate a command for a flow run job.
|
227
234
|
"""
|
228
|
-
|
229
|
-
if (
|
230
|
-
PREFECT_EXPERIMENTAL_WARN
|
231
|
-
and PREFECT_EXPERIMENTAL_WARN_ENHANCED_CANCELLATION
|
232
|
-
):
|
233
|
-
warnings.warn(
|
234
|
-
EXPERIMENTAL_WARNING.format(
|
235
|
-
feature="Enhanced flow run cancellation",
|
236
|
-
group="enhanced_cancellation",
|
237
|
-
help="",
|
238
|
-
),
|
239
|
-
ExperimentalFeature,
|
240
|
-
stacklevel=3,
|
241
|
-
)
|
242
|
-
return "prefect flow-run execute"
|
243
|
-
return "python -m prefect.engine"
|
235
|
+
return "prefect flow-run execute"
|
244
236
|
|
245
237
|
@staticmethod
|
246
238
|
def _base_flow_run_labels(flow_run: "FlowRun") -> Dict[str, str]:
|
@@ -325,6 +317,33 @@ class BaseVariables(BaseModel):
|
|
325
317
|
),
|
326
318
|
)
|
327
319
|
|
320
|
+
@classmethod
|
321
|
+
def model_json_schema(
|
322
|
+
cls,
|
323
|
+
by_alias: bool = True,
|
324
|
+
ref_template: str = "#/definitions/{model}",
|
325
|
+
schema_generator: Type[GenerateJsonSchema] = GenerateJsonSchema,
|
326
|
+
mode: Literal["validation", "serialization"] = "validation",
|
327
|
+
) -> Dict[str, Any]:
|
328
|
+
"""TODO: stop overriding this method - use GenerateSchema in ConfigDict instead?"""
|
329
|
+
schema = super().model_json_schema(
|
330
|
+
by_alias, ref_template, schema_generator, mode
|
331
|
+
)
|
332
|
+
|
333
|
+
# ensure backwards compatibility by copying $defs into definitions
|
334
|
+
if "$defs" in schema:
|
335
|
+
schema["definitions"] = schema.pop("$defs")
|
336
|
+
|
337
|
+
# we aren't expecting these additional fields in the schema
|
338
|
+
if "additionalProperties" in schema:
|
339
|
+
schema.pop("additionalProperties")
|
340
|
+
|
341
|
+
for _, definition in schema.get("definitions", {}).items():
|
342
|
+
if "additionalProperties" in definition:
|
343
|
+
definition.pop("additionalProperties")
|
344
|
+
|
345
|
+
return schema
|
346
|
+
|
328
347
|
|
329
348
|
class BaseWorkerResult(BaseModel, abc.ABC):
|
330
349
|
identifier: str
|
@@ -374,12 +393,14 @@ class BaseWorker(abc.ABC):
|
|
374
393
|
ensure that work pools are not created accidentally.
|
375
394
|
limit: The maximum number of flow runs this worker should be running at
|
376
395
|
a given time.
|
396
|
+
heartbeat_interval_seconds: The number of seconds between worker heartbeats.
|
377
397
|
base_job_template: If creating the work pool, provide the base job
|
378
398
|
template to use. Logs a warning if the pool already exists.
|
379
399
|
"""
|
380
400
|
if name and ("/" in name or "%" in name):
|
381
401
|
raise ValueError("Worker name cannot contain '/' or '%'")
|
382
402
|
self.name = name or f"{self.__class__.__name__} {uuid4()}"
|
403
|
+
self._started_event: Optional[Event] = None
|
383
404
|
self._logger = get_logger(f"worker.{self.__class__.type}.{self.name.lower()}")
|
384
405
|
|
385
406
|
self.is_setup = False
|
@@ -396,6 +417,7 @@ class BaseWorker(abc.ABC):
|
|
396
417
|
)
|
397
418
|
|
398
419
|
self._work_pool: Optional[WorkPool] = None
|
420
|
+
self._exit_stack: AsyncExitStack = AsyncExitStack()
|
399
421
|
self._runs_task_group: Optional[anyio.abc.TaskGroup] = None
|
400
422
|
self._client: Optional[PrefectClient] = None
|
401
423
|
self._last_polled_time: pendulum.DateTime = pendulum.now("utc")
|
@@ -420,7 +442,7 @@ class BaseWorker(abc.ABC):
|
|
420
442
|
@classmethod
|
421
443
|
def get_default_base_job_template(cls) -> Dict:
|
422
444
|
if cls.job_configuration_variables is None:
|
423
|
-
schema = cls.job_configuration.
|
445
|
+
schema = cls.job_configuration.model_json_schema()
|
424
446
|
# remove "template" key from all dicts in schema['properties'] because it is not a
|
425
447
|
# relevant field
|
426
448
|
for key, value in schema["properties"].items():
|
@@ -428,7 +450,7 @@ class BaseWorker(abc.ABC):
|
|
428
450
|
schema["properties"][key].pop("template", None)
|
429
451
|
variables_schema = schema
|
430
452
|
else:
|
431
|
-
variables_schema = cls.job_configuration_variables.
|
453
|
+
variables_schema = cls.job_configuration_variables.model_json_schema()
|
432
454
|
variables_schema.pop("title", None)
|
433
455
|
return {
|
434
456
|
"job_configuration": cls.job_configuration.json_template(),
|
@@ -472,6 +494,86 @@ class BaseWorker(abc.ABC):
|
|
472
494
|
},
|
473
495
|
)
|
474
496
|
|
497
|
+
async def start(
|
498
|
+
self,
|
499
|
+
run_once: bool = False,
|
500
|
+
with_healthcheck: bool = False,
|
501
|
+
printer: Callable[..., None] = print,
|
502
|
+
):
|
503
|
+
"""
|
504
|
+
Starts the worker and runs the main worker loops.
|
505
|
+
|
506
|
+
By default, the worker will run loops to poll for scheduled/cancelled flow
|
507
|
+
runs and sync with the Prefect API server.
|
508
|
+
|
509
|
+
If `run_once` is set, the worker will only run each loop once and then return.
|
510
|
+
|
511
|
+
If `with_healthcheck` is set, the worker will start a healthcheck server which
|
512
|
+
can be used to determine if the worker is still polling for flow runs and restart
|
513
|
+
the worker if necessary.
|
514
|
+
|
515
|
+
Args:
|
516
|
+
run_once: If set, the worker will only run each loop once then return.
|
517
|
+
with_healthcheck: If set, the worker will start a healthcheck server.
|
518
|
+
printer: A `print`-like function where logs will be reported.
|
519
|
+
"""
|
520
|
+
healthcheck_server = None
|
521
|
+
healthcheck_thread = None
|
522
|
+
try:
|
523
|
+
async with self as worker:
|
524
|
+
# wait for an initial heartbeat to configure the worker
|
525
|
+
await worker.sync_with_backend()
|
526
|
+
# schedule the scheduled flow run polling loop
|
527
|
+
async with anyio.create_task_group() as loops_task_group:
|
528
|
+
loops_task_group.start_soon(
|
529
|
+
partial(
|
530
|
+
critical_service_loop,
|
531
|
+
workload=self.get_and_submit_flow_runs,
|
532
|
+
interval=PREFECT_WORKER_QUERY_SECONDS.value(),
|
533
|
+
run_once=run_once,
|
534
|
+
jitter_range=0.3,
|
535
|
+
backoff=4, # Up to ~1 minute interval during backoff
|
536
|
+
)
|
537
|
+
)
|
538
|
+
# schedule the sync loop
|
539
|
+
loops_task_group.start_soon(
|
540
|
+
partial(
|
541
|
+
critical_service_loop,
|
542
|
+
workload=self.sync_with_backend,
|
543
|
+
interval=self.heartbeat_interval_seconds,
|
544
|
+
run_once=run_once,
|
545
|
+
jitter_range=0.3,
|
546
|
+
backoff=4,
|
547
|
+
)
|
548
|
+
)
|
549
|
+
|
550
|
+
self._started_event = await self._emit_worker_started_event()
|
551
|
+
|
552
|
+
if with_healthcheck:
|
553
|
+
from prefect.workers.server import build_healthcheck_server
|
554
|
+
|
555
|
+
# we'll start the ASGI server in a separate thread so that
|
556
|
+
# uvicorn does not block the main thread
|
557
|
+
healthcheck_server = build_healthcheck_server(
|
558
|
+
worker=worker,
|
559
|
+
query_interval_seconds=PREFECT_WORKER_QUERY_SECONDS.value(),
|
560
|
+
)
|
561
|
+
healthcheck_thread = threading.Thread(
|
562
|
+
name="healthcheck-server-thread",
|
563
|
+
target=healthcheck_server.run,
|
564
|
+
daemon=True,
|
565
|
+
)
|
566
|
+
healthcheck_thread.start()
|
567
|
+
printer(f"Worker {worker.name!r} started!")
|
568
|
+
finally:
|
569
|
+
if healthcheck_server and healthcheck_thread:
|
570
|
+
self._logger.debug("Stopping healthcheck server...")
|
571
|
+
healthcheck_server.should_exit = True
|
572
|
+
healthcheck_thread.join()
|
573
|
+
self._logger.debug("Healthcheck server stopped.")
|
574
|
+
|
575
|
+
printer(f"Worker {worker.name!r} stopped!")
|
576
|
+
|
475
577
|
@abc.abstractmethod
|
476
578
|
async def run(
|
477
579
|
self,
|
@@ -486,20 +588,6 @@ class BaseWorker(abc.ABC):
|
|
486
588
|
"Workers must implement a method for running submitted flow runs"
|
487
589
|
)
|
488
590
|
|
489
|
-
async def kill_infrastructure(
|
490
|
-
self,
|
491
|
-
infrastructure_pid: str,
|
492
|
-
configuration: BaseJobConfiguration,
|
493
|
-
grace_seconds: int = 30,
|
494
|
-
):
|
495
|
-
"""
|
496
|
-
Method for killing infrastructure created by a worker. Should be implemented by
|
497
|
-
individual workers if they support killing infrastructure.
|
498
|
-
"""
|
499
|
-
raise NotImplementedError(
|
500
|
-
"This worker does not support killing infrastructure."
|
501
|
-
)
|
502
|
-
|
503
591
|
@classmethod
|
504
592
|
def __dispatch_key__(cls):
|
505
593
|
if cls.__name__ == "BaseWorker":
|
@@ -513,9 +601,13 @@ class BaseWorker(abc.ABC):
|
|
513
601
|
self._limiter = (
|
514
602
|
anyio.CapacityLimiter(self._limit) if self._limit is not None else None
|
515
603
|
)
|
604
|
+
|
605
|
+
if not PREFECT_TEST_MODE and not PREFECT_API_URL.value():
|
606
|
+
raise ValueError("`PREFECT_API_URL` must be set to start a Worker.")
|
607
|
+
|
516
608
|
self._client = get_client()
|
517
|
-
await self.
|
518
|
-
await self.
|
609
|
+
await self._exit_stack.enter_async_context(self._client)
|
610
|
+
await self._exit_stack.enter_async_context(self._runs_task_group)
|
519
611
|
|
520
612
|
self.is_setup = True
|
521
613
|
|
@@ -525,14 +617,14 @@ class BaseWorker(abc.ABC):
|
|
525
617
|
self.is_setup = False
|
526
618
|
for scope in self._scheduled_task_scopes:
|
527
619
|
scope.cancel()
|
528
|
-
|
529
|
-
|
530
|
-
if self.
|
531
|
-
await self.
|
620
|
+
|
621
|
+
await self._exit_stack.__aexit__(*exc_info)
|
622
|
+
if self._started_event:
|
623
|
+
await self._emit_worker_stopped_event(self._started_event)
|
532
624
|
self._runs_task_group = None
|
533
625
|
self._client = None
|
534
626
|
|
535
|
-
def is_worker_still_polling(self, query_interval_seconds:
|
627
|
+
def is_worker_still_polling(self, query_interval_seconds: float) -> bool:
|
536
628
|
"""
|
537
629
|
This method is invoked by a webserver healthcheck handler
|
538
630
|
and returns a boolean indicating if the worker has recorded a
|
@@ -568,143 +660,12 @@ class BaseWorker(abc.ABC):
|
|
568
660
|
|
569
661
|
return await self._submit_scheduled_flow_runs(flow_run_response=runs_response)
|
570
662
|
|
571
|
-
async def check_for_cancelled_flow_runs(self):
|
572
|
-
if not self.is_setup:
|
573
|
-
raise RuntimeError(
|
574
|
-
"Worker is not set up. Please make sure you are running this worker "
|
575
|
-
"as an async context manager."
|
576
|
-
)
|
577
|
-
|
578
|
-
self._logger.debug("Checking for cancelled flow runs...")
|
579
|
-
|
580
|
-
work_queue_filter = (
|
581
|
-
WorkQueueFilter(name=WorkQueueFilterName(any_=list(self._work_queues)))
|
582
|
-
if self._work_queues
|
583
|
-
else None
|
584
|
-
)
|
585
|
-
|
586
|
-
named_cancelling_flow_runs = await self._client.read_flow_runs(
|
587
|
-
flow_run_filter=FlowRunFilter(
|
588
|
-
state=FlowRunFilterState(
|
589
|
-
type=FlowRunFilterStateType(any_=[StateType.CANCELLED]),
|
590
|
-
name=FlowRunFilterStateName(any_=["Cancelling"]),
|
591
|
-
),
|
592
|
-
# Avoid duplicate cancellation calls
|
593
|
-
id=FlowRunFilterId(not_any_=list(self._cancelling_flow_run_ids)),
|
594
|
-
),
|
595
|
-
work_pool_filter=WorkPoolFilter(
|
596
|
-
name=WorkPoolFilterName(any_=[self._work_pool_name])
|
597
|
-
),
|
598
|
-
work_queue_filter=work_queue_filter,
|
599
|
-
)
|
600
|
-
|
601
|
-
typed_cancelling_flow_runs = await self._client.read_flow_runs(
|
602
|
-
flow_run_filter=FlowRunFilter(
|
603
|
-
state=FlowRunFilterState(
|
604
|
-
type=FlowRunFilterStateType(any_=[StateType.CANCELLING]),
|
605
|
-
),
|
606
|
-
# Avoid duplicate cancellation calls
|
607
|
-
id=FlowRunFilterId(not_any_=list(self._cancelling_flow_run_ids)),
|
608
|
-
),
|
609
|
-
work_pool_filter=WorkPoolFilter(
|
610
|
-
name=WorkPoolFilterName(any_=[self._work_pool_name])
|
611
|
-
),
|
612
|
-
work_queue_filter=work_queue_filter,
|
613
|
-
)
|
614
|
-
|
615
|
-
cancelling_flow_runs = named_cancelling_flow_runs + typed_cancelling_flow_runs
|
616
|
-
|
617
|
-
if cancelling_flow_runs:
|
618
|
-
self._logger.info(
|
619
|
-
f"Found {len(cancelling_flow_runs)} flow runs awaiting cancellation."
|
620
|
-
)
|
621
|
-
|
622
|
-
for flow_run in cancelling_flow_runs:
|
623
|
-
self._cancelling_flow_run_ids.add(flow_run.id)
|
624
|
-
self._runs_task_group.start_soon(self.cancel_run, flow_run)
|
625
|
-
|
626
|
-
return cancelling_flow_runs
|
627
|
-
|
628
|
-
async def cancel_run(self, flow_run: "FlowRun"):
|
629
|
-
run_logger = self.get_flow_run_logger(flow_run)
|
630
|
-
|
631
|
-
try:
|
632
|
-
configuration = await self._get_configuration(flow_run)
|
633
|
-
except ObjectNotFound:
|
634
|
-
self._logger.warning(
|
635
|
-
f"Flow run {flow_run.id!r} cannot be cancelled by this worker:"
|
636
|
-
f" associated deployment {flow_run.deployment_id!r} does not exist."
|
637
|
-
)
|
638
|
-
await self._mark_flow_run_as_cancelled(
|
639
|
-
flow_run,
|
640
|
-
state_updates={
|
641
|
-
"message": (
|
642
|
-
"This flow run is missing infrastructure configuration information"
|
643
|
-
" and cancellation cannot be guaranteed."
|
644
|
-
)
|
645
|
-
},
|
646
|
-
)
|
647
|
-
return
|
648
|
-
else:
|
649
|
-
if configuration.is_using_a_runner:
|
650
|
-
self._logger.info(
|
651
|
-
f"Skipping cancellation because flow run {str(flow_run.id)!r} is"
|
652
|
-
" using enhanced cancellation. A dedicated runner will handle"
|
653
|
-
" cancellation."
|
654
|
-
)
|
655
|
-
return
|
656
|
-
|
657
|
-
if not flow_run.infrastructure_pid:
|
658
|
-
run_logger.error(
|
659
|
-
f"Flow run '{flow_run.id}' does not have an infrastructure pid"
|
660
|
-
" attached. Cancellation cannot be guaranteed."
|
661
|
-
)
|
662
|
-
await self._mark_flow_run_as_cancelled(
|
663
|
-
flow_run,
|
664
|
-
state_updates={
|
665
|
-
"message": (
|
666
|
-
"This flow run is missing infrastructure tracking information"
|
667
|
-
" and cancellation cannot be guaranteed."
|
668
|
-
)
|
669
|
-
},
|
670
|
-
)
|
671
|
-
return
|
672
|
-
|
673
|
-
try:
|
674
|
-
await self.kill_infrastructure(
|
675
|
-
infrastructure_pid=flow_run.infrastructure_pid,
|
676
|
-
configuration=configuration,
|
677
|
-
)
|
678
|
-
except NotImplementedError:
|
679
|
-
self._logger.error(
|
680
|
-
f"Worker type {self.type!r} does not support killing created "
|
681
|
-
"infrastructure. Cancellation cannot be guaranteed."
|
682
|
-
)
|
683
|
-
except InfrastructureNotFound as exc:
|
684
|
-
self._logger.warning(f"{exc} Marking flow run as cancelled.")
|
685
|
-
await self._mark_flow_run_as_cancelled(flow_run)
|
686
|
-
except InfrastructureNotAvailable as exc:
|
687
|
-
self._logger.warning(f"{exc} Flow run cannot be cancelled by this worker.")
|
688
|
-
except Exception:
|
689
|
-
run_logger.exception(
|
690
|
-
"Encountered exception while killing infrastructure for flow run "
|
691
|
-
f"'{flow_run.id}'. Flow run may not be cancelled."
|
692
|
-
)
|
693
|
-
# We will try again on generic exceptions
|
694
|
-
self._cancelling_flow_run_ids.remove(flow_run.id)
|
695
|
-
return
|
696
|
-
else:
|
697
|
-
self._emit_flow_run_cancelled_event(
|
698
|
-
flow_run=flow_run, configuration=configuration
|
699
|
-
)
|
700
|
-
await self._mark_flow_run_as_cancelled(flow_run)
|
701
|
-
run_logger.info(f"Cancelled flow run '{flow_run.id}'!")
|
702
|
-
|
703
663
|
async def _update_local_work_pool_info(self):
|
704
664
|
try:
|
705
665
|
work_pool = await self._client.read_work_pool(
|
706
666
|
work_pool_name=self._work_pool_name
|
707
667
|
)
|
668
|
+
|
708
669
|
except ObjectNotFound:
|
709
670
|
if self._create_pool_if_not_found:
|
710
671
|
wp = WorkPoolCreate(
|
@@ -798,11 +759,10 @@ class BaseWorker(abc.ABC):
|
|
798
759
|
for execution by the worker.
|
799
760
|
"""
|
800
761
|
submittable_flow_runs = [entry.flow_run for entry in flow_run_response]
|
801
|
-
|
762
|
+
|
802
763
|
for flow_run in submittable_flow_runs:
|
803
764
|
if flow_run.id in self._submitting_flow_run_ids:
|
804
765
|
continue
|
805
|
-
|
806
766
|
try:
|
807
767
|
if self._limiter:
|
808
768
|
self._limiter.acquire_on_behalf_of_nowait(flow_run.id)
|
@@ -847,8 +807,6 @@ class BaseWorker(abc.ABC):
|
|
847
807
|
" Please use an agent to execute this flow run."
|
848
808
|
)
|
849
809
|
|
850
|
-
#
|
851
|
-
|
852
810
|
async def _submit_run(self, flow_run: "FlowRun") -> None:
|
853
811
|
"""
|
854
812
|
Submits a given flow run for execution by the worker.
|
@@ -888,28 +846,59 @@ class BaseWorker(abc.ABC):
|
|
888
846
|
"not be cancellable."
|
889
847
|
)
|
890
848
|
|
891
|
-
|
849
|
+
run_logger.info(f"Completed submission of flow run '{flow_run.id}'")
|
892
850
|
|
893
|
-
|
894
|
-
|
895
|
-
|
896
|
-
|
851
|
+
else:
|
852
|
+
# If the run is not ready to submit, release the concurrency slot
|
853
|
+
if self._limiter:
|
854
|
+
self._limiter.release_on_behalf_of(flow_run.id)
|
897
855
|
|
898
|
-
|
856
|
+
self._submitting_flow_run_ids.remove(flow_run.id)
|
899
857
|
|
900
858
|
async def _submit_run_and_capture_errors(
|
901
859
|
self, flow_run: "FlowRun", task_status: Optional[anyio.abc.TaskStatus] = None
|
902
860
|
) -> Union[BaseWorkerResult, Exception]:
|
903
861
|
run_logger = self.get_flow_run_logger(flow_run)
|
862
|
+
deployment = None
|
863
|
+
|
864
|
+
if flow_run.deployment_id:
|
865
|
+
deployment = await self._client.read_deployment(flow_run.deployment_id)
|
866
|
+
if deployment and deployment.concurrency_limit:
|
867
|
+
limit_name = f"deployment:{deployment.id}"
|
868
|
+
concurrency_limit = deployment.concurrency_limit
|
869
|
+
concurrency_ctx = concurrency
|
870
|
+
else:
|
871
|
+
limit_name = None
|
872
|
+
concurrency_limit = None
|
873
|
+
concurrency_ctx = asyncnullcontext
|
904
874
|
|
905
875
|
try:
|
906
|
-
|
907
|
-
|
908
|
-
|
909
|
-
|
910
|
-
|
911
|
-
|
876
|
+
async with concurrency_ctx(
|
877
|
+
limit_name, occupy=concurrency_limit, max_retries=0
|
878
|
+
):
|
879
|
+
configuration = await self._get_configuration(flow_run, deployment)
|
880
|
+
submitted_event = self._emit_flow_run_submitted_event(configuration)
|
881
|
+
result = await self.run(
|
882
|
+
flow_run=flow_run,
|
883
|
+
task_status=task_status,
|
884
|
+
configuration=configuration,
|
885
|
+
)
|
886
|
+
except (
|
887
|
+
AcquireConcurrencySlotTimeoutError,
|
888
|
+
ConcurrencySlotAcquisitionError,
|
889
|
+
) as exc:
|
890
|
+
self._logger.info(
|
891
|
+
(
|
892
|
+
"Deployment %s has reached its concurrency limit when submitting flow run %s"
|
893
|
+
),
|
894
|
+
flow_run.deployment_id,
|
895
|
+
flow_run.name,
|
912
896
|
)
|
897
|
+
await self._propose_scheduled_state(flow_run)
|
898
|
+
|
899
|
+
if not task_status._future.done():
|
900
|
+
task_status.started(exc)
|
901
|
+
return exc
|
913
902
|
except Exception as exc:
|
914
903
|
if not task_status._future.done():
|
915
904
|
# This flow run was being submitted and did not start successfully
|
@@ -963,7 +952,7 @@ class BaseWorker(abc.ABC):
|
|
963
952
|
return {
|
964
953
|
"name": self.name,
|
965
954
|
"work_pool": (
|
966
|
-
self._work_pool.
|
955
|
+
self._work_pool.model_dump(mode="json")
|
967
956
|
if self._work_pool is not None
|
968
957
|
else None
|
969
958
|
),
|
@@ -975,8 +964,13 @@ class BaseWorker(abc.ABC):
|
|
975
964
|
async def _get_configuration(
|
976
965
|
self,
|
977
966
|
flow_run: "FlowRun",
|
967
|
+
deployment: Optional["DeploymentResponse"] = None,
|
978
968
|
) -> BaseJobConfiguration:
|
979
|
-
deployment =
|
969
|
+
deployment = (
|
970
|
+
deployment
|
971
|
+
if deployment
|
972
|
+
else await self._client.read_deployment(flow_run.deployment_id)
|
973
|
+
)
|
980
974
|
flow = await self._client.read_flow(flow_run.flow_id)
|
981
975
|
|
982
976
|
deployment_vars = deployment.job_variables or {}
|
@@ -1030,6 +1024,21 @@ class BaseWorker(abc.ABC):
|
|
1030
1024
|
|
1031
1025
|
return True
|
1032
1026
|
|
1027
|
+
async def _propose_scheduled_state(self, flow_run: "FlowRun") -> None:
|
1028
|
+
run_logger = self.get_flow_run_logger(flow_run)
|
1029
|
+
try:
|
1030
|
+
state = await propose_state(
|
1031
|
+
self._client,
|
1032
|
+
AwaitingConcurrencySlot(),
|
1033
|
+
flow_run_id=flow_run.id,
|
1034
|
+
)
|
1035
|
+
self._logger.info(f"Flow run {flow_run.id} now has state {state.name}")
|
1036
|
+
except Abort:
|
1037
|
+
# Flow run already marked as failed
|
1038
|
+
pass
|
1039
|
+
except Exception:
|
1040
|
+
run_logger.exception(f"Failed to update state of flow run '{flow_run.id}'")
|
1041
|
+
|
1033
1042
|
async def _propose_failed_state(self, flow_run: "FlowRun", exc: Exception) -> None:
|
1034
1043
|
run_logger = self.get_flow_run_logger(flow_run)
|
1035
1044
|
try:
|
@@ -1073,7 +1082,7 @@ class BaseWorker(abc.ABC):
|
|
1073
1082
|
state_updates = state_updates or {}
|
1074
1083
|
state_updates.setdefault("name", "Cancelled")
|
1075
1084
|
state_updates.setdefault("type", StateType.CANCELLED)
|
1076
|
-
state = flow_run.state.
|
1085
|
+
state = flow_run.state.model_copy(update=state_updates)
|
1077
1086
|
|
1078
1087
|
await self._client.set_flow_run_state(flow_run.id, state, force=True)
|
1079
1088
|
|
@@ -1124,6 +1133,7 @@ class BaseWorker(abc.ABC):
|
|
1124
1133
|
async def __aenter__(self):
|
1125
1134
|
self._logger.debug("Entering worker context...")
|
1126
1135
|
await self.setup()
|
1136
|
+
|
1127
1137
|
return self
|
1128
1138
|
|
1129
1139
|
async def __aexit__(self, *exc_info):
|
@@ -1160,7 +1170,7 @@ class BaseWorker(abc.ABC):
|
|
1160
1170
|
if include_self:
|
1161
1171
|
worker_resource = self._event_resource()
|
1162
1172
|
worker_resource["prefect.resource.role"] = "worker"
|
1163
|
-
related.append(RelatedResource.
|
1173
|
+
related.append(RelatedResource.model_validate(worker_resource))
|
1164
1174
|
|
1165
1175
|
return related
|
1166
1176
|
|
@@ -1207,20 +1217,3 @@ class BaseWorker(abc.ABC):
|
|
1207
1217
|
related=self._event_related_resources(),
|
1208
1218
|
follows=started_event,
|
1209
1219
|
)
|
1210
|
-
|
1211
|
-
def _emit_flow_run_cancelled_event(
|
1212
|
-
self, flow_run: "FlowRun", configuration: BaseJobConfiguration
|
1213
|
-
):
|
1214
|
-
related = self._event_related_resources(configuration=configuration)
|
1215
|
-
|
1216
|
-
for resource in related:
|
1217
|
-
if resource.role == "flow-run":
|
1218
|
-
resource["prefect.infrastructure.identifier"] = str(
|
1219
|
-
flow_run.infrastructure_pid
|
1220
|
-
)
|
1221
|
-
|
1222
|
-
emit_event(
|
1223
|
-
event="prefect.worker.cancelled-flow-run",
|
1224
|
-
resource=self._event_resource(),
|
1225
|
-
related=related,
|
1226
|
-
)
|