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

prefect/task_engine.py

@@ -4,7 +4,7 @@ import logging
 import threading
 import time
 from asyncio import CancelledError
-from contextlib import ExitStack, asynccontextmanager, contextmanager
+from contextlib import ExitStack, asynccontextmanager, contextmanager, nullcontext
 from dataclasses import dataclass, field
 from functools import partial
 from textwrap import dedent
@@ -523,7 +523,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self.set_state(terminal_state)
         self._return_value = result
 
-        self._telemetry.end_span_on_success(terminal_state.message)
+        self._telemetry.end_span_on_success()
         return result
 
     def handle_retry(self, exc: Exception) -> bool:
@@ -586,7 +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)
+            self._telemetry.end_span_on_failure(state.message if state else None)
 
     def handle_timeout(self, exc: TimeoutError) -> None:
         if not self.handle_retry(exc):
@@ -600,6 +600,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 message=message,
                 name="TimedOut",
             )
+            self.record_terminal_state_timing(state)
             self.set_state(state)
             self._raised = exc
 
@@ -611,7 +612,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self.set_state(state, force=True)
         self._raised = exc
         self._telemetry.record_exception(exc)
-        self._telemetry.end_span_on_failure(state.message)
+        self._telemetry.end_span_on_failure(state.message if state else None)
 
     @contextmanager
     def setup_run_context(self, client: Optional[SyncPrefectClient] = None):
@@ -669,7 +670,7 @@ 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_flow_run_context = FlowRunContext.get()
                 parent_task_run_context = TaskRunContext.get()
 
                 try:
@@ -678,7 +679,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                             self.task.create_local_run(
                                 id=task_run_id,
                                 parameters=self.parameters,
-                                flow_run_context=flow_run_context,
+                                flow_run_context=parent_flow_run_context,
                                 parent_task_run_context=parent_task_run_context,
                                 wait_for=self.wait_for,
                                 extra_task_inputs=dependencies,
@@ -696,11 +697,12 @@ 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
+                            run=self.task_run,
+                            name=self.task.name,
+                            client=self.client,
+                            parameters=self.parameters,
                         )
 
                         yield self
@@ -754,7 +756,9 @@ 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):
-            with trace.use_span(self._telemetry._span):
+            with trace.use_span(
+                self._telemetry.span
+            ) if self._telemetry.span else nullcontext():
                 self.begin_run()
                 try:
                     yield
@@ -1057,7 +1061,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         await self.set_state(terminal_state)
         self._return_value = result
 
-        self._telemetry.end_span_on_success(terminal_state.message)
+        self._telemetry.end_span_on_success()
 
         return result
 
@@ -1134,6 +1138,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 message=message,
                 name="TimedOut",
             )
+            self.record_terminal_state_timing(state)
             await self.set_state(state)
             self._raised = exc
             self._telemetry.end_span_on_failure(state.message)
@@ -1204,15 +1209,16 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             async with AsyncClientContext.get_or_create():
                 self._client = get_client()
                 self._is_started = True
-                flow_run_context = FlowRunContext.get()
+                parent_flow_run_context = FlowRunContext.get()
+                parent_task_run_context = TaskRunContext.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=flow_run_context,
-                            parent_task_run_context=TaskRunContext.get(),
+                            flow_run_context=parent_flow_run_context,
+                            parent_task_run_context=parent_task_run_context,
                             wait_for=self.wait_for,
                             extra_task_inputs=dependencies,
                         )
@@ -1229,11 +1235,11 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                             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
+                        await self._telemetry.async_start_span(
+                            run=self.task_run,
+                            name=self.task.name,
+                            client=self.client,
+                            parameters=self.parameters,
                         )
 
                         yield self
@@ -1289,7 +1295,9 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         async with self.initialize_run(
             task_run_id=task_run_id, dependencies=dependencies
         ):
-            with trace.use_span(self._telemetry._span):
+            with trace.use_span(
+                self._telemetry.span
+            ) if self._telemetry.span else nullcontext():
                 await self.begin_run()
                 try:
                     yield
