prefect-client 2.20.4__py3-none-any.whl → 3.0.0__py3-none-any.whl

prefect/utilities/dockerutils.py

@@ -9,6 +9,7 @@ from tempfile import TemporaryDirectory
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
+    Any,
     Generator,
     Iterable,
     List,
@@ -27,6 +28,10 @@ import prefect
 from prefect.utilities.importtools import lazy_import
 from prefect.utilities.slugify import slugify
 
+CONTAINER_LABELS = {
+    "io.prefect.version": prefect.__version__,
+}
+
 
 def python_version_minor() -> str:
     return f"{sys.version_info.major}.{sys.version_info.minor}"
@@ -37,7 +42,9 @@ def python_version_micro() -> str:
 
 
 def get_prefect_image_name(
-    prefect_version: str = None, python_version: str = None, flavor: str = None
+    prefect_version: Optional[str] = None,
+    python_version: Optional[str] = None,
+    flavor: Optional[str] = None,
 ) -> str:
     """
     Get the Prefect image name matching the current Prefect and Python versions.
@@ -134,9 +141,9 @@ def build_image(
     dockerfile: str = "Dockerfile",
     tag: Optional[str] = None,
     pull: bool = False,
-    platform: str = None,
+    platform: Optional[str] = None,
     stream_progress_to: Optional[TextIO] = None,
-    **kwargs,
+    **kwargs: Any,
 ) -> str:
     """Builds a Docker image, returning the image ID
 
@@ -205,7 +212,7 @@ class ImageBuilder:
         self,
         base_image: str,
         base_directory: Path = None,
