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

prefect/__init__.py

@@ -43,7 +43,6 @@ import prefect.runtime
 
 # Import modules that register types
 import prefect.serializers
-import prefect.blocks.kubernetes
 import prefect.blocks.notifications
 import prefect.blocks.system
 

prefect/client/subscriptions.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Any, Dict, Generic, List, Optional, Type, TypeVar
+from typing import Any, Dict, Generic, Iterable, Optional, Type, TypeVar
 
 import orjson
 import websockets
@@ -21,7 +21,7 @@ class Subscription(Generic[S]):
         self,
         model: Type[S],
         path: str,
-        keys: List[str],
+        keys: Iterable[str],
         client_id: Optional[str] = None,
         base_url: Optional[str] = None,
     ):
@@ -30,7 +30,7 @@ class Subscription(Generic[S]):
         base_url = base_url.replace("http", "ws", 1)
         self.subscription_url = f"{base_url}{path}"
 
-        self.keys = keys
+        self.keys = list(keys)
 
         self._connect = websockets.connect(
             self.subscription_url,

prefect/flow_engine.py

@@ -6,6 +6,7 @@ from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass, field
 from typing import (
     Any,
+    AsyncGenerator,
     Callable,
     Coroutine,
     Dict,
@@ -50,12 +51,13 @@ from prefect.states import (
     return_value_to_state,
 )
 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,
     _resolve_custom_flow_run_name,
     capture_sigterm,
+    link_state_to_result,
     propose_state_sync,
     resolve_to_final_result,
 )
@@ -632,6 +634,80 @@ async def run_flow_async(
     return engine.state if return_type == "state" else engine.result()
 
 
+def run_generator_flow_sync(
+    flow: Flow[P, R],
+    flow_run: Optional[FlowRun] = None,
+    parameters: Optional[Dict[str, Any]] = None,
+    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    return_type: Literal["state", "result"] = "result",
+) -> Generator[R, None, None]:
+    if return_type != "result":
+        raise ValueError("The return_type for a generator flow must be 'result'")
+
+    engine = FlowRunEngine[P, R](
+        flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
+    )
+
+    with engine.start():
+        while engine.is_running():
+            with engine.run_context():
+                call_args, call_kwargs = parameters_to_args_kwargs(
+                    flow.fn, engine.parameters or {}
+                )
+                gen = flow.fn(*call_args, **call_kwargs)
+                try:
+                    while True:
+                        gen_result = next(gen)
+                        # link the current state to the result for dependency tracking
+                        link_state_to_result(engine.state, gen_result)
+                        yield gen_result
+                except StopIteration as exc:
+                    engine.handle_success(exc.value)
+                except GeneratorExit as exc:
+                    engine.handle_success(None)
+                    gen.throw(exc)
+
+    return engine.result()
+
+
+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,
+    return_type: Literal["state", "result"] = "result",
+) -> AsyncGenerator[R, None]:
+    if return_type != "result":
+        raise ValueError("The return_type for a generator flow must be 'result'")
+
+    engine = FlowRunEngine[P, R](
+        flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
+    )
+
+    with engine.start():
+        while engine.is_running():
+            with engine.run_context():
+                call_args, call_kwargs = parameters_to_args_kwargs(
+                    flow.fn, engine.parameters or {}
+                )
+                gen = flow.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
+                        link_state_to_result(engine.state, gen_result)
+                        yield gen_result
+                except (StopAsyncIteration, GeneratorExit) as exc:
+                    engine.handle_success(None)
+                    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_flow(
     flow: Flow[P, R],
     flow_run: Optional[FlowRun] = None,
@@ -646,7 +722,11 @@ def run_flow(
         wait_for=wait_for,
         return_type=return_type,
     )
-    if flow.isasync:
+    if flow.isasync and flow.isgenerator:
+        return run_generator_flow_async(**kwargs)
+    elif flow.isgenerator:
+        return run_generator_flow_sync(**kwargs)
+    elif flow.isasync:
         return run_flow_async(**kwargs)
     else:
         return run_flow_sync(**kwargs)

prefect/flows.py

@@ -89,7 +89,6 @@ from prefect.task_runners import TaskRunner, ThreadPoolTaskRunner
 from prefect.types import BANNED_CHARACTERS, WITHOUT_BANNED_CHARACTERS
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import (
-    is_async_fn,
     run_sync_in_worker_thread,
     sync_compatible,
 )
@@ -289,7 +288,18 @@ class Flow(Generic[P, R]):
         self.description = description or inspect.getdoc(fn)
         update_wrapper(self, fn)
         self.fn = fn
-        self.isasync = is_async_fn(self.fn)
+
+        # the flow is considered async if its function is async or an async
+        # generator
+        self.isasync = inspect.iscoroutinefunction(
+            self.fn
+        ) or inspect.isasyncgenfunction(self.fn)
+
+        # the flow is considered a generator if its function is a generator or
+        # an async generator
+        self.isgenerator = inspect.isgeneratorfunction(
+            self.fn
+        ) or inspect.isasyncgenfunction(self.fn)
 
         raise_for_reserved_arguments(self.fn, ["return_state", "wait_for"])
 

prefect/futures.py

@@ -56,7 +56,7 @@ class PrefectFuture(abc.ABC):
     def wait(self, timeout: Optional[float] = None) -> None:
         ...
         """
-        Wait for the task run to complete. 
+        Wait for the task run to complete.
 
         If the task run has already completed, this method will return immediately.
 
@@ -163,6 +163,10 @@ class PrefectDistributedFuture(PrefectFuture):
             )
             return
 
+        # Ask for the instance of TaskRunWaiter _now_ so that it's already running and
+        # can catch the completion event if it happens before we start listening for it.
+        TaskRunWaiter.instance()
+
         # Read task run to see if it is still running
         async with get_client() as client:
             task_run = await client.read_task_run(task_run_id=self._task_run_id)
@@ -245,6 +249,10 @@ def resolve_futures_to_states(
         context={},
     )
 
+    # if no futures were found, return the original expression
+    if not futures:
+        return expr
+
     # Get final states for each future
     states = []
     for future in futures:

prefect/results.py

@@ -38,7 +38,6 @@ from prefect.settings import (
     PREFECT_RESULTS_DEFAULT_SERIALIZER,
     PREFECT_RESULTS_PERSIST_BY_DEFAULT,
     PREFECT_TASK_SCHEDULING_DEFAULT_STORAGE_BLOCK,
-    default_result_storage_block_name,
 )
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import sync_compatible
@@ -62,35 +61,15 @@ logger = get_logger("results")
 P = ParamSpec("P")
 R = TypeVar("R")
 
-
-@sync_compatible
-async def get_default_result_storage() -> ResultStorage:
-    """
-    Generate a default file system for result storage.
-    """
-    try:
-        return await Block.load(PREFECT_DEFAULT_RESULT_STORAGE_BLOCK.value())
-    except ValueError as e:
-        if "Unable to find" not in str(e):
-            raise e
-        elif (
-            PREFECT_DEFAULT_RESULT_STORAGE_BLOCK.value()
-            == default_result_storage_block_name()
-        ):
-            return LocalFileSystem(basepath=PREFECT_LOCAL_STORAGE_PATH.value())
-        else:
-            raise
+_default_storages: Dict[Tuple[str, str], WritableFileSystem] = {}
 
 
-_default_task_scheduling_storages: Dict[Tuple[str, str], WritableFileSystem] = {}
-
-
-async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
+async def _get_or_create_default_storage(block_document_slug: str) -> ResultStorage:
     """
-    Generate a default file system for background task parameter/result storage.
+    Generate a default file system for storage.
     """
     default_storage_name, storage_path = cache_key = (
-        PREFECT_TASK_SCHEDULING_DEFAULT_STORAGE_BLOCK.value(),
+        block_document_slug,
         PREFECT_LOCAL_STORAGE_PATH.value(),
     )
 
@@ -105,8 +84,8 @@ async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
         if block_type_slug == "local-file-system":
             block = LocalFileSystem(basepath=storage_path)
         else:
-            raise Exception(
-                "The default task storage block does not exist, but it is of type "
+            raise ValueError(
+                "The default storage block does not exist, but it is of type "
                 f"'{block_type_slug}' which cannot be created implicitly.  Please create "
                 "the block manually."
             )
@@ -123,13 +102,32 @@ async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
         return block
 
     try:
-        return _default_task_scheduling_storages[cache_key]
+        return _default_storages[cache_key]
     except KeyError:
         storage = await get_storage()
-        _default_task_scheduling_storages[cache_key] = storage
+        _default_storages[cache_key] = storage
         return storage
 
 
+@sync_compatible
+async def get_or_create_default_result_storage() -> ResultStorage:
+    """
+    Generate a default file system for result storage.
+    """
+    return await _get_or_create_default_storage(
+        PREFECT_DEFAULT_RESULT_STORAGE_BLOCK.value()
+    )
+
+
+async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
+    """
+    Generate a default file system for background task parameter/result storage.
+    """
+    return await _get_or_create_default_storage(
+        PREFECT_TASK_SCHEDULING_DEFAULT_STORAGE_BLOCK.value()
+    )
+
+
 def get_default_result_serializer() -> ResultSerializer:
     """
     Generate a default file system for result storage.
@@ -210,7 +208,9 @@ class ResultFactory(BaseModel):
                 kwargs.pop(key)
 
         # Apply defaults
-        kwargs.setdefault("result_storage", await get_default_result_storage())
+        kwargs.setdefault(
+            "result_storage", await get_or_create_default_result_storage()
+        )
         kwargs.setdefault("result_serializer", get_default_result_serializer())
         kwargs.setdefault("persist_result", get_default_persist_setting())
         kwargs.setdefault("cache_result_in_memory", True)
@@ -280,7 +280,9 @@ class ResultFactory(BaseModel):
         """
         Create a new result factory for a task.
         """
-        return await cls._from_task(task, get_default_result_storage, client=client)
+        return await cls._from_task(
+            task, get_or_create_default_result_storage, client=client
+        )
 
     @classmethod
     @inject_client

prefect/settings.py

@@ -1541,7 +1541,7 @@ The maximum number of retries to queue for submission.
 
 PREFECT_TASK_SCHEDULING_PENDING_TASK_TIMEOUT = Setting(
     timedelta,
-    default=timedelta(seconds=30),
+    default=timedelta(0),
 )
 """
 How long before a PENDING task are made available to another task worker.  In practice,

prefect/task_engine.py

@@ -417,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,
                 )
             )
@@ -467,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,
@@ -478,6 +473,10 @@ class TaskRunEngine(Generic[P, R]):
                     )
 
                     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:

prefect/task_runs.py

@@ -92,13 +92,18 @@ class TaskRunWaiter:
             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,6 +114,7 @@ class TaskRunWaiter:
                 )
             )
         ) as subscriber:
+            consumer_started.set()
             async for event in subscriber:
                 try:
                     self.logger.debug(
@@ -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):

prefect/task_worker.py

@@ -8,10 +8,14 @@ from concurrent.futures import ThreadPoolExecutor
 from contextlib import AsyncExitStack
 from contextvars import copy_context
 from typing import List, Optional
+from uuid import UUID
 
 import anyio
 import anyio.abc
+import pendulum
+import uvicorn
 from exceptiongroup import BaseExceptionGroup  # novermin
+from fastapi import FastAPI
 from websockets.exceptions import InvalidStatusCode
 
 from prefect import Task
@@ -73,8 +77,9 @@ class TaskWorker:
         limit: Optional[int] = 10,
     ):
         self.tasks: List[Task] = list(tasks)
+        self.task_keys = set(t.task_key for t in tasks if isinstance(t, Task))
 
-        self.started: bool = False
+        self._started_at: Optional[pendulum.DateTime] = None
         self.stopping: bool = False
 
         self._client = get_client()
@@ -89,10 +94,41 @@ class TaskWorker:
         self._executor = ThreadPoolExecutor(max_workers=limit if limit else None)
         self._limiter = anyio.CapacityLimiter(limit) if limit else None
 
+        self.in_flight_task_runs: dict[str, dict[UUID, pendulum.DateTime]] = {
+            task_key: {} for task_key in self.task_keys
+        }
+        self.finished_task_runs: dict[str, int] = {
+            task_key: 0 for task_key in self.task_keys
+        }
+
     @property
-    def _client_id(self) -> str:
+    def client_id(self) -> str:
         return f"{socket.gethostname()}-{os.getpid()}"
 
+    @property
+    def started_at(self) -> Optional[pendulum.DateTime]:
+        return self._started_at
+
+    @property
+    def started(self) -> bool:
+        return self._started_at is not None
+
+    @property
+    def limit(self) -> Optional[int]:
+        return int(self._limiter.total_tokens) if self._limiter else None
+
+    @property
+    def current_tasks(self) -> Optional[int]:
+        return (
+            int(self._limiter.borrowed_tokens)
+            if self._limiter
+            else sum(len(runs) for runs in self.in_flight_task_runs.values())
+        )
+
+    @property
+    def available_tasks(self) -> Optional[int]:
+        return int(self._limiter.available_tokens) if self._limiter else None
+
     def handle_sigterm(self, signum, frame):
         """
         Shuts down the task worker when a SIGTERM is received.
@@ -133,11 +169,31 @@ class TaskWorker:
                 " calling .start()"
             )
 
-        self.started = False
+        self._started_at = None
         self.stopping = True
 
         raise StopTaskWorker
 
+    async def _acquire_token(self, task_run_id: UUID) -> bool:
+        try:
+            if self._limiter:
+                await self._limiter.acquire_on_behalf_of(task_run_id)
+        except RuntimeError:
+            logger.debug(f"Token already acquired for task run: {task_run_id!r}")
+            return False
+
+        return True
+
+    def _release_token(self, task_run_id: UUID) -> bool:
+        try:
+            if self._limiter:
+                self._limiter.release_on_behalf_of(task_run_id)
+        except RuntimeError:
+            logger.debug(f"No token to release for task run: {task_run_id!r}")
+            return False
+
+        return True
+
     async def _subscribe_to_task_scheduling(self):
         base_url = PREFECT_API_URL.value()
         if base_url is None:
@@ -146,24 +202,26 @@ class TaskWorker:
                 "Task workers are not compatible with the ephemeral API."
             )
         task_keys_repr = " | ".join(
-            t.task_key.split(".")[-1].split("-")[0] for t in self.tasks
+            task_key.split(".")[-1].split("-")[0] for task_key in sorted(self.task_keys)
         )
         logger.info(f"Subscribing to runs of task(s): {task_keys_repr}")
         async for task_run in Subscription(
             model=TaskRun,
             path="/task_runs/subscriptions/scheduled",
-            keys=[task.task_key for task in self.tasks],
-            client_id=self._client_id,
+            keys=self.task_keys,
+            client_id=self.client_id,
             base_url=base_url,
         ):
             logger.info(f"Received task run: {task_run.id} - {task_run.name}")
