prefect-client 3.0.0rc1__py3-none-any.whl → 3.0.0rc2__py3-none-any.whl

prefect/blocks/redis.py

@@ -0,0 +1,168 @@
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import AsyncGenerator, Optional, Union
+
+try:
+    import redis.asyncio as redis
+except ImportError:
+    raise ImportError(
+        "`redis-py` must be installed to use the `RedisStorageContainer` block. "
+        "You can install it with `pip install redis>=5.0.1"
+    )
+
+from pydantic import Field
+from pydantic.types import SecretStr
+from typing_extensions import Self
+
+from prefect.filesystems import WritableFileSystem
+from prefect.utilities.asyncutils import sync_compatible
+
+
+class RedisStorageContainer(WritableFileSystem):
+    """
+    Block used to interact with Redis as a filesystem
+
+    Attributes:
+        host (str): The value to store.
+        port (int): The value to store.
+        db (int): The value to store.
+        username (str): The value to store.
+        password (str): The value to store.
+        connection_string (str): The value to store.
+
+    Example:
+        Create a new block from hostname, username and password:
+        ```python
+        from prefect.blocks.redis import RedisStorageContainer
+
+        block = RedisStorageContainer.from_host(
+            host="myredishost.com", username="redis", password="SuperSecret")
+        block.save("BLOCK_NAME")
+        ```
+
+        Create a new block from a connection string
+        ```python
+        from prefect.blocks.redis import RedisStorageContainer
+        block = RedisStorageContainer.from_url(""redis://redis:SuperSecret@myredishost.com:6379")
+        block.save("BLOCK_NAME")
+        ```
+    """
+
+    _logo_url = "https://stprododpcmscdnendpoint.azureedge.net/assets/icons/redis.png"
+
+    host: Optional[str] = Field(default=None, description="Redis hostname")
+    port: int = Field(default=6379, description="Redis port")
+    db: int = Field(default=0, description="Redis DB index")
+    username: Optional[SecretStr] = Field(default=None, description="Redis username")
+    password: Optional[SecretStr] = Field(default=None, description="Redis password")
+    connection_string: Optional[SecretStr] = Field(
+        default=None, description="Redis connection string"
+    )
+
+    def block_initialization(self) -> None:
+        if self.connection_string:
+            return
+        if not self.host:
+            raise ValueError("Initialization error: 'host' is required but missing.")
+        if self.username and not self.password:
+            raise ValueError(
+                "Initialization error: 'username' is provided, but 'password' is missing. Both are required."
+            )
+
+    @sync_compatible
+    async def read_path(self, path: Union[Path, str]):
+        """Read the redis content at `path`
+
+        Args:
+            path: Redis key to read from
+
+        Returns:
+            Contents at key as bytes
+        """
+        async with self._client() as client:
+            return await client.get(str(path))
+
+    @sync_compatible
+    async def write_path(self, path: Union[Path, str], content: bytes):
+        """Write `content` to the redis at `path`
+
+        Args:
+            path: Redis key to write to
+            content: Binary object to write
+        """
+
+        async with self._client() as client:
+            return await client.set(str(path), content)
+
+    @asynccontextmanager
+    async def _client(self) -> AsyncGenerator[redis.Redis, None]:
+        if self.connection_string:
+            client = redis.Redis.from_url(self.connection_string.get_secret_value())
+        else:
+            assert self.host
+            client = redis.Redis(
+                host=self.host,
+                port=self.port,
+                username=self.username.get_secret_value() if self.username else None,
+                password=self.password.get_secret_value() if self.password else None,
+                db=self.db,
+            )
+
+        try:
+            yield client
+        finally:
+            await client.aclose()
+
+    @classmethod
+    def from_host(
+        cls,
+        host: str,
+        port: int = 6379,
+        db: int = 0,
+        username: Union[None, str, SecretStr] = None,
+        password: Union[None, str, SecretStr] = None,
+    ) -> Self:
+        """Create block from hostname, username and password
+
+        Args:
+            host: Redis hostname
+            username: Redis username
+            password: Redis password
+            port: Redis port
+
+        Returns:
+            `RedisStorageContainer` instance
+        """
+
+        username = SecretStr(username) if isinstance(username, str) else username
+        password = SecretStr(password) if isinstance(password, str) else password
+
+        return cls(host=host, port=port, db=db, username=username, password=password)
+
+    @classmethod
+    def from_connection_string(cls, connection_string: Union[str, SecretStr]) -> Self:
+        """Create block from a Redis connection string
+
+        Supports the following URL schemes:
+        - `redis://` creates a TCP socket connection
+        - `rediss://` creates a SSL wrapped TCP socket connection
+        - `unix://` creates a Unix Domain Socket connection
+
+        See [Redis docs](https://redis.readthedocs.io/en/stable/examples
+        /connection_examples.html#Connecting-to-Redis-instances-by-specifying-a-URL
+        -scheme.) for more info.
+
+        Args:
+            connection_string: Redis connection string
+
+        Returns:
+            `RedisStorageContainer` instance
+        """
+
+        connection_string = (
+            SecretStr(connection_string)
+            if isinstance(connection_string, str)
+            else connection_string
+        )
+
+        return cls(connection_string=connection_string)

