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

prefect/task_engine.py

@@ -29,6 +29,7 @@ from uuid import UUID
 
 import anyio
 import pendulum
+from opentelemetry import trace
 from typing_extensions import ParamSpec
 
 from prefect import Task
@@ -79,13 +80,14 @@ from prefect.states import (
     exception_to_failed_state,
     return_value_to_state,
 )
+from prefect.telemetry.run_telemetry import RunTelemetry
 from prefect.transactions import IsolationLevel, Transaction, transaction
+from prefect.utilities._engine import get_hook_name
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.callables import call_with_parameters, parameters_to_args_kwargs
 from prefect.utilities.collections import visit_collection
 from prefect.utilities.engine import (
-    _get_hook_name,
     emit_task_run_state_change_event,
     link_state_to_result,
     resolve_to_final_result,
@@ -120,6 +122,7 @@ class BaseTaskRunEngine(Generic[P, R]):
     _is_started: bool = False
     _task_name_set: bool = False
     _last_event: Optional[PrefectEvent] = None
+    _telemetry: RunTelemetry = field(default_factory=RunTelemetry)
 
     def __post_init__(self):
         if self.parameters is None:
@@ -193,11 +196,11 @@ class BaseTaskRunEngine(Generic[P, R]):
         self.parameters = resolved_parameters
 
     def _set_custom_task_run_name(self):
-        from prefect.utilities.engine import _resolve_custom_task_run_name
+        from prefect.utilities._engine import resolve_custom_task_run_name
 
         # update the task run name if necessary
         if not self._task_name_set and self.task.task_run_name:
-            task_run_name = _resolve_custom_task_run_name(
+            task_run_name = resolve_custom_task_run_name(
                 task=self.task, parameters=self.parameters or {}
             )
 
@@ -351,7 +354,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[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(
@@ -427,6 +430,11 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
         self.task_run.state = new_state = state
 
+        if last_state.timestamp == new_state.timestamp:
+            # Ensure that the state timestamp is unique, or at least not equal to the last state.
+            # This might occur especially on Windows where the timestamp resolution is limited.
+            new_state.timestamp += pendulum.duration(microseconds=1)
+
         # Ensure that the state_details are populated with the current run IDs
         new_state.state_details.task_run_id = self.task_run.id
         new_state.state_details.flow_run_id = self.task_run.flow_run_id
@@ -460,7 +468,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             validated_state=self.task_run.state,
             follows=self._last_event,
         )
-
+        self._telemetry.update_state(new_state)
         return new_state
 
     def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
@@ -514,6 +522,8 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self.record_terminal_state_timing(terminal_state)
         self.set_state(terminal_state)
         self._return_value = result
+
+        self._telemetry.end_span_on_success(terminal_state.message)
         return result
 
     def handle_retry(self, exc: Exception) -> bool:
@@ -562,6 +572,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
     def handle_exception(self, exc: Exception) -> None:
         # If the task fails, and we have retries left, set the task to retrying.
+        self._telemetry.record_exception(exc)
         if not self.handle_retry(exc):
             # If the task has no retries left, or the retry condition is not met, set the task to failed.
             state = run_coro_as_sync(
@@ -575,6 +586,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             self.record_terminal_state_timing(state)
             self.set_state(state)
             self._raised = exc
+            self._telemetry.end_span_on_failure(state.message)
 
     def handle_timeout(self, exc: TimeoutError) -> None:
         if not self.handle_retry(exc):
@@ -598,6 +610,8 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self.record_terminal_state_timing(state)
         self.set_state(state, force=True)
         self._raised = exc
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(state.message)
 
     @contextmanager
     def setup_run_context(self, client: Optional[SyncPrefectClient] = None):
@@ -655,14 +669,17 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             with SyncClientContext.get_or_create() as client_ctx:
                 self._client = client_ctx.client
                 self._is_started = True
+                flow_run_context = FlowRunContext.get()
+                parent_task_run_context = TaskRunContext.get()
+
                 try:
                     if not self.task_run:
                         self.task_run = run_coro_as_sync(
                             self.task.create_local_run(
                                 id=task_run_id,
                                 parameters=self.parameters,
-                                flow_run_context=FlowRunContext.get(),
-                                parent_task_run_context=TaskRunContext.get(),
+                                flow_run_context=flow_run_context,
+                                parent_task_run_context=parent_task_run_context,
                                 wait_for=self.wait_for,
                                 extra_task_inputs=dependencies,
                             )
@@ -679,6 +696,13 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                         self.logger.debug(
                             f"Created task run {self.task_run.name!r} for task {self.task.name!r}"
                         )
+                        labels = (
+                            flow_run_context.flow_run.labels if flow_run_context else {}
+                        )
+                        self._telemetry.start_span(
+                            self.task_run, self.parameters, labels
+                        )
+
                         yield self
 
                 except TerminationSignal as exc:
@@ -730,11 +754,12 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     ) -> Generator[None, None, None]:
         with self.initialize_run(task_run_id=task_run_id, dependencies=dependencies):
-            self.begin_run()
-            try:
-                yield
-            finally:
-                self.call_hooks()
+            with trace.use_span(self._telemetry._span):
+                self.begin_run()
+                try:
+                    yield
+                finally:
+                    self.call_hooks()
 
     @contextmanager
     def transaction_context(self) -> Generator[Transaction, None, None]:
@@ -866,7 +891,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[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(
@@ -942,6 +967,11 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
         self.task_run.state = new_state = state
 
+        if last_state.timestamp == new_state.timestamp:
+            # Ensure that the state timestamp is unique, or at least not equal to the last state.
+            # This might occur especially on Windows where the timestamp resolution is limited.
+            new_state.timestamp += pendulum.duration(microseconds=1)
+
         # Ensure that the state_details are populated with the current run IDs
         new_state.state_details.task_run_id = self.task_run.id
         new_state.state_details.flow_run_id = self.task_run.flow_run_id
@@ -977,6 +1007,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             follows=self._last_event,
         )
 
+        self._telemetry.update_state(new_state)
         return new_state
 
     async def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
@@ -1025,6 +1056,9 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self.record_terminal_state_timing(terminal_state)
         await self.set_state(terminal_state)
         self._return_value = result
+
+        self._telemetry.end_span_on_success(terminal_state.message)
+
         return result
 
     async def handle_retry(self, exc: Exception) -> bool:
@@ -1073,6 +1107,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
     async def handle_exception(self, exc: Exception) -> None:
         # If the task fails, and we have retries left, set the task to retrying.
+        self._telemetry.record_exception(exc)
         if not await self.handle_retry(exc):
             # If the task has no retries left, or the retry condition is not met, set the task to failed.
             state = await exception_to_failed_state(
@@ -1084,7 +1119,10 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             await self.set_state(state)
             self._raised = exc
 
+            self._telemetry.end_span_on_failure(state.message)
+
     async def handle_timeout(self, exc: TimeoutError) -> None:
+        self._telemetry.record_exception(exc)
         if not await self.handle_retry(exc):
             if isinstance(exc, TaskRunTimeoutError):
                 message = f"Task run exceeded timeout of {self.task.timeout_seconds} second(s)"
@@ -1098,6 +1136,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             )
             await self.set_state(state)
             self._raised = exc
+            self._telemetry.end_span_on_failure(state.message)
 
     async def handle_crash(self, exc: BaseException) -> None:
         state = await exception_to_crashed_state(exc)
@@ -1107,6 +1146,9 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         await self.set_state(state, force=True)
         self._raised = exc
 
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(state.message)
+
     @asynccontextmanager
     async def setup_run_context(self, client: Optional[PrefectClient] = None):
         from prefect.utilities.engine import (
@@ -1162,12 +1204,14 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             async with AsyncClientContext.get_or_create():
                 self._client = get_client()
                 self._is_started = True
+                flow_run_context = FlowRunContext.get()
+
                 try:
                     if not self.task_run:
                         self.task_run = await self.task.create_local_run(
                             id=task_run_id,
                             parameters=self.parameters,
-                            flow_run_context=FlowRunContext.get(),
+                            flow_run_context=flow_run_context,
                             parent_task_run_context=TaskRunContext.get(),
                             wait_for=self.wait_for,
                             extra_task_inputs=dependencies,
@@ -1184,6 +1228,14 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                         self.logger.debug(
                             f"Created task run {self.task_run.name!r} for task {self.task.name!r}"
                         )
+
+                        labels = (
+                            flow_run_context.flow_run.labels if flow_run_context else {}
+                        )
+                        self._telemetry.start_span(
+                            self.task_run, self.parameters, labels
+                        )
+
                         yield self
 
                 except TerminationSignal as exc:
@@ -1237,11 +1289,12 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         async with self.initialize_run(
             task_run_id=task_run_id, dependencies=dependencies
         ):
-            await self.begin_run()
-            try:
-                yield
-            finally:
-                await self.call_hooks()
+            with trace.use_span(self._telemetry._span):
+                await self.begin_run()
+                try:
+                    yield
+                finally:
+                    await self.call_hooks()
 
     @asynccontextmanager
     async def transaction_context(self) -> AsyncGenerator[Transaction, None]:

prefect/task_runners.py

@@ -97,9 +97,9 @@ class TaskRunner(abc.ABC, Generic[F]):
 
     def map(
         self,
-        task: "Task",
+        task: "Task[P, R]",
         parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     ) -> PrefectFutureList[F]:
         """
         Submit multiple tasks to the task run engine.

prefect/tasks.py

@@ -21,6 +21,7 @@ from typing import (
     List,
     NoReturn,
     Optional,
+    Protocol,
     Set,
     Tuple,
     Type,
@@ -31,7 +32,7 @@ from typing import (
 )
 from uuid import UUID, uuid4
 
-from typing_extensions import Literal, ParamSpec
+from typing_extensions import Literal, ParamSpec, Self, TypeIs
 
 import prefect.states
 from prefect.cache_policies import DEFAULT, NONE, CachePolicy
@@ -223,6 +224,16 @@ def _generate_task_key(fn: Callable[..., Any]) -> str:
     return f"{qualname}-{code_hash}"
 
 
+class TaskRunNameCallbackWithParameters(Protocol):
+    @classmethod
+    def is_callback_with_parameters(cls, callable: Callable[..., str]) -> TypeIs[Self]:
+        sig = inspect.signature(callable)
+        return "parameters" in sig.parameters
+
+    def __call__(self, parameters: dict[str, Any]) -> str:
+        ...
+
+
 class Task(Generic[P, R]):
     """
     A Prefect task definition.
@@ -311,7 +322,7 @@ class Task(Generic[P, R]):
         ] = None,
         cache_expiration: Optional[datetime.timedelta] = None,
         task_run_name: Optional[
-            Union[Callable[[], str], Callable[[Dict[str, Any]], str], str]
+            Union[Callable[[], str], TaskRunNameCallbackWithParameters, str]
         ] = None,
         retries: Optional[int] = None,
         retry_delay_seconds: Optional[
@@ -534,7 +545,9 @@ class Task(Generic[P, R]):
             Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
         ] = None,
         task_run_name: Optional[
-            Union[Callable[[], str], Callable[[Dict[str, Any]], str], str, Type[NotSet]]
+            Union[
+                Callable[[], str], TaskRunNameCallbackWithParameters, str, Type[NotSet]
+            ]
         ] = NotSet,
         cache_expiration: Optional[datetime.timedelta] = None,
         retries: Union[int, Type[NotSet]] = NotSet,
@@ -732,10 +745,8 @@ class Task(Generic[P, R]):
         extra_task_inputs: Optional[Dict[str, Set[TaskRunInput]]] = None,
         deferred: bool = False,
     ) -> TaskRun:
-        from prefect.utilities.engine import (
-            _dynamic_key_for_task_run,
-            collect_task_run_inputs_sync,
-        )
+        from prefect.utilities._engine import dynamic_key_for_task_run
+        from prefect.utilities.engine import collect_task_run_inputs_sync
 
         if flow_run_context is None:
             flow_run_context = FlowRunContext.get()
@@ -751,7 +762,7 @@ class Task(Generic[P, R]):
                 dynamic_key = f"{self.task_key}-{str(uuid4().hex)}"
                 task_run_name = self.name
             else:
-                dynamic_key = _dynamic_key_for_task_run(
+                dynamic_key = dynamic_key_for_task_run(
                     context=flow_run_context, task=self
                 )
                 task_run_name = f"{self.name}-{dynamic_key}"
@@ -835,10 +846,8 @@ class Task(Generic[P, R]):
         extra_task_inputs: Optional[Dict[str, Set[TaskRunInput]]] = None,
         deferred: bool = False,
     ) -> TaskRun:
-        from prefect.utilities.engine import (
-            _dynamic_key_for_task_run,
-            collect_task_run_inputs_sync,
-        )
+        from prefect.utilities._engine import dynamic_key_for_task_run
+        from prefect.utilities.engine import collect_task_run_inputs_sync
 
         if flow_run_context is None:
             flow_run_context = FlowRunContext.get()
@@ -854,7 +863,7 @@ class Task(Generic[P, R]):
                 dynamic_key = f"{self.task_key}-{str(uuid4().hex)}"
                 task_run_name = self.name
             else:
-                dynamic_key = _dynamic_key_for_task_run(
+                dynamic_key = dynamic_key_for_task_run(
                     context=flow_run_context, task=self, stable=False
                 )
                 task_run_name = f"{self.name}-{dynamic_key[:3]}"
@@ -1179,7 +1188,7 @@ class Task(Generic[P, R]):
         self: "Task[P, R]",
         *args: Any,
         return_state: Literal[True],
-        wait_for: Optional[Iterable[Union[PrefectFuture[T], T]]] = ...,
+        wait_for: Optional[Iterable[Union[PrefectFuture[R], R]]] = ...,
         deferred: bool = ...,
         **kwargs: Any,
     ) -> List[State[R]]:
@@ -1189,7 +1198,7 @@ class Task(Generic[P, R]):
     def map(
         self: "Task[P, R]",
         *args: Any,
-        wait_for: Optional[Iterable[Union[PrefectFuture[T], T]]] = ...,
+        wait_for: Optional[Iterable[Union[PrefectFuture[R], R]]] = ...,
         deferred: bool = ...,
         **kwargs: Any,
     ) -> PrefectFutureList[R]:
@@ -1200,7 +1209,7 @@ class Task(Generic[P, R]):
         self: "Task[P, R]",
         *args: Any,
         return_state: Literal[True],
-        wait_for: Optional[Iterable[Union[PrefectFuture[T], T]]] = ...,
+        wait_for: Optional[Iterable[Union[PrefectFuture[R], R]]] = ...,
         deferred: bool = ...,
         **kwargs: Any,
     ) -> List[State[R]]:
@@ -1210,7 +1219,7 @@ class Task(Generic[P, R]):
     def map(
         self: "Task[P, R]",
         *args: Any,
-        wait_for: Optional[Iterable[Union[PrefectFuture[T], T]]] = ...,
+        wait_for: Optional[Iterable[Union[PrefectFuture[R], R]]] = ...,
         deferred: bool = ...,
         **kwargs: Any,
     ) -> PrefectFutureList[R]:
@@ -1221,7 +1230,7 @@ class Task(Generic[P, R]):
         self: "Task[P, Coroutine[Any, Any, R]]",
         *args: Any,
         return_state: Literal[True],
-        wait_for: Optional[Iterable[Union[PrefectFuture[T], T]]] = ...,
+        wait_for: Optional[Iterable[Union[PrefectFuture[R], R]]] = ...,
         deferred: bool = ...,
         **kwargs: Any,
     ) -> List[State[R]]:
@@ -1232,7 +1241,7 @@ class Task(Generic[P, R]):
         self: "Task[P, Coroutine[Any, Any, R]]",
         *args: Any,
         return_state: Literal[False],
-        wait_for: Optional[Iterable[Union[PrefectFuture[T], T]]] = ...,
+        wait_for: Optional[Iterable[Union[PrefectFuture[R], R]]] = ...,
         deferred: bool = ...,
         **kwargs: Any,
     ) -> PrefectFutureList[R]:
@@ -1242,10 +1251,10 @@ class Task(Generic[P, R]):
         self,
         *args: Any,
         return_state: bool = False,
-        wait_for: Optional[Iterable[Union[PrefectFuture[T], T]]] = None,
+        wait_for: Optional[Iterable[Union[PrefectFuture[R], R]]] = None,
         deferred: bool = False,
         **kwargs: Any,
-    ):
+    ) -> Union[List[State[R]], PrefectFutureList[R]]:
         """
         Submit a mapped run of the task to a worker.
 
@@ -1394,7 +1403,7 @@ class Task(Generic[P, R]):
                 " execution."
             )
         if return_state:
-            states = []
+            states: list[State[R]] = []
             for future in futures:
                 future.wait()
                 states.append(future.state)
@@ -1406,9 +1415,9 @@ class Task(Generic[P, R]):
         self,
         args: Optional[Tuple[Any, ...]] = None,
         kwargs: Optional[Dict[str, Any]] = None,
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
         dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
-    ) -> PrefectDistributedFuture:
+    ) -> PrefectDistributedFuture[R]:
         """
         Create a pending task run for a task worker to execute.
 