-            if self._limiter:
-                await self._limiter.acquire_on_behalf_of(task_run.id)
-            self._runs_task_group.start_soon(
-                self._safe_submit_scheduled_task_run, task_run
-            )
+
+            token_acquired = await self._acquire_token(task_run.id)
+            if token_acquired:
+                self._runs_task_group.start_soon(
+                    self._safe_submit_scheduled_task_run, task_run
+                )
 
     async def _safe_submit_scheduled_task_run(self, task_run: TaskRun):
+        self.in_flight_task_runs[task_run.task_key][task_run.id] = pendulum.now()
         try:
             await self._submit_scheduled_task_run(task_run)
         except BaseException as exc:
@@ -172,8 +230,9 @@ class TaskWorker:
                 exc_info=exc,
             )
         finally:
-            if self._limiter:
-                self._limiter.release_on_behalf_of(task_run.id)
+            self.in_flight_task_runs[task_run.task_key].pop(task_run.id, None)
+            self.finished_task_runs[task_run.task_key] += 1
+            self._release_token(task_run.id)
 
     async def _submit_scheduled_task_run(self, task_run: TaskRun):
         logger.debug(
@@ -284,9 +343,9 @@ class TaskWorker:
     async def execute_task_run(self, task_run: TaskRun):
         """Execute a task run in the task worker."""
         async with self if not self.started else asyncnullcontext():
-            if self._limiter:
-                await self._limiter.acquire_on_behalf_of(task_run.id)
-            await self._safe_submit_scheduled_task_run(task_run)
+            token_acquired = await self._acquire_token(task_run.id)
+            if token_acquired:
+                await self._safe_submit_scheduled_task_run(task_run)
 
     async def __aenter__(self):
         logger.debug("Starting task worker...")
@@ -298,17 +357,42 @@ class TaskWorker:
         await self._exit_stack.enter_async_context(self._runs_task_group)
         self._exit_stack.enter_context(self._executor)
 
-        self.started = True
+        self._started_at = pendulum.now()
         return self
 
     async def __aexit__(self, *exc_info):
         logger.debug("Stopping task worker...")
-        self.started = False
+        self._started_at = None
         await self._exit_stack.__aexit__(*exc_info)
 
 
+def create_status_server(task_worker: TaskWorker) -> FastAPI:
+    status_app = FastAPI()
+
+    @status_app.get("/status")
+    def status():
+        return {
+            "client_id": task_worker.client_id,
+            "started_at": task_worker.started_at.isoformat(),
+            "stopping": task_worker.stopping,
+            "limit": task_worker.limit,
+            "current": task_worker.current_tasks,
+            "available": task_worker.available_tasks,
+            "tasks": sorted(task_worker.task_keys),
+            "finished": task_worker.finished_task_runs,
+            "in_flight": {
+                key: {str(run): start.isoformat() for run, start in tasks.items()}
+                for key, tasks in task_worker.in_flight_task_runs.items()
+            },
+        }
+
+    return status_app
+
+
 @sync_compatible
-async def serve(*tasks: Task, limit: Optional[int] = 10):
+async def serve(
+    *tasks: Task, 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`.
@@ -318,6 +402,9 @@ async def serve(*tasks: Task, limit: Optional[int] = 10):
             given task, the task run will be submitted to the engine for execution.
         - limit: The maximum number of tasks that can be run concurrently. Defaults to 10.
             Pass `None` to remove the limit.
+        - status_server_port: An optional port on which to start an HTTP server
+            exposing status information about the task worker. If not provided, no
+            status server will run.
 
     Example:
         ```python
@@ -339,6 +426,20 @@ async def serve(*tasks: Task, limit: Optional[int] = 10):
     """
     task_worker = TaskWorker(*tasks, limit=limit)
 
+    status_server_task = None
+    if status_server_port is not None:
+        server = uvicorn.Server(
+            uvicorn.Config(
+                app=create_status_server(task_worker),
+                host="127.0.0.1",
+                port=status_server_port,
+                access_log=False,
+                log_level="warning",
+            )
+        )
+        loop = asyncio.get_event_loop()
+        status_server_task = loop.create_task(server.serve())
+
     try:
         await task_worker.start()
 
@@ -355,3 +456,11 @@ async def serve(*tasks: Task, limit: Optional[int] = 10):
 
     except (asyncio.CancelledError, KeyboardInterrupt):
         logger.info("Task worker interrupted, stopping...")
+
+    finally:
+        if status_server_task:
+            status_server_task.cancel()
+            try:
+                await status_server_task
+            except asyncio.CancelledError:
+                pass