prefect/client/orchestration.py

@@ -8,11 +8,13 @@ from typing import (
     Dict,
     Iterable,
     List,
+    Literal,
     Optional,
     Set,
     Tuple,
     TypeVar,
     Union,
+    overload,
 )
 from uuid import UUID, uuid4
 
@@ -156,9 +158,23 @@ class ServerType(AutoEnum):
     CLOUD = AutoEnum.auto()
 
 
+@overload
+def get_client(
+    httpx_settings: Optional[Dict[str, Any]] = None, sync_client: Literal[False] = False
+) -> "PrefectClient":
+    ...
+
+
+@overload
+def get_client(
+    httpx_settings: Optional[Dict[str, Any]] = None, sync_client: Literal[True] = True
+) -> "SyncPrefectClient":
+    ...
+
+
 def get_client(
     httpx_settings: Optional[Dict[str, Any]] = None, sync_client: bool = False
-) -> Union["PrefectClient", "SyncPrefectClient"]:
+):
     """
     Retrieve a HTTP client for communicating with the Prefect REST API.
 

prefect/client/schemas/objects.py

@@ -94,6 +94,14 @@ class StateType(AutoEnum):
     CANCELLING = AutoEnum.auto()
 
 
+TERMINAL_STATES = {
+    StateType.COMPLETED,
+    StateType.CANCELLED,
+    StateType.FAILED,
+    StateType.CRASHED,
+}
+
+
 class WorkPoolStatus(AutoEnum):
     """Enumeration of work pool statuses."""
 
@@ -280,7 +288,7 @@ class State(ObjectBaseModel, Generic[R]):
     def default_scheduled_start_time(self) -> Self:
         if self.type == StateType.SCHEDULED:
             if not self.state_details.scheduled_time:
-                self.state_details.scheduled_time = pendulum.now("utc")
+                self.state_details.scheduled_time = DateTime.now("utc")
         return self
 
     def is_scheduled(self) -> bool:
@@ -308,12 +316,7 @@ class State(ObjectBaseModel, Generic[R]):
         return self.type == StateType.CANCELLING
 
     def is_final(self) -> bool:
-        return self.type in {
-            StateType.CANCELLED,
-            StateType.FAILED,
-            StateType.COMPLETED,
-            StateType.CRASHED,
-        }
+        return self.type in TERMINAL_STATES
 
     def is_paused(self) -> bool:
         return self.type == StateType.PAUSED
@@ -550,7 +553,8 @@ class FlowRun(ObjectBaseModel):
         examples=["State(type=StateType.COMPLETED)"],
     )
     job_variables: Optional[dict] = Field(
-        default=None, description="Job variables for the flow run."
+        default=None,
+        description="Job variables for the flow run.",
     )
 
     # These are server-side optimizations and should not be present on client models

prefect/concurrency/asyncio.py

@@ -11,7 +11,7 @@ except ImportError:
     # pendulum < 3
     from pendulum.period import Period as Interval  # type: ignore
 
-from prefect import get_client
+from prefect.client.orchestration import get_client
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
 from prefect.utilities.timeout import timeout_async
 

prefect/concurrency/services.py

@@ -10,9 +10,9 @@ from typing import (
 import httpx
 from starlette import status
 
-from prefect import get_client
 from prefect._internal.concurrency import logger
 from prefect._internal.concurrency.services import QueueService
+from prefect.client.orchestration import get_client
 
 if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient

prefect/deployments/base.py

@@ -398,7 +398,13 @@ async def _find_flow_functions_in_file(filename: str) -> List[Dict]:
             return decorated_functions
 
     for node in ast.walk(tree):
-        if isinstance(node, ast.FunctionDef):
+        if isinstance(
+            node,
+            (
+                ast.FunctionDef,
+                ast.AsyncFunctionDef,
+            ),
+        ):
             for decorator in node.decorator_list:
                 # handles @flow
                 is_name_match = (

prefect/events/schemas/events.py

@@ -83,6 +83,8 @@ class RelatedResource(Resource):
 class Event(PrefectBaseModel):
     """The client-side view of an event that has happened to a Resource"""
 
+    model_config = ConfigDict(extra="ignore")
+
     occurred: DateTime = Field(
         default_factory=lambda: pendulum.now("UTC"),
         description="When the event happened from the sender's perspective",

prefect/flow_engine.py

@@ -25,9 +25,9 @@ import anyio._backends._asyncio
 from sniffio import AsyncLibraryNotFoundError
 from typing_extensions import ParamSpec
 
-from prefect import Task, get_client
+from prefect import Task
 from prefect._internal.concurrency.api import create_call, from_sync
-from prefect.client.orchestration import SyncPrefectClient
+from prefect.client.orchestration import SyncPrefectClient, get_client
 from prefect.client.schemas import FlowRun, TaskRun
 from prefect.client.schemas.filters import FlowRunFilter
 from prefect.client.schemas.sorting import FlowRunSort

prefect/flow_runs.py

@@ -76,7 +76,7 @@ async def wait_for_flow_run(
             ```python
             import asyncio
 
-            from prefect import get_client
+            from prefect.client.orchestration import get_client
             from prefect.flow_runs import wait_for_flow_run
 
             async def main():
@@ -94,7 +94,7 @@ async def wait_for_flow_run(
             ```python
             import asyncio
 
-            from prefect import get_client
+            from prefect.client.orchestration import get_client
             from prefect.flow_runs import wait_for_flow_run
 
             async def main(num_runs: int):

prefect/flows.py

@@ -1913,7 +1913,14 @@ def load_flow_argument_from_entrypoint(
         (
             node
             for node in ast.walk(parsed_code)
-            if isinstance(node, ast.FunctionDef) and node.name == func_name
+            if isinstance(
+                node,
+                (
+                    ast.FunctionDef,
+                    ast.AsyncFunctionDef,
+                ),
+            )
+            and node.name == func_name
         ),
         None,
     )

prefect/futures.py

@@ -1,7 +1,6 @@
 import abc
 import concurrent.futures
 import inspect
-import time
 import uuid
 from functools import partial
 from typing import Any, Generic, Optional, Set, Union, cast
@@ -11,13 +10,17 @@ from typing_extensions import TypeVar
 from prefect.client.orchestration import get_client
 from prefect.client.schemas.objects import TaskRun
 from prefect.exceptions import ObjectNotFound
+from prefect.logging.loggers import get_logger
 from prefect.states import Pending, State
+from prefect.task_runs import TaskRunWaiter
 from prefect.utilities.annotations import quote
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.collections import StopVisiting, visit_collection
 
 F = TypeVar("F")
 
+logger = get_logger(__name__)
+
 
 class PrefectFuture(abc.ABC):
     """
