prefect-client 2.19.3__py3-none-any.whl → 3.0.0rc1__py3-none-any.whl

prefect/utilities/collections.py

@@ -4,6 +4,7 @@ Utilities for extensions of and operations on Python collections.
 
 import io
 import itertools
+import warnings
 from collections import OrderedDict, defaultdict
 from collections.abc import Iterator as IteratorABC
 from collections.abc import Sequence
@@ -28,12 +29,7 @@ from typing import (
 )
 from unittest.mock import Mock
 
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-
-if HAS_PYDANTIC_V2:
-    import pydantic.v1 as pydantic
-else:
-    import pydantic
+import pydantic
 
 # Quote moved to `prefect.utilities.annotations` but preserved here for compatibility
 from prefect.utilities.annotations import BaseAnnotation, Quote, quote  # noqa
@@ -225,7 +221,7 @@ class StopVisiting(BaseException):
 
 def visit_collection(
     expr,
-    visit_fn: Callable[[Any], Any],
+    visit_fn: Union[Callable[[Any, dict], Any], Callable[[Any], Any]],
     return_data: bool = False,
     max_depth: int = -1,
     context: Optional[dict] = None,
@@ -252,8 +248,10 @@ def visit_collection(
 
     Args:
         expr (Any): a Python object or expression
-        visit_fn (Callable[[Any], Awaitable[Any]]): an async function that
-            will be applied to every non-collection element of expr.
+        visit_fn (Callable[[Any, Optional[dict]], Awaitable[Any]]): a function that
+            will be applied to every non-collection element of expr. The function can
+            accept one or two arguments. If two arguments are accepted, the second
+            argument will be the context dictionary.
         return_data (bool): if `True`, a copy of `expr` containing data modified
             by `visit_fn` will be returned. This is slower than `return_data=False`
             (the default).
@@ -343,39 +341,26 @@ def visit_collection(
         result = typ(**items) if return_data else None
 
     elif isinstance(expr, pydantic.BaseModel):
-        # NOTE: This implementation *does not* traverse private attributes
-        # Pydantic does not expose extras in `__fields__` so we use `__fields_set__`
-        # as well to get all of the relevant attributes
-        # Check for presence of attrs even if they're in the field set due to pydantic#4916
-        model_fields = {
-            f for f in expr.__fields_set__.union(expr.__fields__) if hasattr(expr, f)
-        }
-        items = [visit_nested(getattr(expr, key)) for key in model_fields]
+        typ = cast(Type[pydantic.BaseModel], typ)
 
-        if return_data:
-            # Collect fields with aliases so reconstruction can use the correct field name
-            aliases = {
-                key: value.alias
-                for key, value in expr.__fields__.items()
-                if value.has_alias
-            }
+        # when extra=allow, fields not in model_fields may be in model_fields_set
+        model_fields = expr.model_fields_set.union(expr.model_fields.keys())
 
-            model_instance = typ(
-                **{
-                    aliases.get(key) or key: value
-                    for key, value in zip(model_fields, items)
-                }
-            )
+        # We may encounter a deprecated field here, but this isn't the caller's fault
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=DeprecationWarning)
 
-            # Private attributes are not included in `__fields_set__` but we do not want
-            # to drop them from the model so we restore them after constructing a new
-            # model
-            for attr in expr.__private_attributes__:
-                # Use `object.__setattr__` to avoid errors on immutable models
-                object.__setattr__(model_instance, attr, getattr(expr, attr))
+            updated_data = {
+                field: visit_nested(getattr(expr, field)) for field in model_fields
+            }
 
-            # Preserve data about which fields were explicitly set on the original model
-            object.__setattr__(model_instance, "__fields_set__", expr.__fields_set__)
+        if return_data:
+            # Use construct to avoid validation and handle immutability
+            model_instance = typ.model_construct(
+                _fields_set=expr.model_fields_set, **updated_data
+            )
+            for private_attr in expr.__private_attributes__:
+                setattr(model_instance, private_attr, getattr(expr, private_attr))
             result = model_instance
         else:
             result = None

prefect/utilities/dispatch.py

@@ -90,8 +90,10 @@ def get_dispatch_key(
 
 @classmethod
 def _register_subclass_of_base_type(cls, **kwargs):
-    if cls.__init_subclass_original__:
+    if hasattr(cls, "__init_subclass_original__"):
         cls.__init_subclass_original__(**kwargs)
+    elif hasattr(cls, "__pydantic_init_subclass_original__"):
+        cls.__pydantic_init_subclass_original__(**kwargs)
 
     # Do not register abstract base classes
     if abc.ABC in getattr(cls, "__bases__", []):
@@ -114,8 +116,14 @@ def register_base_type(cls: T) -> T:
         registry[base_key] = cls
 
     # Add automatic subtype registration
-    cls.__init_subclass_original__ = getattr(cls, "__init_subclass__")
-    cls.__init_subclass__ = _register_subclass_of_base_type
+    if hasattr(cls, "__pydantic_init_subclass__"):
+        cls.__pydantic_init_subclass_original__ = getattr(
+            cls, "__pydantic_init_subclass__"
+        )
+        cls.__pydantic_init_subclass__ = _register_subclass_of_base_type
+    else:
+        cls.__init_subclass_original__ = getattr(cls, "__init_subclass__")
+        cls.__init_subclass__ = _register_subclass_of_base_type
 
     return cls
 

prefect/utilities/dockerutils.py

@@ -27,6 +27,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}"

prefect/utilities/engine.py

@@ -6,6 +6,7 @@ import signal
 import time
 from functools import partial
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Dict,
@@ -24,7 +25,6 @@ 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,
@@ -44,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,
@@ -55,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")
@@ -93,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
@@ -118,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 = []
 
@@ -225,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
@@ -264,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
@@ -311,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,
@@ -412,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,
@@ -459,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
 
@@ -515,7 +559,9 @@ 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
+    if 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())
         )
@@ -528,14 +574,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.
@@ -664,7 +702,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:
@@ -734,3 +772,85 @@ def emit_task_run_state_change_event(
         },
         follows=follows,
     )
+
+
+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):
+        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`
+    """
+
+    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/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