-        platform: str = None,
+        platform: Optional[str] = None,
         context: Path = None,
     ):
         """Create an ImageBuilder

prefect/utilities/engine.py

@@ -4,13 +4,12 @@ import inspect
 import os
 import signal
 import time
-from contextlib import contextmanager
 from functools import partial
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Dict,
-    Generator,
     Iterable,
     Optional,
     Set,
@@ -20,14 +19,12 @@ from typing import (
 from uuid import UUID, uuid4
 
 import anyio
-from exceptiongroup import BaseExceptionGroup  # novermin
 from typing_extensions import Literal
 
 import prefect
 import prefect.context
 import prefect.plugins
 from prefect._internal.concurrency.cancellation import get_deadline
-from prefect.client.orchestration import PrefectClient, SyncPrefectClient
 from prefect.client.schemas import OrchestrationResult, TaskRun
 from prefect.client.schemas.objects import (
     StateType,
@@ -47,6 +44,7 @@ from prefect.exceptions import (
 )
 from prefect.flows import Flow
 from prefect.futures import PrefectFuture
+from prefect.futures import PrefectFuture as NewPrefectFuture
 from prefect.logging.loggers import (
     get_logger,
     task_run_logger,
@@ -58,17 +56,19 @@ from prefect.settings import (
 from prefect.states import (
     State,
     get_state_exception,
-    is_state,
 )
 from prefect.tasks import Task
 from prefect.utilities.annotations import allow_failure, quote
 from prefect.utilities.asyncutils import (
     gather,
-    run_sync,
+    run_coro_as_sync,
 )
 from prefect.utilities.collections import StopVisiting, visit_collection
 from prefect.utilities.text import truncated_to
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient, SyncPrefectClient
+
 API_HEALTHCHECKS = {}
 UNTRACKABLE_TYPES = {bool, type(None), type(...), type(NotImplemented)}
 engine_logger = get_logger("engine")
@@ -96,7 +96,7 @@ async def collect_task_run_inputs(expr: Any, max_depth: int = -1) -> Set[TaskRun
             # We need to wait for futures to be submitted before we can get the task
             # run id but we want to do so asynchronously
             futures.add(obj)
-        elif is_state(obj):
+        elif isinstance(obj, State):
             if obj.state_details.task_run_id:
                 inputs.add(TaskRunResult(id=obj.state_details.task_run_id))
         # Expressions inside quotes should not be traversed
@@ -121,8 +121,49 @@ async def collect_task_run_inputs(expr: Any, max_depth: int = -1) -> Set[TaskRun
     return inputs
 
 
+def collect_task_run_inputs_sync(
+    expr: Any, future_cls: Any = NewPrefectFuture, max_depth: int = -1
+) -> Set[TaskRunInput]:
+    """
+    This function recurses through an expression to generate a set of any discernible
+    task run inputs it finds in the data structure. It produces a set of all inputs
+    found.
+
+    Examples:
+        >>> task_inputs = {
+        >>>    k: collect_task_run_inputs(v) for k, v in parameters.items()
+        >>> }
+    """
+    # TODO: This function needs to be updated to detect parameters and constants
+
+    inputs = set()
+
+    def add_futures_and_states_to_inputs(obj):
+        if isinstance(obj, future_cls) and hasattr(obj, "task_run_id"):
+            inputs.add(TaskRunResult(id=obj.task_run_id))
+        elif isinstance(obj, State):
+            if obj.state_details.task_run_id:
+                inputs.add(TaskRunResult(id=obj.state_details.task_run_id))
+        # Expressions inside quotes should not be traversed
+        elif isinstance(obj, quote):
+            raise StopVisiting
+        else:
+            state = get_state_for_result(obj)
+            if state and state.state_details.task_run_id:
+                inputs.add(TaskRunResult(id=state.state_details.task_run_id))
+
+    visit_collection(
+        expr,
+        visit_fn=add_futures_and_states_to_inputs,
+        return_data=False,
+        max_depth=max_depth,
+    )
+
+    return inputs
+
+
 async def wait_for_task_runs_and_report_crashes(
-    task_run_futures: Iterable[PrefectFuture], client: PrefectClient
+    task_run_futures: Iterable[PrefectFuture], client: "PrefectClient"
 ) -> Literal[True]:
     crash_exceptions = []
 
@@ -228,7 +269,7 @@ async def resolve_inputs(
 
         if isinstance(expr, PrefectFuture):
             futures.add(expr)
-        if is_state(expr):
+        if isinstance(expr, State):
             states.add(expr)
 
         return expr
@@ -267,7 +308,7 @@ async def resolve_inputs(
 
         if isinstance(expr, PrefectFuture):
             state = expr._final_state
-        elif is_state(expr):
+        elif isinstance(expr, State):
             state = expr
         else:
             return expr
@@ -314,7 +355,7 @@ async def resolve_inputs(
 
 
 async def propose_state(
-    client: PrefectClient,
+    client: "PrefectClient",
     state: State[object],
     force: bool = False,
     task_run_id: Optional[UUID] = None,
@@ -415,7 +456,7 @@ async def propose_state(
 
 
 def propose_state_sync(
-    client: SyncPrefectClient,
+    client: "SyncPrefectClient",
     state: State[object],
     force: bool = False,
     task_run_id: Optional[UUID] = None,
@@ -462,7 +503,7 @@ def propose_state_sync(
             # the purpose of disabling `cache_result_in_memory`
             result = state.result(raise_on_failure=False, fetch=True)
             if inspect.isawaitable(result):
-                result = run_sync(result)
+                result = run_coro_as_sync(result)
         else:
             result = state.data
 
@@ -517,8 +558,14 @@ def propose_state_sync(
         )
 
 
-def _dynamic_key_for_task_run(context: FlowRunContext, task: Task) -> int:
-    if context.flow_run is None:  # this is an autonomous task run
+def _dynamic_key_for_task_run(
+    context: FlowRunContext, task: Task, stable: bool = True
+) -> Union[int, str]:
+    if (
+        stable is False or context.detached
+    ):  # this task is running on remote infrastructure
+        return str(uuid4())
+    elif context.flow_run is None:  # this is an autonomous task run
         context.task_run_dynamic_keys[task.task_key] = getattr(
             task, "dynamic_key", str(uuid4())
         )
@@ -531,14 +578,6 @@ def _dynamic_key_for_task_run(context: FlowRunContext, task: Task) -> int:
     return context.task_run_dynamic_keys[task.task_key]
 
 
-def _observed_flow_pauses(context: FlowRunContext) -> int:
-    if "counter" not in context.observed_flow_pauses:
-        context.observed_flow_pauses["counter"] = 1
-    else:
-        context.observed_flow_pauses["counter"] += 1
-    return context.observed_flow_pauses["counter"]
-
-
 def get_state_for_result(obj: Any) -> Optional[State]:
     """
     Get the state related to a result object.
@@ -667,7 +706,7 @@ def _get_hook_name(hook: Callable) -> str:
     )
 
 
-async def check_api_reachable(client: PrefectClient, fail_message: str):
+async def check_api_reachable(client: "PrefectClient", fail_message: str):
     # Do not perform a healthcheck if it exists and is not expired
     api_url = str(client.api_url)
     if api_url in API_HEALTHCHECKS:
@@ -709,6 +748,12 @@ def emit_task_run_state_change_event(
                     "message": truncated_to(
                         state_message_truncation_length, initial_state.message
                     ),
+                    "state_details": initial_state.state_details.model_dump(
+                        mode="json",
+                        exclude_none=True,
+                        exclude_unset=True,
+                        exclude={"flow_run_id", "task_run_id"},
+                    ),
                 }
                 if initial_state
                 else None
