prefect-client 3.0.0rc2__py3-none-any.whl → 3.0.0rc4__py3-none-any.whl

prefect/task_engine.py

@@ -3,8 +3,10 @@ import logging
 import time
 from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass, field
+from textwrap import dedent
 from typing import (
     Any,
+    AsyncGenerator,
     Callable,
     Coroutine,
     Dict,
@@ -25,7 +27,6 @@ import pendulum
 from typing_extensions import ParamSpec
 
 from prefect import Task
-from prefect._internal.concurrency.api import create_call, from_sync
 from prefect.client.orchestration import SyncPrefectClient
 from prefect.client.schemas import TaskRun
 from prefect.client.schemas.objects import State, TaskRunInput
@@ -43,7 +44,6 @@ from prefect.exceptions import (
     UpstreamTaskError,
 )
 from prefect.futures import PrefectFuture
-from prefect.logging.handlers import APILogHandler
 from prefect.logging.loggers import get_logger, patch_print, task_run_logger
 from prefect.records.result_store import ResultFactoryStore
 from prefect.results import ResultFactory, _format_user_supplied_storage_key
@@ -64,11 +64,12 @@ from prefect.states import (
 )
 from prefect.transactions import Transaction, transaction
 from prefect.utilities.asyncutils import run_coro_as_sync
-from prefect.utilities.callables import call_with_parameters
+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,
     propose_state_sync,
     resolve_to_final_result,
 )
@@ -219,7 +220,7 @@ class TaskRunEngine(Generic[P, R]):
             return_data=False,
             max_depth=-1,
             remove_annotations=True,
-            context={},
+            context={"current_task_run": self.task_run, "current_task": self.task},
         )
 
     def begin_run(self):
@@ -298,9 +299,17 @@ class TaskRunEngine(Generic[P, R]):
         if result_factory is None:
             raise ValueError("Result factory is not set")
 
+        if self.task.cache_expiration is not None:
+            expiration = pendulum.now("utc") + self.task.cache_expiration
+        else:
+            expiration = None
+
         terminal_state = run_coro_as_sync(
             return_value_to_state(
-                result, result_factory=result_factory, key=transaction.key
+                result,
+                result_factory=result_factory,
+                key=transaction.key,
+                expiration=expiration,
             )
         )
         transaction.stage(
@@ -333,10 +342,24 @@ class TaskRunEngine(Generic[P, R]):
                     scheduled_time=pendulum.now("utc").add(seconds=delay)
                 )
             else:
+                delay = None
                 new_state = Retrying()
+
+            self.logger.info(
+                f"Task run failed with exception {exc!r} - "
+                f"Retry {self.retries + 1}/{self.task.retries} will start "
+                f"{str(delay) + ' second(s) from now' if delay else 'immediately'}"
+            )
+
             self.set_state(new_state, force=True)
             self.retries = self.retries + 1
             return True
+        elif self.retries >= self.task.retries:
+            self.logger.error(
+                f"Task run failed with exception {exc!r} - Retries are exhausted"
+            )
+            return False
+
         return False
 
     def handle_exception(self, exc: Exception) -> None:
@@ -373,7 +396,7 @@ class TaskRunEngine(Generic[P, R]):
         self.set_state(state, force=True)
 
     @contextmanager
-    def enter_run_context(self, client: Optional[SyncPrefectClient] = None):
+    def setup_run_context(self, client: Optional[SyncPrefectClient] = None):
         from prefect.utilities.engine import (
             _resolve_custom_task_run_name,
             should_log_prints,
@@ -394,9 +417,7 @@ class TaskRunEngine(Generic[P, R]):
                     log_prints=log_prints,
                     task_run=self.task_run,
                     parameters=self.parameters,
-                    result_factory=run_coro_as_sync(
-                        ResultFactory.from_autonomous_task(self.task)
-                    ),  # type: ignore
+                    result_factory=run_coro_as_sync(ResultFactory.from_task(self.task)),  # type: ignore
                     client=client,
                 )
             )
@@ -444,9 +465,6 @@ class TaskRunEngine(Generic[P, R]):
                                 extra_task_inputs=dependencies,
                             )
                         )
-                        self.logger.info(
-                            f"Created task run {self.task_run.name!r} for task {self.task.name!r}"
-                        )
                     # Emit an event to capture that the task run was in the `PENDING` state.
                     self._last_event = emit_task_run_state_change_event(
                         task_run=self.task_run,
@@ -454,13 +472,20 @@ class TaskRunEngine(Generic[P, R]):
                         validated_state=self.task_run.state,
                     )
 
