prefect-client 3.0.0rc9__py3-none-any.whl → 3.0.0rc11__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 (49) hide show
  1. prefect/_internal/compatibility/migration.py +48 -8
  2. prefect/_internal/concurrency/api.py +1 -1
  3. prefect/_internal/retries.py +61 -0
  4. prefect/agent.py +6 -0
  5. prefect/client/cloud.py +1 -1
  6. prefect/client/schemas/objects.py +3 -4
  7. prefect/concurrency/asyncio.py +3 -3
  8. prefect/concurrency/events.py +1 -1
  9. prefect/concurrency/services.py +3 -2
  10. prefect/concurrency/sync.py +19 -5
  11. prefect/context.py +14 -2
  12. prefect/deployments/__init__.py +28 -15
  13. prefect/deployments/schedules.py +5 -2
  14. prefect/deployments/steps/pull.py +7 -0
  15. prefect/events/schemas/automations.py +3 -3
  16. prefect/exceptions.py +4 -1
  17. prefect/filesystems.py +4 -3
  18. prefect/flow_engine.py +76 -14
  19. prefect/flows.py +222 -64
  20. prefect/futures.py +53 -7
  21. prefect/infrastructure/__init__.py +6 -0
  22. prefect/infrastructure/base.py +6 -0
  23. prefect/logging/loggers.py +1 -1
  24. prefect/results.py +50 -67
  25. prefect/runner/runner.py +93 -20
  26. prefect/runner/server.py +20 -22
  27. prefect/runner/submit.py +0 -8
  28. prefect/runtime/flow_run.py +38 -3
  29. prefect/serializers.py +3 -3
  30. prefect/settings.py +15 -45
  31. prefect/task_engine.py +77 -21
  32. prefect/task_runners.py +28 -16
  33. prefect/task_worker.py +6 -4
  34. prefect/tasks.py +30 -5
  35. prefect/transactions.py +18 -2
  36. prefect/utilities/asyncutils.py +9 -3
  37. prefect/utilities/engine.py +34 -1
  38. prefect/utilities/importtools.py +1 -1
  39. prefect/utilities/timeout.py +20 -5
  40. prefect/workers/base.py +98 -208
  41. prefect/workers/block.py +6 -0
  42. prefect/workers/cloud.py +6 -0
  43. prefect/workers/process.py +262 -4
  44. prefect/workers/server.py +27 -9
  45. {prefect_client-3.0.0rc9.dist-info → prefect_client-3.0.0rc11.dist-info}/METADATA +4 -4
  46. {prefect_client-3.0.0rc9.dist-info → prefect_client-3.0.0rc11.dist-info}/RECORD +49 -44
  47. {prefect_client-3.0.0rc9.dist-info → prefect_client-3.0.0rc11.dist-info}/LICENSE +0 -0
  48. {prefect_client-3.0.0rc9.dist-info → prefect_client-3.0.0rc11.dist-info}/WHEEL +0 -0
  49. {prefect_client-3.0.0rc9.dist-info → prefect_client-3.0.0rc11.dist-info}/top_level.txt +0 -0
