prefect-client 3.0.0rc20__py3-none-any.whl → 3.0.2__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 (57) hide show
  1. prefect/_internal/compatibility/deprecated.py +1 -1
  2. prefect/_internal/compatibility/migration.py +1 -1
  3. prefect/artifacts.py +1 -1
  4. prefect/blocks/core.py +3 -4
  5. prefect/blocks/notifications.py +31 -10
  6. prefect/blocks/system.py +4 -4
  7. prefect/blocks/webhook.py +11 -1
  8. prefect/client/cloud.py +2 -1
  9. prefect/client/orchestration.py +93 -21
  10. prefect/client/schemas/actions.py +2 -2
  11. prefect/client/schemas/objects.py +24 -6
  12. prefect/client/types/flexible_schedule_list.py +1 -1
  13. prefect/concurrency/asyncio.py +45 -6
  14. prefect/concurrency/services.py +1 -1
  15. prefect/concurrency/sync.py +21 -27
  16. prefect/concurrency/v1/asyncio.py +3 -0
  17. prefect/concurrency/v1/sync.py +4 -5
  18. prefect/context.py +11 -9
  19. prefect/deployments/runner.py +4 -3
  20. prefect/events/actions.py +6 -0
  21. prefect/exceptions.py +6 -0
  22. prefect/filesystems.py +5 -3
  23. prefect/flow_engine.py +22 -11
  24. prefect/flows.py +0 -2
  25. prefect/futures.py +2 -1
  26. prefect/locking/__init__.py +0 -0
  27. prefect/locking/filesystem.py +243 -0
  28. prefect/locking/memory.py +213 -0
  29. prefect/locking/protocol.py +122 -0
  30. prefect/logging/handlers.py +0 -2
  31. prefect/logging/loggers.py +0 -18
  32. prefect/logging/logging.yml +1 -0
  33. prefect/main.py +19 -5
  34. prefect/records/base.py +12 -0
  35. prefect/records/filesystem.py +10 -4
  36. prefect/records/memory.py +6 -0
  37. prefect/records/result_store.py +18 -6
  38. prefect/results.py +702 -205
  39. prefect/runner/runner.py +74 -5
  40. prefect/settings.py +11 -4
  41. prefect/states.py +40 -23
  42. prefect/task_engine.py +39 -37
  43. prefect/task_worker.py +6 -4
  44. prefect/tasks.py +24 -6
  45. prefect/transactions.py +116 -54
  46. prefect/utilities/callables.py +1 -3
  47. prefect/utilities/engine.py +16 -8
  48. prefect/utilities/importtools.py +1 -0
  49. prefect/utilities/urls.py +70 -12
  50. prefect/variables.py +34 -24
  51. prefect/workers/base.py +14 -6
  52. prefect/workers/process.py +1 -3
  53. {prefect_client-3.0.0rc20.dist-info → prefect_client-3.0.2.dist-info}/METADATA +2 -2
  54. {prefect_client-3.0.0rc20.dist-info → prefect_client-3.0.2.dist-info}/RECORD +57 -53
  55. {prefect_client-3.0.0rc20.dist-info → prefect_client-3.0.2.dist-info}/LICENSE +0 -0
  56. {prefect_client-3.0.0rc20.dist-info → prefect_client-3.0.2.dist-info}/WHEEL +0 -0
  57. {prefect_client-3.0.0rc20.dist-info → prefect_client-3.0.2.dist-info}/top_level.txt +0 -0
prefect/results.py CHANGED
@@ -1,11 +1,14 @@
1
1
  import abc
2
2
  import inspect
3
+ import os
4
+ import socket
5
+ import threading
3
6
  import uuid
4
7
  from functools import partial