@@ -1370,7 +1378,7 @@ def run_task_sync(
     task_run_id: Optional[UUID] = None,
     task_run: Optional[TaskRun] = 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",
     dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     context: Optional[Dict[str, Any]] = None,
@@ -1397,7 +1405,7 @@ async def run_task_async(
     task_run_id: Optional[UUID] = None,
     task_run: Optional[TaskRun] = 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",
     dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     context: Optional[Dict[str, Any]] = None,
@@ -1424,7 +1432,7 @@ def run_generator_task_sync(
     task_run_id: Optional[UUID] = None,
     task_run: Optional[TaskRun] = 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",
     dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     context: Optional[Dict[str, Any]] = None,
@@ -1479,7 +1487,7 @@ async def run_generator_task_async(
     task_run_id: Optional[UUID] = None,
     task_run: Optional[TaskRun] = 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",
     dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     context: Optional[Dict[str, Any]] = None,
@@ -1535,7 +1543,7 @@ def run_task(
     task_run_id: Optional[UUID] = None,
     task_run: Optional[TaskRun] = 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",
     dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     context: Optional[Dict[str, Any]] = None,

prefect/task_worker.py

@@ -7,7 +7,7 @@ import sys
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import AsyncExitStack
 from contextvars import copy_context
-from typing import Optional
+from typing import TYPE_CHECKING, Any, Optional
 from uuid import UUID
 
 import anyio
@@ -16,6 +16,7 @@ import pendulum
 import uvicorn
 from exceptiongroup import BaseExceptionGroup  # novermin
 from fastapi import FastAPI
+from typing_extensions import ParamSpec, TypeVar
 from websockets.exceptions import InvalidStatusCode
 
 from prefect import Task
@@ -35,12 +36,17 @@ from prefect.task_engine import run_task_async, run_task_sync
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import asyncnullcontext, sync_compatible
 from prefect.utilities.engine import emit_task_run_state_change_event
-from prefect.utilities.processutils import _register_signal
+from prefect.utilities.processutils import (
+    _register_signal,  # pyright: ignore[reportPrivateUsage]
+)
 from prefect.utilities.services import start_client_metrics_server
 from prefect.utilities.urls import url_for
 
 logger = get_logger("task_worker")
 
+P = ParamSpec("P")
+R = TypeVar("R", infer_variance=True)
+
 
 class StopTaskWorker(Exception):
     """Raised when the task worker is stopped."""
@@ -48,8 +54,10 @@ class StopTaskWorker(Exception):
     pass
 
 
-def should_try_to_read_parameters(task: Task, task_run: TaskRun) -> bool:
+def should_try_to_read_parameters(task: Task[P, R], task_run: TaskRun) -> bool:
     """Determines whether a task run should read parameters from the result store."""
+    if TYPE_CHECKING:
+        assert task_run.state is not None
     new_enough_state_details = hasattr(
         task_run.state.state_details, "task_parameters_id"
     )
@@ -76,20 +84,23 @@ class TaskWorker:
 
     def __init__(
         self,
-        *tasks: Task,
+        *tasks: Task[P, R],
         limit: Optional[int] = 10,
     ):
-        self.tasks = []
+        self.tasks: list["Task[..., Any]"] = []
         for t in tasks:
-            if isinstance(t, Task):
-                if t.cache_policy in [None, NONE, NotSet]:
-                    self.tasks.append(
-                        t.with_options(persist_result=True, cache_policy=DEFAULT)
-                    )
-                else:
-                    self.tasks.append(t.with_options(persist_result=True))
+            if not TYPE_CHECKING:
+                if not isinstance(t, Task):
+                    continue
 
-        self.task_keys = set(t.task_key for t in tasks if isinstance(t, Task))
+            if t.cache_policy in [None, NONE, NotSet]:
+                self.tasks.append(
+                    t.with_options(persist_result=True, cache_policy=DEFAULT)
+                )
+            else:
+                self.tasks.append(t.with_options(persist_result=True))
+
+        self.task_keys = set(t.task_key for t in tasks if isinstance(t, Task))  # pyright: ignore[reportUnnecessaryIsInstance]
 
         self._started_at: Optional[pendulum.DateTime] = None
         self.stopping: bool = False
@@ -97,7 +108,9 @@ class TaskWorker:
         self._client = get_client()
         self._exit_stack = AsyncExitStack()
 
-        if not asyncio.get_event_loop().is_running():
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
             raise RuntimeError(
                 "TaskWorker must be initialized within an async context."
             )
@@ -141,7 +154,7 @@ class TaskWorker:
     def available_tasks(self) -> Optional[int]:
         return int(self._limiter.available_tokens) if self._limiter else None
 
-    def handle_sigterm(self, signum, frame):
+    def handle_sigterm(self, signum: int, frame: object):
         """
         Shuts down the task worker when a SIGTERM is received.
         """
@@ -252,6 +265,8 @@ class TaskWorker:
             self._release_token(task_run.id)
 
     async def _submit_scheduled_task_run(self, task_run: TaskRun):
+        if TYPE_CHECKING:
+            assert task_run.state is not None
         logger.debug(
             f"Found task run: {task_run.name!r} in state: {task_run.state.name!r}"
         )
@@ -280,7 +295,7 @@ class TaskWorker:
                 result_storage=await get_or_create_default_task_scheduling_storage()
             ).update_for_task(task)
             try:
-                run_data = await store.read_parameters(parameters_id)
+                run_data: dict[str, Any] = await store.read_parameters(parameters_id)
                 parameters = run_data.get("parameters", {})
                 wait_for = run_data.get("wait_for", [])
                 run_context = run_data.get("context", None)
@@ -350,7 +365,7 @@ class TaskWorker:
     async def __aenter__(self):
         logger.debug("Starting task worker...")
 
-        if self._client._closed:
+        if self._client._closed:  # pyright: ignore[reportPrivateUsage]
             self._client = get_client()
         self._runs_task_group = anyio.create_task_group()
 
@@ -362,7 +377,7 @@ class TaskWorker:
         self._started_at = pendulum.now()
         return self
 
-    async def __aexit__(self, *exc_info):
+    async def __aexit__(self, *exc_info: Any) -> None:
         logger.debug("Stopping task worker...")
         self._started_at = None
         await self._exit_stack.__aexit__(*exc_info)
@@ -372,7 +387,9 @@ def create_status_server(task_worker: TaskWorker) -> FastAPI:
     status_app = FastAPI()
 
     @status_app.get("/status")
-    def status():
+    def status():  # pyright: ignore[reportUnusedFunction]
+        if TYPE_CHECKING:
+            assert task_worker.started_at is not None
         return {
             "client_id": task_worker.client_id,
             "started_at": task_worker.started_at.isoformat(),
@@ -393,11 +410,13 @@ def create_status_server(task_worker: TaskWorker) -> FastAPI:
 
 @sync_compatible
 async def serve(
-    *tasks: Task, limit: Optional[int] = 10, status_server_port: Optional[int] = None
+    *tasks: Task[P, R],
+    limit: Optional[int] = 10,
+    status_server_port: Optional[int] = None,
 ):
-    """Serve the provided tasks so that their runs may be submitted to and executed.
-    in the engine. Tasks do not need to be within a flow run context to be submitted.
-    You must `.submit` the same task object that you pass to `serve`.
+    """Serve the provided tasks so that their runs may be submitted to
+    and executed in the engine. Tasks do not need to be within a flow run context to be
+    submitted. You must `.submit` the same task object that you pass to `serve`.
 
     Args:
         - tasks: A list of tasks to serve. When a scheduled task run is found for a
@@ -422,8 +441,7 @@ async def serve(
             print(message.upper())
 
         # starts a long-lived process that listens for scheduled runs of these tasks
-        if __name__ == "__main__":
-            serve(say, yell)
+        serve(say, yell)
         ```
     """
     task_worker = TaskWorker(*tasks, limit=limit)