prefect-client 3.0.0rc3__tar.gz → 3.0.0rc4__tar.gz

{prefect-client-3.0.0rc3/src/prefect_client.egg-info → prefect-client-3.0.0rc4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.0.0rc3
+Version: 3.0.0rc4
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.

{prefect-client-3.0.0rc3 → prefect-client-3.0.0rc4}/requirements-client.txt

@@ -21,8 +21,7 @@ pathspec >= 0.8.0
 pendulum >= 3.0.0, <4
 pydantic >= 2.7, < 3.0.0
 pydantic_core >= 2.12.0, < 3.0.0
-# Avoiding 2.8.1 due to https://github.com/pydantic/pydantic-extra-types/issues/188
-pydantic_extra_types > 2.7, != 2.8.1, < 3.0.0
+pydantic_extra_types >= 2.8.2, < 3.0.0
 pydantic_settings
 python_dateutil >= 2.8.2, < 3.0.0
 python-slugify >= 5.0, < 9.0

{prefect-client-3.0.0rc3 → prefect-client-3.0.0rc4}/requirements-dev.txt

@@ -22,7 +22,7 @@ pytest-benchmark
 pytest-env
 pytest-flakefinder
 pytest-timeout
-pytest-xdist < 3.4.0 # 3.5.0 introduces some flakiness. Need to investigate and resolve.
+pytest-xdist >= 3.6.1
 pytkdocs >= 0.14.2
 pyyaml
 redis>=5.0.1

{prefect-client-3.0.0rc3 → prefect-client-3.0.0rc4}/requirements.txt

@@ -13,7 +13,6 @@ griffe >= 0.20.0
 jinja2 >= 3.0.0, < 4.0.0
 jinja2-humanize-extension >= 0.4.0
 humanize >= 4.9.0
-kubernetes >= 24.2.0, < 30.0.0
 pytz >= 2021.1, < 2025
 readchar >= 4.0.0, < 5.0.0
 sqlalchemy[asyncio] >= 2.0, < 3.0.0

{prefect-client-3.0.0rc3 → prefect-client-3.0.0rc4}/src/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-3.0.0rc3 → prefect-client-3.0.0rc4}/src/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-client-3.0.0rc3 → prefect-client-3.0.0rc4}/src/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-client-3.0.0rc3 → prefect-client-3.0.0rc4}/src/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-client-3.0.0rc3 → prefect-client-3.0.0rc4}/src/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-client-3.0.0rc3 → prefect-client-3.0.0rc4}/src/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-client-3.0.0rc3 → prefect-client-3.0.0rc4}/src/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-client-3.0.0rc3 → prefect-client-3.0.0rc4}/src/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-client-3.0.0rc3 → prefect-client-3.0.0rc4}/src/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):