8
+ from pathlib import Path
5
9
  from typing import (
6
10
  TYPE_CHECKING,
7
11
  Any,
8
- Awaitable,
9
12
  Callable,
10
13
  Dict,
11
14
  Generic,
@@ -17,6 +20,8 @@ from typing import (
17
20
  )
18
21
  from uuid import UUID
19
22
 
23
+ import pendulum
24
+ from cachetools import LRUCache
20
25
  from pydantic import (
21
26
  BaseModel,
22
27
  ConfigDict,
@@ -31,13 +36,20 @@ from pydantic_extra_types.pendulum_dt import DateTime
31
36
  from typing_extensions import ParamSpec, Self
32
37
 
33
38
  import prefect
39
+ from prefect._internal.compatibility import deprecated
40
+ from prefect._internal.compatibility.deprecated import deprecated_field
34
41
  from prefect.blocks.core import Block
35
42
  from prefect.client.utilities import inject_client
36
- from prefect.exceptions import SerializationError
43
+ from prefect.exceptions import (
44
+ ConfigurationError,
45
+ MissingContextError,
46
+ SerializationError,
47
+ )
37
48
  from prefect.filesystems import (
38
49
  LocalFileSystem,
39
50
  WritableFileSystem,
40
51
  )
52
+ from prefect.locking.protocol import LockManager
41
53
  from prefect.logging import get_logger
42
54
  from prefect.serializers import PickleSerializer, Serializer
43
55
  from prefect.settings import (
@@ -54,6 +66,7 @@ from prefect.utilities.pydantic import get_dispatch_key, lookup_type, register_b
54
66
  if TYPE_CHECKING:
55
67
  from prefect import Flow, Task
56
68
  from prefect.client.orchestration import PrefectClient
69
+ from prefect.transactions import IsolationLevel
57
70
 
58
71
 
59
72
  ResultStorage = Union[WritableFileSystem, str]
@@ -73,18 +86,69 @@ _default_storages: Dict[Tuple[str, str], WritableFileSystem] = {}
73
86
 
74
87
 
75
88
  @sync_compatible
76
- async def get_default_result_storage() -> ResultStorage:
89
+ async def get_default_result_storage() -> WritableFileSystem:
77
90
  """
78
91
  Generate a default file system for result storage.
79
92
  """
80
93
  default_block = PREFECT_DEFAULT_RESULT_STORAGE_BLOCK.value()
81
94
 
82
95
  if default_block is not None:
83
- return await Block.load(default_block)
96
+ return await resolve_result_storage(default_block)
84
97
 
85
98
  # otherwise, use the local file system
86
99
  basepath = PREFECT_LOCAL_STORAGE_PATH.value()
87
- return LocalFileSystem(basepath=basepath)
100
+ return LocalFileSystem(basepath=str(basepath))
101
+
102
+
103
+ @sync_compatible
104
+ async def resolve_result_storage(
105
+ result_storage: Union[ResultStorage, UUID],
106
+ ) -> WritableFileSystem:
107
+ """
108
+ Resolve one of the valid `ResultStorage` input types into a saved block
109
+ document id and an instance of the block.
110
+ """
111
+ from prefect.client.orchestration import get_client
112
+
113
+ client = get_client()
114
+ if isinstance(result_storage, Block):
115
+ storage_block = result_storage
116
+
117
+ if storage_block._block_document_id is not None:
118
+ # Avoid saving the block if it already has an identifier assigned
119
+ storage_block_id = storage_block._block_document_id
120
+ else:
121
+ storage_block_id = None
122
+ elif isinstance(result_storage, str):
123
+ storage_block = await Block.load(result_storage, client=client)
124
+ storage_block_id = storage_block._block_document_id
125
+ assert storage_block_id is not None, "Loaded storage blocks must have ids"
126
+ elif isinstance(result_storage, UUID):
127
+ block_document = await client.read_block_document(result_storage)
128
+ storage_block = Block._from_block_document(block_document)
129
+ else:
130
+ raise TypeError(
131
+ "Result storage must be one of the following types: 'UUID', 'Block', "
132
+ f"'str'. Got unsupported type {type(result_storage).__name__!r}."
133
+ )
134
+
135
+ return storage_block
136
+
137
+
138
+ def resolve_serializer(serializer: ResultSerializer) -> Serializer:
139
+ """
140
+ Resolve one of the valid `ResultSerializer` input types into a serializer
141
+ instance.
142
+ """
143
+ if isinstance(serializer, Serializer):
144
+ return serializer
145
+ elif isinstance(serializer, str):
146
+ return Serializer(type=serializer)
147
+ else:
148
+ raise TypeError(
149
+ "Result serializer must be one of the following types: 'Serializer', "
150
+ f"'str'. Got unsupported type {type(serializer).__name__!r}."
151
+ )
88
152
 
89
153
 
90
154
  async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
@@ -101,11 +165,11 @@ async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
101
165
  return LocalFileSystem(basepath=basepath)
102
166
 
103
167
 
104
- def get_default_result_serializer() -> ResultSerializer:
168
+ def get_default_result_serializer() -> Serializer:
105
169
  """
106
170
  Generate a default file system for result storage.
107
171
  """
108
- return PREFECT_RESULTS_DEFAULT_SERIALIZER.value()
172
+ return resolve_serializer(PREFECT_RESULTS_DEFAULT_SERIALIZER.value())
109
173
 
110
174
 
111
175
  def get_default_persist_setting() -> bool:
@@ -115,6 +179,25 @@ def get_default_persist_setting() -> bool:
115
179
  return PREFECT_RESULTS_PERSIST_BY_DEFAULT.value()
116
180
 
117
181
 
182
+ def should_persist_result() -> bool:
183
+ """
184
+ Return the default option for result persistence determined by the current run context.
185
+
186
+ If there is no current run context, the default value set by
187
+ `PREFECT_RESULTS_PERSIST_BY_DEFAULT` will be returned.
188
+ """
189
+ from prefect.context import FlowRunContext, TaskRunContext
190
+
191
+ task_run_context = TaskRunContext.get()
192
+ if task_run_context is not None:
193
+ return task_run_context.persist_result
194
+ flow_run_context = FlowRunContext.get()
195
+ if flow_run_context is not None:
196
+ return flow_run_context.persist_result
197
+
198
+ return PREFECT_RESULTS_PERSIST_BY_DEFAULT.value()
199
+
200
+
118
201
  def _format_user_supplied_storage_key(key: str) -> str:
119
202
  # Note here we are pinning to task runs since flow runs do not support storage keys
120
203
  # yet; we'll need to split logic in the future or have two separate functions
@@ -122,218 +205,562 @@ def _format_user_supplied_storage_key(key: str) -> str:
122
205
  return key.format(**runtime_vars, parameters=prefect.runtime.task_run.parameters)
123
206
 
124
207
 
125
- class ResultFactory(BaseModel):
208
+ T = TypeVar("T")
209
+
210
+
211
+ @deprecated_field(
212
+ "persist_result",
213
+ when=lambda x: x is not None,
214
+ when_message="use the `should_persist_result` utility function instead",
215
+ start_date="Sep 2024",
216
+ end_date="Nov 2024",
217
+ )
218
+ class ResultStore(BaseModel):
126
219
  """
127
- A utility to generate `Result` types.
220
+ Manages the storage and retrieval of results.
221
+
222
+ Attributes:
223
+ result_storage: The storage for result records. If not provided, the default
224
+ result storage will be used.
225
+ metadata_storage: The storage for result record metadata. If not provided,
226
+ the metadata will be stored alongside the results.
227
+ lock_manager: The lock manager to use for locking result records. If not provided,
228
+ the store cannot be used in transactions with the SERIALIZABLE isolation level.
229
+ persist_result: Whether to persist results.
230
+ cache_result_in_memory: Whether to cache results in memory.
231
+ serializer: The serializer to use for results.
232
+ storage_key_fn: The function to generate storage keys.
128
233
  """
129
234
 
130
- persist_result: bool
131
- cache_result_in_memory: bool
132
- serializer: Serializer
133
- storage_block_id: Optional[uuid.UUID] = None
134
- storage_block: WritableFileSystem
135
- storage_key_fn: Callable[[], str]
235
+ model_config = ConfigDict(arbitrary_types_allowed=True)
136
236
 
137
- @classmethod
138
- @inject_client
139
- async def default_factory(cls, client: "PrefectClient" = None, **kwargs):
237
+ result_storage: Optional[WritableFileSystem] = Field(default=None)
238
+ metadata_storage: Optional[WritableFileSystem] = Field(default=None)
239
+ lock_manager: Optional[LockManager] = Field(default=None)
240
+ cache_result_in_memory: bool = Field(default=True)
241
+ serializer: Serializer = Field(default_factory=get_default_result_serializer)
242
+ storage_key_fn: Callable[[], str] = Field(default=DEFAULT_STORAGE_KEY_FN)
243
+ cache: LRUCache = Field(default_factory=lambda: LRUCache(maxsize=1000))
244
+
245
+ # Deprecated fields
246
+ persist_result: Optional[bool] = Field(default=None)
247
+
248
+ @property
249
+ def result_storage_block_id(self) -> Optional[UUID]:
250
+ if self.result_storage is None:
251
+ return None
252
+ return self.result_storage._block_document_id
253
+
254
+ @sync_compatible
255
+ async def update_for_flow(self, flow: "Flow") -> Self:
140
256
  """
141
- Create a new result factory with default options.
257
+ Create a new result store for a flow with updated settings.
142
258
 
143
- Keyword arguments may be provided to override defaults. Null keys will be
144
- ignored.
259
+ Args:
260
+ flow: The flow to update the result store for.
261
+
262
+ Returns:
263
+ An updated result store.
145
264
  """
146
- # Remove any null keys so `setdefault` can do its magic
147
- for key, value in tuple(kwargs.items()):
148
- if value is None:
149
- kwargs.pop(key)
265
+ update = {}
266
+ if flow.result_storage is not None:
267
+ update["result_storage"] = await resolve_result_storage(flow.result_storage)
268
+ if flow.result_serializer is not None:
269
+ update["serializer"] = resolve_serializer(flow.result_serializer)
270
+ if flow.cache_result_in_memory is not None:
271
+ update["cache_result_in_memory"] = flow.cache_result_in_memory
272
+ if self.result_storage is None and update.get("result_storage") is None:
273
+ update["result_storage"] = await get_default_result_storage()
274
+ return self.model_copy(update=update)
150
275
 
151
- # Apply defaults
152
- kwargs.setdefault("result_storage", await get_default_result_storage())
153
- kwargs.setdefault("result_serializer", get_default_result_serializer())
154
- kwargs.setdefault("persist_result", get_default_persist_setting())
155
- kwargs.setdefault("cache_result_in_memory", True)
156
- kwargs.setdefault("storage_key_fn", DEFAULT_STORAGE_KEY_FN)
276
+ @sync_compatible
277
+ async def update_for_task(self: Self, task: "Task") -> Self:
278
+ """
279
+ Create a new result store for a task.
157
280
 
158
- return await cls.from_settings(**kwargs, client=client)
281
+ Args:
282
+ task: The task to update the result store for.
159
283
 
160
- @classmethod
161
- @inject_client
162
- async def from_flow(
163
- cls: Type[Self], flow: "Flow", client: "PrefectClient" = None
164
- ) -> Self:
165
- """
166
- Create a new result factory for a flow.
167
- """
168
- from prefect.context import FlowRunContext
169
-
170
- ctx = FlowRunContext.get()
171
- if ctx:
172
- # This is a child flow run
173
- return await cls.from_settings(
174
- result_storage=flow.result_storage or ctx.result_factory.storage_block,
175
- result_serializer=flow.result_serializer
176
- or ctx.result_factory.serializer,
177
- persist_result=flow.persist_result,
178
- cache_result_in_memory=flow.cache_result_in_memory,
179
- storage_key_fn=DEFAULT_STORAGE_KEY_FN,
180
- client=client,
284
+ Returns:
285
+ An updated result store.
286
+ """
287
+ update = {}
288
+ if task.result_storage is not None:
289
+ update["result_storage"] = await resolve_result_storage(task.result_storage)
290
+ if task.result_serializer is not None:
291
+ update["serializer"] = resolve_serializer(task.result_serializer)
292
+ if task.cache_result_in_memory is not None:
293
+ update["cache_result_in_memory"] = task.cache_result_in_memory
294
+ if task.result_storage_key is not None:
295
+ update["storage_key_fn"] = partial(
296
+ _format_user_supplied_storage_key, task.result_storage_key
181
297
  )
298
+ if self.result_storage is None and update.get("result_storage") is None:
299
+ update["result_storage"] = await get_default_result_storage()
300
+ return self.model_copy(update=update)
301
+
302
+ @staticmethod
303
+ def generate_default_holder() -> str:
304
+ """
305
+ Generate a default holder string using hostname, PID, and thread ID.
306
+
307
+ Returns:
308
+ str: A unique identifier string.
309
+ """
310
+ hostname = socket.gethostname()
311
+ pid = os.getpid()
312
+ thread_name = threading.current_thread().name
313
+ thread_id = threading.get_ident()
314
+ return f"{hostname}:{pid}:{thread_id}:{thread_name}"
315
+
316
+ @sync_compatible
317
+ async def _exists(self, key: str) -> bool:
318
+ """
319
+ Check if a result record exists in storage.
320
+
321
+ Args:
322
+ key: The key to check for the existence of a result record.
323
+
324
+ Returns:
325
+ bool: True if the result record exists, False otherwise.
326
+ """
327
+ if self.metadata_storage is not None:
328
+ # TODO: Add an `exists` method to commonly used storage blocks
329
+ # so the entire payload doesn't need to be read
330
+ try:
331
+ metadata_content = await self.metadata_storage.read_path(key)
332
+ if metadata_content is None:
333
+ return False
334
+ metadata = ResultRecordMetadata.load_bytes(metadata_content)
335
+
336
+ except Exception:
337
+ return False
338
+ else:
339
+ try:
340
+ content = await self.result_storage.read_path(key)
341
+ if content is None:
342
+ return False
343
+ record = ResultRecord.deserialize(content)
344
+ metadata = record.metadata
345
+ except Exception:
346
+ return False
347
+
348
+ if metadata.expiration:
349
+ # if the result has an expiration,
350
+ # check if it is still in the future
351
+ exists = metadata.expiration > pendulum.now("utc")
182
352
  else:
183
- # This is a root flow run
184
- # Pass the flow settings up to the default which will replace nulls with
185
- # our default options
186
- return await cls.default_factory(
187
- client=client,
188
- result_storage=flow.result_storage,
189
- result_serializer=flow.result_serializer,
190
- persist_result=flow.persist_result,
191
- cache_result_in_memory=flow.cache_result_in_memory,
192
- storage_key_fn=DEFAULT_STORAGE_KEY_FN,
353
+ exists = True
354
+ return exists
355
+
356
+ def exists(self, key: str) -> bool:
357
+ """
358
+ Check if a result record exists in storage.
359
+
360
+ Args:
361
+ key: The key to check for the existence of a result record.
362
+
363
+ Returns:
364
+ bool: True if the result record exists, False otherwise.
365
+ """
366
+ return self._exists(key=key, _sync=True)
367
+
368
+ async def aexists(self, key: str) -> bool:
369
+ """
370
+ Check if a result record exists in storage.
371
+
372
+ Args:
373
+ key: The key to check for the existence of a result record.
374
+
375
+ Returns:
376
+ bool: True if the result record exists, False otherwise.
377
+ """
378
+ return await self._exists(key=key, _sync=False)
379
+
380
+ @sync_compatible
381
+ async def _read(self, key: str, holder: str) -> "ResultRecord":
382
+ """
383
+ Read a result record from storage.
384
+
385
+ This is the internal implementation. Use `read` or `aread` for synchronous and
386
+ asynchronous result reading respectively.
387
+
388
+ Args:
389
+ key: The key to read the result record from.
390
+ holder: The holder of the lock if a lock was set on the record.
391
+
392
+ Returns:
393
+ A result record.
394
+ """
395
+
396
+ if self.lock_manager is not None and not self.is_lock_holder(key, holder):
397
+ await self.await_for_lock(key)
398
+
399
+ if key in self.cache:
400
+ return self.cache[key]
401
+
402
+ if self.result_storage is None:
403
+ self.result_storage = await get_default_result_storage()
404
+
405
+ if self.metadata_storage is not None:
406
+ metadata_content = await self.metadata_storage.read_path(key)
407
+ metadata = ResultRecordMetadata.load_bytes(metadata_content)
408
+ assert (
409
+ metadata.storage_key is not None
410
+ ), "Did not find storage key in metadata"
411
+ result_content = await self.result_storage.read_path(metadata.storage_key)
412
+ result_record = ResultRecord.deserialize_from_result_and_metadata(
413
+ result=result_content, metadata=metadata_content
193
414
  )
415
+ else:
416
+ content = await self.result_storage.read_path(key)
417
+ result_record = ResultRecord.deserialize(content)
194
418
 
195
- @classmethod
196
- @inject_client
197
- async def from_task(
198
- cls: Type[Self], task: "Task", client: "PrefectClient" = None
199
- ) -> Self:
419
+ if self.cache_result_in_memory:
420
+ if self.result_storage_block_id is None and hasattr(
421
+ self.result_storage, "_resolve_path"
422
+ ):
423
+ cache_key = str(self.result_storage._resolve_path(key))
424
+ else:
425
+ cache_key = key
426
+
427
+ self.cache[cache_key] = result_record
428
+ return result_record
429
+
430
+ def read(self, key: str, holder: Optional[str] = None) -> "ResultRecord":
200
431
  """
201
- Create a new result factory for a task.
432
+ Read a result record from storage.
433
+
434
+ Args:
435
+ key: The key to read the result record from.
436
+ holder: The holder of the lock if a lock was set on the record.
437
+ Returns:
438
+ A result record.
202
439
  """
203
- return await cls._from_task(task, get_default_result_storage, client=client)
440
+ holder = holder or self.generate_default_holder()
441
+ return self._read(key=key, holder=holder, _sync=True)
204
442
 
205
- @classmethod
206
- @inject_client
207
- async def from_autonomous_task(
208
- cls: Type[Self], task: "Task[P, R]", client: "PrefectClient" = None
209
- ) -> Self:
443
+ async def aread(self, key: str, holder: Optional[str] = None) -> "ResultRecord":
444
+ """
445
+ Read a result record from storage.
446
+
447
+ Args:
448
+ key: The key to read the result record from.
449
+ holder: The holder of the lock if a lock was set on the record.
450
+ Returns:
451
+ A result record.
210
452
  """
211
- Create a new result factory for an autonomous task.
453
+ holder = holder or self.generate_default_holder()
454
+ return await self._read(key=key, holder=holder, _sync=False)
455
+
456
+ def create_result_record(
457
+ self,
458
+ obj: Any,
459
+ key: Optional[str] = None,
460
+ expiration: Optional[DateTime] = None,
461
+ ) -> "ResultRecord":
462
+ """
463
+ Create a result record.
464
+
465
+ Args:
466
+ key: The key to create the result record for.
467
+ obj: The object to create the result record for.
468
+ expiration: The expiration time for the result record.
212
469
  """
213
- return await cls._from_task(
214
- task, get_or_create_default_task_scheduling_storage, client=client
470
+ key = key or self.storage_key_fn()
471
+
472
+ if self.result_storage is None:
473
+ self.result_storage = get_default_result_storage(_sync=True)
474
+
475
+ if self.result_storage_block_id is None:
476
+ if hasattr(self.result_storage, "_resolve_path"):
477
+ key = str(self.result_storage._resolve_path(key))
478
+
479
+ return ResultRecord(
480
+ result=obj,
481
+ metadata=ResultRecordMetadata(
482
+ serializer=self.serializer,
483
+ expiration=expiration,
484
+ storage_key=key,
485
+ storage_block_id=self.result_storage_block_id,
486
+ ),
215
487
  )
216
488
 
217
- @classmethod
218
- @inject_client
219
- async def _from_task(
220
- cls: Type[Self],
221
- task: "Task",
222
- default_storage_getter: Callable[[], Awaitable[ResultStorage]],
223
- client: "PrefectClient" = None,
224
- ) -> Self:
225
- from prefect.context import FlowRunContext
226
-
227
- ctx = FlowRunContext.get()
228
-
229
- result_storage = task.result_storage or (
230
- ctx.result_factory.storage_block
231
- if ctx and ctx.result_factory
232
- else await default_storage_getter()
489
+ def write(
490
+ self,
491
+ obj: Any,
492
+ key: Optional[str] = None,
493
+ expiration: Optional[DateTime] = None,
494
+ holder: Optional[str] = None,
495
+ ):
496
+ """
497
+ Write a result to storage.
498
+
499
+ Handles the creation of a `ResultRecord` and its serialization to storage.
500
+
501
+ Args:
502
+ key: The key to write the result record to.
503
+ obj: The object to write to storage.
504
+ expiration: The expiration time for the result record.
505
+ holder: The holder of the lock if a lock was set on the record.
506
+ """
507
+ holder = holder or self.generate_default_holder()
508
+ result_record = self.create_result_record(
509
+ key=key, obj=obj, expiration=expiration
233
510
  )
234
- result_serializer = task.result_serializer or (
235
- ctx.result_factory.serializer
236
- if ctx and ctx.result_factory
237
- else get_default_result_serializer()
511
+ return self.persist_result_record(
512
+ result_record=result_record,
513
+ holder=holder,
238
514
  )
239
- if task.persist_result is None:
240
- persist_result = (
241
- ctx.result_factory.persist_result
242
- if ctx and ctx.result_factory
243
- else get_default_persist_setting()
244
- )
245
- else:
246
- persist_result = task.persist_result
247
-
248
- cache_result_in_memory = task.cache_result_in_memory
249
-
250
- return await cls.from_settings(
251
- result_storage=result_storage,
252
- result_serializer=result_serializer,
253
- persist_result=persist_result,
254
- cache_result_in_memory=cache_result_in_memory,
255
- client=client,
256
- storage_key_fn=(
257
- partial(_format_user_supplied_storage_key, task.result_storage_key)
258
- if task.result_storage_key is not None
259
- else DEFAULT_STORAGE_KEY_FN
515
+
516
+ async def awrite(
517
+ self,
518
+ obj: Any,
519
+ key: Optional[str] = None,
520
+ expiration: Optional[DateTime] = None,
521
+ holder: Optional[str] = None,
522
+ ):
523
+ """
524
+ Write a result to storage.
525
+
526
+ Args:
527
+ key: The key to write the result record to.
528
+ obj: The object to write to storage.
529
+ expiration: The expiration time for the result record.
530
+ holder: The holder of the lock if a lock was set on the record.
531
+ """
532
+ holder = holder or self.generate_default_holder()
533
+ return await self.apersist_result_record(
534
+ result_record=self.create_result_record(
535
+ key=key, obj=obj, expiration=expiration
260
536
  ),
537
+ holder=holder,
261
538
  )
262
539
 
263
- @classmethod
264
- @inject_client
265
- async def from_settings(
266
- cls: Type[Self],
267
- result_storage: ResultStorage,
268
- result_serializer: ResultSerializer,
269
- persist_result: Optional[bool],
270
- cache_result_in_memory: bool,
271
- storage_key_fn: Callable[[], str],
272
- client: "PrefectClient",
273
- ) -> Self:
274
- if persist_result is None:
275
- persist_result = get_default_persist_setting()
540
+ @sync_compatible
541
+ async def _persist_result_record(self, result_record: "ResultRecord", holder: str):
542
+ """
543
+ Persist a result record to storage.
276
544
 
277
- storage_block_id, storage_block = await cls.resolve_storage_block(
278
- result_storage, client=client, persist_result=persist_result
545
+ Args:
546
+ result_record: The result record to persist.
547
+ holder: The holder of the lock if a lock was set on the record.
548
+ """
549
+ assert (
550
+ result_record.metadata.storage_key is not None
551
+ ), "Storage key is required on result record"
552
+
553
+ key = result_record.metadata.storage_key
554
+ if result_record.metadata.storage_block_id is None:
555
+ basepath = (
556
+ self.result_storage._resolve_path("")
557
+ if hasattr(self.result_storage, "_resolve_path")
558
+ else Path(".").resolve()
559
+ )
560
+ base_key = str(Path(key).relative_to(basepath))
561
+ else:
562
+ base_key = key
563
+ if (
564
+ self.lock_manager is not None
565
+ and self.is_locked(base_key)
566
+ and not self.is_lock_holder(base_key, holder)
567
+ ):
568
+ raise RuntimeError(
569
+ f"Cannot write to result record with key {base_key} because it is locked by "
570
+ f"another holder."
571
+ )
572
+ if self.result_storage is None:
573
+ self.result_storage = await get_default_result_storage()
574
+
575
+ # If metadata storage is configured, write result and metadata separately
576
+ if self.metadata_storage is not None:
577
+ await self.result_storage.write_path(
578
+ result_record.metadata.storage_key,
579
+ content=result_record.serialize_result(),
580
+ )
581
+ await self.metadata_storage.write_path(
582
+ base_key,
583
+ content=result_record.serialize_metadata(),
584
+ )
585
+ # Otherwise, write the result metadata and result together
586
+ else:
587
+ await self.result_storage.write_path(
588
+ result_record.metadata.storage_key, content=result_record.serialize()
589
+ )
590
+
591
+ if self.cache_result_in_memory:
592
+ self.cache[key] = result_record
593
+
594
+ def persist_result_record(
595
+ self, result_record: "ResultRecord", holder: Optional[str] = None
596
+ ):
597
+ """
598
+ Persist a result record to storage.
599
+
600
+ Args:
601
+ result_record: The result record to persist.
602
+ """
603
+ holder = holder or self.generate_default_holder()
604
+ return self._persist_result_record(
605
+ result_record=result_record, holder=holder, _sync=True
279
606
  )
280
- serializer = cls.resolve_serializer(result_serializer)
281
607
 
282
- return cls(
283
- storage_block=storage_block,
284
- storage_block_id=storage_block_id,
285
- serializer=serializer,
286
- persist_result=persist_result,
287
- cache_result_in_memory=cache_result_in_memory,
288
- storage_key_fn=storage_key_fn,
608
+ async def apersist_result_record(
609
+ self, result_record: "ResultRecord", holder: Optional[str] = None
610
+ ):
611
+ """
612
+ Persist a result record to storage.
613
+
614
+ Args:
615
+ result_record: The result record to persist.
616
+ """
617
+ holder = holder or self.generate_default_holder()
618
+ return await self._persist_result_record(
619
+ result_record=result_record, holder=holder, _sync=False
289
620
  )
290
621
 
291
- @staticmethod
292
- async def resolve_storage_block(
293
- result_storage: ResultStorage,
294
- client: "PrefectClient",
295
- persist_result: bool = True,
296
- ) -> Tuple[Optional[uuid.UUID], WritableFileSystem]:
622
+ def supports_isolation_level(self, level: "IsolationLevel") -> bool:
297
623
  """
298
- Resolve one of the valid `ResultStorage` input types into a saved block
299
- document id and an instance of the block.
624
+ Check if the result store supports a given isolation level.
625
+
626
+ Args:
627
+ level: The isolation level to check.
628
+
629
+ Returns:
630
+ bool: True if the isolation level is supported, False otherwise.
300
631
  """
301
- if isinstance(result_storage, Block):
302
- storage_block = result_storage
632
+ from prefect.transactions import IsolationLevel
303
633
 
304
- if storage_block._block_document_id is not None:
305
- # Avoid saving the block if it already has an identifier assigned
306
- storage_block_id = storage_block._block_document_id
307
- else:
308
- storage_block_id = None
309
- elif isinstance(result_storage, str):
310
- storage_block = await Block.load(result_storage, client=client)
311
- storage_block_id = storage_block._block_document_id
312
- assert storage_block_id is not None, "Loaded storage blocks must have ids"
634
+ if level == IsolationLevel.READ_COMMITTED:
635
+ return True
636
+ elif level == IsolationLevel.SERIALIZABLE:
637
+ return self.lock_manager is not None
313
638
  else:
314
- raise TypeError(
315
- "Result storage must be one of the following types: 'UUID', 'Block', "
316
- f"'str'. Got unsupported type {type(result_storage).__name__!r}."
639
+ raise ValueError(f"Unsupported isolation level: {level}")
640
+
641
+ def acquire_lock(
642
+ self, key: str, holder: Optional[str] = None, timeout: Optional[float] = None
643
+ ) -> bool:
644
+ """
645
+ Acquire a lock for a result record.
646
+
647
+ Args:
648
+ key: The key to acquire the lock for.
649
+ holder: The holder of the lock. If not provided, a default holder based on the
650
+ current host, process, and thread will be used.
651
+ timeout: The timeout for the lock.
652
+
653
+ Returns:
654
+ bool: True if the lock was successfully acquired; False otherwise.
655
+ """
656
+ holder = holder or self.generate_default_holder()
657
+ if self.lock_manager is None:
658
+ raise ConfigurationError(
659
+ "Result store is not configured with a lock manager. Please set"
660
+ " a lock manager when creating the result store to enable locking."
317
661
  )
662
+ return self.lock_manager.acquire_lock(key, holder, timeout)
318
663
 
319
- return storage_block_id, storage_block
664
+ async def aacquire_lock(
665
+ self, key: str, holder: Optional[str] = None, timeout: Optional[float] = None
666
+ ) -> bool:
667
+ """
668
+ Acquire a lock for a result record.
320
669
 
321
- @staticmethod
322
- def resolve_serializer(serializer: ResultSerializer) -> Serializer:
670
+ Args:
671
+ key: The key to acquire the lock for.
672
+ holder: The holder of the lock. If not provided, a default holder based on the
673
+ current host, process, and thread will be used.
674
+ timeout: The timeout for the lock.
675
+
676
+ Returns:
677
+ bool: True if the lock was successfully acquired; False otherwise.
323
678
  """
324
- Resolve one of the valid `ResultSerializer` input types into a serializer
325
- instance.
679
+ holder = holder or self.generate_default_holder()
680
+ if self.lock_manager is None:
681
+ raise ConfigurationError(
682
+ "Result store is not configured with a lock manager. Please set"
683
+ " a lock manager when creating the result store to enable locking."
684
+ )
685
+
686
+ return await self.lock_manager.aacquire_lock(key, holder, timeout)
687
+
688
+ def release_lock(self, key: str, holder: Optional[str] = None):
326
689
  """
327
- if isinstance(serializer, Serializer):
328
- return serializer
329
- elif isinstance(serializer, str):
330
- return Serializer(type=serializer)
331
- else:
332
- raise TypeError(
333
- "Result serializer must be one of the following types: 'Serializer', "
334
- f"'str'. Got unsupported type {type(serializer).__name__!r}."
690
+ Release a lock for a result record.
691
+
692
+ Args:
693
+ key: The key to release the lock for.
694
+ holder: The holder of the lock. Must match the holder that acquired the lock.
695
+ If not provided, a default holder based on the current host, process, and
696
+ thread will be used.
697
+ """
698
+ holder = holder or self.generate_default_holder()
699
+ if self.lock_manager is None:
700
+ raise ConfigurationError(
701
+ "Result store is not configured with a lock manager. Please set"
702
+ " a lock manager when creating the result store to enable locking."
703
+ )
704
+ return self.lock_manager.release_lock(key, holder)
705
+
706
+ def is_locked(self, key: str) -> bool:
707
+ """
708
+ Check if a result record is locked.
709
+ """
710
+ if self.lock_manager is None:
711
+ raise ConfigurationError(
712
+ "Result store is not configured with a lock manager. Please set"
713
+ " a lock manager when creating the result store to enable locking."
714
+ )
715
+ return self.lock_manager.is_locked(key)
716
+
717
+ def is_lock_holder(self, key: str, holder: Optional[str] = None) -> bool:
718
+ """
719
+ Check if the current holder is the lock holder for the result record.
720
+
721
+ Args:
722
+ key: The key to check the lock for.
723
+ holder: The holder of the lock. If not provided, a default holder based on the
724
+ current host, process, and thread will be used.
725
+
726
+ Returns:
727
+ bool: True if the current holder is the lock holder; False otherwise.
728
+ """
729
+ holder = holder or self.generate_default_holder()
730
+ if self.lock_manager is None:
731
+ raise ConfigurationError(
732
+ "Result store is not configured with a lock manager. Please set"
733
+ " a lock manager when creating the result store to enable locking."
734
+ )
735
+ return self.lock_manager.is_lock_holder(key, holder)
736
+
737
+ def wait_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
738
+ """
739
+ Wait for the corresponding transaction record to become free.
740
+ """
741
+ if self.lock_manager is None:
742
+ raise ConfigurationError(
743
+ "Result store is not configured with a lock manager. Please set"
744
+ " a lock manager when creating the result store to enable locking."
335
745
  )
746
+ return self.lock_manager.wait_for_lock(key, timeout)
336
747
 
748
+ async def await_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
749
+ """
750
+ Wait for the corresponding transaction record to become free.
751
+ """
752
+ if self.lock_manager is None:
753
+ raise ConfigurationError(
754
+ "Result store is not configured with a lock manager. Please set"
755
+ " a lock manager when creating the result store to enable locking."
756
+ )
757
+ return await self.lock_manager.await_for_lock(key, timeout)
758
+
759
+ @deprecated.deprecated_callable(
760
+ start_date="Sep 2024",
761
+ end_date="Nov 2024",
762
+ help="Use `create_result_record` instead.",
763
+ )
337
764
  @sync_compatible
338
765
  async def create_result(
339
766
  self,
@@ -342,12 +769,15 @@ class ResultFactory(BaseModel):
342
769
  expiration: Optional[DateTime] = None,
343
770
  ) -> Union[R, "BaseResult[R]"]:
344
771
  """
345
- Create a result type for the given object.
346
-
347
- If persistence is enabled the object is serialized, persisted to storage, and a reference is returned.
772
+ Create a `PersistedResult` for the given object.
348
773
  """
349
774
  # Null objects are "cached" in memory at no cost
350
775
  should_cache_object = self.cache_result_in_memory or obj is None
776
+ should_persist_result = (
777
+ self.persist_result
778
+ if self.persist_result is not None
779
+ else not should_cache_object
780
+ )
351
781
 
352
782
  if key:
353
783
 
@@ -358,15 +788,18 @@ class ResultFactory(BaseModel):
358
788
  else:
359
789
  storage_key_fn = self.storage_key_fn
360
790
 
791
+ if self.result_storage is None:
792
+ self.result_storage = await get_default_result_storage()
793
+
361
794
  return await PersistedResult.create(
362
795
  obj,
363
- storage_block=self.storage_block,
364
- storage_block_id=self.storage_block_id,
796
+ storage_block=self.result_storage,
797
+ storage_block_id=self.result_storage_block_id,
365
798
  storage_key_fn=storage_key_fn,
366
799
  serializer=self.serializer,
367
800
  cache_object=should_cache_object,
368
801
  expiration=expiration,
369
- serialize_to_none=not self.persist_result,
802
+ serialize_to_none=not should_persist_result,
370
803
  )
371
804
 
372
805
  # TODO: These two methods need to find a new home
@@ -379,18 +812,33 @@ class ResultFactory(BaseModel):
379
812
  serializer=self.serializer, storage_key=str(identifier)
380
813
  ),
381
814
  )
382
- await self.storage_block.write_path(
815
+ await self.result_storage.write_path(
383
816
  f"parameters/{identifier}", content=record.serialize()
384
817
  )
385
818
 
386
819
  @sync_compatible
387
820
  async def read_parameters(self, identifier: UUID) -> Dict[str, Any]:
388
821
  record = ResultRecord.deserialize(
389
- await self.storage_block.read_path(f"parameters/{identifier}")
822
+ await self.result_storage.read_path(f"parameters/{identifier}")
390
823
  )
391
824
  return record.result
392
825
 
393
826
 
827
+ def get_result_store() -> ResultStore:
828
+ """
829
+ Get the current result store.
830
+ """
831
+ from prefect.context import get_run_context
832
+
833
+ try:
834
+ run_context = get_run_context()
835
+ except MissingContextError:
836
+ result_store = ResultStore()
837
+ else:
838
+ result_store = run_context.result_store
839
+ return result_store
840
+
841
+
394
842
  class ResultRecordMetadata(BaseModel):
395
843
  """
396
844
  Metadata for a result record.
@@ -402,6 +850,7 @@ class ResultRecordMetadata(BaseModel):
402
850
  expiration: Optional[DateTime] = Field(default=None)
403
851
  serializer: Serializer = Field(default_factory=PickleSerializer)
404
852
  prefect_version: str = Field(default=prefect.__version__)
853
+ storage_block_id: Optional[uuid.UUID] = Field(default=None)
405
854
 
406
855
  def dump_bytes(self) -> bytes:
407
856
  """
@@ -425,6 +874,17 @@ class ResultRecordMetadata(BaseModel):
425
874
  """
426
875
  return cls.model_validate_json(data)
427
876
 
877
+ def __eq__(self, other):
878
+ if not isinstance(other, ResultRecordMetadata):
879
+ return False
880
+ return (
881
+ self.storage_key == other.storage_key
882
+ and self.expiration == other.expiration
883
+ and self.serializer == other.serializer
884
+ and self.prefect_version == other.prefect_version
885
+ and self.storage_block_id == other.storage_block_id
886
+ )
887
+
428
888
 
429
889
  class ResultRecord(BaseModel, Generic[R]):
430
890
  """
@@ -502,6 +962,31 @@ class ResultRecord(BaseModel, Generic[R]):
502
962
  value["metadata"]["prefect_version"] = value.pop("prefect_version")
503
963
  return value
504
964
 
965
+ @classmethod
966
+ async def _from_metadata(cls, metadata: ResultRecordMetadata) -> "ResultRecord[R]":
967
+ """
968
+ Create a result record from metadata.
969
+
970
+ Will use the result record metadata to fetch data via a result store.
971
+
972
+ Args:
973
+ metadata: The metadata to create the result record from.
974
+
975
+ Returns:
976
+ ResultRecord: The result record.
977
+ """
978
+ if metadata.storage_block_id is None:
979
+ storage_block = None
980
+ else:
981
+ storage_block = await resolve_result_storage(
982
+ metadata.storage_block_id, _sync=False
983
+ )
984
+ store = ResultStore(
985
+ result_storage=storage_block, serializer=metadata.serializer
986
+ )
987
+ result = await store.aread(metadata.storage_key)
988
+ return result
989
+
505
990
  def serialize_metadata(self) -> bytes:
506
991
  return self.metadata.dump_bytes()
507
992
 
@@ -559,14 +1044,24 @@ class ResultRecord(BaseModel, Generic[R]):
559
1044
  result=result_record_metadata.serializer.loads(result),
560
1045
  )
561
1046
 
1047
+ def __eq__(self, other):
1048
+ if not isinstance(other, ResultRecord):
1049
+ return False
1050
+ return self.metadata == other.metadata and self.result == other.result
562
1051
 
1052
+
1053
+ @deprecated.deprecated_class(
1054
+ start_date="Sep 2024", end_date="Nov 2024", help="Use `ResultRecord` instead."
1055
+ )
563
1056
  @register_base_type
564
1057
  class BaseResult(BaseModel, abc.ABC, Generic[R]):
565
1058
  model_config = ConfigDict(extra="forbid")
566
1059
  type: str
567
1060
 
568
1061
  def __init__(self, **data: Any) -> None:
569
- type_string = get_dispatch_key(self) if type(self) != BaseResult else "__base__"
1062
+ type_string = (
1063
+ get_dispatch_key(self) if type(self) is not BaseResult else "__base__"
1064
+ )
570
1065
  data.setdefault("type", type_string)
571
1066
  super().__init__(**data)
572
1067
 
@@ -608,6 +1103,9 @@ class BaseResult(BaseModel, abc.ABC, Generic[R]):
608
1103
  return cls.__name__ if isinstance(default, PydanticUndefinedType) else default
609
1104
 
610
1105
 
1106
+ @deprecated.deprecated_class(
1107
+ start_date="Sep 2024", end_date="Nov 2024", help="Use `ResultRecord` instead."
1108
+ )
611
1109
  class PersistedResult(BaseResult):
612
1110
  """
613
1111
  Result type which stores a reference to a persisted result.
@@ -660,14 +1158,22 @@ class PersistedResult(BaseResult):
660
1158
 
661
1159
  @sync_compatible
662
1160
  @inject_client
663
- async def get(self, client: "PrefectClient") -> R:
1161
+ async def get(
1162
+ self, ignore_cache: bool = False, client: "PrefectClient" = None
1163
+ ) -> R:
664
1164
  """
665
1165
  Retrieve the data and deserialize it into the original object.
666
1166
  """
667
- if self.has_cached_object():
1167
+ if self.has_cached_object() and not ignore_cache:
668
1168
  return self._cache
669
1169
 
670
- record = await self._read_result_record(client=client)
1170
+ result_store_kwargs = {}
1171
+ if self._serializer:
1172
+ result_store_kwargs["serializer"] = resolve_serializer(self._serializer)
1173
+ storage_block = await self._get_storage_block(client=client)
1174
+ result_store = ResultStore(result_storage=storage_block, **result_store_kwargs)
1175
+
1176
+ record = await result_store.aread(self.storage_key)
671
1177
  self.expiration = record.expiration
672
1178
 
673
1179
  if self._should_cache_object:
@@ -675,13 +1181,6 @@ class PersistedResult(BaseResult):
675
1181
 
676
1182
  return record.result
677
1183
 
678
- @inject_client
679
- async def _read_result_record(self, client: "PrefectClient") -> "ResultRecord":
680
- block = await self._get_storage_block(client=client)
681
- content = await block.read_path(self.storage_key)
682
- record = ResultRecord.deserialize(content)
683
- return record
684
-
685
1184
  @staticmethod
686
1185
  def _infer_path(storage_block, key) -> str:
687
1186
  """
@@ -700,7 +1199,6 @@ class PersistedResult(BaseResult):
700
1199
  """
701
1200
  Write the result to the storage block.
702
1201
  """
703
-
704
1202
  if self._persisted or self.serialize_to_none:
705
1203
  # don't double write or overwrite
706
1204
  return
@@ -721,15 +1219,14 @@ class PersistedResult(BaseResult):
721
1219
  # this could error if the serializer requires kwargs
722
1220
  serializer = Serializer(type=self.serializer_type)
723
1221
 
724
- record = ResultRecord(
725
- result=obj,
726
- metadata=ResultRecordMetadata(
727
- storage_key=self.storage_key,
728
- expiration=self.expiration,
729
- serializer=serializer,
730
- ),
1222
+ result_store = ResultStore(
1223
+ result_storage=storage_block,
1224
+ serializer=serializer,
731
1225
  )
732
- await storage_block.write_path(self.storage_key, content=record.serialize())
1226
+ await result_store.awrite(
1227
+ obj=obj, key=self.storage_key, expiration=self.expiration
1228
+ )
1229
+
733
1230
  self._persisted = True
734
1231
 
735
1232
  if not self._should_cache_object: