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

prefect/flow_engine.py

@@ -2,7 +2,7 @@ import asyncio
 import logging
 import os
 import time
-from contextlib import ExitStack, asynccontextmanager, contextmanager
+from contextlib import ExitStack, asynccontextmanager, contextmanager, nullcontext
 from dataclasses import dataclass, field
 from typing import (
     Any,
@@ -23,7 +23,7 @@ from typing import (
 from uuid import UUID
 
 from anyio import CancelScope
-from opentelemetry import trace
+from opentelemetry import propagate, trace
 from opentelemetry.trace import Tracer, get_tracer
 from typing_extensions import ParamSpec
 
@@ -72,6 +72,9 @@ from prefect.states import (
     exception_to_failed_state,
     return_value_to_state,
 )
+from prefect.telemetry.run_telemetry import OTELSetter
+from prefect.types import KeyValueLabels
+from prefect.utilities._engine import get_hook_name, resolve_custom_flow_run_name
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.callables import (
@@ -81,8 +84,6 @@ from prefect.utilities.callables import (
 )
 from prefect.utilities.collections import visit_collection
 from prefect.utilities.engine import (
-    _get_hook_name,
-    _resolve_custom_flow_run_name,
     capture_sigterm,
     link_state_to_result,
     propose_state,
@@ -94,6 +95,8 @@ from prefect.utilities.urls import url_for
 
 P = ParamSpec("P")
 R = TypeVar("R")
+LABELS_TRACEPARENT_KEY = "__OTEL_TRACEPARENT"
+TRACEPARENT_KEY = "traceparent"
 
 
 class FlowRunTimeoutError(TimeoutError):
@@ -178,6 +181,37 @@ class BaseFlowRunEngine(Generic[P, R]):
         if hasattr(self.flow.task_runner, "cancel_all"):
             self.flow.task_runner.cancel_all()  # type: ignore
 
+    def _update_otel_labels(
+        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
+            ):
+                carrier: KeyValueLabels = {TRACEPARENT_KEY: traceparent}
+                propagate.get_global_textmap().inject(
+                    carrier={TRACEPARENT_KEY: traceparent},
+                    setter=OTELSetter(),
+                )
+            else:
+                carrier: KeyValueLabels = {}
+                propagate.get_global_textmap().inject(
+                    carrier,
+                    context=trace.set_span_in_context(span),
+                    setter=OTELSetter(),
+                )
+            if carrier.get(TRACEPARENT_KEY):
+                if self.flow_run:
+                    client.update_flow_run_labels(
+                        flow_run_id=self.flow_run.id,
+                        labels={LABELS_TRACEPARENT_KEY: carrier[TRACEPARENT_KEY]},
+                    )
+                else:
+                    self.logger.info(
+                        f"Tried to set traceparent {carrier[TRACEPARENT_KEY]} for flow run, but None was found"
+                    )
+
 
 @dataclass
 class FlowRunEngine(BaseFlowRunEngine[P, R]):
@@ -283,7 +317,7 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
 
         if self._span:
             self._span.add_event(
-                state.name,
+                state.name or state.type,
                 {
                     "prefect.state.message": state.message or "",
                     "prefect.state.type": state.type,
@@ -402,7 +436,7 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
         self.set_state(state, force=True)
         self._raised = exc
 
-        self._end_span_on_error(exc, state.message)
+        self._end_span_on_error(exc, state.message if state else "")
 
     def load_subflow_run(
         self,
@@ -537,7 +571,7 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
             hooks = None
 
         for hook in hooks or []:
-            hook_name = _get_hook_name(hook)
+            hook_name = get_hook_name(hook)
 
             try:
                 self.logger.info(
@@ -600,7 +634,7 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
 
             # update the flow run name if necessary
             if not self._flow_run_name_set and self.flow.flow_run_name:
-                flow_run_name = _resolve_custom_flow_run_name(
+                flow_run_name = resolve_custom_flow_run_name(
                     flow=self.flow, parameters=self.parameters
                 )
                 self.client.set_flow_run_name(
@@ -647,7 +681,7 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
                     empirical_policy=self.flow_run.empirical_policy,
                 )
 
-            self._span = self._tracer.start_span(
+            span = self._tracer.start_span(
                 name=self.flow_run.name,
                 attributes={
                     **self.flow_run.labels,
@@ -657,6 +691,9 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
                     "prefect.flow.name": self.flow.name,
                 },
             )
+            self._update_otel_labels(span, self.client)
+
+            self._span = span
 
             try:
                 yield self
@@ -698,12 +735,13 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
 
     @contextmanager
     def start(self) -> Generator[None, None, None]:
-        with self.initialize_run(), trace.use_span(self._span):
-            self.begin_run()
+        with self.initialize_run():
+            with trace.use_span(self._span) if self._span else nullcontext():
+                self.begin_run()
 
-            if self.state.is_running():
-                self.call_hooks()
-            yield
+                if self.state.is_running():
+                    self.call_hooks()
+                yield
 
     @contextmanager
     def run_context(self):
@@ -856,7 +894,7 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
 
         if self._span:
             self._span.add_event(
-                state.name,
+                state.name or state.type,
                 {
                     "prefect.state.message": state.message or "",
                     "prefect.state.type": state.type,
@@ -1107,7 +1145,7 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
             hooks = None
 
         for hook in hooks or []:
-            hook_name = _get_hook_name(hook)
+            hook_name = get_hook_name(hook)
 
             try:
                 self.logger.info(
@@ -1170,7 +1208,7 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
 
             # update the flow run name if necessary
             if not self._flow_run_name_set and self.flow.flow_run_name:
-                flow_run_name = _resolve_custom_flow_run_name(
+                flow_run_name = resolve_custom_flow_run_name(
                     flow=self.flow, parameters=self.parameters
                 )
                 await self.client.set_flow_run_name(
@@ -1217,7 +1255,7 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
                     empirical_policy=self.flow_run.empirical_policy,
                 )
 
-            self._span = self._tracer.start_span(
+            span = self._tracer.start_span(
                 name=self.flow_run.name,
                 attributes={
                     **self.flow_run.labels,
@@ -1227,6 +1265,8 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
                     "prefect.flow.name": self.flow.name,
                 },
             )
+            self._update_otel_labels(span, self.client)
+            self._span = span
 
             try:
                 yield self
@@ -1269,7 +1309,7 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
     @asynccontextmanager
     async def start(self) -> AsyncGenerator[None, None]:
         async with self.initialize_run():
-            with trace.use_span(self._span):
+            with trace.use_span(self._span) if self._span else nullcontext():
                 await self.begin_run()
 
                 if self.state.is_running():
@@ -1392,7 +1432,7 @@ async def run_generator_flow_async(
     flow: Flow[P, R],
     flow_run: Optional[FlowRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     return_type: Literal["state", "result"] = "result",
 ) -> AsyncGenerator[R, None]:
     if return_type != "result":
@@ -1430,7 +1470,7 @@ def run_flow(
     flow: Flow[P, R],
     flow_run: Optional[FlowRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     return_type: Literal["state", "result"] = "result",
 ) -> Union[R, State, None]:
     kwargs = dict(

prefect/flow_runs.py

@@ -307,7 +307,7 @@ async def suspend_flow_run(
     flow_run_id: Optional[UUID] = None,
     timeout: Optional[int] = 3600,
     key: Optional[str] = None,
-    client: PrefectClient = None,
+    client: Optional[PrefectClient] = None,
 ) -> None:
     ...
 
@@ -318,7 +318,7 @@ async def suspend_flow_run(
     flow_run_id: Optional[UUID] = None,
     timeout: Optional[int] = 3600,
     key: Optional[str] = None,
-    client: PrefectClient = None,
+    client: Optional[PrefectClient] = None,
 ) -> T:
     ...
 
@@ -330,7 +330,7 @@ async def suspend_flow_run(
     flow_run_id: Optional[UUID] = None,
     timeout: Optional[int] = 3600,
     key: Optional[str] = None,
-    client: PrefectClient = None,
+    client: Optional[PrefectClient] = None,
 ) -> Optional[T]:
     """
     Suspends a flow run by stopping code execution until resumed.