@@ -1588,7 +1597,7 @@ def task(
     ] = None,
     cache_expiration: Optional[datetime.timedelta] = None,
     task_run_name: Optional[
-        Union[Callable[[], str], Callable[[Dict[str, Any]], str], str]
+        Union[Callable[[], str], TaskRunNameCallbackWithParameters, str]
     ] = None,
     retries: int = 0,
     retry_delay_seconds: Union[
@@ -1606,16 +1615,18 @@ def task(
     timeout_seconds: Union[int, float, None] = None,
     log_prints: Optional[bool] = None,
     refresh_cache: Optional[bool] = None,
-    on_completion: Optional[List[Callable[["Task", TaskRun, State], None]]] = None,
-    on_failure: Optional[List[Callable[["Task", TaskRun, State], None]]] = None,
-    retry_condition_fn: Optional[Callable[["Task", TaskRun, State], bool]] = None,
+    on_completion: Optional[
+        List[Callable[["Task[P, R]", TaskRun, State], None]]
+    ] = None,
+    on_failure: Optional[List[Callable[["Task[P, R]", TaskRun, State], None]]] = None,
+    retry_condition_fn: Optional[Callable[["Task[P, R]", TaskRun, State], bool]] = None,
     viz_return_value: Any = None,
 ) -> Callable[[Callable[P, R]], Task[P, R]]:
     ...
 
 
 def task(
-    __fn=None,
+    __fn: Optional[Callable[P, R]] = None,
     *,
     name: Optional[str] = None,
     description: Optional[str] = None,
@@ -1627,7 +1638,7 @@ def task(
     ] = None,
     cache_expiration: Optional[datetime.timedelta] = None,
     task_run_name: Optional[
-        Union[Callable[[], str], Callable[[Dict[str, Any]], str], str]
+        Union[Callable[[], str], TaskRunNameCallbackWithParameters, str]
     ] = None,
     retries: Optional[int] = None,
     retry_delay_seconds: Union[
@@ -1642,9 +1653,11 @@ def task(
     timeout_seconds: Union[int, float, None] = None,
     log_prints: Optional[bool] = None,
     refresh_cache: Optional[bool] = None,
-    on_completion: Optional[List[Callable[["Task", TaskRun, State], None]]] = None,
-    on_failure: Optional[List[Callable[["Task", TaskRun, State], None]]] = None,
-    retry_condition_fn: Optional[Callable[["Task", TaskRun, State], bool]] = None,
+    on_completion: Optional[
+        List[Callable[["Task[P, R]", TaskRun, State], None]]
+    ] = None,
+    on_failure: Optional[List[Callable[["Task[P, R]", TaskRun, State], None]]] = None,
+    retry_condition_fn: Optional[Callable[["Task[P, R]", TaskRun, State], bool]] = None,
     viz_return_value: Any = None,
 ):
     """

prefect/telemetry/bootstrap.py

@@ -23,10 +23,23 @@ def setup_telemetry() -> (
     if server_type != ServerType.CLOUD:
         return None, None, None
 
-    assert settings.api.key
+    if not settings.api.key:
+        raise ValueError(
+            "A Prefect Cloud API key is required to enable telemetry. Please set "
+            "the `PREFECT_API_KEY` environment variable or authenticate with "
+            "Prefect Cloud via the `prefect cloud login` command."
+        )
+
     assert settings.api.url
 
     # This import is here to defer importing of the `opentelemetry` packages.
-    from .instrumentation import setup_exporters
+    try:
+        from .instrumentation import setup_exporters
+    except ImportError as exc:
+        raise ValueError(
+            "Unable to import OpenTelemetry instrumentation libraries. Please "
+            "ensure you have installed the `otel` extra when installing Prefect: "
+            "`pip install 'prefect[otel]'`"
+        ) from exc
 
     return setup_exporters(settings.api.url, settings.api.key.get_secret_value())

prefect/telemetry/run_telemetry.py

@@ -0,0 +1,107 @@
+import time
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+from opentelemetry.propagators.textmap import Setter
+from opentelemetry.trace import (
+    Status,
+    StatusCode,
+    get_tracer,
+)
+
+import prefect
+from prefect.client.schemas import TaskRun
+from prefect.client.schemas.objects import State
+from prefect.types import KeyValueLabels
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import Tracer
+
+
+class OTELSetter(Setter[KeyValueLabels]):
+    """
+    A setter for OpenTelemetry that supports Prefect's custom labels.
+    """
+
+    def set(self, carrier: KeyValueLabels, key: str, value: str) -> None:
+        carrier[key] = value
+
+
+@dataclass
+class RunTelemetry:
+    """
+    A class for managing the telemetry of runs.
+    """
+
+    _tracer: "Tracer" = field(
+        default_factory=lambda: get_tracer("prefect", prefect.__version__)
+    )
+    _span = None
+
+    def start_span(
+        self,
+        task_run: TaskRun,
+        parameters: Optional[Dict[str, Any]] = None,
+        labels: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Start a span for a task run.
+        """
+        if parameters is None:
+            parameters = {}
+        if labels is None:
+            labels = {}
+        parameter_attributes = {
+            f"prefect.run.parameter.{k}": type(v).__name__
+            for k, v in parameters.items()
+        }
+        self._span = self._tracer.start_span(
+            name=task_run.name,
+            attributes={
+                "prefect.run.type": "task",
+                "prefect.run.id": str(task_run.id),
+                "prefect.tags": task_run.tags,
+                **parameter_attributes,
+                **labels,
+            },
+        )
+
+    def end_span_on_success(self, terminal_message: str) -> None:
+        """
+        End a span for a task run on success.
+        """
+        if self._span:
+            self._span.set_status(Status(StatusCode.OK), terminal_message)
+            self._span.end(time.time_ns())
+            self._span = None
+
+    def end_span_on_failure(self, terminal_message: str) -> None:
+        """
+        End a span for a task run on failure.
+        """
+        if self._span:
+            self._span.set_status(Status(StatusCode.ERROR, terminal_message))
+            self._span.end(time.time_ns())
+            self._span = None
+
+    def record_exception(self, exc: Exception) -> None:
+        """
+        Record an exception on a span.
+        """
+        if self._span:
+            self._span.record_exception(exc)
+
+    def update_state(self, new_state: State) -> None:
+        """
+        Update a span with the state of a task run.
+        """
+        if self._span:
+            self._span.add_event(
+                new_state.name or new_state.type,
+                {
+                    "prefect.state.message": new_state.message or "",
+                    "prefect.state.type": new_state.type,
+                    "prefect.state.name": new_state.name or new_state.type,
+                    "prefect.state.id": str(new_state.id),
+                },
+            )