prefect-client 2.18.3__py3-none-any.whl → 2.19.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (42) hide show
  1. prefect/__init__.py +1 -15
  2. prefect/_internal/compatibility/experimental.py +11 -2
  3. prefect/_internal/concurrency/cancellation.py +2 -0
  4. prefect/_internal/schemas/validators.py +10 -0
  5. prefect/_vendor/starlette/testclient.py +1 -1
  6. prefect/blocks/notifications.py +6 -6
  7. prefect/client/base.py +244 -1
  8. prefect/client/cloud.py +4 -2
  9. prefect/client/orchestration.py +515 -106
  10. prefect/client/schemas/actions.py +58 -8
  11. prefect/client/schemas/objects.py +15 -1
  12. prefect/client/schemas/responses.py +19 -0
  13. prefect/client/schemas/schedules.py +1 -1
  14. prefect/client/utilities.py +2 -2
  15. prefect/concurrency/asyncio.py +34 -4
  16. prefect/concurrency/sync.py +40 -6
  17. prefect/context.py +2 -2
  18. prefect/engine.py +2 -2
  19. prefect/events/clients.py +2 -2
  20. prefect/flows.py +91 -17
  21. prefect/infrastructure/process.py +0 -17
  22. prefect/logging/formatters.py +1 -4
  23. prefect/new_flow_engine.py +137 -168
  24. prefect/new_task_engine.py +137 -202
  25. prefect/runner/__init__.py +1 -1
  26. prefect/runner/runner.py +2 -107
  27. prefect/settings.py +21 -0
  28. prefect/tasks.py +76 -57
  29. prefect/types/__init__.py +27 -5
  30. prefect/utilities/annotations.py +1 -8
  31. prefect/utilities/asyncutils.py +4 -0
  32. prefect/utilities/engine.py +106 -1
  33. prefect/utilities/schema_tools/__init__.py +6 -1
  34. prefect/utilities/schema_tools/validation.py +25 -8
  35. prefect/utilities/timeout.py +34 -0
  36. prefect/workers/base.py +7 -3
  37. prefect/workers/process.py +0 -17
  38. {prefect_client-2.18.3.dist-info → prefect_client-2.19.1.dist-info}/METADATA +1 -1
  39. {prefect_client-2.18.3.dist-info → prefect_client-2.19.1.dist-info}/RECORD +42 -41
  40. {prefect_client-2.18.3.dist-info → prefect_client-2.19.1.dist-info}/LICENSE +0 -0
  41. {prefect_client-2.18.3.dist-info → prefect_client-2.19.1.dist-info}/WHEEL +0 -0
  42. {prefect_client-2.18.3.dist-info → prefect_client-2.19.1.dist-info}/top_level.txt +0 -0
prefect/runner/runner.py CHANGED
@@ -51,9 +51,6 @@ from uuid import UUID, uuid4
51
51
  import anyio
52
52
  import anyio.abc
53
53
  import pendulum
54
- import sniffio
55
- from rich.console import Console, Group
56
- from rich.table import Table
57
54
 