-                    yield self
+                    with self.setup_run_context():
+                        # setup_run_context might update the task run name, so log creation here
+                        self.logger.info(
+                            f"Created task run {self.task_run.name!r} for task {self.task.name!r}"
+                        )
+                        yield self
 
                 except Exception:
                     # regular exceptions are caught and re-raised to the user
                     raise
-                except (Pause, Abort):
+                except (Pause, Abort) as exc:
                     # Do not capture internal signals as crashes
+                    if isinstance(exc, Abort):
+                        self.logger.error("Task run was aborted: %s", exc)
                     raise
                 except GeneratorExit:
                     # Do not capture generator exits as crashes
@@ -474,18 +499,37 @@ class TaskRunEngine(Generic[P, R]):
                     display_state = (
                         repr(self.state) if PREFECT_DEBUG_MODE else str(self.state)
                     )
-                    self.logger.log(
-                        level=(
-                            logging.INFO if self.state.is_completed() else logging.ERROR
-                        ),
-                        msg=f"Finished in state {display_state}",
-                    )
+                    level = logging.INFO if self.state.is_completed() else logging.ERROR
+                    msg = f"Finished in state {display_state}"
+                    if self.state.is_pending():
+                        msg += (
+                            "\nPlease wait for all submitted tasks to complete"
+                            " before exiting your flow by calling `.wait()` on the "
+                            "`PrefectFuture` returned from your `.submit()` calls."
+                        )
+                        msg += dedent(
+                            """
+
+                            Example:
 
-                    # flush all logs if this is not a "top" level run
-                    if not (FlowRunContext.get() or TaskRunContext.get()):
-                        from_sync.call_soon_in_loop_thread(
-                            create_call(APILogHandler.aflush)
+                            from prefect import flow, task
+
+                            @task
+                            def say_hello(name):
+                                print f"Hello, {name}!"
+
+                            @flow
+                            def example_flow():
+                                say_hello.submit(name="Marvin)
+                                say_hello.wait()
+
+                            example_flow()
+                                      """
                         )
+                    self.logger.log(
+                        level=level,
+                        msg=msg,
+                    )
 
                     self._is_started = False
                     self._client = None
@@ -499,10 +543,8 @@ class TaskRunEngine(Generic[P, R]):
     async def wait_until_ready(self):
         """Waits until the scheduled time (if its the future), then enters Running."""
         if scheduled_time := self.state.state_details.scheduled_time:
-            self.logger.info(
-                f"Waiting for scheduled time {scheduled_time} for task {self.task.name!r}"
-            )
-            await anyio.sleep((scheduled_time - pendulum.now("utc")).total_seconds())
+            sleep_time = (scheduled_time - pendulum.now("utc")).total_seconds()
+            await anyio.sleep(sleep_time if sleep_time > 0 else 0)
             self.set_state(
                 Retrying() if self.state.name == "AwaitingRetry" else Running(),
                 force=True,
@@ -521,15 +563,11 @@ class TaskRunEngine(Generic[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 self.enter_run_context():
-                self.logger.debug(
-                    f"Executing task {self.task.name!r} for task run {self.task_run.name!r}..."
-                )
-                self.begin_run()
-                try:
-                    yield
-                finally:
-                    self.call_hooks()
+            self.begin_run()
+            try:
+                yield
+            finally:
+                self.call_hooks()
 
     @contextmanager
     def transaction_context(self) -> Generator[Transaction, None, None]:
@@ -552,9 +590,12 @@ class TaskRunEngine(Generic[P, R]):
     def run_context(self):
         timeout_context = timeout_async if self.task.isasync else timeout
         # reenter the run context to ensure it is up to date for every run
-        with self.enter_run_context():
+        with self.setup_run_context():
             try:
                 with timeout_context(seconds=self.task.timeout_seconds):
+                    self.logger.debug(
+                        f"Executing task {self.task.name!r} for task run {self.task_run.name!r}..."
+                    )
                     yield self
             except TimeoutError as exc:
                 self.handle_timeout(exc)
@@ -641,6 +682,117 @@ async def run_task_async(
     return engine.state if return_type == "state" else engine.result()
 
 
+def run_generator_task_sync(
+    task: Task[P, R],
+    task_run_id: Optional[UUID] = None,
+    task_run: Optional[TaskRun] = None,
+    parameters: Optional[Dict[str, Any]] = None,
+    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    return_type: Literal["state", "result"] = "result",
+    dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+    context: Optional[Dict[str, Any]] = None,
+) -> Generator[R, None, None]:
+    if return_type != "result":
+        raise ValueError("The return_type for a generator task must be 'result'")
+
+    engine = TaskRunEngine[P, R](
+        task=task,
+        parameters=parameters,
+        task_run=task_run,
+        wait_for=wait_for,
+        context=context,
+    )
+
+    with engine.start(task_run_id=task_run_id, dependencies=dependencies):
+        while engine.is_running():
+            run_coro_as_sync(engine.wait_until_ready())
+            with engine.run_context(), engine.transaction_context() as txn:
+                # TODO: generators should default to commit_mode=OFF
+                # because they are dynamic by definition
+                # for now we just prevent this branch explicitly
+                if False and txn.is_committed():
+                    txn.read()
+                else:
+                    call_args, call_kwargs = parameters_to_args_kwargs(
+                        task.fn, engine.parameters or {}
+                    )
+                    gen = task.fn(*call_args, **call_kwargs)
+                    try:
+                        while True:
+                            gen_result = next(gen)
+                            # link the current state to the result for dependency tracking
+                            #
+                            # TODO: this could grow the task_run_result
+                            # dictionary in an unbounded way, so finding a
+                            # way to periodically clean it up (using
+                            # weakrefs or similar) would be good
+                            link_state_to_result(engine.state, gen_result)
+                            yield gen_result
+                    except StopIteration as exc:
+                        engine.handle_success(exc.value, transaction=txn)
+                    except GeneratorExit as exc:
+                        engine.handle_success(None, transaction=txn)
+                        gen.throw(exc)
+
+    return engine.result()
+
+
+async def run_generator_task_async(
+    task: Task[P, R],
+    task_run_id: Optional[UUID] = None,
+    task_run: Optional[TaskRun] = None,
+    parameters: Optional[Dict[str, Any]] = None,
+    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    return_type: Literal["state", "result"] = "result",
+    dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+    context: Optional[Dict[str, Any]] = None,
+) -> AsyncGenerator[R, None]:
+    if return_type != "result":
+        raise ValueError("The return_type for a generator task must be 'result'")
+    engine = TaskRunEngine[P, R](
+        task=task,
+        parameters=parameters,
+        task_run=task_run,
+        wait_for=wait_for,
+        context=context,
+    )
+
+    with engine.start(task_run_id=task_run_id, dependencies=dependencies):
+        while engine.is_running():
+            await engine.wait_until_ready()
+            with engine.run_context(), engine.transaction_context() as txn:
+                # TODO: generators should default to commit_mode=OFF
+                # because they are dynamic by definition
+                # for now we just prevent this branch explicitly
+                if False and txn.is_committed():
+                    txn.read()
+                else:
+                    call_args, call_kwargs = parameters_to_args_kwargs(
+                        task.fn, engine.parameters or {}
+                    )
+                    gen = task.fn(*call_args, **call_kwargs)
+                    try:
+                        while True:
+                            # can't use anext in Python < 3.10
+                            gen_result = await gen.__anext__()
+                            # link the current state to the result for dependency tracking
+                            #
+                            # TODO: this could grow the task_run_result
+                            # dictionary in an unbounded way, so finding a
+                            # way to periodically clean it up (using
+                            # weakrefs or similar) would be good
+                            link_state_to_result(engine.state, gen_result)
+                            yield gen_result
+                    except (StopAsyncIteration, GeneratorExit) as exc:
+                        engine.handle_success(None, transaction=txn)
+                        if isinstance(exc, GeneratorExit):
+                            gen.throw(exc)
+
+    # async generators can't return, but we can raise failures here
+    if engine.state.is_failed():
+        engine.result()
+
+
 def run_task(
     task: Task[P, Union[R, Coroutine[Any, Any, R]]],
     task_run_id: Optional[UUID] = None,
@@ -680,7 +832,11 @@ def run_task(
         dependencies=dependencies,
         context=context,
     )
-    if task.isasync:
+    if task.isasync and task.isgenerator:
+        return run_generator_task_async(**kwargs)
+    elif task.isgenerator:
+        return run_generator_task_sync(**kwargs)
+    elif task.isasync:
         return run_task_async(**kwargs)
     else:
         return run_task_sync(**kwargs)

prefect/task_runners.py

@@ -202,12 +202,13 @@ class TaskRunner(abc.ABC, Generic[F]):
 
 
 class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
-    def __init__(self):
+    def __init__(self, max_workers: Optional[int] = None):
         super().__init__()
         self._executor: Optional[ThreadPoolExecutor] = None
+        self._max_workers = max_workers
 
     def duplicate(self) -> "ThreadPoolTaskRunner":
-        return type(self)()
+        return type(self)(max_workers=self._max_workers)
 
     def submit(
         self,
@@ -278,7 +279,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
 
     def __enter__(self):
         super().__enter__()
-        self._executor = ThreadPoolExecutor()
+        self._executor = ThreadPoolExecutor(max_workers=self._max_workers)
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
@@ -287,6 +288,11 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
             self._executor = None
         super().__exit__(exc_type, exc_value, traceback)
 
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, ThreadPoolTaskRunner):
+            return False
+        return self._max_workers == value._max_workers
+
 
 # Here, we alias ConcurrentTaskRunner to ThreadPoolTaskRunner for backwards compatibility
 ConcurrentTaskRunner = ThreadPoolTaskRunner

prefect/task_runs.py

@@ -71,7 +71,7 @@ class TaskRunWaiter:
         self.logger = get_logger("TaskRunWaiter")
         self._consumer_task: Optional[asyncio.Task] = None
         self._observed_completed_task_runs: TTLCache[uuid.UUID, bool] = TTLCache(
-            maxsize=100, ttl=60
+            maxsize=10000, ttl=600
         )
         self._completion_events: Dict[uuid.UUID, asyncio.Event] = {}
         self._loop: Optional[asyncio.AbstractEventLoop] = None
@@ -85,20 +85,25 @@ class TaskRunWaiter:
         """
         if self._started:
             return
-        self.logger.info("Starting TaskRunWaiter")
+        self.logger.debug("Starting TaskRunWaiter")
         loop_thread = get_global_loop()
 
         if not asyncio.get_running_loop() == loop_thread._loop:
             raise RuntimeError("TaskRunWaiter must run on the global loop thread.")
 
         self._loop = loop_thread._loop
-        self._consumer_task = self._loop.create_task(self._consume_events())
+
+        consumer_started = asyncio.Event()
+        self._consumer_task = self._loop.create_task(
+            self._consume_events(consumer_started)
+        )
+        asyncio.run_coroutine_threadsafe(consumer_started.wait(), self._loop)
 
         loop_thread.add_shutdown_call(create_call(self.stop))
         atexit.register(self.stop)
         self._started = True
 
-    async def _consume_events(self):
+    async def _consume_events(self, consumer_started: asyncio.Event):
         async with get_events_subscriber(
             filter=EventFilter(
                 event=EventNameFilter(
@@ -109,9 +114,10 @@ class TaskRunWaiter:
                 )
             )
         ) as subscriber:
+            consumer_started.set()
             async for event in subscriber:
                 try:
-                    self.logger.info(
+                    self.logger.debug(
                         f"Received event: {event.resource['prefect.resource.id']}"
                     )
                     task_run_id = uuid.UUID(
@@ -119,6 +125,7 @@ class TaskRunWaiter:
                             "prefect.task-run.", ""
                         )
                     )
+
                     with self._observed_completed_task_runs_lock:
                         # Cache the task run ID for a short period of time to avoid
                         # unnecessary waits
@@ -172,14 +179,21 @@ class TaskRunWaiter:
             # when the event is received
             instance._completion_events[task_run_id] = finished_event
 
-        with anyio.move_on_after(delay=timeout):
-            await from_async.wait_for_call_in_loop_thread(
-                create_call(finished_event.wait)
-            )
+        try:
+            # Now check one more time whether the task run arrived before we start to
+            # wait on it, in case it came in while we were setting up the event above.
+            with instance._observed_completed_task_runs_lock:
+                if task_run_id in instance._observed_completed_task_runs:
+                    return
 
-        with instance._completion_events_lock:
-            # Remove the event from the cache after it has been waited on
-            instance._completion_events.pop(task_run_id, None)
+            with anyio.move_on_after(delay=timeout):
+                await from_async.wait_for_call_in_loop_thread(
+                    create_call(finished_event.wait)
+                )
+        finally:
+            with instance._completion_events_lock:
+                # Remove the event from the cache after it has been waited on
+                instance._completion_events.pop(task_run_id, None)
 
     @classmethod
     def instance(cls):