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

prefect/docker/docker_image.py

@@ -0,0 +1,82 @@
+from pathlib import Path
+from typing import Optional
+
+from pendulum import now as pendulum_now
+
+from prefect.settings import (
+    PREFECT_DEFAULT_DOCKER_BUILD_NAMESPACE,
+)
+from prefect.utilities.dockerutils import (
+    PushError,
+    build_image,
+    docker_client,
+    generate_default_dockerfile,
+    parse_image_tag,
+    split_repository_path,
+)
+from prefect.utilities.slugify import slugify
+
+
+class DockerImage:
+    """
+    Configuration used to build and push a Docker image for a deployment.
+
+    Attributes:
+        name: The name of the Docker image to build, including the registry and
+            repository.
+        tag: The tag to apply to the built image.
+        dockerfile: The path to the Dockerfile to use for building the image. If
+            not provided, a default Dockerfile will be generated.
+        **build_kwargs: Additional keyword arguments to pass to the Docker build request.
+            See the [`docker-py` documentation](https://docker-py.readthedocs.io/en/stable/images.html#docker.models.images.ImageCollection.build)
+            for more information.
+
+    """
+
+    def __init__(
+        self, name: str, tag: Optional[str] = None, dockerfile="auto", **build_kwargs
+    ):
+        image_name, image_tag = parse_image_tag(name)
+        if tag and image_tag:
+            raise ValueError(
+                f"Only one tag can be provided - both {image_tag!r} and {tag!r} were"
+                " provided as tags."
+            )
+        namespace, repository = split_repository_path(image_name)
+        # if the provided image name does not include a namespace (registry URL or user/org name),
+        # use the default namespace
+        if not namespace:
+            namespace = PREFECT_DEFAULT_DOCKER_BUILD_NAMESPACE.value()
+        # join the namespace and repository to create the full image name
+        # ignore namespace if it is None
+        self.name = "/".join(filter(None, [namespace, repository]))
+        self.tag = tag or image_tag or slugify(pendulum_now("utc").isoformat())
+        self.dockerfile = dockerfile
+        self.build_kwargs = build_kwargs
+
+    @property
+    def reference(self):
+        return f"{self.name}:{self.tag}"
+
+    def build(self):
+        full_image_name = self.reference
+        build_kwargs = self.build_kwargs.copy()
+        build_kwargs["context"] = Path.cwd()
+        build_kwargs["tag"] = full_image_name
+        build_kwargs["pull"] = build_kwargs.get("pull", True)
+
+        if self.dockerfile == "auto":
+            with generate_default_dockerfile():
+                build_image(**build_kwargs)
+        else:
+            build_kwargs["dockerfile"] = self.dockerfile
+            build_image(**build_kwargs)
+
+    def push(self):
+        with docker_client() as client:
+            events = client.api.push(
+                repository=self.name, tag=self.tag, stream=True, decode=True
+            )
+            for event in events:
+                if "error" in event:
+                    raise PushError(event["error"])

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,
 )
@@ -166,6 +168,20 @@ class FlowRunEngine(Generic[P, R]):
             )
             return state
 
+        # 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:
+                message = "Validation of flow parameters failed with error:"
+                self.logger.error("%s %s", message, exc)
+                self.handle_exception(
+                    exc,
+                    msg=message,
+                    result_factory=run_coro_as_sync(ResultFactory.from_flow(self.flow)),
+                )
+                self.short_circuit = True
+
         new_state = Running()
         state = self.set_state(new_state)
         while state.is_pending():
@@ -482,24 +498,6 @@ class FlowRunEngine(Generic[P, R]):
                     flow_version=self.flow.version,
                     empirical_policy=self.flow_run.empirical_policy,
                 )
-
-            # 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:
-                    message = "Validation of flow parameters failed with error:"
-                    self.logger.error("%s %s", message, exc)
-                    self.handle_exception(
-                        exc,
-                        msg=message,
-                        result_factory=run_coro_as_sync(
-                            ResultFactory.from_flow(self.flow)
-                        ),
-                    )
-                    self.short_circuit = True
             try:
                 yield self
             except Exception:
@@ -632,6 +630,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 +718,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

@@ -17,11 +17,9 @@ import warnings
 from copy import copy
 from functools import partial, update_wrapper
 from pathlib import Path
