prefect-client 3.1.6__py3-none-any.whl → 3.1.8__py3-none-any.whl

prefect/flow_engine.py

@@ -24,10 +24,8 @@ from uuid import UUID
 
 from anyio import CancelScope
 from opentelemetry import propagate, trace
-from opentelemetry.trace import Tracer, get_tracer
 from typing_extensions import ParamSpec
 
-import prefect
 from prefect import Task
 from prefect.client.orchestration import PrefectClient, SyncPrefectClient, get_client
 from prefect.client.schemas import FlowRun, TaskRun
@@ -72,7 +70,12 @@ from prefect.states import (
     exception_to_failed_state,
     return_value_to_state,
 )
-from prefect.telemetry.run_telemetry import OTELSetter
+from prefect.telemetry.run_telemetry import (
+    LABELS_TRACEPARENT_KEY,
+    TRACEPARENT_KEY,
+    OTELSetter,
+    RunTelemetry,
+)
 from prefect.types import KeyValueLabels
 from prefect.utilities._engine import get_hook_name, resolve_custom_flow_run_name
 from prefect.utilities.annotations import NotSet
@@ -95,8 +98,6 @@ from prefect.utilities.urls import url_for
 
 P = ParamSpec("P")
 R = TypeVar("R")
-LABELS_TRACEPARENT_KEY = "__OTEL_TRACEPARENT"
-TRACEPARENT_KEY = "traceparent"
 
 
 class FlowRunTimeoutError(TimeoutError):
@@ -136,10 +137,7 @@ class BaseFlowRunEngine(Generic[P, R]):
     _is_started: bool = False
     short_circuit: bool = False
     _flow_run_name_set: bool = False
-    _tracer: Tracer = field(
-        default_factory=lambda: get_tracer("prefect", prefect.__version__)
-    )
-    _span: Optional[trace.Span] = None
+    _telemetry: RunTelemetry = field(default_factory=RunTelemetry)
 
     def __post_init__(self):
         if self.flow is None and self.flow_run_id is None:
@@ -152,21 +150,6 @@ class BaseFlowRunEngine(Generic[P, R]):
     def state(self) -> State:
         return self.flow_run.state  # type: ignore
 
-    def _end_span_on_success(self):
-        if not self._span:
-            return
-        self._span.set_status(trace.Status(trace.StatusCode.OK))
-        self._span.end(time.time_ns())
-        self._span = None
-
-    def _end_span_on_error(self, exc: BaseException, description: Optional[str]):
-        if not self._span:
-            return
-        self._span.record_exception(exc)
-        self._span.set_status(trace.Status(trace.StatusCode.ERROR, description))
-        self._span.end(time.time_ns())
-        self._span = None
-
     def is_running(self) -> bool:
         if getattr(self, "flow_run", None) is None:
             return False
@@ -185,6 +168,7 @@ class BaseFlowRunEngine(Generic[P, R]):
         self, span: trace.Span, client: Union[SyncPrefectClient, PrefectClient]
     ):
         parent_flow_run_ctx = FlowRunContext.get()
+
         if parent_flow_run_ctx and parent_flow_run_ctx.flow_run:
             if traceparent := parent_flow_run_ctx.flow_run.labels.get(
                 LABELS_TRACEPARENT_KEY
@@ -194,6 +178,7 @@ class BaseFlowRunEngine(Generic[P, R]):
                     carrier={TRACEPARENT_KEY: traceparent},
                     setter=OTELSetter(),
                 )
+
             else:
                 carrier: KeyValueLabels = {}
                 propagate.get_global_textmap().inject(
@@ -315,16 +300,7 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
         self.flow_run.state_name = state.name  # type: ignore
         self.flow_run.state_type = state.type  # type: ignore
 
-        if self._span:
-            self._span.add_event(
-                state.name or state.type,
-                {
-                    "prefect.state.message": state.message or "",
-                    "prefect.state.type": state.type,
-                    "prefect.state.name": state.name or state.type,
-                    "prefect.state.id": str(state.id),
-                },
-            )
+        self._telemetry.update_state(state)
         return state
 
     def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
@@ -374,7 +350,7 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
         self.set_state(terminal_state)
         self._return_value = resolved_result
 
-        self._end_span_on_success()
+        self._telemetry.end_span_on_success()
 
         return result
 
@@ -406,8 +382,8 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
             )
             state = self.set_state(Running())
         self._raised = exc