58
55
  from prefect._internal.concurrency.api import (
59
56
  create_call,
@@ -80,7 +77,6 @@ from prefect.deployments.runner import (
80
77
  RunnerDeployment,
81
78
  )
82
79
  from prefect.deployments.schedules import FlexibleScheduleList
83
- from prefect.engine import propose_state
84
80
  from prefect.events import DeploymentTriggerTypes, TriggerTypes
85
81
  from prefect.exceptions import (
86
82
  Abort,
@@ -94,7 +90,6 @@ from prefect.settings import (
94
90
  PREFECT_RUNNER_POLL_FREQUENCY,
95
91
  PREFECT_RUNNER_PROCESS_LIMIT,
96
92
  PREFECT_RUNNER_SERVER_ENABLE,
97
- PREFECT_UI_URL,
98
93
  get_current_settings,
99
94
  )
100
95
  from prefect.states import Crashed, Pending, exception_to_failed_state
@@ -103,10 +98,11 @@ from prefect.utilities.asyncutils import (
103
98
  is_async_fn,
104
99
  sync_compatible,
105
100
  )
101
+ from prefect.utilities.engine import propose_state
106
102
  from prefect.utilities.processutils import _register_signal, run_process
107
103
  from prefect.utilities.services import critical_service_loop
108
104
 
109
- __all__ = ["Runner", "serve"]
105
+ __all__ = ["Runner"]
110
106
 
111
107
 
112
108
  class Runner:
@@ -550,7 +546,6 @@ class Runner:
550
546
  if sys.platform == "win32":
551
547
  kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
552
548
 
553
- _use_threaded_child_watcher()
554
549
  flow_run_logger.info("Opening process...")
555
550
 
556
551
  env = get_current_settings().to_environment_variables(exclude_unset=True)
@@ -1190,106 +1185,6 @@ if sys.platform == "win32":
1190
1185
  STATUS_CONTROL_C_EXIT = 0xC000013A
1191
1186
 
1192
1187
 
1193
- def _use_threaded_child_watcher():
1194
- if (
1195
- sys.version_info < (3, 8)
1196
- and sniffio.current_async_library() == "asyncio"
1197
- and sys.platform != "win32"
1198
- ):
1199
- from prefect.utilities.compat import ThreadedChildWatcher
1200
-
1201
- # Python < 3.8 does not use a `ThreadedChildWatcher` by default which can
1202
- # lead to errors in tests on unix as the previous default `SafeChildWatcher`
1203
- # is not compatible with threaded event loops.
1204
- asyncio.get_event_loop_policy().set_child_watcher(ThreadedChildWatcher())
1205
-
1206
-
1207
- @sync_compatible
1208
- async def serve(
1209
- *args: RunnerDeployment,
1210
- pause_on_shutdown: bool = True,
1211
- print_starting_message: bool = True,
1212
- limit: Optional[int] = None,
1213
- **kwargs,
1214
- ):
1215
- """
1216
- Serve the provided list of deployments.
1217
-
1218
- Args:
1219
- *args: A list of deployments to serve.
1220
- pause_on_shutdown: A boolean for whether or not to automatically pause
1221
- deployment schedules on shutdown.
1222
- print_starting_message: Whether or not to print message to the console
1223
- on startup.
1224
- limit: The maximum number of runs that can be executed concurrently.
1225
- **kwargs: Additional keyword arguments to pass to the runner.
1226
-
1227
- Examples:
1228
- Prepare two deployments and serve them:
1229
-
1230
- ```python
1231
- import datetime
1232
-
1233
- from prefect import flow, serve
1234
-
1235
- @flow
1236
- def my_flow(name):
1237
- print(f"hello {name}")
1238
-
1239
- @flow
1240
- def my_other_flow(name):
1241
- print(f"goodbye {name}")
1242
-
1243
- if __name__ == "__main__":
1244
- # Run once a day
1245
- hello_deploy = my_flow.to_deployment(
1246
- "hello", tags=["dev"], interval=datetime.timedelta(days=1)
1247
- )
1248
-
1249
- # Run every Sunday at 4:00 AM
1250
- bye_deploy = my_other_flow.to_deployment(
1251
- "goodbye", tags=["dev"], cron="0 4 * * sun"
1252
- )
1253
-
1254
- serve(hello_deploy, bye_deploy)
1255
- ```
1256
- """
1257
- runner = Runner(pause_on_shutdown=pause_on_shutdown, limit=limit, **kwargs)
1258
- for deployment in args:
1259
- await runner.add_deployment(deployment)
1260
-
1261
- if print_starting_message:
1262
- help_message_top = (
1263
- "[green]Your deployments are being served and polling for"
1264
- " scheduled runs!\n[/]"
1265
- )
1266
-
1267
- table = Table(title="Deployments", show_header=False)
1268
-
1269
- table.add_column(style="blue", no_wrap=True)
1270
-
1271
- for deployment in args:
1272
- table.add_row(f"{deployment.flow_name}/{deployment.name}")
1273
-
1274
- help_message_bottom = (
1275
- "\nTo trigger any of these deployments, use the"
1276
- " following command:\n[blue]\n\t$ prefect deployment run"
1277
- " [DEPLOYMENT_NAME]\n[/]"
1278
- )
1279
- if PREFECT_UI_URL:
1280
- help_message_bottom += (
1281
- "\nYou can also trigger your deployments via the Prefect UI:"
1282
- f" [blue]{PREFECT_UI_URL.value()}/deployments[/]\n"
1283
- )
1284
-
1285
- console = Console()
1286
- console.print(
1287
- Group(help_message_top, table, help_message_bottom), soft_wrap=True
1288
- )
1289
-
1290
- await runner.start()
1291
-
1292
-
1293
1188
  async def _run_hooks(
1294
1189
  hooks: List[Callable[[Flow, "FlowRun", State], None]], flow_run, flow, state
1295
1190
  ):
prefect/settings.py CHANGED
@@ -1501,6 +1501,11 @@ PREFECT_RUNNER_SERVER_ENABLE = Setting(bool, default=False)
1501
1501
  Whether or not to enable the runner's webserver.
1502
1502
  """
1503
1503
 
1504
+ PREFECT_DEPLOYMENT_SCHEDULE_MAX_SCHEDULED_RUNS = Setting(int, default=50)
1505
+ """
1506
+ The maximum number of scheduled runs to create for a deployment.
1507
+ """
1508
+
1504
1509
  PREFECT_WORKER_HEARTBEAT_SECONDS = Setting(float, default=30)
1505
1510
  """
1506
1511
  Number of seconds a worker should wait between sending a heartbeat.
@@ -1614,6 +1619,11 @@ PREFECT_EXPERIMENTAL_ENABLE_NEW_ENGINE = Setting(bool, default=False)
1614
1619
  Whether or not to enable experimental new engine.
1615
1620
  """
1616
1621
 
1622
+ PREFECT_EXPERIMENTAL_DISABLE_SYNC_COMPAT = Setting(bool, default=False)
1623
+ """
1624
+ Whether or not to disable the sync_compatible decorator utility.
1625
+ """
1626
+
1617
1627
 
1618
1628
  # Defaults -----------------------------------------------------------------------------
1619
1629
 
@@ -1745,6 +1755,17 @@ PREFECT_API_EVENTS_RELATED_RESOURCE_CACHE_TTL = Setting(
1745
1755
  How long to cache related resource data for emitting server-side vents
1746
1756
  """
1747
1757
 
1758
+
1759
+ def automation_settings_enabled() -> bool:
1760
+ """
1761
+ Whether or not automations are enabled.
1762
+ """
1763
+ return (
1764
+ PREFECT_EXPERIMENTAL_EVENTS.value()
1765
+ and PREFECT_API_SERVICES_TRIGGERS_ENABLED.value()
1766
+ )
1767
+
1768
+
1748
1769
  # Deprecated settings ------------------------------------------------------------------
1749
1770
 
1750
1771
 
prefect/tasks.py CHANGED
@@ -7,7 +7,6 @@ Module containing the base workflow task class and decorator - for most use case
7
7
  import datetime
8
8
  import inspect
9
9
  import os
10
- import warnings
11
10
  from copy import copy
12
11
  from functools import partial, update_wrapper
13
12
  from typing import (
@@ -33,9 +32,15 @@ from uuid import uuid4
33
32
  from typing_extensions import Literal, ParamSpec
34
33
 
35
34
  from prefect._internal.concurrency.api import create_call, from_async, from_sync
35
+ from prefect.client.orchestration import PrefectClient, SyncPrefectClient
36
36
  from prefect.client.schemas import TaskRun
37
- from prefect.client.schemas.objects import TaskRunInput
38
- from prefect.context import FlowRunContext, PrefectObjectRegistry, TagsContext
37
+ from prefect.client.schemas.objects import TaskRunInput, TaskRunResult
38
+ from prefect.context import (
39
+ FlowRunContext,
40
+ PrefectObjectRegistry,
41
+ TagsContext,
42
+ TaskRunContext,
43
+ )
39
44
  from prefect.futures import PrefectFuture
40
45
  from prefect.logging.loggers import get_logger, get_run_logger
41
46
  from prefect.results import ResultSerializer, ResultStorage
@@ -332,33 +337,7 @@ class Task(Generic[P, R]):
332
337
  self.result_serializer = result_serializer
333
338
  self.result_storage_key = result_storage_key
334
339
  self.cache_result_in_memory = cache_result_in_memory
335
-
336
340
  self.timeout_seconds = float(timeout_seconds) if timeout_seconds else None
337
- # Warn if this task's `name` conflicts with another task while having a
338
- # different function. This is to detect the case where two or more tasks
339
- # share a name or are lambdas, which should result in a warning, and to
340
- # differentiate it from the case where the task was 'copied' via
341
- # `with_options`, which should not result in a warning.
342
- registry = PrefectObjectRegistry.get()
343
-
344
- if registry and any(
345
- other
346
- for other in registry.get_instances(Task)
347
- if other.name == self.name and id(other.fn) != id(self.fn)
348
- ):
349
- try:
350
- file = inspect.getsourcefile(self.fn)
351
- line_number = inspect.getsourcelines(self.fn)[1]
352
- except TypeError:
353
- file = "unknown"
354
- line_number = "unknown"
355
-
356
- warnings.warn(
357
- f"A task named {self.name!r} and defined at '{file}:{line_number}' "
358
- "conflicts with another task. Consider specifying a unique `name` "
359
- "parameter in the task definition:\n\n "
360
- "`@task(name='my_unique_name', ...)`"
361
- )
362
341
  self.on_completion = on_completion
363
342
  self.on_failure = on_failure
364
343
 
@@ -539,48 +518,93 @@ class Task(Generic[P, R]):
539
518
 
540
519
  async def create_run(
541
520
  self,
542
- flow_run_context: FlowRunContext,
543
- parameters: Dict[str, Any],
544
- wait_for: Optional[Iterable[PrefectFuture]],
521
+ client: Optional[Union[PrefectClient, SyncPrefectClient]],
522
+ parameters: Dict[str, Any] = None,
523
+ flow_run_context: Optional[FlowRunContext] = None,
524
+ parent_task_run_context: Optional[TaskRunContext] = None,
525
+ wait_for: Optional[Iterable[PrefectFuture]] = None,
545
526
  extra_task_inputs: Optional[Dict[str, Set[TaskRunInput]]] = None,
546
527
  ) -> TaskRun:
547
- # TODO: Investigate if we can replace create_task_run on the task run engine
548
- # with this method. Would require updating to work without the flow run context.
549
528
  from prefect.utilities.engine import (
550
529
  _dynamic_key_for_task_run,
530
+ _resolve_custom_task_run_name,
551
531
  collect_task_run_inputs,
552
532
  )
553
533
 
554
- dynamic_key = _dynamic_key_for_task_run(flow_run_context, self)
534
+ if flow_run_context is None:
535
+ flow_run_context = FlowRunContext.get()
536
+ if parent_task_run_context is None:
537
+ parent_task_run_context = TaskRunContext.get()
538
+ if parameters is None:
539
+ parameters = {}
540
+
541
+ try:
542
+ task_run_name = _resolve_custom_task_run_name(self, parameters)
543
+ except TypeError:
544
+ task_run_name = None
545
+
546
+ if flow_run_context:
547
+ dynamic_key = _dynamic_key_for_task_run(context=flow_run_context, task=self)
548
+ else:
549
+ dynamic_key = uuid4().hex
550
+
551
+ # collect task inputs
555
552
  task_inputs = {
556
553
  k: await collect_task_run_inputs(v) for k, v in parameters.items()
557
554
  }
555
+
556
+ # check if this task has a parent task run based on running in another
557
+ # task run's existing context. A task run is only considered a parent if
558
+ # it is in the same flow run (because otherwise presumably the child is
559
+ # in a subflow, so the subflow serves as the parent) or if there is no
560
+ # flow run
561
+ if parent_task_run_context:
562
+ # there is no flow run
563
+ if not flow_run_context:
564
+ task_inputs["__parents__"] = [
565
+ TaskRunResult(id=parent_task_run_context.task_run.id)
566
+ ]
567
+ # there is a flow run and the task run is in the same flow run
568
+ elif (
569
+ flow_run_context
570
+ and parent_task_run_context.task_run.flow_run_id
571
+ == flow_run_context.flow_run.id
572
+ ):
573
+ task_inputs["__parents__"] = [
574
+ TaskRunResult(id=parent_task_run_context.task_run.id)
575
+ ]
576
+
558
577
  if wait_for:
559
578
  task_inputs["wait_for"] = await collect_task_run_inputs(wait_for)
560
579
 
561
580
  # Join extra task inputs
562
- extra_task_inputs = extra_task_inputs or {}
563
- for k, extras in extra_task_inputs.items():
581
+ for k, extras in (extra_task_inputs or {}).items():
564
582
  task_inputs[k] = task_inputs[k].union(extras)
565
583
 
566
- flow_run_logger = get_run_logger(flow_run_context)
567
-
568
- task_run = await flow_run_context.client.create_task_run(
584
+ # create the task run
585
+ task_run = client.create_task_run(
569
586
  task=self,
570
- name=f"{self.name} - {dynamic_key}",
571
- flow_run_id=flow_run_context.flow_run.id,
572
- dynamic_key=dynamic_key,
587
+ name=task_run_name,
588
+ flow_run_id=(
589
+ getattr(flow_run_context.flow_run, "id", None)
590
+ if flow_run_context and flow_run_context.flow_run
591
+ else None
592
+ ),
593
+ dynamic_key=str(dynamic_key),
573
594
  state=Pending(),
574
- extra_tags=TagsContext.get().current_tags,
575
595
  task_inputs=task_inputs,
596
+ extra_tags=TagsContext.get().current_tags,
576
597
  )
598
+ # the new engine uses sync clients but old engines use async clients
599
+ if inspect.isawaitable(task_run):
600
+ task_run = await task_run
577
601
 
578
- if flow_run_context.flow_run:
579
- flow_run_logger.info(
602
+ if flow_run_context and flow_run_context.flow_run:
603
+ get_run_logger(flow_run_context).debug(
580
604
  f"Created task run {task_run.name!r} for task {self.name!r}"
581
605
  )
582
606
  else:
583
- logger.info(f"Created task run {task_run.name!r} for task {self.name!r}")
607
+ logger.debug(f"Created task run {task_run.name!r} for task {self.name!r}")
584
608
 
585
609
  return task_run
586
610
 
@@ -637,21 +661,15 @@ class Task(Generic[P, R]):
637
661
  self.isasync, self.name, parameters, self.viz_return_value
638
662
  )
639
663
 
640
- # new engine currently only compatible with async tasks
641
664
  if PREFECT_EXPERIMENTAL_ENABLE_NEW_ENGINE.value():
642
- from prefect.new_task_engine import run_task, run_task_sync
665
+ from prefect.new_task_engine import run_task
643
666
 
644
- run_kwargs = dict(
667
+ return run_task(
645
668
  task=self,
646
669
  parameters=parameters,
647
670
  wait_for=wait_for,
648
671
  return_type=return_type,
649
672
  )
650
- if self.isasync:
651
- # this returns an awaitable coroutine
652
- return run_task(**run_kwargs)
653
- else:
654
- return run_task_sync(**run_kwargs)
655
673
 
656
674
  if (
657
675
  PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING.value()
@@ -931,11 +949,12 @@ class Task(Generic[P, R]):
931
949
  wait_for: Optional[Iterable[PrefectFuture]],
932
950
  return_state: bool,
933
951
  ):
934
- from prefect.new_task_engine import run_task
952
+ from prefect.new_task_engine import run_task_async
935
953
 
936
954
  task_runner = flow_run_context.task_runner
937
955
 
938
956
  task_run = await self.create_run(
957
+ client=flow_run_context.client,
939
958
  flow_run_context=flow_run_context,
940
959
  parameters=parameters,
941
960
  wait_for=wait_for,
@@ -952,7 +971,7 @@ class Task(Generic[P, R]):
952
971
  await task_runner.submit(
953
972
  key=future.key,
954
973
  call=partial(
955
- run_task,
974
+ run_task_async,
956
975
  task=self,
957
976
  task_run=task_run,
958
977
  parameters=parameters,
prefect/types/__init__.py CHANGED
@@ -1,4 +1,3 @@
1
- from dataclasses import dataclass
2
1
  from typing import Any, Callable, ClassVar, Generator
3
2
 
4
3
  from pydantic_core import core_schema, CoreSchema, SchemaValidator
@@ -6,8 +5,9 @@ from typing_extensions import Self
6
5
  from datetime import timedelta
7
6
 
8
7
 
9
- @dataclass
10
8
  class NonNegativeInteger(int):
9
+ """An integer that must be greater than or equal to 0."""
10
+
11
11
  schema: ClassVar[CoreSchema] = core_schema.int_schema(ge=0)
12
12
 
13
13
  @classmethod
@@ -25,8 +25,9 @@ class NonNegativeInteger(int):
25
25
  return SchemaValidator(schema=cls.schema).validate_python(v)
26
26
 
27
27
 
28
- @dataclass
29
28
  class PositiveInteger(int):
29
+ """An integer that must be greater than 0."""
30
+
30
31
  schema: ClassVar[CoreSchema] = core_schema.int_schema(gt=0)
31
32
 
32
33
  @classmethod
@@ -44,8 +45,27 @@ class PositiveInteger(int):
44
45
  return SchemaValidator(schema=cls.schema).validate_python(v)
45
46
 
46
47
 
47
- @dataclass
48
+ class NonNegativeFloat(float):
49
+ schema: ClassVar[CoreSchema] = core_schema.float_schema(ge=0)
50
+
51
+ @classmethod
52
+ def __get_validators__(cls) -> Generator[Callable[..., Any], None, None]:
53
+ yield cls.validate
54
+
55
+ @classmethod
56
+ def __get_pydantic_core_schema__(
57
+ cls, source_type: Any, handler: Callable[..., Any]
58
+ ) -> CoreSchema:
59
+ return cls.schema
60
+
61
+ @classmethod
62
+ def validate(cls, v: Any) -> Self:
63
+ return SchemaValidator(schema=cls.schema).validate_python(v)
64
+
65
+
48
66
  class NonNegativeDuration(timedelta):
67
+ """A timedelta that must be greater than or equal to 0."""
68
+
49
69
  schema: ClassVar = core_schema.timedelta_schema(ge=timedelta(seconds=0))
50
70
 
51
71
  @classmethod
@@ -63,8 +83,9 @@ class NonNegativeDuration(timedelta):
63
83
  return SchemaValidator(schema=cls.schema).validate_python(v)
64
84
 
65
85
 
66
- @dataclass
67
86
  class PositiveDuration(timedelta):
87
+ """A timedelta that must be greater than 0."""
88
+
68
89
  schema: ClassVar = core_schema.timedelta_schema(gt=timedelta(seconds=0))
69
90
 
70
91
  @classmethod
@@ -85,6 +106,7 @@ class PositiveDuration(timedelta):
85
106
  __all__ = [
86
107
  "NonNegativeInteger",
87
108
  "PositiveInteger",
109
+ "NonNegativeFloat",
88
110
  "NonNegativeDuration",
89
111
  "PositiveDuration",
90
112
  ]
@@ -1,4 +1,3 @@
1
- import sys
2
1
  import warnings
3
2
  from abc import ABC
4
3
  from collections import namedtuple
@@ -17,13 +16,7 @@ class BaseAnnotation(
17
16
  """
18
17
 
19
18
  def unwrap(self) -> T:
20
- if sys.version_info < (3, 8):
21
- # cannot simply return self.value due to recursion error in Python 3.7
22
- # also _asdict does not follow convention; it's not an internal method
23
- # https://stackoverflow.com/a/26180604
24
- return self._asdict()["value"]
25
- else:
26
- return self.value
19
+ return self.value
27
20
 
28
21
  def rewrap(self, value: T) -> "BaseAnnotation[T]":
29
22
  return type(self)(value)
@@ -266,6 +266,10 @@ def sync_compatible(async_fn: T) -> T:
266
266
  from prefect._internal.concurrency.calls import get_current_call, logger
267
267
  from prefect._internal.concurrency.event_loop import get_running_loop
268
268
  from prefect._internal.concurrency.threads import get_global_loop
269
+ from prefect.settings import PREFECT_EXPERIMENTAL_DISABLE_SYNC_COMPAT
270
+
271
+ if PREFECT_EXPERIMENTAL_DISABLE_SYNC_COMPAT:
272
+ return async_fn(*args, **kwargs)
269
273
 
270
274
  global_thread_portal = get_global_loop()
271
275
  current_thread = threading.current_thread()
@@ -1,5 +1,6 @@
1
1
  import asyncio
2
2
  import contextlib
3
+ import inspect
3
4
  import os
4
5
  import signal
5
6
  import time
@@ -23,7 +24,7 @@ import prefect
23
24
  import prefect.context
24
25
  import prefect.plugins
25
26
  from prefect._internal.concurrency.cancellation import get_deadline
26
- from prefect.client.orchestration import PrefectClient
27
+ from prefect.client.orchestration import PrefectClient, SyncPrefectClient
27
28
  from prefect.client.schemas import OrchestrationResult, TaskRun
28
29
  from prefect.client.schemas.objects import (
29
30
  StateType,
@@ -60,6 +61,7 @@ from prefect.tasks import Task
60
61
  from prefect.utilities.annotations import allow_failure, quote
61
62
  from prefect.utilities.asyncutils import (
62
63
  gather,
64
+ run_sync,
63
65
  )
64
66
  from prefect.utilities.collections import StopVisiting, visit_collection
65
67
  from prefect.utilities.text import truncated_to
@@ -409,6 +411,109 @@ async def propose_state(
409
411
  )
410
412
 
411
413
 
414
+ def propose_state_sync(
415
+ client: SyncPrefectClient,
416
+ state: State[object],
417
+ force: bool = False,
418
+ task_run_id: Optional[UUID] = None,
419
+ flow_run_id: Optional[UUID] = None,
420
+ ) -> State[object]:
421
+ """
422
+ Propose a new state for a flow run or task run, invoking Prefect orchestration logic.
423
+
424
+ If the proposed state is accepted, the provided `state` will be augmented with
425
+ details and returned.
426
+
427
+ If the proposed state is rejected, a new state returned by the Prefect API will be
428
+ returned.
429
+
430
+ If the proposed state results in a WAIT instruction from the Prefect API, the
431
+ function will sleep and attempt to propose the state again.
432
+
433
+ If the proposed state results in an ABORT instruction from the Prefect API, an
434
+ error will be raised.
435
+
436
+ Args:
437
+ state: a new state for the task or flow run
438
+ task_run_id: an optional task run id, used when proposing task run states
439
+ flow_run_id: an optional flow run id, used when proposing flow run states
440
+
441
+ Returns:
442
+ a [State model][prefect.client.schemas.objects.State] representation of the
443
+ flow or task run state
444
+
445
+ Raises:
446
+ ValueError: if neither task_run_id or flow_run_id is provided
447
+ prefect.exceptions.Abort: if an ABORT instruction is received from
448
+ the Prefect API
449
+ """
450
+
451
+ # Determine if working with a task run or flow run
452
+ if not task_run_id and not flow_run_id:
453
+ raise ValueError("You must provide either a `task_run_id` or `flow_run_id`")
454
+
455
+ # Handle task and sub-flow tracing
456
+ if state.is_final():
457
+ if isinstance(state.data, BaseResult) and state.data.has_cached_object():
458
+ # Avoid fetching the result unless it is cached, otherwise we defeat
459
+ # the purpose of disabling `cache_result_in_memory`
460
+ result = state.result(raise_on_failure=False, fetch=True)
461
+ if inspect.isawaitable(result):
462
+ result = run_sync(result)
463
+ else:
464
+ result = state.data
465
+
466
+ link_state_to_result(state, result)
467
+
468
+ # Handle repeated WAITs in a loop instead of recursively, to avoid
469
+ # reaching max recursion depth in extreme cases.
470
+ def set_state_and_handle_waits(set_state_func) -> OrchestrationResult:
471
+ response = set_state_func()
472
+ while response.status == SetStateStatus.WAIT:
473
+ engine_logger.debug(
474
+ f"Received wait instruction for {response.details.delay_seconds}s: "
475
+ f"{response.details.reason}"
476
+ )
477
+ time.sleep(response.details.delay_seconds)
478
+ response = set_state_func()
479
+ return response
480
+
481
+ # Attempt to set the state
482
+ if task_run_id:
483
+ set_state = partial(client.set_task_run_state, task_run_id, state, force=force)
484
+ response = set_state_and_handle_waits(set_state)
485
+ elif flow_run_id:
486
+ set_state = partial(client.set_flow_run_state, flow_run_id, state, force=force)
487
+ response = set_state_and_handle_waits(set_state)
488
+ else:
489
+ raise ValueError(
490
+ "Neither flow run id or task run id were provided. At least one must "
491
+ "be given."
492
+ )
493
+
494
+ # Parse the response to return the new state
495
+ if response.status == SetStateStatus.ACCEPT:
496
+ # Update the state with the details if provided
497
+ state.id = response.state.id
498
+ state.timestamp = response.state.timestamp
499
+ if response.state.state_details:
500
+ state.state_details = response.state.state_details
501
+ return state
502
+
503
+ elif response.status == SetStateStatus.ABORT:
504
+ raise prefect.exceptions.Abort(response.details.reason)
505
+
506
+ elif response.status == SetStateStatus.REJECT:
507
+ if response.state.is_paused():
508
+ raise Pause(response.details.reason, state=response.state)
509
+ return response.state
510
+
511
+ else:
512
+ raise ValueError(
513
+ f"Received unexpected `SetStateStatus` from server: {response.status!r}"
514
+ )
515
+
516
+
412
517
  def _dynamic_key_for_task_run(context: FlowRunContext, task: Task) -> int:
413
518
  if context.flow_run is None: # this is an autonomous task run
414
519
  context.task_run_dynamic_keys[task.task_key] = getattr(
@@ -1,5 +1,10 @@
1
1
  from .hydration import HydrationContext, HydrationError, hydrate
2
- from .validation import CircularSchemaRefError, ValidationError, validate
2
+ from .validation import (
3
+ CircularSchemaRefError,
4
+ ValidationError,
5
+ validate,
6
+ is_valid_schema,
7
+ )
3
8
 
4
9
  __all__ = [
5
10
  "CircularSchemaRefError",