@@ -38,6 +38,7 @@ __all__ = [
38
38
  "parameters",
39
39
  "parent_flow_run_id",
40
40
  "parent_deployment_id",
41
+ "root_flow_run_id",
41
42
  "run_count",
42
43
  "api_url",
43
44
  "ui_url",
@@ -237,11 +238,12 @@ def get_parent_flow_run_id() -> Optional[str]:
237
238
  parent_task_run = from_sync.call_soon_in_loop_thread(
238
239
  create_call(_get_task_run, parent_task_run_id)
239
240
  ).result()
240
- return parent_task_run.flow_run_id
241
+ return str(parent_task_run.flow_run_id) if parent_task_run.flow_run_id else None
242
+
241
243
  return None
242
244
 
243
245
 
244
- def get_parent_deployment_id() -> Dict[str, Any]:
246
+ def get_parent_deployment_id() -> Optional[str]:
245
247
  parent_flow_run_id = get_parent_flow_run_id()
246
248
  if parent_flow_run_id is None:
247
249
  return None
@@ -249,7 +251,39 @@ def get_parent_deployment_id() -> Dict[str, Any]:
249
251
  parent_flow_run = from_sync.call_soon_in_loop_thread(
250
252
  create_call(_get_flow_run, parent_flow_run_id)
251
253
  ).result()
252
- return parent_flow_run.deployment_id if parent_flow_run else None
254
+
255
+ if parent_flow_run:
256
+ return (
257
+ str(parent_flow_run.deployment_id)
258
+ if parent_flow_run.deployment_id
259
+ else None
260
+ )
261
+
262
+ return None
263
+
264
+
265
+ def get_root_flow_run_id() -> str:
266
+ run_id = get_id()
267
+ parent_flow_run_id = get_parent_flow_run_id()
268
+ if parent_flow_run_id is None:
269
+ return run_id
270
+
271
+ def _get_root_flow_run_id(flow_run_id):
272
+ flow_run = from_sync.call_soon_in_loop_thread(
273
+ create_call(_get_flow_run, flow_run_id)
274
+ ).result()
275
+
276
+ if flow_run.parent_task_run_id is None:
277
+ return str(flow_run_id)
278
+ else:
279
+ parent_task_run = from_sync.call_soon_in_loop_thread(
280
+ create_call(_get_task_run, flow_run.parent_task_run_id)
281
+ ).result()
282
+ return _get_root_flow_run_id(parent_task_run.flow_run_id)
283
+
284
+ root_flow_run_id = _get_root_flow_run_id(parent_flow_run_id)
285
+
286
+ return root_flow_run_id
253
287
 
254
288
 
255
289
  def get_flow_run_api_url() -> Optional[str]:
@@ -275,6 +309,7 @@ FIELDS = {
275
309
  "parameters": get_parameters,
276
310
  "parent_flow_run_id": get_parent_flow_run_id,
277
311
  "parent_deployment_id": get_parent_deployment_id,
312
+ "root_flow_run_id": get_root_flow_run_id,
278
313
  "run_count": get_run_count,
279
314
  "api_url": get_flow_run_api_url,
280
315
  "ui_url": get_flow_run_ui_url,
prefect/serializers.py CHANGED
@@ -13,7 +13,7 @@ bytes to an object respectively.
13
13
 
14
14
  import abc
15
15
  import base64
16
- from typing import Any, Dict, Generic, Optional, Type, TypeVar
16
+ from typing import Any, Dict, Generic, Optional, Type
17
17
 
18
18
  from pydantic import (
19
19
  BaseModel,
@@ -23,7 +23,7 @@ from pydantic import (
23
23
  ValidationError,
24
24
  field_validator,
25
25
  )
26
- from typing_extensions import Literal, Self
26
+ from typing_extensions import Literal, Self, TypeVar
27
27
 
28
28
  from prefect._internal.schemas.validators import (
29
29
  cast_type_names_to_serializers,
@@ -36,7 +36,7 @@ from prefect.utilities.dispatch import get_dispatch_key, lookup_type, register_b
36
36
  from prefect.utilities.importtools import from_qualified_name, to_qualified_name
37
37
  from prefect.utilities.pydantic import custom_pydantic_encoder
38
38
 
39
- D = TypeVar("D")
39
+ D = TypeVar("D", default=Any)
40
40
 
41
41
 
42
42
  def prefect_json_object_encoder(obj: Any) -> Any:
prefect/settings.py CHANGED
@@ -42,7 +42,7 @@ dependent on the value of other settings or perform other dynamic effects.
42
42
 
43
43
  import logging
44
44
  import os
45
- import socket
45
+ import re
46
46
  import string
47
47
  import warnings
48
48
  from contextlib import contextmanager
@@ -85,7 +85,6 @@ from prefect._internal.schemas.validators import validate_settings
85
85
  from prefect.exceptions import MissingProfileError
86
86
  from prefect.utilities.names import OBFUSCATED_PREFIX, obfuscate
87
87
  from prefect.utilities.pydantic import add_cloudpickle_reduction
88
- from prefect.utilities.slugify import slugify
89
88
 
90
89
  T = TypeVar("T")
91
90
 
@@ -404,18 +403,6 @@ def warn_on_misconfigured_api_url(values):
404
403
  return values
405
404
 
406
405
 
407
- def default_result_storage_block_name(
408
- settings: Optional["Settings"] = None, value: Optional[str] = None
409
- ):
410
- """
411
- `value_callback` for `PREFECT_DEFAULT_RESULT_STORAGE_BLOCK` that sets the default
412
- value to the hostname of the machine.
413
- """
414
- if value is None:
415
- return f"local-file-system/{slugify(socket.gethostname())}-storage"
416
- return value
417
-
418
-
419
406
  def default_database_connection_url(settings, value):
420
407
  templater = template_with_settings(PREFECT_HOME, PREFECT_API_DATABASE_PASSWORD)
421
408
 
@@ -474,10 +461,8 @@ def default_cloud_ui_url(settings, value):
474
461
  # Otherwise, infer a value from the API URL
475
462
  ui_url = api_url = PREFECT_CLOUD_API_URL.value_from(settings)
476
463
 
477
- if api_url.startswith("https://api.prefect.cloud"):
478
- ui_url = ui_url.replace(
479
- "https://api.prefect.cloud", "https://app.prefect.cloud", 1
480
- )
464
+ if re.match(r"^https://api[\.\w]*.prefect.[^\.]+/", api_url):
465
+ ui_url = ui_url.replace("https://api", "https://app", 1)
481
466
 
482
467
  if ui_url.endswith("/api"):
483
468
  ui_url = ui_url[:-4]
@@ -1175,6 +1160,11 @@ polled."""
1175
1160
  PREFECT_API_LOG_RETRYABLE_ERRORS = Setting(bool, default=False)
1176
1161
  """If `True`, log retryable errors in the API and it's services."""
1177
1162
 
1163
+ PREFECT_API_SERVICES_TASK_RUN_RECORDER_ENABLED = Setting(bool, default=True)
1164
+ """
1165
+ Whether or not to start the task run recorder service in the server application.
1166
+ """
1167
+
1178
1168
 
1179
1169
  PREFECT_API_DEFAULT_LIMIT = Setting(
1180
1170
  int,
@@ -1323,27 +1313,14 @@ PREFECT_API_MAX_FLOW_RUN_GRAPH_ARTIFACTS = Setting(int, default=10000)
1323
1313
  The maximum number of artifacts to show on a flow run graph on the v2 API
1324
1314
  """
1325
1315
 
1326
- PREFECT_EXPERIMENTAL_ENABLE_WORKERS = Setting(bool, default=True)
1327
- """
1328
- Whether or not to enable experimental Prefect workers.
1329
- """
1330
-
1331
- PREFECT_EXPERIMENTAL_WARN_WORKERS = Setting(bool, default=False)
1332
- """
1333
- Whether or not to warn when experimental Prefect workers are used.
1334
- """
1335
-
1336
- PREFECT_EXPERIMENTAL_ENABLE_ENHANCED_CANCELLATION = Setting(bool, default=True)
1337
- """
1338
- Whether or not to enable experimental enhanced flow run cancellation.
1339
- """
1340
1316
 
1341
- PREFECT_EXPERIMENTAL_WARN_ENHANCED_CANCELLATION = Setting(bool, default=False)
1317
+ PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION = Setting(
1318
+ bool, default=False
1319
+ )
1342
1320
  """
1343
- Whether or not to warn when experimental enhanced flow run cancellation is used.
1321
+ Whether or not to enable experimental client side task run orchestration.
1344
1322
  """
1345
1323
 
1346
-
1347
1324
  # Prefect Events feature flags
1348
1325
 
1349
1326
  PREFECT_RUNNER_PROCESS_LIMIT = Setting(int, default=5)
@@ -1423,10 +1400,7 @@ PREFECT_API_SERVICES_TASK_SCHEDULING_ENABLED = Setting(bool, default=True)
1423
1400
  Whether or not to start the task scheduling service in the server application.
1424
1401
  """
1425
1402
 
1426
- PREFECT_TASK_SCHEDULING_DEFAULT_STORAGE_BLOCK = Setting(
1427
- str,
1428
- default="local-file-system/prefect-task-scheduling",
1429
- )
1403
+ PREFECT_TASK_SCHEDULING_DEFAULT_STORAGE_BLOCK = Setting(Optional[str], default=None)
1430
1404
  """The `block-type/block-document` slug of a block to use as the default storage
1431
1405
  for autonomous tasks."""
1432
1406
 
@@ -1464,11 +1438,6 @@ a task worker should move a task from PENDING to RUNNING very quickly, so runs s
1464
1438
  PENDING for a while is a sign that the task worker may have crashed.
1465
1439
  """
1466
1440
 
1467
- PREFECT_EXPERIMENTAL_ENABLE_EXTRA_RUNNER_ENDPOINTS = Setting(bool, default=False)
1468
- """
1469
- Whether or not to enable experimental worker webserver endpoints.
1470
- """
1471
-
1472
1441
  PREFECT_EXPERIMENTAL_DISABLE_SYNC_COMPAT = Setting(bool, default=False)
1473
1442
  """
1474
1443
  Whether or not to disable the sync_compatible decorator utility.
@@ -1479,7 +1448,8 @@ PREFECT_EXPERIMENTAL_ENABLE_SCHEDULE_CONCURRENCY = Setting(bool, default=False)
1479
1448
  # Defaults -----------------------------------------------------------------------------
1480
1449
 
1481
1450
  PREFECT_DEFAULT_RESULT_STORAGE_BLOCK = Setting(
1482
- Optional[str], default=None, value_callback=default_result_storage_block_name
1451
+ Optional[str],
1452
+ default=None,
1483
1453
  )
1484
1454
  """The `block-type/block-document` slug of a block to use as the default result storage."""
1485
1455
 
prefect/task_engine.py CHANGED
@@ -1,6 +1,8 @@
1
1
  import inspect
2
2
  import logging
3
+ import threading
3
4
  import time
5
+ from asyncio import CancelledError
4
6
  from contextlib import ExitStack, contextmanager
5
7
  from dataclasses import dataclass, field
6
8
  from textwrap import dedent
@@ -17,6 +19,7 @@ from typing import (
17
19
  Optional,
18
20
  Sequence,
19
21
  Set,
22
+ Type,
20
23
  TypeVar,
21
24
  Union,
22
25
  )
@@ -36,17 +39,18 @@ from prefect.context import (
36
39
  TaskRunContext,
37
40
  hydrated_context,
38
41
  )
39
- from prefect.events.schemas.events import Event
42
+ from prefect.events.schemas.events import Event as PrefectEvent
40
43
  from prefect.exceptions import (
41
44
  Abort,
42
45
  Pause,
43
46
  PrefectException,
47
+ TerminationSignal,
44
48
  UpstreamTaskError,
45
49
  )
46
50
  from prefect.futures import PrefectFuture
47
51
  from prefect.logging.loggers import get_logger, patch_print, task_run_logger
48
52
  from prefect.records.result_store import ResultFactoryStore
49
- from prefect.results import ResultFactory, _format_user_supplied_storage_key
53
+ from prefect.results import BaseResult, ResultFactory, _format_user_supplied_storage_key
50
54
  from prefect.settings import (
51
55
  PREFECT_DEBUG_MODE,
52
56
  PREFECT_TASKS_REFRESH_CACHE,
@@ -63,6 +67,7 @@ from prefect.states import (
63
67
  return_value_to_state,
64
68
  )
65
69
  from prefect.transactions import Transaction, transaction
70
+ from prefect.utilities.annotations import NotSet
66
71
  from prefect.utilities.asyncutils import run_coro_as_sync
67
72
  from prefect.utilities.callables import call_with_parameters, parameters_to_args_kwargs
68
73
  from prefect.utilities.collections import visit_collection
@@ -80,6 +85,10 @@ P = ParamSpec("P")
80
85
  R = TypeVar("R")
81
86
 
82
87
 
88
+ class TaskRunTimeoutError(TimeoutError):
89
+ """Raised when a task run exceeds its timeout."""
90
+
91
+
83
92
  @dataclass
84
93
  class TaskRunEngine(Generic[P, R]):
85
94
  task: Union[Task[P, R], Task[P, Coroutine[Any, Any, R]]]
@@ -89,11 +98,15 @@ class TaskRunEngine(Generic[P, R]):
89
98
  retries: int = 0
90
99
  wait_for: Optional[Iterable[PrefectFuture]] = None
91
100
  context: Optional[Dict[str, Any]] = None
101
+ # holds the return value from the user code
102
+ _return_value: Union[R, Type[NotSet]] = NotSet
103
+ # holds the exception raised by the user code, if any
104
+ _raised: Union[Exception, Type[NotSet]] = NotSet
92
105
  _initial_run_context: Optional[TaskRunContext] = None
93
106
  _is_started: bool = False
94
107
  _client: Optional[SyncPrefectClient] = None
95
108
  _task_name_set: bool = False
96
- _last_event: Optional[Event] = None
109
+ _last_event: Optional[PrefectEvent] = None
97
110
 
98
111
  def __post_init__(self):
99
112
  if self.parameters is None:
@@ -136,7 +149,16 @@ class TaskRunEngine(Generic[P, R]):
136
149
  )
137
150
  return False
138
151
 
139
- def call_hooks(self, state: State = None) -> Iterable[Callable]:
152
+ def is_cancelled(self) -> bool:
153
+ if (
154
+ self.context
155
+ and "cancel_event" in self.context
156
+ and isinstance(self.context["cancel_event"], threading.Event)
157
+ ):
158
+ return self.context["cancel_event"].is_set()
159
+ return False
160
+
161
+ def call_hooks(self, state: Optional[State] = None):
140
162
  if state is None:
141
163
  state = self.state
142
164
  task = self.task
@@ -171,7 +193,7 @@ class TaskRunEngine(Generic[P, R]):
171
193
  else:
172
194
  self.logger.info(f"Hook {hook_name!r} finished running successfully")
173
195
 
174
- def compute_transaction_key(self) -> str:
196
+ def compute_transaction_key(self) -> Optional[str]:
175
197
  key = None
176
198
  if self.task.cache_policy:
177
199
  flow_run_context = FlowRunContext.get()
@@ -304,12 +326,24 @@ class TaskRunEngine(Generic[P, R]):
304
326
  return new_state
305
327
 
306
328
  def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
307
- _result = self.state.result(raise_on_failure=raise_on_failure, fetch=True)
308
- # state.result is a `sync_compatible` function that may or may not return an awaitable
309
- # depending on whether the parent frame is sync or not
310
- if inspect.isawaitable(_result):
311
- _result = run_coro_as_sync(_result)
312
- return _result
329
+ if self._return_value is not NotSet:
330
+ # if the return value is a BaseResult, we need to fetch it
331
+ if isinstance(self._return_value, BaseResult):
332
+ _result = self._return_value.get()
333
+ if inspect.isawaitable(_result):
334
+ _result = run_coro_as_sync(_result)
335
+ return _result
336
+
337
+ # otherwise, return the value as is
338
+ return self._return_value
339
+
340
+ if self._raised is not NotSet:
341
+ # if the task raised an exception, raise it
342
+ if raise_on_failure:
343
+ raise self._raised
344
+
345
+ # otherwise, return the exception
346
+ return self._raised
313
347
 
314
348
  def handle_success(self, result: R, transaction: Transaction) -> R:
315
349
  result_factory = getattr(TaskRunContext.get(), "result_factory", None)
@@ -339,6 +373,7 @@ class TaskRunEngine(Generic[P, R]):
339
373
  if transaction.is_committed():
340
374
  terminal_state.name = "Cached"
341
375
  self.set_state(terminal_state)
376
+ self._return_value = result
342
377
  return result
343
378
 
344
379
  def handle_retry(self, exc: Exception) -> bool:
@@ -365,9 +400,11 @@ class TaskRunEngine(Generic[P, R]):
365
400
  new_state = Retrying()
366
401
 
367
402
  self.logger.info(
368
- f"Task run failed with exception {exc!r} - "
369
- f"Retry {self.retries + 1}/{self.task.retries} will start "
370
- f"{str(delay) + ' second(s) from now' if delay else 'immediately'}"
403
+ "Task run failed with exception: %r - " "Retry %s/%s will start %s",
404
+ exc,
405
+ self.retries + 1,
406
+ self.task.retries,
407
+ str(delay) + " second(s) from now" if delay else "immediately",
371
408
  )
372
409
 
373
410
  self.set_state(new_state, force=True)
@@ -375,7 +412,9 @@ class TaskRunEngine(Generic[P, R]):
375
412
  return True
376
413
  elif self.retries >= self.task.retries:
377
414
  self.logger.error(
378
- f"Task run failed with exception {exc!r} - Retries are exhausted"
415
+ "Task run failed with exception: %r - Retries are exhausted",
416
+ exc,
417
+ exc_info=True,
379
418
  )
380
419
  return False
381
420
 
@@ -394,12 +433,14 @@ class TaskRunEngine(Generic[P, R]):
394
433
  )
395
434
  )
396
435
  self.set_state(state)
436
+ self._raised = exc
397
437
 
398
438
  def handle_timeout(self, exc: TimeoutError) -> None:
399
439
  if not self.handle_retry(exc):
400
- message = (
401
- f"Task run exceeded timeout of {self.task.timeout_seconds} seconds"
402
- )
440
+ if isinstance(exc, TaskRunTimeoutError):
441
+ message = f"Task run exceeded timeout of {self.task.timeout_seconds} second(s)"
442
+ else:
443
+ message = f"Task run failed due to timeout: {exc!r}"
403
444
  self.logger.error(message)
404
445
  state = Failed(
405
446
  data=exc,
@@ -407,12 +448,14 @@ class TaskRunEngine(Generic[P, R]):
407
448
  name="TimedOut",
408
449
  )
409
450
  self.set_state(state)
451
+ self._raised = exc
410
452
 
411
453
  def handle_crash(self, exc: BaseException) -> None:
412
454
  state = run_coro_as_sync(exception_to_crashed_state(exc))
413
455
  self.logger.error(f"Crash detected! {state.message}")
414
456
  self.logger.debug("Crash details:", exc_info=exc)
415
457
  self.set_state(state, force=True)
458
+ self._raised = exc
416
459
 
417
460
  @contextmanager
418
461
  def setup_run_context(self, client: Optional[SyncPrefectClient] = None):
@@ -498,6 +541,11 @@ class TaskRunEngine(Generic[P, R]):
498
541
  )
499
542
  yield self
500
543
 
544
+ except TerminationSignal as exc:
545
+ # TerminationSignals are caught and handled as crashes
546
+ self.handle_crash(exc)
547
+ raise exc
548
+
501
549
  except Exception:
502
550
  # regular exceptions are caught and re-raised to the user
503
551
  raise
@@ -539,8 +587,8 @@ class TaskRunEngine(Generic[P, R]):
539
587
 
540
588
  @flow
541
589
  def example_flow():
542
- say_hello.submit(name="Marvin)
543
- say_hello.wait()
590
+ future = say_hello.submit(name="Marvin)
591
+ future.wait()
544
592
 
545
593
  example_flow()
546
594
  """
@@ -612,10 +660,16 @@ class TaskRunEngine(Generic[P, R]):
612
660
  # reenter the run context to ensure it is up to date for every run
613
661
  with self.setup_run_context():
614
662
  try:
615
- with timeout_context(seconds=self.task.timeout_seconds):
663
+ with timeout_context(
664
+ seconds=self.task.timeout_seconds,
665
+ timeout_exc_type=TaskRunTimeoutError,
666
+ ):
616
667
  self.logger.debug(
617
668
  f"Executing task {self.task.name!r} for task run {self.task_run.name!r}..."
618
669
  )
670
+ if self.is_cancelled():
671
+ raise CancelledError("Task run cancelled by the task runner")
672
+
619
673
  yield self
620
674
  except TimeoutError as exc:
621
675
  self.handle_timeout(exc)
@@ -638,6 +692,7 @@ class TaskRunEngine(Generic[P, R]):
638
692
  else:
639
693
  result = await call_with_parameters(self.task.fn, parameters)
640
694
  self.handle_success(result, transaction=transaction)
695
+ return result
641
696
 
642
697
  return _call_task_fn()
643
698
  else:
@@ -646,6 +701,7 @@ class TaskRunEngine(Generic[P, R]):
646
701
  else:
647
702
  result = call_with_parameters(self.task.fn, parameters)
648
703
  self.handle_success(result, transaction=transaction)
704
+ return result
649
705
 
650
706
 
651
707
  def run_task_sync(
prefect/task_runners.py CHANGED
@@ -1,6 +1,7 @@
1
1
  import abc
2
2
  import asyncio
3
3
  import sys
4
+ import threading
4
5
  import uuid
5
6
  from concurrent.futures import ThreadPoolExecutor
6
7
  from contextvars import copy_context
@@ -41,8 +42,8 @@ if TYPE_CHECKING:
41
42
 
42
43
  P = ParamSpec("P")
43
44
  T = TypeVar("T")
44
- F = TypeVar("F", bound=PrefectFuture)
45
45
  R = TypeVar("R")
46
+ F = TypeVar("F", bound=PrefectFuture, default=PrefectConcurrentFuture)
46
47
 
47
48
 
48
49
  class TaskRunner(abc.ABC, Generic[F]):
@@ -220,6 +221,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
220
221
  super().__init__()
221
222
  self._executor: Optional[ThreadPoolExecutor] = None
222
223
  self._max_workers = sys.maxsize if max_workers is None else max_workers
224
+ self._cancel_events: Dict[uuid.UUID, threading.Event] = {}
223
225
 
224
226
  def duplicate(self) -> "ThreadPoolTaskRunner":
225
227
  return type(self)(max_workers=self._max_workers)
@@ -270,6 +272,8 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
270
272
  from prefect.task_engine import run_task_async, run_task_sync
271
273
 
272
274
  task_run_id = uuid.uuid4()
275
+ cancel_event = threading.Event()
276
+ self._cancel_events[task_run_id] = cancel_event
273
277
  context = copy_context()
274
278
 
275
279
  flow_run_ctx = FlowRunContext.get()
@@ -280,31 +284,29 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
280
284
  else:
281
285
  self.logger.info(f"Submitting task {task.name} to thread pool executor...")
282
286
 
287
+ submit_kwargs = dict(
288
+ task=task,
289
+ task_run_id=task_run_id,
290
+ parameters=parameters,
291
+ wait_for=wait_for,
292
+ return_type="state",
293
+ dependencies=dependencies,
294
+ context=dict(cancel_event=cancel_event),
295
+ )
296
+
283
297
  if task.isasync:
284
298
  # TODO: Explore possibly using a long-lived thread with an event loop
285
299
  # for better performance
286
300
  future = self._executor.submit(
287
301
  context.run,
288
302
  asyncio.run,
289
- run_task_async(
290
- task=task,
291
- task_run_id=task_run_id,
292
- parameters=parameters,
293
- wait_for=wait_for,
294
- return_type="state",
295
- dependencies=dependencies,
296
- ),
303
+ run_task_async(**submit_kwargs),
297
304
  )
298
305
  else:
299
306
  future = self._executor.submit(
300
307
  context.run,
301
308
  run_task_sync,
302
- task=task,
303
- task_run_id=task_run_id,
304
- parameters=parameters,
305
- wait_for=wait_for,
306
- return_type="state",
307
- dependencies=dependencies,
309
+ **submit_kwargs,
308
310
  )
309
311
  prefect_future = PrefectConcurrentFuture(
310
312
  task_run_id=task_run_id, wrapped_future=future
@@ -337,14 +339,24 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
337
339
  ):
338
340
  return super().map(task, parameters, wait_for)
339
341
 
342
+ def cancel_all(self):
343
+ for event in self._cancel_events.values():
344
+ event.set()
345
+ self.logger.debug("Set cancel event")
346
+
347
+ if self._executor is not None:
348
+ self._executor.shutdown(cancel_futures=True)
349
+ self._executor = None
350
+
340
351
  def __enter__(self):
341
352
  super().__enter__()
342
353
  self._executor = ThreadPoolExecutor(max_workers=self._max_workers)
343
354
  return self
344
355
 
345
356
  def __exit__(self, exc_type, exc_value, traceback):
357
+ self.cancel_all()
346
358
  if self._executor is not None:
347
- self._executor.shutdown()
359
+ self._executor.shutdown(cancel_futures=True)
348
360
  self._executor = None
349
361
  super().__exit__(exc_type, exc_value, traceback)
350
362
 
prefect/task_worker.py CHANGED
@@ -37,6 +37,7 @@ from prefect.utilities.annotations import NotSet
37
37
  from prefect.utilities.asyncutils import asyncnullcontext, sync_compatible
38
38
  from prefect.utilities.engine import emit_task_run_state_change_event, propose_state
39
39
  from prefect.utilities.processutils import _register_signal
40
+ from prefect.utilities.urls import url_for
40
41
 
41
42
  logger = get_logger("task_worker")
42
43
 
@@ -288,10 +289,6 @@ class TaskWorker:
288
289
  await self._client._client.delete(f"/task_runs/{task_run.id}")
289
290
  return
290
291
 
291
- logger.debug(
292
- f"Submitting run {task_run.name!r} of task {task.name!r} to engine"
293
- )
294
-
295
292
  try:
296
293
  new_state = Pending()
297
294
  new_state.state_details.deferred = True
@@ -326,6 +323,11 @@ class TaskWorker:
326
323
  validated_state=state,
327
324
  )
328
325
 
326
+ if task_run_url := url_for(task_run):
327
+ logger.info(
328
+ f"Submitting task run {task_run.name!r} to engine. View in the UI: {task_run_url}"
329
+ )
330
+
329
331
  if task.isasync:
330
332
  await run_task_async(
331
333
  task=task,
prefect/tasks.py CHANGED
@@ -50,7 +50,6 @@ from prefect.futures import PrefectDistributedFuture, PrefectFuture, PrefectFutu
50
50
  from prefect.logging.loggers import get_logger
51
51
  from prefect.results import ResultFactory, ResultSerializer, ResultStorage
52
52
  from prefect.settings import (
53
- PREFECT_RESULTS_PERSIST_BY_DEFAULT,
54
53
  PREFECT_TASK_DEFAULT_RETRIES,
55
54
  PREFECT_TASK_DEFAULT_RETRY_DELAY_SECONDS,
56
55
  )
@@ -67,6 +66,7 @@ from prefect.utilities.callables import (
67
66
  )
68
67
  from prefect.utilities.hashing import hash_objects
69
68
  from prefect.utilities.importtools import to_qualified_name
69
+ from prefect.utilities.urls import url_for
70
70
 
71
71
  if TYPE_CHECKING:
72
72
  from prefect.client.orchestration import PrefectClient
@@ -387,9 +387,20 @@ class Task(Generic[P, R]):
387
387
  self.cache_expiration = cache_expiration
388
388
  self.refresh_cache = refresh_cache
389
389
 
390
+ # result persistence settings
390
391
  if persist_result is None:
391
- persist_result = PREFECT_RESULTS_PERSIST_BY_DEFAULT.value()
392
- if not persist_result:
392
+ if any(
393
+ [
394
+ cache_policy and cache_policy != NONE and cache_policy != NotSet,
395
+ cache_key_fn is not None,
396
+ result_storage_key is not None,
397
+ result_storage is not None,
398
+ result_serializer is not None,
399
+ ]
400
+ ):
401
+ persist_result = True
402
+
403
+ if persist_result is False:
393
404
  self.cache_policy = None if cache_policy is None else NONE
394
405
  if cache_policy and cache_policy is not NotSet and cache_policy != NONE:
395
406
  logger.warning(
@@ -428,6 +439,14 @@ class Task(Generic[P, R]):
428
439
 
429
440
  self.retry_jitter_factor = retry_jitter_factor
430
441
  self.persist_result = persist_result
442
+
443
+ if result_storage and not isinstance(result_storage, str):
444
+ if getattr(result_storage, "_block_document_id", None) is None:
445
+ raise TypeError(
446
+ "Result storage configuration must be persisted server-side."
447
+ " Please call `.save()` on your block before passing it in."
448
+ )
449
+
431
450
  self.result_storage = result_storage
432
451
  self.result_serializer = result_serializer
433
452
  self.result_storage_key = result_storage_key
@@ -1282,14 +1301,20 @@ class Task(Generic[P, R]):
1282
1301
  # Convert the call args/kwargs to a parameter dict
1283
1302
  parameters = get_call_parameters(self.fn, args, kwargs)
1284
1303
 
1285
- task_run = run_coro_as_sync(
1304
+ task_run: TaskRun = run_coro_as_sync(
1286
1305
  self.create_run(
1287
1306
  parameters=parameters,
1288
1307
  deferred=True,
1289
1308
  wait_for=wait_for,
1290
1309
  extra_task_inputs=dependencies,
1291
1310
  )
1292
- )
1311
+ ) # type: ignore
1312
+
1313
+ if task_run_url := url_for(task_run):
1314
+ logger.info(
1315
+ f"Created task run {task_run.name!r}. View it in the UI at {task_run_url!r}"
1316
+ )
1317
+
1293
1318
  return PrefectDistributedFuture(task_run_id=task_run.id)
1294
1319
 
1295
1320
  def delay(self, *args: P.args, **kwargs: P.kwargs) -> PrefectDistributedFuture: