prefect-client 2.18.1__py3-none-any.whl → 2.18.3__py3-none-any.whl

prefect/flows.py

@@ -733,7 +733,7 @@ class Flow(Generic[P, R]):
     @sync_compatible
     async def serve(
         self,
-        name: str,
+        name: Optional[str] = None,
         interval: Optional[
             Union[
                 Iterable[Union[int, float, datetime.timedelta]],
@@ -764,7 +764,7 @@ class Flow(Generic[P, R]):
         Creates a deployment for this flow and starts a runner to monitor for scheduled work.
 
         Args:
-            name: The name to give the created deployment.
+            name: The name to give the created deployment. Defaults to the name of the flow.
             interval: An interval on which to execute the deployment. Accepts a number or a
                 timedelta object to create a single schedule. If a number is given, it will be
                 interpreted as seconds. Also accepts an iterable of numbers or timedelta to create
@@ -827,10 +827,13 @@ class Flow(Generic[P, R]):
         """
         from prefect.runner import Runner
 
-        # Handling for my_flow.serve(__file__)
-        # Will set name to name of file where my_flow.serve() without the extension
-        # Non filepath strings will pass through unchanged
-        name = Path(name).stem
+        if not name:
+            name = self.name
+        else:
+            # Handling for my_flow.serve(__file__)
+            # Will set name to name of file where my_flow.serve() without the extension
+            # Non filepath strings will pass through unchanged
+            name = Path(name).stem
 
         runner = Runner(name=name, pause_on_shutdown=pause_on_shutdown, limit=limit)
         deployment_id = await runner.add_flow(
@@ -1226,19 +1229,19 @@ class Flow(Generic[P, R]):
             return track_viz_task(self.isasync, self.name, parameters)
 
         if PREFECT_EXPERIMENTAL_ENABLE_NEW_ENGINE.value():
-            from prefect.new_flow_engine import run_flow
-            from prefect.utilities.asyncutils import run_sync
+            from prefect.new_flow_engine import run_flow, run_flow_sync
 
-            awaitable = run_flow(
+            run_kwargs = dict(
                 flow=self,
                 parameters=parameters,
                 wait_for=wait_for,
                 return_type=return_type,
             )
             if self.isasync:
-                return awaitable
+                # this returns an awaitable coroutine
+                return run_flow(**run_kwargs)
             else:
-                return run_sync(awaitable)
+                return run_flow_sync(**run_kwargs)
 
         return enter_flow_run_engine_from_flow_call(
             self,

prefect/input/run_input.py

@@ -582,7 +582,9 @@ def receive_input(
     # the signature is the same as here:
     #   Union[Type[R], Type[T], pydantic.BaseModel],
     # Seems like a possible mypy bug, so we'll ignore the type check here.
-    input_cls = run_input_subclass_from_type(input_type)  # type: ignore[arg-type]
+    input_cls: Union[
+        Type[AutomaticRunInput[T]], Type[R]
+    ] = run_input_subclass_from_type(input_type)  # type: ignore[arg-type]
 
     if issubclass(input_cls, AutomaticRunInput):
         return input_cls.receive(

prefect/new_flow_engine.py

@@ -1,5 +1,7 @@
 import asyncio
-from contextlib import asynccontextmanager
+import inspect
+import os
+from contextlib import AsyncExitStack, asynccontextmanager, contextmanager
 from dataclasses import dataclass
 from typing import (
     Any,
@@ -9,20 +11,26 @@ from typing import (
     Iterable,
     Literal,
     Optional,
+    Tuple,
     TypeVar,
     Union,
     cast,
 )
+from uuid import UUID
 
 import anyio
+import anyio._backends._asyncio
+from sniffio import AsyncLibraryNotFoundError
 from typing_extensions import ParamSpec
 
-from prefect import Flow, Task, get_client
+from prefect import Task, get_client
 from prefect.client.orchestration import PrefectClient
 from prefect.client.schemas import FlowRun, TaskRun
 from prefect.client.schemas.filters import FlowRunFilter
 from prefect.client.schemas.sorting import FlowRunSort
 from prefect.context import FlowRunContext
+from prefect.deployments import load_flow_from_flow_run
+from prefect.flows import Flow, load_flow_from_entrypoint
 from prefect.futures import PrefectFuture, resolve_futures_to_states
 from prefect.logging.loggers import flow_run_logger
 from prefect.results import ResultFactory
@@ -33,7 +41,8 @@ from prefect.states import (
     exception_to_failed_state,
     return_value_to_state,
 )
-from prefect.utilities.asyncutils import A, Async
+from prefect.utilities.asyncutils import A, Async, run_sync
+from prefect.utilities.callables import parameters_to_args_kwargs
 from prefect.utilities.engine import (
     _dynamic_key_for_task_run,
     _resolve_custom_flow_run_name,
@@ -45,16 +54,34 @@ P = ParamSpec("P")
 R = TypeVar("R")
 
 
+async def load_flow_and_flow_run(flow_run_id: UUID) -> Tuple[FlowRun, Flow]:
+    ## TODO: add error handling to update state and log tracebacks
+    entrypoint = os.environ.get("PREFECT__FLOW_ENTRYPOINT")
+
+    async with get_client() as client:
+        flow_run = await client.read_flow_run(flow_run_id)
+        flow = (
+            load_flow_from_entrypoint(entrypoint)
+            if entrypoint
+            else await load_flow_from_flow_run(flow_run, client=client)
+        )
+    return flow_run, flow
+
+
 @dataclass
 class FlowRunEngine(Generic[P, R]):
-    flow: Flow[P, Coroutine[Any, Any, R]]
+    flow: Optional[Union[Flow[P, R], Flow[P, Coroutine[Any, Any, R]]]] = None
     parameters: Optional[Dict[str, Any]] = None
     flow_run: Optional[FlowRun] = None
+    flow_run_id: Optional[UUID] = None
     _is_started: bool = False
     _client: Optional[PrefectClient] = None
     short_circuit: bool = False
 
     def __post_init__(self):
+        if self.flow is None and self.flow_run_id is None:
+            raise ValueError("Either a flow or a flow_run_id must be provided.")
+
         if self.parameters is None:
             self.parameters = {}
 
@@ -89,10 +116,17 @@ class FlowRunEngine(Generic[P, R]):
         return state
 
     async def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
-        return await self.state.result(raise_on_failure=raise_on_failure, fetch=True)
+        _result = self.state.result(raise_on_failure=raise_on_failure, fetch=True)  # type: ignore
+        # state.result is a `sync_compatible` function that may or may not return an awaitable
+        # depending on whether the parent frame is sync or not
+        if inspect.isawaitable(_result):
+            _result = await _result
+        return _result
 
     async def handle_success(self, result: R) -> R:
         result_factory = getattr(FlowRunContext.get(), "result_factory", None)
+        if result_factory is None:
+            raise ValueError("Result factory is not set")
         terminal_state = await return_value_to_state(
             await resolve_futures_to_states(result),
             result_factory=result_factory,
@@ -117,71 +151,102 @@ class FlowRunEngine(Generic[P, R]):
             state = await self.set_state(Running())
         return state
 
+    async def load_subflow_run(
+        self, parent_task_run: TaskRun, client: PrefectClient, context: FlowRunContext
+    ) -> Union[FlowRun, None]:
+        """
+        This method attempts to load an existing flow run for a subflow task
+        run, if appropriate.
+
+        If the parent task run is in a final but not COMPLETED state, and not
+        being rerun, then we attempt to load an existing flow run instead of
+        creating a new one. This will prevent the engine from running the
+        subflow again.
+
+        If no existing flow run is found, or if the subflow should be rerun,
+        then no flow run is returned.
+        """
+
+        # check if the parent flow run is rerunning
+        rerunning = (
+            context.flow_run.run_count > 1
+            if getattr(context, "flow_run", None)
+            and isinstance(context.flow_run, FlowRun)
+            else False
+        )
+
+        # if the parent task run is in a final but not completed state, and
+        # not rerunning, then retrieve the most recent flow run instead of
+        # creating a new one. This effectively loads a cached flow run for
+        # situations where we are confident the flow should not be run
+        # again.
+        assert isinstance(parent_task_run.state, State)
+        if parent_task_run.state.is_final() and not (
+            rerunning and not parent_task_run.state.is_completed()
+        ):
+            # return the most recent flow run, if it exists
+            flow_runs = await client.read_flow_runs(
+                flow_run_filter=FlowRunFilter(
+                    parent_task_run_id={"any_": [parent_task_run.id]}
+                ),
+                sort=FlowRunSort.EXPECTED_START_TIME_ASC,
+                limit=1,
+            )
+            if flow_runs:
+                return flow_runs[-1]
+
     async def create_subflow_task_run(
         self, client: PrefectClient, context: FlowRunContext
     ) -> TaskRun:
+        """
+        Adds a task to a parent flow run that represents the execution of a subflow run.
+
+        The task run is referred to as the "parent task run" of the subflow and will be kept
+        in sync with the subflow run's state by the orchestration engine.
+        """
         dummy_task = Task(
             name=self.flow.name, fn=self.flow.fn, version=self.flow.version
         )
         task_inputs = {
-            k: await collect_task_run_inputs(v) for k, v in self.parameters.items()
+            k: await collect_task_run_inputs(v)
+            for k, v in (self.parameters or {}).items()
         }
         parent_task_run = await client.create_task_run(
             task=dummy_task,
             flow_run_id=(
-                context.flow_run.id if getattr(context, "flow_run", None) else None
+                context.flow_run.id
+                if getattr(context, "flow_run", None)
+                and isinstance(context.flow_run, FlowRun)
+                else None
             ),
-            dynamic_key=_dynamic_key_for_task_run(context, dummy_task),
-            task_inputs=task_inputs,
+            dynamic_key=_dynamic_key_for_task_run(context, dummy_task),  # type: ignore
+            task_inputs=task_inputs,  # type: ignore
             state=Pending(),
         )
         return parent_task_run
 
-    async def get_most_recent_flow_run_for_parent_task_run(
-        self, client: PrefectClient, parent_task_run: TaskRun
-    ) -> "Union[FlowRun, None]":
-        """
-        Get the most recent flow run associated with the provided parent task run.
-
-        Args:
-            - An orchestration client
-            - The parent task run to get the most recent flow run for
-
-        Returns:
-            The most recent flow run associated with the parent task run or `None` if
-            no flow runs are found
-        """
-        flow_runs = await client.read_flow_runs(
-            flow_run_filter=FlowRunFilter(
-                parent_task_run_id={"any_": [parent_task_run.id]}
-            ),
-            sort=FlowRunSort.EXPECTED_START_TIME_ASC,
-        )
-        return flow_runs[-1] if flow_runs else None
-
     async def create_flow_run(self, client: PrefectClient) -> FlowRun:
         flow_run_ctx = FlowRunContext.get()
+        parameters = self.parameters or {}
 
         parent_task_run = None
+
         # this is a subflow run
         if flow_run_ctx:
+            # get the parent task run
             parent_task_run = await self.create_subflow_task_run(
                 client=client, context=flow_run_ctx
             )
-            # If the parent task run already completed, return the last flow run
-            # associated with the parent task run. This prevents rerunning a completed
-            # flow run when the parent task run is rerun.
-            most_recent_flow_run = (
-                await self.get_most_recent_flow_run_for_parent_task_run(
-                    client=client, parent_task_run=parent_task_run
-                )
-            )
-            if most_recent_flow_run:
-                return most_recent_flow_run
+
+            # check if there is already a flow run for this subflow
+            if subflow_run := await self.load_subflow_run(
+                parent_task_run=parent_task_run, client=client, context=flow_run_ctx
+            ):
+                return subflow_run
 
         try:
             flow_run_name = _resolve_custom_flow_run_name(
-                flow=self.flow, parameters=self.parameters
+                flow=self.flow, parameters=parameters
             )
         except TypeError:
             flow_run_name = None
@@ -189,7 +254,7 @@ class FlowRunEngine(Generic[P, R]):
         flow_run = await client.create_flow_run(
             flow=self.flow,
             name=flow_run_name,
-            parameters=self.flow.serialize_parameters(self.parameters),
+            parameters=self.flow.serialize_parameters(parameters),
             state=Pending(),
             parent_task_run_id=getattr(parent_task_run, "id", None),
         )
@@ -199,8 +264,46 @@ class FlowRunEngine(Generic[P, R]):
     async def enter_run_context(self, client: Optional[PrefectClient] = None):
         if client is None:
             client = self.client
+        if not self.flow_run:
+            raise ValueError("Flow run not set")
 
         self.flow_run = await client.read_flow_run(self.flow_run.id)
+        task_runner = self.flow.task_runner.duplicate()
+
+        async with AsyncExitStack() as stack:
+            task_runner = await stack.enter_async_context(
+                self.flow.task_runner.duplicate().start()
+            )
+            stack.enter_context(
+                FlowRunContext(
+                    flow=self.flow,
+                    log_prints=self.flow.log_prints or False,
+                    flow_run=self.flow_run,
+                    parameters=self.parameters,
+                    client=client,
+                    background_tasks=anyio.create_task_group(),
+                    result_factory=await ResultFactory.from_flow(self.flow),
+                    task_runner=task_runner,
+                )
+            )
+            self.logger = flow_run_logger(flow_run=self.flow_run, flow=self.flow)
+            yield
+
+    @contextmanager
+    def enter_run_context_sync(self, client: Optional[PrefectClient] = None):
+        if client is None:
+            client = self.client
+        if not self.flow_run:
+            raise ValueError("Flow run not set")
+
+        self.flow_run = run_sync(client.read_flow_run(self.flow_run.id))
+
+        # if running in a completely synchronous frame, anyio will not detect the
+        # backend to use for the task group
+        try:
+            task_group = anyio.create_task_group()
+        except AsyncLibraryNotFoundError:
+            task_group = anyio._backends._asyncio.TaskGroup()
 
         with FlowRunContext(
             flow=self.flow,
@@ -208,8 +311,8 @@ class FlowRunEngine(Generic[P, R]):
             flow_run=self.flow_run,
             parameters=self.parameters,
             client=client,
-            background_tasks=anyio.create_task_group(),
-            result_factory=await ResultFactory.from_flow(self.flow),
+            background_tasks=task_group,
+            result_factory=run_sync(ResultFactory.from_flow(self.flow)),
             task_runner=self.flow.task_runner,
         ):
             self.logger = flow_run_logger(flow_run=self.flow_run, flow=self.flow)
@@ -224,13 +327,27 @@ class FlowRunEngine(Generic[P, R]):
             self._client = client
             self._is_started = True
 
+            # this conditional is engaged whenever a run is triggered via deployment
+            if self.flow_run_id and not self.flow:
+                self.flow_run = await client.read_flow_run(self.flow_run_id)
+                try:
+                    self.flow = await self.load_flow(client)
+                except Exception as exc:
+                    await self.handle_exception(
+                        exc,
+                        msg="Failed to load flow from entrypoint.",
+                    )
+                    self.short_circuit = True
+
             if not self.flow_run:
                 self.flow_run = await self.create_flow_run(client)
 
             # validate prior to context so that context receives validated params
             if self.flow.should_validate_parameters:
                 try:
-                    self.parameters = self.flow.validate_parameters(self.parameters)
+                    self.parameters = self.flow.validate_parameters(
+                        self.parameters or {}
+                    )
                 except Exception as exc:
                     await self.handle_exception(
                         exc,
@@ -238,11 +355,46 @@ class FlowRunEngine(Generic[P, R]):
                         result_factory=await ResultFactory.from_flow(self.flow),
                     )
                     self.short_circuit = True
+            try:
+                yield self
+            finally:
+                self._is_started = False
+                self._client = None
+
+    @contextmanager
+    def start_sync(self):
+        """
+        Enters a client context and creates a flow run if needed.
+        """
 
+        client = get_client()
+        run_sync(client.__aenter__())
+        self._client = client
+        self._is_started = True
+
+        if not self.flow_run:
+            self.flow_run = run_sync(self.create_flow_run(client))
+
+        # validate prior to context so that context receives validated params
+        if self.flow.should_validate_parameters:
+            try:
+                self.parameters = self.flow.validate_parameters(self.parameters or {})
+            except Exception as exc:
+                run_sync(
+                    self.handle_exception(
+                        exc,
+                        msg="Validation of flow parameters failed with error",
+                        result_factory=run_sync(ResultFactory.from_flow(self.flow)),
+                    )
+                )
+                self.short_circuit = True
+        try:
             yield self
-
-        self._is_started = False
-        self._client = None
+        finally:
+            # quickly close client
+            run_sync(client.__aexit__(None, None, None))
+            self._is_started = False
+            self._client = None
 
     def is_running(self) -> bool:
         if getattr(self, "flow_run", None) is None:
@@ -256,31 +408,33 @@ class FlowRunEngine(Generic[P, R]):
 
 
 async def run_flow(
-    flow: Task[P, Coroutine[Any, Any, R]],
+    flow: Optional[Flow[P, Coroutine[Any, Any, R]]] = None,
     flow_run: Optional[FlowRun] = None,
+    flow_run_id: Optional[UUID] = None,
     parameters: Optional[Dict[str, Any]] = None,
     wait_for: Optional[Iterable[PrefectFuture[A, Async]]] = None,
     return_type: Literal["state", "result"] = "result",
-) -> "Union[R, None]":
+) -> Union[R, None]:
     """
     Runs a flow against the API.
 
     We will most likely want to use this logic as a wrapper and return a coroutine for type inference.
     """
 
-    engine = FlowRunEngine[P, R](flow, parameters, flow_run)
+    engine = FlowRunEngine[P, R](flow, parameters, flow_run, flow_run_id)
+
+    # This is a context manager that keeps track of the state of the flow run.
     async with engine.start() as run:
-        # This is a context manager that keeps track of the state of the flow run.
         await run.begin_run()
 
         while run.is_running():
             async with run.enter_run_context():
                 try:
                     # This is where the flow is actually run.
-                    if flow.isasync:
-                        result = cast(R, await flow.fn(**(run.parameters or {})))  # type: ignore
-                    else:
-                        result = cast(R, flow.fn(**(run.parameters or {})))  # type: ignore
+                    call_args, call_kwargs = parameters_to_args_kwargs(
+                        flow.fn, run.parameters or {}
+                    )
+                    result = cast(R, await flow.fn(*call_args, **call_kwargs))  # type: ignore
                     # If the flow run is successful, finalize it.
                     await run.handle_success(result)
 
@@ -291,3 +445,36 @@ async def run_flow(
         if return_type == "state":
             return run.state
         return await run.result()
+
+
+def run_flow_sync(
+    flow: Flow[P, R],
+    flow_run: Optional[FlowRun] = None,
+    parameters: Optional[Dict[str, Any]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[A, Async]]] = None,
+    return_type: Literal["state", "result"] = "result",
+) -> Union[R, State, None]:
+    engine = FlowRunEngine[P, R](flow, parameters, flow_run)
+
+    # This is a context manager that keeps track of the state of the flow run.
+    with engine.start_sync() as run:
+        run_sync(run.begin_run())
+
+        while run.is_running():
+            with run.enter_run_context_sync():
+                try:
+                    # This is where the flow is actually run.
+                    call_args, call_kwargs = parameters_to_args_kwargs(
+                        flow.fn, run.parameters or {}
+                    )
+                    result = cast(R, flow.fn(*call_args, **call_kwargs))  # type: ignore
+                    # If the flow run is successful, finalize it.
+                    run_sync(run.handle_success(result))
+
+                except Exception as exc:
+                    # If the flow fails, and we have retries left, set the flow to retrying.
+                    run_sync(run.handle_exception(exc))
+
+        if return_type == "state":
+            return run.state
+        return run_sync(run.result())