@@ -146,68 +149,66 @@ class PrefectDistributedFuture(PrefectFuture):
     Represents the result of a computation happening anywhere.
 
     This class is typically used to interact with the result of a task run
-    scheduled to run in a Prefect task server but can be used to interact with
+    scheduled to run in a Prefect task worker but can be used to interact with
     any task run scheduled in Prefect's API.
     """
 
-    def __init__(self, task_run_id: uuid.UUID):
-        self._task_run = None
-        self._client = None
-        super().__init__(task_run_id=task_run_id)
-
-    @property
-    def client(self):
-        if self._client is None:
-            self._client = get_client(sync_client=True)
-        return self._client
+    def wait(self, timeout: Optional[float] = None) -> None:
+        return run_coro_as_sync(self.wait_async(timeout=timeout))
 
-    @property
-    def task_run(self):
-        if self._task_run is None:
-            self._task_run = self.client.read_task_run(task_run_id=self.task_run_id)
-        return self._task_run
-
-    @task_run.setter
-    def task_run(self, task_run):
-        self._task_run = task_run
-
-    def wait(
-        self, timeout: Optional[float] = None, polling_interval: Optional[float] = 0.2
-    ) -> None:
-        start_time = time.time()
-        # TODO: Websocket implementation?
-        while True:
-            self.task_run = cast(
-                TaskRun, self.client.read_task_run(task_run_id=self.task_run_id)
+    async def wait_async(self, timeout: Optional[float] = None):
+        if self._final_state:
+            logger.debug(
+                "Final state already set for %s. Returning...", self.task_run_id
             )
-            if self.task_run.state and self.task_run.state.is_final():
-                self._final_state = self.task_run.state
-                return
-            if timeout is not None and (time.time() - start_time) > timeout:
+            return
+
+        # Read task run to see if it is still running
+        async with get_client() as client:
+            task_run = await client.read_task_run(task_run_id=self._task_run_id)
+            if task_run.state.is_final():
+                logger.debug(
+                    "Task run %s already finished. Returning...",
+                    self.task_run_id,
+                )
+                self._final_state = task_run.state
                 return
-            time.sleep(polling_interval)
+
+            # If still running, wait for a completed event from the server
+            logger.debug(
+                "Waiting for completed event for task run %s...",
+                self.task_run_id,
+            )
+            await TaskRunWaiter.wait_for_task_run(self._task_run_id, timeout=timeout)
+            task_run = await client.read_task_run(task_run_id=self._task_run_id)
+            if task_run.state.is_final():
+                self._final_state = task_run.state
+            return
 
     def result(
         self,
         timeout: Optional[float] = None,
         raise_on_failure: bool = True,
-        polling_interval: Optional[float] = 0.2,
     ) -> Any:
+        return run_coro_as_sync(
+            self.result_async(timeout=timeout, raise_on_failure=raise_on_failure)
+        )
+
+    async def result_async(
+        self,
+        timeout: Optional[float] = None,
+        raise_on_failure: bool = True,
+    ):
         if not self._final_state:
-            self.wait(timeout=timeout)
+            await self.wait_async(timeout=timeout)
             if not self._final_state:
                 raise TimeoutError(
                     f"Task run {self.task_run_id} did not complete within {timeout} seconds"
                 )
 
-        _result = self._final_state.result(
+        return await self._final_state.result(
             raise_on_failure=raise_on_failure, fetch=True
         )
-        # state.result is a `sync_compatible` function that may or may not return an awaitable
-        # depending on whether the parent frame is sync or not
-        if inspect.isawaitable(_result):
-            _result = run_coro_as_sync(_result)
-        return _result
 
     def __eq__(self, other):
         if not isinstance(other, PrefectDistributedFuture):

prefect/input/run_input.py

@@ -18,7 +18,8 @@ Sender flow:
 ```python
 import random
 from uuid import UUID
-from prefect import flow, get_run_logger
+from prefect import flow
+from prefect.logging import get_run_logger
 from prefect.input import RunInput
 
 class NumberData(RunInput):
@@ -43,7 +44,8 @@ Receiver flow:
 ```python
 import random
 from uuid import UUID
-from prefect import flow, get_run_logger
+from prefect import flow
+from prefect.logging import get_run_logger
 from prefect.input import RunInput
 
 class NumberData(RunInput):

prefect/records/cache_policies.py

@@ -0,0 +1,179 @@
+import inspect
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Optional
+
+from prefect.context import TaskRunContext
+from prefect.utilities.hashing import hash_objects
+
+
+@dataclass
+class CachePolicy:
+    @classmethod
+    def from_cache_key_fn(
+        cls, cache_key_fn: Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
+    ) -> "CacheKeyFnPolicy":
+        """
+        Given a function generates a key policy.
+        """
+        return CacheKeyFnPolicy(cache_key_fn=cache_key_fn)
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        raise NotImplementedError
+
+    def __sub__(self, other: str) -> "CompoundCachePolicy":
+        if not isinstance(other, str):
+            raise TypeError("Can only subtract strings from key policies.")
+        if isinstance(self, Inputs):
+            exclude = self.exclude or []
+            return Inputs(exclude=exclude + [other])
+        elif isinstance(self, CompoundCachePolicy):
+            new = Inputs(exclude=[other])
+            policies = self.policies or []
+            return CompoundCachePolicy(policies=policies + [new])
+        else:
+            new = Inputs(exclude=[other])
+            return CompoundCachePolicy(policies=[self, new])
+
+    def __add__(self, other: "CachePolicy") -> "CompoundCachePolicy":
+        # adding _None is a no-op
+        if isinstance(other, _None):
+            return self
+        elif isinstance(self, _None):
+            return other
+
+        if isinstance(self, CompoundCachePolicy):
+            policies = self.policies or []
+            return CompoundCachePolicy(policies=policies + [other])
+        elif isinstance(other, CompoundCachePolicy):
+            policies = other.policies or []
+            return CompoundCachePolicy(policies=policies + [self])
+        else:
+            return CompoundCachePolicy(policies=[self, other])
+
+
+@dataclass
+class CacheKeyFnPolicy(CachePolicy):
+    # making it optional for tests
+    cache_key_fn: Optional[
+        Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
+    ] = None
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        if self.cache_key_fn:
+            return self.cache_key_fn(task_ctx, inputs)
+
+
+@dataclass
+class CompoundCachePolicy(CachePolicy):
+    policies: list = None
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        keys = []
+        for policy in self.policies:
+            keys.append(
+                policy.compute_key(
+                    task_ctx=task_ctx,
+                    inputs=inputs,
+                    flow_parameters=flow_parameters,
+                    **kwargs,
+                )
+            )
+        return hash_objects(*keys)
+
+
+@dataclass
+class Default(CachePolicy):
+    "Execution run ID only"
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        return str(task_ctx.task_run.id)
+
+
+@dataclass
+class _None(CachePolicy):
+    "ignore key policies altogether, always run - prevents persistence"
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        return None
+
+
+@dataclass
+class TaskDef(CachePolicy):
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        lines = inspect.getsource(task_ctx.task)
+        return hash_objects(lines)
+
+
+@dataclass
+class FlowParameters(CachePolicy):
+    pass
+
+
+@dataclass
+class Inputs(CachePolicy):
+    """
+    Exposes flag for whether to include flow parameters as well.
+
+    And exclude/include config.
+    """
+
+    exclude: list = None
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        hashed_inputs = {}
+        inputs = inputs or {}
+        exclude = self.exclude or []
+
+        for key, val in inputs.items():
+            if key not in exclude:
+                hashed_inputs[key] = val
+
+        return hash_objects(hashed_inputs)
+
+
+DEFAULT = Default()
+INPUTS = Inputs()
+NONE = _None()
+TASKDEF = TaskDef()