-
-        self._end_span_on_error(exc, state.message)
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(state.message)
 
         return state
 
@@ -426,8 +402,8 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
         )
         self.set_state(state)
         self._raised = exc
-
-        self._end_span_on_error(exc, message)
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(message)
 
     def handle_crash(self, exc: BaseException) -> None:
         state = run_coro_as_sync(exception_to_crashed_state(exc))
@@ -435,8 +411,8 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
         self.logger.debug("Crash details:", exc_info=exc)
         self.set_state(state, force=True)
         self._raised = exc
-
-        self._end_span_on_error(exc, state.message if state else "")
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(state.message if state else None)
 
     def load_subflow_run(
         self,
@@ -681,19 +657,12 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
                     empirical_policy=self.flow_run.empirical_policy,
                 )
 
-            span = self._tracer.start_span(
-                name=self.flow_run.name,
-                attributes={
-                    **self.flow_run.labels,
-                    "prefect.run.type": "flow",
-                    "prefect.run.id": str(self.flow_run.id),
-                    "prefect.tags": self.flow_run.tags,
-                    "prefect.flow.name": self.flow.name,
-                },
+            self._telemetry.start_span(
+                name=self.flow.name,
+                run=self.flow_run,
+                client=self.client,
+                parameters=self.parameters,
             )
-            self._update_otel_labels(span, self.client)
-
-            self._span = span
 
             try:
                 yield self
@@ -736,7 +705,9 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
     @contextmanager
     def start(self) -> Generator[None, None, None]:
         with self.initialize_run():
-            with trace.use_span(self._span) if self._span else nullcontext():
+            with trace.use_span(
+                self._telemetry.span
+            ) if self._telemetry.span else nullcontext():
                 self.begin_run()
 
                 if self.state.is_running():
@@ -892,16 +863,7 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
         self.flow_run.state_name = state.name  # type: ignore
         self.flow_run.state_type = state.type  # type: ignore
 
-        if self._span:
-            self._span.add_event(
-                state.name or state.type,
-                {
-                    "prefect.state.message": state.message or "",
-                    "prefect.state.type": state.type,
-                    "prefect.state.name": state.name or state.type,
-                    "prefect.state.id": str(state.id),
-                },
-            )
+        self._telemetry.update_state(state)
         return state
 
     async def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
@@ -949,7 +911,7 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
         await self.set_state(terminal_state)
         self._return_value = resolved_result
 
-        self._end_span_on_success()
+        self._telemetry.end_span_on_success()
 
         return result
 
@@ -979,8 +941,8 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
             )
             state = await self.set_state(Running())
         self._raised = exc
-
-        self._end_span_on_error(exc, state.message)
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(state.message)
 
         return state
 
@@ -1000,7 +962,8 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
         await self.set_state(state)
         self._raised = exc
 
-        self._end_span_on_error(exc, message)
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(message)
 
     async def handle_crash(self, exc: BaseException) -> None:
         # need to shield from asyncio cancellation to ensure we update the state
@@ -1012,7 +975,8 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
             await self.set_state(state, force=True)
             self._raised = exc
 
-            self._end_span_on_error(exc, state.message)
+            self._telemetry.record_exception(exc)
+            self._telemetry.end_span_on_failure(state.message)
 
     async def load_subflow_run(
         self,
@@ -1255,18 +1219,12 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
                     empirical_policy=self.flow_run.empirical_policy,
                 )
 
