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

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,9 +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
@@ -14,8 +15,11 @@ from typing import Any, Callable, Dict, Iterable, NamedTuple, Optional, Union
 import fsspec
 
 from prefect.exceptions import ScriptError
+from prefect.logging.loggers import get_logger
 from prefect.utilities.filesystem import filename, is_local_path, tmpchdir
 
+logger = get_logger(__name__)
+
 
 def to_qualified_name(obj: Any) -> str:
     """
@@ -224,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(
@@ -252,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
@@ -263,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
@@ -356,3 +359,70 @@ class AliasedModuleLoader(Loader):
         if self.callback is not None:
             self.callback(self.alias)
         sys.modules[self.alias] = root_module
+
+
+def safe_load_namespace(source_code: str):
+    """
+    Safely load a namespace from source code.
+
+    This function will attempt to import all modules and classes defined in the source
+    code. If an import fails, the error is caught and the import is skipped. This function
+    will also attempt to compile and evaluate class and function definitions locally.
+
+    Args:
+        source_code: The source code to load
+
+    Returns:
+        The namespace loaded from the source code. Can be used when evaluating source
+        code.
+    """
+    parsed_code = ast.parse(source_code)
+
+    namespace = {}
+
+    # Walk through the AST and find all import statements
+    for node in ast.walk(parsed_code):
+        if isinstance(node, ast.Import):
+            for alias in node.names:
+                module_name = alias.name
+                as_name = alias.asname if alias.asname else module_name
+                try:
+                    # Attempt to import the module
+                    namespace[as_name] = importlib.import_module(module_name)
+                    logger.debug("Successfully imported %s", module_name)
+                except ImportError as e:
+                    logger.debug(f"Failed to import {module_name}: {e}")
+        elif isinstance(node, ast.ImportFrom):
+            module_name = node.module
+            if module_name is None:
+                continue
+            try:
+                module = importlib.import_module(module_name)
+                for alias in node.names:
+                    name = alias.name
+                    asname = alias.asname if alias.asname else name
+                    try:
+                        # Get the specific attribute from the module
+                        attribute = getattr(module, name)
+                        namespace[asname] = attribute
+                    except AttributeError as e:
+                        logger.debug(
+                            "Failed to retrieve %s from %s: %s", name, module_name, e
+                        )
+            except ImportError as e:
+                logger.debug("Failed to import from %s: %s", node.module, e)
+
+    # Handle local class definitions
+    for node in ast.walk(parsed_code):
+        if isinstance(node, (ast.ClassDef, ast.FunctionDef)):
+            try:
+                # Compile and execute each class and function definition locally
+                code = compile(
+                    ast.Module(body=[node], type_ignores=[]),
+                    filename="<ast>",
+                    mode="exec",
+                )
+                exec(code, namespace)
+            except Exception as e:
+                logger.debug("Failed to compile class definition: %s", e)
+    return namespace