prefect-client 3.1.6__tar.gz → 3.1.8__tar.gz

{prefect-client-3.1.6/src/prefect_client.egg-info → prefect-client-3.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.1.6
+Version: 3.1.8
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.

{prefect-client-3.1.6 → prefect-client-3.1.8}/requirements-client.txt

@@ -26,6 +26,7 @@ pydantic_extra_types >= 2.8.2, < 3.0.0
 pydantic_settings > 2.2.1
 python_dateutil >= 2.8.2, < 3.0.0
 python-slugify >= 5.0, < 9.0
+python-socks[asyncio] >= 2.5.3, < 3.0
 pyyaml >= 5.4.1, < 7.0.0
 rfc3339-validator >= 0.1.4, < 0.2.0
 rich >= 11.0, < 14.0

prefect-client-3.1.8/src/prefect/_experimental/lineage.py

@@ -0,0 +1,181 @@
+from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Sequence, Union
+
+from prefect.events.related import related_resources_from_run_context
+from prefect.events.schemas.events import RelatedResource, Resource
+from prefect.events.utilities import emit_event
+from prefect.settings import get_current_settings
+
+if TYPE_CHECKING:
+    from prefect.results import ResultStore
+
+UpstreamResources = Sequence[Union[RelatedResource, dict[str, str]]]
+DownstreamResources = Sequence[Union[Resource, dict[str, str]]]
+
+# Map block types to their URI schemes
+STORAGE_URI_SCHEMES = {
+    "local-file-system": "file://{path}",
+    "s3-bucket": "s3://{storage.bucket_name}/{path}",
+    "gcs-bucket": "gs://{storage.bucket}/{path}",
+    "azure-blob-storage": "azure-blob://{storage.container_name}/{path}",
+}
+
+
+def get_result_resource_uri(
+    store: "ResultStore",
+    key: str,
+) -> Optional[str]:
+    """
+    Generate a URI for a result based on its storage backend.
+
+    Args:
+        store: A `ResultStore` instance.
+        key: The key of the result to generate a URI for.
+    """
+    storage = store.result_storage
+    if storage is None:
+        return
+
+    path = store._resolved_key_path(key)
+
+    block_type = storage.get_block_type_slug()
+    if block_type and block_type in STORAGE_URI_SCHEMES:
+        return STORAGE_URI_SCHEMES[block_type].format(storage=storage, path=path)
+
+    # Generic fallback
+    return f"prefect://{block_type}/{path}"
+
+
+async def emit_lineage_event(
+    event_name: str,
+    upstream_resources: Optional[UpstreamResources] = None,
+    downstream_resources: Optional[DownstreamResources] = None,
+    direction_of_run_from_event: Literal["upstream", "downstream"] = "downstream",
+) -> None:
+    """Emit lineage events showing relationships between resources.
+
+    Args:
+        event_name: The name of the event to emit
+        upstream_resources: Optional list of RelatedResources that were upstream of
+            the event
+        downstream_resources: Optional list of Resources that were downstream
+            of the event
+        direction_of_run_from_event: The direction of the current run from
+            the event. E.g., if we're in a flow run and
+            `direction_of_run_from_event` is "downstream", then the flow run is
+            considered downstream of the resource's event.
+    """
+    from prefect.client.orchestration import get_client  # Avoid a circular import
+
+    if not get_current_settings().experiments.lineage_events_enabled:
+        return
+
+    upstream_resources = list(upstream_resources) if upstream_resources else []
+    downstream_resources = list(downstream_resources) if downstream_resources else []
+
+    async with get_client() as client:
+        related_resources = await related_resources_from_run_context(client)
+
+    # NOTE: We handle adding run-related resources to the event here instead of in
+    # the EventsWorker because not all run-related resources are upstream from
+    # every lineage event (they might be downstream). The EventsWorker only adds
+    # related resources to the "related" field in the event, which, for
+    # lineage-related events, tracks upstream resources only. For downstream
+    # resources, we need to emit an event for each downstream resource.
+    if direction_of_run_from_event == "downstream":
+        downstream_resources.extend(related_resources)
+    else:
+        upstream_resources.extend(related_resources)
+
+    # Emit an event for each downstream resource. This is necessary because
+    # our event schema allows one primary resource and many related resources,
+    # and for the purposes of lineage, related resources can only represent
+    # upstream resources.
+    for resource in downstream_resources:
+        # Downstream lineage resources need to have the
+        # prefect.resource.lineage-group label. All upstram resources from a
+        # downstream resource with this label will be considered lineage-related
+        # resources.
+        if "prefect.resource.lineage-group" not in resource:
+            resource["prefect.resource.lineage-group"] = "global"
+
+        emit_kwargs: Dict[str, Any] = {
+            "event": event_name,
+            "resource": resource,
+            "related": upstream_resources,
+        }
+
+        emit_event(**emit_kwargs)
+
+
+async def emit_result_read_event(
+    store: "ResultStore",
+    result_key: str,
+    downstream_resources: Optional[DownstreamResources] = None,
+    cached: bool = False,
+) -> None:
+    """
+    Emit a lineage event showing a task or flow result was read.
+
+    Args:
+        store: A `ResultStore` instance.
+        result_key: The key of the result to generate a URI for.
+        downstream_resources: List of resources that were
+            downstream of the event's resource.
+    """
+    if not get_current_settings().experiments.lineage_events_enabled:
+        return
+
+    result_resource_uri = get_result_resource_uri(store, result_key)
+    if result_resource_uri:
+        upstream_resources = [
+            RelatedResource(
+                root={
+                    "prefect.resource.id": result_resource_uri,
+                    "prefect.resource.role": "result",
+                }
+            )
+        ]
+        event_name = "prefect.result.read"
+        if cached:
+            event_name += ".cached"
+
+        await emit_lineage_event(
+            event_name=event_name,
+            upstream_resources=upstream_resources,
+            downstream_resources=downstream_resources,
+            direction_of_run_from_event="downstream",
+        )
+
+
+async def emit_result_write_event(
+    store: "ResultStore",
+    result_key: str,
+    upstream_resources: Optional[UpstreamResources] = None,
+) -> None:
+    """
+    Emit a lineage event showing a task or flow result was written.
+
+    Args:
+        store: A `ResultStore` instance.
+        result_key: The key of the result to generate a URI for.
+        upstream_resources: Optional list of resources that were
+            upstream of the event's resource.
+    """
+    if not get_current_settings().experiments.lineage_events_enabled:
+        return
+
+    result_resource_uri = get_result_resource_uri(store, result_key)
+    if result_resource_uri:
+        downstream_resources = [
+            {
+                "prefect.resource.id": result_resource_uri,
+                "prefect.resource.role": "result",
+                "prefect.resource.lineage-group": "global",
+            }
+        ]
+        await emit_lineage_event(
+            event_name="prefect.result.write",
+            upstream_resources=upstream_resources,
+            downstream_resources=downstream_resources,
+            direction_of_run_from_event="upstream",
+        )

{prefect-client-3.1.6 → prefect-client-3.1.8}/src/prefect/_internal/compatibility/async_dispatch.py

@@ -1,11 +1,12 @@
 import asyncio
 import inspect
 from functools import wraps
-from typing import Any, Callable, Coroutine, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Callable, Coroutine, Optional, TypeVar, Union
 
 from typing_extensions import ParamSpec
 
-from prefect.tasks import Task
+if TYPE_CHECKING:
+    from prefect.tasks import Task
 
 R = TypeVar("R")
 P = ParamSpec("P")
@@ -13,20 +14,46 @@ P = ParamSpec("P")
 
 def is_in_async_context() -> bool:
     """
-    Returns True if called from within an async context (coroutine or running event loop)
+    Returns True if called from within an async context.
+
+    An async context is one of:
+        - a coroutine
+        - a running event loop
+        - a task or flow that is async
     """
+    from prefect.context import get_run_context
+    from prefect.exceptions import MissingContextError
+
     try:
-        asyncio.get_running_loop()
-        return True
-    except RuntimeError:
-        return False
+        run_ctx = get_run_context()
+        parent_obj = getattr(run_ctx, "task", None)
+        if not parent_obj:
+            parent_obj = getattr(run_ctx, "flow", None)
+        return getattr(parent_obj, "isasync", True)
+    except MissingContextError:
+        # not in an execution context, make best effort to
+        # decide whether to syncify
+        try:
+            asyncio.get_running_loop()
+            return True
+        except RuntimeError:
+            return False
 
 
-def _is_acceptable_callable(obj: Union[Callable, Task]) -> bool:
+def _is_acceptable_callable(
+    obj: Union[Callable[P, R], "Task[P, R]", classmethod],
+) -> bool:
     if inspect.iscoroutinefunction(obj):
         return True
-    if isinstance(obj, Task) and inspect.iscoroutinefunction(obj.fn):
+
+    # Check if a task or flow. Need to avoid importing `Task` or `Flow` here
+    # due to circular imports.
+    if (fn := getattr(obj, "fn", None)) and inspect.iscoroutinefunction(fn):
         return True
+
+    if isinstance(obj, classmethod) and inspect.iscoroutinefunction(obj.__func__):
+        return True
+
     return False
 
 
@@ -56,6 +83,8 @@ def async_dispatch(
 
             if should_run_sync:
                 return sync_fn(*args, **kwargs)
+            if isinstance(async_impl, classmethod):
+                return async_impl.__func__(*args, **kwargs)
             return async_impl(*args, **kwargs)
 
         return wrapper  # type: ignore

{prefect-client-3.1.6 → prefect-client-3.1.8}/src/prefect/_internal/pydantic/v2_validated_func.py

@@ -6,16 +6,17 @@ Specifically it allows for us to validate v2 models used as flow/task
 arguments.
 """
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
 
 # importing directly from v2 to be able to create a v2 model
 from pydantic import BaseModel, ConfigDict, create_model, field_validator
 from pydantic.v1.decorator import ValidatedFunction
 from pydantic.v1.errors import ConfigError
 from pydantic.v1.utils import to_camel
+from typing_extensions import TypeAlias
 
 if TYPE_CHECKING:
-    ConfigType = Union[None, Type[Any], Dict[str, Any]]
+    ConfigType: TypeAlias = Union[None, type[Any], dict[str, Any]]
 
 V_POSITIONAL_ONLY_NAME = "v__positional_only"
 V_DUPLICATE_KWARGS = "v__duplicate_kwargs"
@@ -24,13 +25,17 @@ V_DUPLICATE_KWARGS = "v__duplicate_kwargs"
 class V2ValidatedFunction(ValidatedFunction):
     def create_model(
         self,
-        fields: Dict[str, Any],
+        fields: dict[str, Any],
         takes_args: bool,
         takes_kwargs: bool,
-        config: ConfigDict,
+        config: "ConfigType",
     ) -> None:
         pos_args = len(self.arg_mapping)
 
+        config = {} if config is None else config
+        if not isinstance(config, dict):
+            raise TypeError(f"config must be None or a dict, got {type(config)}")
+
         if config.get("fields") or config.get("alias_generator"):
             raise ConfigError(
                 'Setting the "fields" and "alias_generator" property on custom Config'
@@ -42,11 +47,11 @@ class V2ValidatedFunction(ValidatedFunction):
 
         # This is the key change -- inheriting the BaseModel class from v2
         class DecoratorBaseModel(BaseModel):
-            model_config = config
+            model_config: ClassVar[ConfigDict] = ConfigDict(**config)
 
             @field_validator(self.v_args_name, check_fields=False)
             @classmethod
-            def check_args(cls, v: Optional[List[Any]]) -> Optional[List[Any]]:
+            def check_args(cls, v: Optional[list[Any]]) -> Optional[list[Any]]:
                 if takes_args or v is None:
                     return v
 
@@ -58,8 +63,8 @@ class V2ValidatedFunction(ValidatedFunction):
             @field_validator(self.v_kwargs_name, check_fields=False)
             @classmethod
             def check_kwargs(
-                cls, v: Optional[Dict[str, Any]]
-            ) -> Optional[Dict[str, Any]]:
+                cls, v: Optional[dict[str, Any]]
+            ) -> Optional[dict[str, Any]]:
                 if takes_kwargs or v is None:
                     return v
 
@@ -69,7 +74,7 @@ class V2ValidatedFunction(ValidatedFunction):
 
             @field_validator(V_POSITIONAL_ONLY_NAME, check_fields=False)
             @classmethod
-            def check_positional_only(cls, v: Optional[List[str]]) -> None:
+            def check_positional_only(cls, v: Optional[list[str]]) -> None:
                 if v is None:
                     return
 
@@ -82,7 +87,7 @@ class V2ValidatedFunction(ValidatedFunction):
 
             @field_validator(V_DUPLICATE_KWARGS, check_fields=False)
             @classmethod
-            def check_duplicate_kwargs(cls, v: Optional[List[str]]) -> None:
+            def check_duplicate_kwargs(cls, v: Optional[list[str]]) -> None:
                 if v is None:
                     return
 

{prefect-client-3.1.6 → prefect-client-3.1.8}/src/prefect/_internal/retries.py

@@ -1,10 +1,15 @@
 import asyncio
 from functools import wraps
-from typing import Any, Callable, Tuple, Type
+from typing import Callable, Optional, Tuple, Type, TypeVar
+
+from typing_extensions import ParamSpec
 
 from prefect._internal._logging import logger
 from prefect.utilities.math import clamped_poisson_interval
 
+P = ParamSpec("P")
+R = TypeVar("R")
+
 
 def exponential_backoff_with_jitter(
     attempt: int, base_delay: float, max_delay: float
@@ -21,7 +26,8 @@ def retry_async_fn(
     base_delay: float = 1,
     max_delay: float = 10,
     retry_on_exceptions: Tuple[Type[Exception], ...] = (Exception,),
-):
+    operation_name: Optional[str] = None,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
     """A decorator for retrying an async function.
 
     Args:
@@ -33,23 +39,26 @@ def retry_async_fn(
         max_delay: The maximum delay to use for the last attempt.
         retry_on_exceptions: A tuple of exception types to retry on. Defaults to
             retrying on all exceptions.
+        operation_name: Optional name to use for logging the operation instead of
+            the function name. If None, uses the function name.
     """
 
-    def decorator(func):
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
         @wraps(func)
-        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            name = operation_name or func.__name__
             for attempt in range(max_attempts):
                 try:
                     return await func(*args, **kwargs)
                 except retry_on_exceptions as e:
                     if attempt == max_attempts - 1:
                         logger.exception(
-                            f"Function {func.__name__!r} failed after {max_attempts} attempts"
+                            f"Function {name!r} failed after {max_attempts} attempts"
                         )
                         raise
                     delay = backoff_strategy(attempt, base_delay, max_delay)
                     logger.warning(
-                        f"Attempt {attempt + 1} of function {func.__name__!r} failed with {type(e).__name__}. "
+                        f"Attempt {attempt + 1} of function {name!r} failed with {type(e).__name__}: {str(e)}. "
                         f"Retrying in {delay:.2f} seconds..."
                     )
                     await asyncio.sleep(delay)

{prefect-client-3.1.6 → prefect-client-3.1.8}/src/prefect/_internal/schemas/bases.py

@@ -13,9 +13,10 @@ from pydantic import (
     ConfigDict,
     Field,
 )
-from pydantic_extra_types.pendulum_dt import DateTime
 from typing_extensions import Self
 
+from prefect.types import DateTime
+
 T = TypeVar("T")
 
 

{prefect-client-3.1.6 → prefect-client-3.1.8}/src/prefect/_internal/schemas/validators.py

@@ -18,9 +18,9 @@ from uuid import UUID
 import jsonschema
 import pendulum
 import yaml
-from pydantic_extra_types.pendulum_dt import DateTime
 
 from prefect.exceptions import InvalidRepositoryURLError
+from prefect.types import DateTime
 from prefect.utilities.collections import isiterable
 from prefect.utilities.dockerutils import get_prefect_image_name
 from prefect.utilities.filesystem import relative_path_to_current_platform
@@ -33,6 +33,7 @@ LOWERCASE_LETTERS_NUMBERS_AND_UNDERSCORES_REGEX = "^[a-z0-9_]*$"
 
 if TYPE_CHECKING:
     from prefect.blocks.core import Block
+    from prefect.serializers import Serializer
     from prefect.utilities.callables import ParameterSchema
 
 
@@ -578,7 +579,7 @@ def validate_picklelib_and_modules(values: dict) -> dict:
     return values
 
 
-def validate_dump_kwargs(value: dict) -> dict:
+def validate_dump_kwargs(value: dict[str, Any]) -> dict[str, Any]:
     # `default` is set by `object_encoder`. A user provided callable would make this
     # class unserializable anyway.
     if "default" in value:
@@ -586,7 +587,7 @@ def validate_dump_kwargs(value: dict) -> dict:
     return value
 
 
-def validate_load_kwargs(value: dict) -> dict:
+def validate_load_kwargs(value: dict[str, Any]) -> dict[str, Any]:
     # `object_hook` is set by `object_decoder`. A user provided callable would make
     # this class unserializable anyway.
     if "object_hook" in value:
@@ -596,7 +597,7 @@ def validate_load_kwargs(value: dict) -> dict:
     return value
 
 
-def cast_type_names_to_serializers(value):
+def cast_type_names_to_serializers(value: Union[str, "Serializer"]) -> "Serializer":
     from prefect.serializers import Serializer
 
     if isinstance(value, str):

{prefect-client-3.1.6 → prefect-client-3.1.8}/src/prefect/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-12-11T10:46:23-0500",
+ "date": "2024-12-17T10:20:43-0800",
  "dirty": true,
  "error": null,
- "full-revisionid": "6bfce9e87c6e7e3e44c65d8b5ecc6ab156f2c7c8",
- "version": "3.1.6"
+ "full-revisionid": "53a83ebc4a9a26a1d7cfd91d6be7cf6bd30e7f21",
+ "version": "3.1.8"
 }
 '''  # END VERSION_JSON