@@ -719,7 +764,33 @@ def emit_task_run_state_change_event(
                 "message": truncated_to(
                     state_message_truncation_length, validated_state.message
                 ),
+                "state_details": validated_state.state_details.model_dump(
+                    mode="json",
+                    exclude_none=True,
+                    exclude_unset=True,
+                    exclude={"flow_run_id", "task_run_id"},
+                ),
+                "data": validated_state.data.model_dump(mode="json")
+                if isinstance(validated_state.data, BaseResult)
+                else None,
             },
+            "task_run": task_run.model_dump(
+                mode="json",
+                exclude_none=True,
+                exclude={
+                    "id",
+                    "created",
+                    "updated",
+                    "flow_run_id",
+                    "state_id",
+                    "state_type",
+                    "state_name",
+                    "state",
+                    # server materialized fields
+                    "estimated_start_time_delta",
+                    "estimated_run_time",
+                },
+            ),
         },
         resource={
             "prefect.resource.id": f"prefect.task-run.{task_run.id}",
@@ -734,17 +805,100 @@ def emit_task_run_state_change_event(
                 else ""
             ),
             "prefect.state-type": str(validated_state.type.value),
+            "prefect.orchestration": "client",
         },
         follows=follows,
     )
 
 
-@contextmanager
-def collapse_excgroups() -> Generator[None, None, None]:
-    try:
-        yield
-    except BaseException as exc:
-        while isinstance(exc, BaseExceptionGroup) and len(exc.exceptions) == 1:
-            exc = exc.exceptions[0]
+def resolve_to_final_result(expr, context):
+    """
+    Resolve any `PrefectFuture`, or `State` types nested in parameters into
+    data. Designed to be use with `visit_collection`.
+    """
+    state = None
+
+    # Expressions inside quotes should not be modified
+    if isinstance(context.get("annotation"), quote):
+        raise StopVisiting()
+
+    if isinstance(expr, NewPrefectFuture):
+        upstream_task_run = context.get("current_task_run")
+        upstream_task = context.get("current_task")
+        if (
+            upstream_task
+            and upstream_task_run
+            and expr.task_run_id == upstream_task_run.id
+        ):
+            raise ValueError(
+                f"Discovered a task depending on itself. Raising to avoid a deadlock. Please inspect the inputs and dependencies of {upstream_task.name}."
+            )
+
+        expr.wait()
+        state = expr.state
+    elif isinstance(expr, State):
+        state = expr
+    else:
+        return expr
+
+    assert state
+
+    # Do not allow uncompleted upstreams except failures when `allow_failure` has
+    # been used
+    if not state.is_completed() and not (
+        # TODO: Note that the contextual annotation here is only at the current level
+        #       if `allow_failure` is used then another annotation is used, this will
+        #       incorrectly evaluate to false — to resolve this, we must track all
+        #       annotations wrapping the current expression but this is not yet
+        #       implemented.
+        isinstance(context.get("annotation"), allow_failure) and state.is_failed()
+    ):
+        raise UpstreamTaskError(
+            f"Upstream task run '{state.state_details.task_run_id}' did not reach a"
+            " 'COMPLETED' state."
+        )
+
+    _result = state.result(raise_on_failure=False, fetch=True)
+    if inspect.isawaitable(_result):
+        _result = run_coro_as_sync(_result)
+    return _result
+
+
+def resolve_inputs_sync(
+    parameters: Dict[str, Any], return_data: bool = True, max_depth: int = -1
+) -> Dict[str, Any]:
+    """
+    Resolve any `Quote`, `PrefectFuture`, or `State` types nested in parameters into
+    data.
+
+    Returns:
+        A copy of the parameters with resolved data
+
+    Raises:
+        UpstreamTaskError: If any of the upstream states are not `COMPLETED`
+    """
 
-        raise exc
+    if not parameters:
+        return {}
+
+    resolved_parameters = {}
+    for parameter, value in parameters.items():
+        try:
+            resolved_parameters[parameter] = visit_collection(
+                value,
+                visit_fn=resolve_to_final_result,
+                return_data=return_data,
+                max_depth=max_depth,
+                remove_annotations=True,
+                context={},
+            )
+        except UpstreamTaskError:
+            raise
+        except Exception as exc:
+            raise PrefectException(
+                f"Failed to resolve inputs in parameter {parameter!r}. If your"
+                " parameter type is not supported, consider using the `quote`"
+                " annotation to skip resolution of inputs."
+            ) from exc
+
+    return resolved_parameters

prefect/utilities/filesystem.py