-from tempfile import NamedTemporaryFile
 from typing import (
     TYPE_CHECKING,
     Any,
-    AnyStr,
     Awaitable,
     Callable,
     Coroutine,
@@ -56,9 +54,9 @@ from prefect.client.schemas.objects import Flow as FlowSchema
 from prefect.client.schemas.objects import FlowRun
 from prefect.client.schemas.schedules import SCHEDULE_TYPES
 from prefect.client.utilities import client_injector
-from prefect.context import PrefectObjectRegistry, registry_from_script
-from prefect.deployments.runner import DeploymentImage, EntrypointType, deploy
+from prefect.deployments.runner import deploy
 from prefect.deployments.steps.core import run_steps
+from prefect.docker.docker_image import DockerImage
 from prefect.events import DeploymentTriggerTypes, TriggerTypes
 from prefect.exceptions import (
     InvalidNameError,
@@ -87,9 +85,9 @@ from prefect.settings import (
 from prefect.states import State
 from prefect.task_runners import TaskRunner, ThreadPoolTaskRunner
 from prefect.types import BANNED_CHARACTERS, WITHOUT_BANNED_CHARACTERS
+from prefect.types.entrypoint import EntrypointType
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import (
-    is_async_fn,
     run_sync_in_worker_thread,
     sync_compatible,
 )
@@ -119,11 +117,11 @@ logger = get_logger("flows")
 
 if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient
-    from prefect.deployments.runner import FlexibleScheduleList, RunnerDeployment
+    from prefect.client.types.flexible_schedule_list import FlexibleScheduleList
+    from prefect.deployments.runner import RunnerDeployment
     from prefect.flows import FlowRun
 
 
-@PrefectObjectRegistry.register_instances
 class Flow(Generic[P, R]):
     """
     A Prefect workflow definition.
@@ -146,7 +144,7 @@ class Flow(Generic[P, R]):
             be provided as a string template with the flow's parameters as variables,
             or a function that returns a string.
         task_runner: An optional task runner to use for task execution within the flow;
-            if not provided, a `ConcurrentTaskRunner` will be used.
+            if not provided, a `ThreadPoolTaskRunner` will be used.
         description: An optional string description for the flow; if not provided, the
             description will be pulled from the docstring for the decorated function.
         timeout_seconds: An optional number of seconds indicating a maximum runtime for
@@ -289,7 +287,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"])
 
@@ -988,7 +997,7 @@ class Flow(Generic[P, R]):
         self,
         name: str,
         work_pool_name: Optional[str] = None,
-        image: Optional[Union[str, DeploymentImage]] = None,
+        image: Optional[Union[str, DockerImage]] = None,
         build: bool = True,
         push: bool = True,
         work_queue_name: Optional[str] = None,
@@ -1024,7 +1033,7 @@ class Flow(Generic[P, R]):
             work_pool_name: The name of the work pool to use for this deployment. Defaults to
                 the value of `PREFECT_DEFAULT_WORK_POOL_NAME`.
             image: The name of the Docker image to build, including the registry and
-                repository. Pass a DeploymentImage instance to customize the Dockerfile used
+                repository. Pass a DockerImage instance to customize the Dockerfile used
                 and build arguments.
             build: Whether or not to build a new image for the flow. If False, the provided
                 image will be used as-is and pulled at runtime.
@@ -1628,47 +1637,6 @@ def select_flow(
         return list(flows_dict.values())[0]
 
 
-def load_flows_from_script(path: str) -> List[Flow]:
-    """
-    Load all flow objects from the given python script. All of the code in the file
-    will be executed.
-
-    Returns:
-        A list of flows
-
-    Raises:
-        FlowScriptError: If an exception is encountered while running the script
-    """
-    return registry_from_script(path).get_instances(Flow)
-
-
-def load_flow_from_script(path: str, flow_name: Optional[str] = None) -> Flow:
-    """
-    Extract a flow object from a script by running all of the code in the file.
-
-    If the script has multiple flows in it, a flow name must be provided to specify
-    the flow to return.
-
-    Args:
-        path: A path to a Python script containing flows
-        flow_name: An optional flow name to look for in the script
-
-    Returns:
-        The flow object from the script
-
-    Raises:
-        FlowScriptError: If an exception is encountered while running the script
-        MissingFlowError: If no flows exist in the iterable
-        MissingFlowError: If a flow name is provided and that flow does not exist
-        UnspecifiedFlowError: If multiple flows exist but no flow name was provided
-    """
-    return select_flow(
-        load_flows_from_script(path),
-        flow_name=flow_name,
-        from_message=f"in script '{path}'",
-    )
-
-
 def load_flow_from_entrypoint(
     entrypoint: str,
 ) -> Flow:
@@ -1686,52 +1654,25 @@ def load_flow_from_entrypoint(
         FlowScriptError: If an exception is encountered while running the script
         MissingFlowError: If the flow function specified in the entrypoint does not exist
     """
-    with PrefectObjectRegistry(  # type: ignore
-        block_code_execution=True,
-        capture_failures=True,
-    ):
-        if ":" in entrypoint:
-            # split by the last colon once to handle Windows paths with drive letters i.e C:\path\to\file.py:do_stuff
-            path, func_name = entrypoint.rsplit(":", maxsplit=1)
-        else:
-            path, func_name = entrypoint.rsplit(".", maxsplit=1)
-        try:
-            flow = import_object(entrypoint)
-        except AttributeError as exc:
-            raise MissingFlowError(
-                f"Flow function with name {func_name!r} not found in {path!r}. "
-            ) from exc
-
-        if not isinstance(flow, Flow):
-            raise MissingFlowError(
-                f"Function with name {func_name!r} is not a flow. Make sure that it is "
-                "decorated with '@flow'."
-            )
-
-        return flow
 
+    if ":" in entrypoint:
+        # split by the last colon once to handle Windows paths with drive letters i.e C:\path\to\file.py:do_stuff
+        path, func_name = entrypoint.rsplit(":", maxsplit=1)
+    else:
+        path, func_name = entrypoint.rsplit(".", maxsplit=1)
+    try:
+        flow = import_object(entrypoint)
+    except AttributeError as exc:
+        raise MissingFlowError(
+            f"Flow function with name {func_name!r} not found in {path!r}. "
+        ) from exc
 
-def load_flow_from_text(script_contents: AnyStr, flow_name: str) -> Flow:
-    """
-    Load a flow from a text script.
+    if not isinstance(flow, Flow):
+        raise MissingFlowError(
+            f"Function with name {func_name!r} is not a flow. Make sure that it is "
+            "decorated with '@flow'."
+        )
 
-    The script will be written to a temporary local file path so errors can refer
-    to line numbers and contextual tracebacks can be provided.
-    """
-    with NamedTemporaryFile(
-        mode="wt" if isinstance(script_contents, str) else "wb",
-        prefix=f"flow-script-{flow_name}",
-        suffix=".py",
-        delete=False,
-    ) as tmpfile:
-        tmpfile.write(script_contents)
-        tmpfile.flush()
-    try:
-        flow = load_flow_from_script(tmpfile.name, flow_name=flow_name)
-    finally:
-        # windows compat
-        tmpfile.close()
-        os.remove(tmpfile.name)
     return flow
 
 

prefect/futures.py

@@ -10,7 +10,7 @@ from typing_extensions import TypeVar
 from prefect.client.orchestration import get_client
 from prefect.client.schemas.objects import TaskRun
 from prefect.exceptions import ObjectNotFound
-from prefect.logging.loggers import get_logger
+from prefect.logging.loggers import get_logger, get_run_logger
 from prefect.states import Pending, State
 from prefect.task_runs import TaskRunWaiter
 from prefect.utilities.annotations import quote
@@ -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.
 
@@ -143,6 +143,18 @@ class PrefectConcurrentFuture(PrefectWrappedFuture[concurrent.futures.Future]):
             _result = run_coro_as_sync(_result)
         return _result
 
+    def __del__(self):
+        if self._final_state or self._wrapped_future.done():
+            return
+        try:
+            local_logger = get_run_logger()
+        except Exception:
+            local_logger = logger
+        local_logger.warning(
+            "A future was garbage collected before it resolved."
+            " Please call `.wait()` or `.result()` on futures to ensure they resolve.",
+        )
+
 
 class PrefectDistributedFuture(PrefectFuture):
     """
@@ -163,6 +175,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 +261,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/infrastructure/provisioners/cloud_run.py

@@ -404,7 +404,7 @@ class CloudRunPushProvisioner:
                         dedent(
                             f"""\
                         from prefect import flow
-                        from prefect.deployments import DeploymentImage
+                        from prefect.docker import DockerImage
 
 
                         @flow(log_prints=True)
@@ -416,7 +416,7 @@ class CloudRunPushProvisioner:
                             my_flow.deploy(
                                 name="my-deployment",
                                 work_pool_name="{work_pool_name}",
-                                image=DeploymentImage(
+                                image=DockerImage(
                                     name="my-image:latest",
                                     platform="linux/amd64",
                                 )

prefect/infrastructure/provisioners/container_instance.py

@@ -1042,7 +1042,7 @@ class ContainerInstancePushProvisioner:
                     dedent(
                         f"""\
                         from prefect import flow
-                        from prefect.deployments import DeploymentImage
+                        from prefect.docker import DockerImage
 
 
                         @flow(log_prints=True)
@@ -1054,7 +1054,7 @@ class ContainerInstancePushProvisioner:
                             my_flow.deploy(
                                 name="my-deployment",
                                 work_pool_name="{work_pool_name}",
-                                image=DeploymentImage(
+                                image=DockerImage(
                                     name="my-image:latest",
                                     platform="linux/amd64",
                                 )

prefect/infrastructure/provisioners/ecs.py

@@ -950,7 +950,7 @@ class ContainerRepositoryResource:
                         dedent(
                             f"""\
                                 from prefect import flow
-                                from prefect.deployments import DeploymentImage
+                                from prefect.docker import DockerImage
 
 
                                 @flow(log_prints=True)
@@ -962,7 +962,7 @@ class ContainerRepositoryResource:
                                     my_flow.deploy(
                                         name="my-deployment",
                                         work_pool_name="{self._work_pool_name}",
-                                        image=DeploymentImage(
+                                        image=DockerImage(
                                             name="{self._repository_name}:latest",
                                             platform="linux/amd64",
                                         )

prefect/records/result_store.py

@@ -44,6 +44,10 @@ class ResultFactoryStore(RecordStore):
             raise ValueError("Result could not be read")
 
     def write(self, key: str, value: Any) -> BaseResult:
-        if isinstance(value, BaseResult):
+        if isinstance(value, PersistedResult):
+            # if the value is already a persisted result, write it
+            value.write(_sync=True)
+            return value
+        elif isinstance(value, BaseResult):
             return value
         return run_coro_as_sync(self.result_factory.create_result(obj=value, key=key))