prefect-client 3.3.5.dev3__py3-none-any.whl → 3.3.6__py3-none-any.whl

prefect/flows.py

@@ -40,6 +40,7 @@ from typing import (
 from uuid import UUID
 
 import pydantic
+from exceptiongroup import BaseExceptionGroup, ExceptionGroup
 from pydantic.v1 import BaseModel as V1BaseModel
 from pydantic.v1.decorator import ValidatedFunction as V1ValidatedFunction
 from pydantic.v1.errors import ConfigError  # TODO
@@ -64,7 +65,7 @@ from prefect.exceptions import (
     UnspecifiedFlowError,
 )
 from prefect.filesystems import LocalFileSystem, ReadableDeploymentStorage
-from prefect.futures import PrefectFuture
+from prefect.futures import PrefectFlowRunFuture, PrefectFuture
 from prefect.logging import get_logger
 from prefect.logging.loggers import flow_run_logger
 from prefect.results import ResultSerializer, ResultStorage
@@ -105,6 +106,9 @@ from ._internal.pydantic.v2_validated_func import (
     V2ValidatedFunction as ValidatedFunction,
 )
 
+if TYPE_CHECKING:
+    from prefect.workers.base import BaseWorker
+
 T = TypeVar("T")  # Generic type var for capturing the inner return type of async funcs
 R = TypeVar("R")  # The return type of the user's function
 P = ParamSpec("P")  # The parameters of the flow
@@ -1984,6 +1988,241 @@ class FlowDecorator:
 flow: FlowDecorator = FlowDecorator()
 
 
+class InfrastructureBoundFlow(Flow[P, R]):
+    """
+    EXPERIMENTAL: This class is experimental and may be removed or changed in future
+        releases.
+
+    A flow that is bound to running on a specific infrastructure.
+
+    Attributes:
+        work_pool: The name of the work pool to run the flow on. The base job
+            configuration of the work pool will determine the configuration of the
+            infrastructure the flow will run on.
+        job_variables: Infrastructure configuration that will override the base job
+            configuration of the work pool.
+        worker_cls: The class of the worker to use to spin up infrastructure and submit
+            the flow to it.
+    """
+
+    def __init__(
+        self,
+        *args: Any,
+        work_pool: str,
+        job_variables: dict[str, Any],
+        worker_cls: type["BaseWorker[Any, Any, Any]"],
+        **kwargs: Any,
+    ):
+        super().__init__(*args, **kwargs)
+        self.work_pool = work_pool
+        self.job_variables = job_variables
+        self.worker_cls = worker_cls
+
+    @overload
+    def __call__(self: "Flow[P, NoReturn]", *args: P.args, **kwargs: P.kwargs) -> None:
+        # `NoReturn` matches if a type can't be inferred for the function which stops a
+        # sync function from matching the `Coroutine` overload
+        ...
+
+    @overload
+    def __call__(
+        self: "Flow[P, Coroutine[Any, Any, T]]",
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> Coroutine[Any, Any, T]: ...
+
+    @overload
+    def __call__(
+        self: "Flow[P, T]",
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> T: ...
+
+    @overload
+    def __call__(
+        self: "Flow[P, Coroutine[Any, Any, T]]",
+        *args: P.args,
+        return_state: Literal[True],
+        **kwargs: P.kwargs,
+    ) -> Awaitable[State[T]]: ...
+
+    @overload
+    def __call__(
+        self: "Flow[P, T]",
+        *args: P.args,
+        return_state: Literal[True],
+        **kwargs: P.kwargs,
+    ) -> State[T]: ...
+
+    def __call__(
+        self,
+        *args: "P.args",
+        return_state: bool = False,
+        wait_for: Optional[Iterable[PrefectFuture[Any]]] = None,
+        **kwargs: "P.kwargs",
+    ):
+        async def modified_call(
+            *args: P.args,
+            return_state: bool = False,
+            # TODO: Handle wait_for once we have an asynchronous way to wait for futures
+            # We should wait locally for futures to resolve before spinning up
+            # infrastructure.
+            wait_for: Optional[Iterable[PrefectFuture[Any]]] = None,
+            **kwargs: P.kwargs,
+        ) -> R | State[R]:
+            try:
+                async with self.worker_cls(work_pool_name=self.work_pool) as worker:
+                    parameters = get_call_parameters(self, args, kwargs)
+                    future = await worker.submit(
+                        flow=self,
+                        parameters=parameters,
+                        job_variables=self.job_variables,
+                    )
+                    if return_state:
+                        await future.wait_async()
+                        return future.state
+                    return await future.aresult()
+            except (ExceptionGroup, BaseExceptionGroup) as exc:
+                # For less verbose tracebacks
+                exceptions = exc.exceptions
+                if len(exceptions) == 1:
+                    raise exceptions[0] from None
+                else:
+                    raise
+
+        if inspect.iscoroutinefunction(self.fn):
+            return modified_call(
+                *args, return_state=return_state, wait_for=wait_for, **kwargs
+            )
+        else:
+            return run_coro_as_sync(
+                modified_call(
+                    *args,
+                    return_state=return_state,
+                    wait_for=wait_for,
+                    **kwargs,
+                )
+            )
+
+    def submit(self, *args: P.args, **kwargs: P.kwargs) -> PrefectFlowRunFuture[R]:
+        """
+        EXPERIMENTAL: This method is experimental and may be removed or changed in future
+            releases.
+
+        Submit the flow to run on remote infrastructure.
+
+        Args:
+            *args: Positional arguments to pass to the flow.
+            **kwargs: Keyword arguments to pass to the flow.
+
+        Returns:
+            A `PrefectFlowRunFuture` that can be used to retrieve the result of the flow run.
+
+        Examples:
+            Submit a flow to run on Kubernetes:
+
+            ```python
+            from prefect import flow
+            from prefect_kubernetes.experimental import kubernetes
+
+            @kubernetes(work_pool="my-kubernetes-work-pool")
+            @flow
+            def my_flow(x: int, y: int):
+                return x + y
+
+            future = my_flow.submit(x=1, y=2)
+            result = future.result()
+            print(result)
+            ```
+        """
+
+        async def submit_func():
+            async with self.worker_cls(work_pool_name=self.work_pool) as worker:
+                parameters = get_call_parameters(self, args, kwargs)
+                return await worker.submit(
+                    flow=self,
+                    parameters=parameters,
+                    job_variables=self.job_variables,
+                )
+
+        return run_coro_as_sync(submit_func())
+
+    def with_options(
+        self,
+        *,
+        name: Optional[str] = None,
+        version: Optional[str] = None,
+        retries: Optional[int] = None,
+        retry_delay_seconds: Optional[Union[int, float]] = None,
+        description: Optional[str] = None,
+        flow_run_name: Optional[Union[Callable[[], str], str]] = None,
+        task_runner: Union[
+            Type[TaskRunner[PrefectFuture[Any]]], TaskRunner[PrefectFuture[Any]], None
+        ] = None,
+        timeout_seconds: Union[int, float, None] = None,
+        validate_parameters: Optional[bool] = None,
+        persist_result: Optional[bool] = NotSet,  # type: ignore
+        result_storage: Optional[ResultStorage] = NotSet,  # type: ignore
+        result_serializer: Optional[ResultSerializer] = NotSet,  # type: ignore
+        cache_result_in_memory: Optional[bool] = None,
+        log_prints: Optional[bool] = NotSet,  # type: ignore
+        on_completion: Optional[list[FlowStateHook[P, R]]] = None,
+        on_failure: Optional[list[FlowStateHook[P, R]]] = None,
+        on_cancellation: Optional[list[FlowStateHook[P, R]]] = None,
+        on_crashed: Optional[list[FlowStateHook[P, R]]] = None,
+        on_running: Optional[list[FlowStateHook[P, R]]] = None,
+        job_variables: Optional[dict[str, Any]] = None,
+    ) -> "InfrastructureBoundFlow[P, R]":
+        new_flow = super().with_options(
+            name=name,
+            version=version,
+            retries=retries,
+            retry_delay_seconds=retry_delay_seconds,
+            description=description,
+            flow_run_name=flow_run_name,
+            task_runner=task_runner,
+            timeout_seconds=timeout_seconds,
+            validate_parameters=validate_parameters,
+            persist_result=persist_result,
+            result_storage=result_storage,
+            result_serializer=result_serializer,
+            cache_result_in_memory=cache_result_in_memory,
+            log_prints=log_prints,
+            on_completion=on_completion,
+            on_failure=on_failure,
+            on_cancellation=on_cancellation,
+            on_crashed=on_crashed,
+            on_running=on_running,
+        )
+        new_infrastructure_bound_flow = bind_flow_to_infrastructure(
+            new_flow,
+            self.work_pool,
+            self.worker_cls,
+            job_variables=job_variables
+            if job_variables is not None
+            else self.job_variables,
+        )
+        return new_infrastructure_bound_flow
+
+
+def bind_flow_to_infrastructure(
+    flow: Flow[P, R],
+    work_pool: str,
+    worker_cls: type["BaseWorker[Any, Any, Any]"],
+    job_variables: dict[str, Any] | None = None,
+) -> InfrastructureBoundFlow[P, R]:
+    new = InfrastructureBoundFlow[P, R](
+        flow.fn,
+        work_pool=work_pool,
+        job_variables=job_variables or {},
+        worker_cls=worker_cls,
+    )
+    # Copy all attributes from the original flow
+    for attr, value in flow.__dict__.items():
+        setattr(new, attr, value)
+    return new
+
+
 def _raise_on_name_with_banned_characters(name: Optional[str]) -> Optional[str]:
     """
     Raise an InvalidNameError if the given name contains any invalid

prefect/futures.py

@@ -612,19 +612,44 @@ def resolve_futures_to_states(
 
     Unsupported object types will be returned without modification.
     """
-    futures: set[PrefectFuture[R]] = set()
 
-    def _collect_futures(
-        futures: set[PrefectFuture[R]], expr: Any | PrefectFuture[R], context: Any
-    ) -> Any | PrefectFuture[R]:
-        # Expressions inside quotes should not be traversed
-        if isinstance(context.get("annotation"), quote):
-            raise StopVisiting()
+    def _resolve_state(future: PrefectFuture[R]):
+        future.wait()
+        return future.state
 
-        if isinstance(expr, PrefectFuture):
-            futures.add(expr)
+    return _resolve_futures(
+        expr,
+        resolve_fn=_resolve_state,
+    )
 
-        return expr
+
+def resolve_futures_to_results(
+    expr: PrefectFuture[R] | Any,
+) -> Any:
+    """
+    Given a Python built-in collection, recursively find `PrefectFutures` and build a
+    new collection with the same structure with futures resolved to their final results.
+    Resolving futures to their final result may wait for execution to complete.
+
+    Unsupported object types will be returned without modification.
+    """
+
+    def _resolve_result(future: PrefectFuture[R]) -> Any:
+        future.wait()
+        if future.state.is_completed():
+            return future.result()
+        else:
+            raise Exception("At least one result did not complete successfully")
+
+    return _resolve_futures(expr, resolve_fn=_resolve_result)
+
+
+def _resolve_futures(
+    expr: PrefectFuture[R] | Any,
+    resolve_fn: Callable[[PrefectFuture[R]], Any],
+) -> Any:
+    """Helper function to resolve PrefectFutures in a collection."""
+    futures: set[PrefectFuture[R]] = set()
 
     visit_collection(
         expr,
@@ -633,31 +658,39 @@ def resolve_futures_to_states(
         context={},
     )
 
-    # if no futures were found, return the original expression
+    # If no futures were found, return the original expression
     if not futures:
         return expr
 
-    # Get final states for each future
-    states: list[State] = []
-    for future in futures:
-        future.wait()
-        states.append(future.state)
-
-    states_by_future = dict(zip(futures, states))
+    # Resolve each future using the provided resolve function
+    resolved_values = {future: resolve_fn(future) for future in futures}
 
-    def replace_futures_with_states(expr: Any, context: Any) -> Any:
+    def replace_futures(expr: Any, context: Any) -> Any:
         # Expressions inside quotes should not be modified
         if isinstance(context.get("annotation"), quote):
             raise StopVisiting()
 
         if isinstance(expr, PrefectFuture):
-            return states_by_future[expr]
+            return resolved_values[expr]
         else:
             return expr
 
     return visit_collection(
         expr,
-        visit_fn=replace_futures_with_states,
+        visit_fn=replace_futures,
         return_data=True,
         context={},
     )
+
+
+def _collect_futures(
+    futures: set[PrefectFuture[R]], expr: Any | PrefectFuture[R], context: Any
+) -> Any | PrefectFuture[R]:
+    # Expressions inside quotes should not be traversed
+    if isinstance(context.get("annotation"), quote):
+        raise StopVisiting()
+
+    if isinstance(expr, PrefectFuture):
+        futures.add(expr)
+
+    return expr

prefect/runner/runner.py

@@ -636,7 +636,12 @@ class Runner:
 
                     return process
 
-    async def execute_bundle(self, bundle: SerializedBundle) -> None:
+    async def execute_bundle(
+        self,
+        bundle: SerializedBundle,
+        cwd: Path | str | None = None,
+        env: dict[str, str | None] | None = None,
+    ) -> None:
         """
         Executes a bundle in a subprocess.
         """
@@ -651,7 +656,7 @@ class Runner:
             if not self._acquire_limit_slot(flow_run.id):
                 return
 
-            process = execute_bundle_in_subprocess(bundle)
+            process = execute_bundle_in_subprocess(bundle, cwd=cwd, env=env)
 
             if process.pid is None:
                 # This shouldn't happen because `execute_bundle_in_subprocess` starts the process
@@ -776,7 +781,7 @@ class Runner:
         if command is None:
             runner_command = [get_sys_executable(), "-m", "prefect.engine"]
         else:
-            runner_command = shlex.split(command)
+            runner_command = shlex.split(command, posix=(os.name != "nt"))
 
         flow_run_logger = self._get_flow_run_logger(flow_run)
 

prefect/server/api/deployments.py

@@ -98,9 +98,31 @@ async def create_deployment(
             )
 
         # hydrate the input model into a full model
-        deployment_dict = deployment.model_dump(
-            exclude={"work_pool_name"}, exclude_unset=True
+        deployment_dict: dict = deployment.model_dump(
+            exclude={"work_pool_name"},
+            exclude_unset=True,
         )
+
+        requested_concurrency_limit = deployment_dict.pop(
+            "global_concurrency_limit_id", "unset"
+        )
+        if requested_concurrency_limit != "unset":
+            if requested_concurrency_limit:
+                concurrency_limit = (
+                    await models.concurrency_limits_v2.read_concurrency_limit(
+                        session=session,
+                        concurrency_limit_id=requested_concurrency_limit,
+                    )
+                )
+
+                if not concurrency_limit:
+                    raise HTTPException(
+                        status_code=status.HTTP_404_NOT_FOUND,
+                        detail="Concurrency limit not found",
+                    )
+
+            deployment_dict["concurrency_limit_id"] = requested_concurrency_limit
+
         if deployment.work_pool_name and deployment.work_queue_name:
             # If a specific pool name/queue name combination was provided, get the
             # ID for that work pool queue.
@@ -300,8 +322,24 @@ async def update_deployment(
                     detail="Invalid schema: Unable to validate schema with circular references.",
                 )
 
+        if deployment.global_concurrency_limit_id:
+            concurrency_limit = (
+                await models.concurrency_limits_v2.read_concurrency_limit(
+                    session=session,
+                    concurrency_limit_id=deployment.global_concurrency_limit_id,
+                )
+            )
+
+            if not concurrency_limit:
+                raise HTTPException(
+                    status_code=status.HTTP_404_NOT_FOUND,
+                    detail="Concurrency limit not found",
+                )
+
         result = await models.deployments.update_deployment(
-            session=session, deployment_id=deployment_id, deployment=deployment
+            session=session,
+            deployment_id=deployment_id,
+            deployment=deployment,
         )
 
         for schedule in schedules_to_patch:

prefect/utilities/_ast.py

@@ -0,0 +1,117 @@
+import ast
+import math
+from typing import TYPE_CHECKING, Literal
+
+import anyio
+from typing_extensions import TypeAlias
+
+from prefect.logging.loggers import get_logger
+from prefect.settings import get_current_settings
+from prefect.utilities.asyncutils import LazySemaphore
+from prefect.utilities.filesystem import get_open_file_limit
+
+# Only allow half of the open file limit to be open at once to allow for other
+# actors to open files.
+OPEN_FILE_SEMAPHORE = LazySemaphore(lambda: math.floor(get_open_file_limit() * 0.5))
+
+# this potentially could be a TypedDict, but you
+# need some way to convince the type checker that
+# Literal["flow_name", "task_name"] are being provided
+DecoratedFnMetadata: TypeAlias = dict[str, str]
+
+
+async def find_prefect_decorated_functions_in_file(
+    path: anyio.Path, decorator_module: str, decorator_name: Literal["flow", "task"]
+) -> list[DecoratedFnMetadata]:
+    logger = get_logger()
+    decorator_name_key = f"{decorator_name}_name"
+    decorated_functions: list[DecoratedFnMetadata] = []
+
+    async with OPEN_FILE_SEMAPHORE:
+        try:
+            async with await anyio.open_file(path) as f:
+                try:
+                    tree = ast.parse(await f.read())
+                except SyntaxError:
+                    if get_current_settings().debug_mode:
+                        logger.debug(
+                            f"Could not parse {path} as a Python file. Skipping."
+                        )
+                    return decorated_functions
+        except Exception as exc:
+            if get_current_settings().debug_mode:
+                logger.debug(f"Could not open {path}: {exc}. Skipping.")
+            return decorated_functions
+
+    for node in ast.walk(tree):
+        if isinstance(
+            node,
+            (
+                ast.FunctionDef,
+                ast.AsyncFunctionDef,
+            ),
+        ):
+            for decorator in node.decorator_list:
+                # handles @decorator_name
+                is_name_match = (
+                    isinstance(decorator, ast.Name) and decorator.id == decorator_name
+                )
+                # handles @decorator_name()
+                is_func_name_match = (
+                    isinstance(decorator, ast.Call)
+                    and isinstance(decorator.func, ast.Name)
+                    and decorator.func.id == decorator_name
+                )
+                # handles @decorator_module.decorator_name
+                is_module_attribute_match = (
+                    isinstance(decorator, ast.Attribute)
+                    and isinstance(decorator.value, ast.Name)
+                    and decorator.value.id == decorator_module
+                    and decorator.attr == decorator_name
+                )
+                # handles @decorator_module.decorator_name()
+                is_module_attribute_func_match = (
+                    isinstance(decorator, ast.Call)
+                    and isinstance(decorator.func, ast.Attribute)
+                    and decorator.func.attr == decorator_name
+                    and isinstance(decorator.func.value, ast.Name)
+                    and decorator.func.value.id == decorator_module
+                )
+                if is_name_match or is_module_attribute_match:
+                    decorated_functions.append(
+                        {
+                            decorator_name_key: node.name,
+                            "function_name": node.name,
+                            "filepath": str(path),
+                        }
+                    )
+                if is_func_name_match or is_module_attribute_func_match:
+                    name_kwarg_node = None
+                    if TYPE_CHECKING:
+                        assert isinstance(decorator, ast.Call)
+                    for kw in decorator.keywords:
+                        if kw.arg == "name":
+                            name_kwarg_node = kw
+                            break
+                    if name_kwarg_node is not None and isinstance(
+                        name_kwarg_node.value, ast.Constant
+                    ):
+                        decorated_fn_name = name_kwarg_node.value.value
+                    else:
+                        decorated_fn_name = node.name
+                    decorated_functions.append(
+                        {
+                            decorator_name_key: decorated_fn_name,
+                            "function_name": node.name,
+                            "filepath": str(path),
+                        }
+                    )
+    return decorated_functions
+
+
+async def find_flow_functions_in_file(path: anyio.Path) -> list[DecoratedFnMetadata]:
+    return await find_prefect_decorated_functions_in_file(path, "prefect", "flow")
+
+
+async def find_task_functions_in_file(path: anyio.Path) -> list[DecoratedFnMetadata]:
+    return await find_prefect_decorated_functions_in_file(path, "prefect", "task")