@@ -1,12 +1,13 @@
 """
 Utilities for working with file systems
 """
+
 import os
 import pathlib
 import threading
 from contextlib import contextmanager
 from pathlib import Path, PureWindowsPath
-from typing import Union
+from typing import Optional, Union
 
 import fsspec
 import pathspec
@@ -32,7 +33,7 @@ def create_default_ignore_file(path: str) -> bool:
 
 
 def filter_files(
-    root: str = ".", ignore_patterns: list = None, include_dirs: bool = True
+    root: str = ".", ignore_patterns: Optional[list] = None, include_dirs: bool = True
 ) -> set:
     """
     This function accepts a root directory path and a list of file patterns to ignore, and returns
@@ -40,9 +41,7 @@ def filter_files(
 
     The specification matches that of [.gitignore files](https://git-scm.com/docs/gitignore).
     """
-    if ignore_patterns is None:
-        ignore_patterns = []
-    spec = pathspec.PathSpec.from_lines("gitwildmatch", ignore_patterns)
+    spec = pathspec.PathSpec.from_lines("gitwildmatch", ignore_patterns or [])
     ignored_files = {p.path for p in spec.match_tree_entries(root)}
     if include_dirs:
         all_files = {p.path for p in pathspec.util.iter_tree_entries(root)}

prefect/utilities/importtools.py

@@ -1,10 +1,10 @@
 import ast
 import importlib
 import importlib.util
-import inspect
 import os
 import runpy
 import sys
+import warnings
 from importlib.abc import Loader, MetaPathFinder
 from importlib.machinery import ModuleSpec
 from pathlib import Path
@@ -228,24 +228,18 @@ class DelayedImportErrorModule(ModuleType):
     [1]: https://github.com/scientific-python/lazy_loader
     """
 
-    def __init__(self, frame_data, help_message, *args, **kwargs):
-        self.__frame_data = frame_data
+    def __init__(self, error_message, help_message, *args, **kwargs):
+        self.__error_message = error_message
         self.__help_message = (
             help_message or "Import errors for this module are only reported when used."
         )
         super().__init__(*args, **kwargs)
 
     def __getattr__(self, attr):
-        if attr in ("__class__", "__file__", "__frame_data", "__help_message"):
+        if attr in ("__class__", "__file__", "__help_message"):
             super().__getattr__(attr)
         else:
-            fd = self.__frame_data
-            raise ModuleNotFoundError(
-                f"No module named '{fd['spec']}'\n\nThis module was originally imported"
-                f" at:\n  File \"{fd['filename']}\", line {fd['lineno']}, in"
-                f" {fd['function']}\n\n    {''.join(fd['code_context']).strip()}\n"
-                + self.__help_message
-            )
+            raise ModuleNotFoundError(self.__error_message)
 
 
 def lazy_import(
@@ -256,6 +250,13 @@ def lazy_import(
     Use this to retain module-level imports for libraries that we don't want to
     actually import until they are needed.
 
+    NOTE: Lazy-loading a subpackage can cause the subpackage to be imported
+    twice if another non-lazy import also imports the subpackage. For example,
+    using both `lazy_import("docker.errors")` and `import docker.errors` in the
+    same codebase will import `docker.errors` twice and can lead to unexpected
+    behavior, e.g. type check failures and import-time side effects running
+    twice.
+
     Adapted from the [Python documentation][1] and [lazy_loader][2]
 
     [1]: https://docs.python.org/3/library/importlib.html#implementing-lazy-imports
@@ -267,25 +268,23 @@ def lazy_import(
     except KeyError:
         pass
 
+    if "." in name:
+        warnings.warn(
+            "Lazy importing subpackages can lead to unexpected behavior.",
+            RuntimeWarning,
+        )
+
     spec = importlib.util.find_spec(name)
+
     if spec is None:
+        import_error_message = f"No module named '{name}'.\n{help_message}"
+
         if error_on_import:
-            raise ModuleNotFoundError(f"No module named '{name}'.\n{help_message}")
-        else:
-            try:
-                parent = inspect.stack()[1]
-                frame_data = {
-                    "spec": name,
-                    "filename": parent.filename,
-                    "lineno": parent.lineno,
-                    "function": parent.function,
-                    "code_context": parent.code_context,
-                }
-                return DelayedImportErrorModule(
-                    frame_data, help_message, "DelayedImportErrorModule"
-                )
-            finally:
-                del parent
+            raise ModuleNotFoundError(import_error_message)
+
+        return DelayedImportErrorModule(
+            import_error_message, help_message, "DelayedImportErrorModule"
+        )
 
     module = importlib.util.module_from_spec(spec)
     sys.modules[name] = module