-            span = self._tracer.start_span(
-                name=self.flow_run.name,
-                attributes={
-                    **self.flow_run.labels,
-                    "prefect.run.type": "flow",
-                    "prefect.run.id": str(self.flow_run.id),
-                    "prefect.tags": self.flow_run.tags,
-                    "prefect.flow.name": self.flow.name,
-                },
+            await self._telemetry.async_start_span(
+                name=self.flow.name,
+                run=self.flow_run,
+                client=self.client,
+                parameters=self.parameters,
             )
-            self._update_otel_labels(span, self.client)
-            self._span = span
 
             try:
                 yield self
@@ -1309,7 +1267,9 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
     @asynccontextmanager
     async def start(self) -> AsyncGenerator[None, None]:
         async with self.initialize_run():
-            with trace.use_span(self._span) if self._span else nullcontext():
+            with trace.use_span(
+                self._telemetry.span
+            ) if self._telemetry.span else nullcontext():
                 await self.begin_run()
 
                 if self.state.is_running():

prefect/flow_runs.py

@@ -1,6 +1,6 @@
 from typing import (
     TYPE_CHECKING,
-    Dict,
+    Any,
     Optional,
     Type,
     TypeVar,
@@ -430,7 +430,9 @@ async def suspend_flow_run(
 
 
 @sync_compatible
-async def resume_flow_run(flow_run_id, run_input: Optional[Dict] = None):
+async def resume_flow_run(
+    flow_run_id: UUID, run_input: Optional[dict[str, Any]] = None
+) -> None:
     """
     Resumes a paused flow.
 

prefect/flows.py

@@ -564,14 +564,12 @@ class Flow(Generic[P, R]):
                 "Cannot mix Pydantic v1 and v2 types as arguments to a flow."
             )
 
+        validated_fn_kwargs = dict(arbitrary_types_allowed=True)
+
         if has_v1_models:
-            validated_fn = V1ValidatedFunction(
-                self.fn, config={"arbitrary_types_allowed": True}
-            )
+            validated_fn = V1ValidatedFunction(self.fn, config=validated_fn_kwargs)
         else:
-            validated_fn = V2ValidatedFunction(
-                self.fn, config=pydantic.ConfigDict(arbitrary_types_allowed=True)
-            )
+            validated_fn = V2ValidatedFunction(self.fn, config=validated_fn_kwargs)
 
         try:
             with warnings.catch_warnings():

prefect/results.py

@@ -35,10 +35,13 @@ from pydantic import (
     model_validator,
 )
 from pydantic_core import PydanticUndefinedType
-from pydantic_extra_types.pendulum_dt import DateTime
 from typing_extensions import ParamSpec, Self
 
 import prefect
+from prefect._experimental.lineage import (
+    emit_result_read_event,
+    emit_result_write_event,
+)
 from prefect._internal.compatibility import deprecated
 from prefect._internal.compatibility.deprecated import deprecated_field
 from prefect.blocks.core import Block
@@ -57,6 +60,7 @@ from prefect.locking.protocol import LockManager
 from prefect.logging import get_logger
 from prefect.serializers import PickleSerializer, Serializer
 from prefect.settings.context import get_current_settings
+from prefect.types import DateTime
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import sync_compatible
 from prefect.utilities.pydantic import get_dispatch_key, lookup_type, register_base_type
@@ -129,7 +133,7 @@ async def resolve_result_storage(
     elif isinstance(result_storage, Path):
         storage_block = LocalFileSystem(basepath=str(result_storage))
     elif isinstance(result_storage, str):
-        storage_block = await Block.load(result_storage, client=client)
+        storage_block = await Block.aload(result_storage, client=client)
         storage_block_id = storage_block._block_document_id
         assert storage_block_id is not None, "Loaded storage blocks must have ids"
     elif isinstance(result_storage, UUID):
@@ -168,7 +172,7 @@ async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
     default_block = settings.tasks.scheduling.default_storage_block
 
     if default_block is not None:
-        return await Block.load(default_block)
+        return await Block.aload(default_block)
 
     # otherwise, use the local file system
     basepath = settings.results.local_storage_path
@@ -232,6 +236,10 @@ def _format_user_supplied_storage_key(key: str) -> str:
 T = TypeVar("T")
 
 
+def default_cache() -> LRUCache[str, "ResultRecord[Any]"]:
+    return LRUCache(maxsize=1000)
+
+
 def result_storage_discriminator(x: Any) -> str:
     if isinstance(x, dict):
         if "block_type_slug" in x:
@@ -284,7 +292,7 @@ class ResultStore(BaseModel):
     cache_result_in_memory: bool = Field(default=True)
     serializer: Serializer = Field(default_factory=get_default_result_serializer)
     storage_key_fn: Callable[[], str] = Field(default=DEFAULT_STORAGE_KEY_FN)
-    cache: LRUCache = Field(default_factory=lambda: LRUCache(maxsize=1000))
+    cache: LRUCache[str, "ResultRecord[Any]"] = Field(default_factory=default_cache)
 
     # Deprecated fields
     persist_result: Optional[bool] = Field(default=None)
@@ -319,7 +327,7 @@ class ResultStore(BaseModel):
         return self.model_copy(update=update)
 
     @sync_compatible
-    async def update_for_task(self: Self, task: "Task") -> Self:
+    async def update_for_task(self: Self, task: "Task[P, R]") -> Self:
         """
         Create a new result store for a task.
 
@@ -446,8 +454,15 @@ class ResultStore(BaseModel):
         """
         return await self._exists(key=key, _sync=False)
 
+    def _resolved_key_path(self, key: str) -> str:
+        if self.result_storage_block_id is None and hasattr(
+            self.result_storage, "_resolve_path"
+        ):
+            return str(self.result_storage._resolve_path(key))
+        return key
+
     @sync_compatible
-    async def _read(self, key: str, holder: str) -> "ResultRecord":
+    async def _read(self, key: str, holder: str) -> "ResultRecord[Any]":
         """
         Read a result record from storage.
 
@@ -465,8 +480,12 @@ class ResultStore(BaseModel):
         if self.lock_manager is not None and not self.is_lock_holder(key, holder):
             await self.await_for_lock(key)
 
-        if key in self.cache:
-            return self.cache[key]
+        resolved_key_path = self._resolved_key_path(key)
+
+        if resolved_key_path in self.cache:
+            cached_result = self.cache[resolved_key_path]
+            await emit_result_read_event(self, resolved_key_path, cached=True)
+            return cached_result
 
         if self.result_storage is None:
             self.result_storage = await get_default_result_storage()
@@ -478,31 +497,28 @@ class ResultStore(BaseModel):
                 metadata.storage_key is not None
             ), "Did not find storage key in metadata"
             result_content = await self.result_storage.read_path(metadata.storage_key)
-            result_record = ResultRecord.deserialize_from_result_and_metadata(
+            result_record: ResultRecord[
+                Any
+            ] = ResultRecord.deserialize_from_result_and_metadata(
                 result=result_content, metadata=metadata_content
             )
+            await emit_result_read_event(self, resolved_key_path)
         else:
             content = await self.result_storage.read_path(key)
-            result_record = ResultRecord.deserialize(
+            result_record: ResultRecord[Any] = ResultRecord.deserialize(
                 content, backup_serializer=self.serializer
             )
+            await emit_result_read_event(self, resolved_key_path)
 
         if self.cache_result_in_memory:
-            if self.result_storage_block_id is None and hasattr(
-                self.result_storage, "_resolve_path"
-            ):
-                cache_key = str(self.result_storage._resolve_path(key))
-            else:
-                cache_key = key
-
-            self.cache[cache_key] = result_record
+            self.cache[resolved_key_path] = result_record
         return result_record
 
     def read(
         self,
         key: str,
         holder: Optional[str] = None,
-    ) -> "ResultRecord":
+    ) -> "ResultRecord[Any]":
         """
         Read a result record from storage.
 
@@ -520,7 +536,7 @@ class ResultStore(BaseModel):
         self,
         key: str,
         holder: Optional[str] = None,
-    ) -> "ResultRecord":
+    ) -> "ResultRecord[Any]":
         """
         Read a result record from storage.
 
@@ -663,12 +679,13 @@ class ResultStore(BaseModel):
                 base_key,
                 content=result_record.serialize_metadata(),
             )
+            await emit_result_write_event(self, result_record.metadata.storage_key)
         # Otherwise, write the result metadata and result together
         else:
             await self.result_storage.write_path(
                 result_record.metadata.storage_key, content=result_record.serialize()
             )
-
+            await emit_result_write_event(self, result_record.metadata.storage_key)
         if self.cache_result_in_memory:
             self.cache[key] = result_record
 
@@ -898,7 +915,11 @@ class ResultStore(BaseModel):
         )
 
     @sync_compatible
-    async def read_parameters(self, identifier: UUID) -> Dict[str, Any]:
+    async def read_parameters(self, identifier: UUID) -> dict[str, Any]:
+        if self.result_storage is None:
+            raise ValueError(
+                "Result store is not configured - must have a result storage block to read parameters"
+            )
         record = ResultRecord.deserialize(
             await self.result_storage.read_path(f"parameters/{identifier}")
         )