prefect-client 3.0.0rc20__py3-none-any.whl → 3.0.2__py3-none-any.whl

prefect/locking/filesystem.py

@@ -0,0 +1,243 @@
+import time
+from pathlib import Path
+from typing import Dict, Optional
+
+import anyio
+import pendulum
+import pydantic_core
+from typing_extensions import TypedDict
+
+from prefect.logging.loggers import get_logger
+
+from .protocol import LockManager
+
+logger = get_logger(__name__)
+
+
+class _LockInfo(TypedDict):
+    """
+    A dictionary containing information about a lock.
+
+    Attributes:
+        holder: The holder of the lock.
+        expiration: Datetime when the lock expires.
+        path: Path to the lock file.
+    """
+
+    holder: str
+    expiration: Optional[pendulum.DateTime]
+    path: Path
+
+
+class FileSystemLockManager(LockManager):
+    """
+    A lock manager that implements locking using local files.
+
+    Attributes:
+        lock_files_directory: the directory where lock files are stored
+    """
+
+    def __init__(self, lock_files_directory: Path):
+        self.lock_files_directory = lock_files_directory
+        self._locks: Dict[str, _LockInfo] = {}
+
+    def _ensure_records_directory_exists(self):
+        self.lock_files_directory.mkdir(parents=True, exist_ok=True)
+
+    def _lock_path_for_key(self, key: str) -> Path:
+        if (lock_info := self._locks.get(key)) is not None:
+            return lock_info["path"]
+        return self.lock_files_directory.joinpath(key).with_suffix(".lock")
+
+    def _get_lock_info(self, key: str, use_cache=True) -> Optional[_LockInfo]:
+        if use_cache:
+            if (lock_info := self._locks.get(key)) is not None:
+                return lock_info
+
+        lock_path = self._lock_path_for_key(key)
+
+        try:
+            with open(lock_path, "rb") as lock_file:
+                lock_info = pydantic_core.from_json(lock_file.read())
+                lock_info["path"] = lock_path
+                expiration = lock_info.get("expiration")
+                lock_info["expiration"] = (
+                    pendulum.parse(expiration) if expiration is not None else None
+                )
+            self._locks[key] = lock_info
+            return lock_info
+        except FileNotFoundError:
+            return None
+
+    async def _aget_lock_info(
+        self, key: str, use_cache: bool = True
+    ) -> Optional[_LockInfo]:
+        if use_cache:
+            if (lock_info := self._locks.get(key)) is not None:
+                return lock_info
+
+        lock_path = self._lock_path_for_key(key)
+
+        try:
+            lock_info_bytes = await anyio.Path(lock_path).read_bytes()
+            lock_info = pydantic_core.from_json(lock_info_bytes)
+            lock_info["path"] = lock_path
+            expiration = lock_info.get("expiration")
+            lock_info["expiration"] = (
+                pendulum.parse(expiration) if expiration is not None else None
+            )
+            self._locks[key] = lock_info
+            return lock_info
+        except FileNotFoundError:
+            return None
+
+    def acquire_lock(
+        self,
+        key: str,
+        holder: str,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ) -> bool:
+        self._ensure_records_directory_exists()
+        lock_path = self._lock_path_for_key(key)
+
+        if self.is_locked(key) and not self.is_lock_holder(key, holder):
+            lock_free = self.wait_for_lock(key, acquire_timeout)
+            if not lock_free:
+                return False
+
+        try:
+            Path(lock_path).touch(exist_ok=False)
+        except FileExistsError:
+            if not self.is_lock_holder(key, holder):
+                logger.debug(
+                    f"Another actor acquired the lock for record with key {key}. Trying again."
+                )
+                return self.acquire_lock(key, holder, acquire_timeout, hold_timeout)
+        expiration = (
+            pendulum.now("utc") + pendulum.duration(seconds=hold_timeout)
+            if hold_timeout is not None
+            else None
+        )
+
+        with open(Path(lock_path), "wb") as lock_file:
+            lock_file.write(
+                pydantic_core.to_json(
+                    {
+                        "holder": holder,
+                        "expiration": str(expiration)
+                        if expiration is not None
+                        else None,
+                    },
+                )
+            )
+
+        self._locks[key] = {
+            "holder": holder,
+            "expiration": expiration,
+            "path": lock_path,
+        }
+
+        return True
+
+    async def aacquire_lock(
+        self,
+        key: str,
+        holder: str,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ) -> bool:
+        await anyio.Path(self.lock_files_directory).mkdir(parents=True, exist_ok=True)
+        lock_path = self._lock_path_for_key(key)
+
+        if self.is_locked(key) and not self.is_lock_holder(key, holder):
+            lock_free = await self.await_for_lock(key, acquire_timeout)
+            if not lock_free:
+                return False
+
+        try:
+            await anyio.Path(lock_path).touch(exist_ok=False)
+        except FileExistsError:
+            if not self.is_lock_holder(key, holder):
+                logger.debug(
+                    f"Another actor acquired the lock for record with key {key}. Trying again."
+                )
+                return self.acquire_lock(key, holder, acquire_timeout, hold_timeout)
+        expiration = (
+            pendulum.now("utc") + pendulum.duration(seconds=hold_timeout)
+            if hold_timeout is not None
+            else None
+        )
+
+        async with await anyio.Path(lock_path).open("wb") as lock_file:
+            await lock_file.write(
+                pydantic_core.to_json(
+                    {
+                        "holder": holder,
+                        "expiration": str(expiration)
+                        if expiration is not None
+                        else None,
+                    },
+                )
+            )
+
+        self._locks[key] = {
+            "holder": holder,
+            "expiration": expiration,
+            "path": lock_path,
+        }
+
+        return True
+
+    def release_lock(self, key: str, holder: str) -> None:
+        lock_path = self._lock_path_for_key(key)
+        if not self.is_locked(key):
+            ValueError(f"No lock for transaction with key {key}")
+        if self.is_lock_holder(key, holder):
+            Path(lock_path).unlink(missing_ok=True)
+            self._locks.pop(key, None)
+        else:
+            raise ValueError(f"No lock held by {holder} for transaction with key {key}")
+
+    def is_locked(self, key: str, use_cache: bool = False) -> bool:
+        if (lock_info := self._get_lock_info(key, use_cache=use_cache)) is None:
+            return False
+
+        if (expiration := lock_info.get("expiration")) is None:
+            return True
+
+        expired = expiration < pendulum.now("utc")
+        if expired:
+            Path(lock_info["path"]).unlink()
+            self._locks.pop(key, None)
+            return False
+        else:
+            return True
+
+    def is_lock_holder(self, key: str, holder: str) -> bool:
+        if not self.is_locked(key):
+            return False
+
+        if not self.is_locked(key):
+            return False
+        if (lock_info := self._get_lock_info(key)) is None:
+            return False
+        return lock_info["holder"] == holder
+
+    def wait_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+        seconds_waited = 0
+        while self.is_locked(key, use_cache=False):
+            if timeout and seconds_waited >= timeout:
+                return False
+            seconds_waited += 0.1
+            time.sleep(0.1)
+        return True
+
+    async def await_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+        seconds_waited = 0
+        while self.is_locked(key, use_cache=False):
+            if timeout and seconds_waited >= timeout:
+                return False
+            seconds_waited += 0.1
+            await anyio.sleep(0.1)
+        return True

prefect/locking/memory.py

@@ -0,0 +1,213 @@
+import asyncio
+import threading
+from typing import Dict, Optional, TypedDict
+
+from .protocol import LockManager
+
+
+class _LockInfo(TypedDict):
+    """
+    A dictionary containing information about a lock.
+
+    Attributes:
+        holder: The holder of the lock.
+        lock: The lock object.
+        expiration_timer: The timer for the lock expiration
+    """
+
+    holder: str
+    lock: threading.Lock
+    expiration_timer: Optional[threading.Timer]
+
+
+class MemoryLockManager(LockManager):
+    """
+    A lock manager that stores lock information in memory.
+
+    Note: because this lock manager stores data in memory, it is not suitable for
+    use in a distributed environment or across different processes.
+    """
+
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        self._locks_dict_lock = threading.Lock()
+        self._locks: Dict[str, _LockInfo] = {}
+
+    def _expire_lock(self, key: str):
+        """
+        Expire the lock for the given key.
+
+        Used as a callback for the expiration timer of a lock.
+
+        Args:
+            key: The key of the lock to expire.
+        """
+        with self._locks_dict_lock:
+            if key in self._locks:
+                lock_info = self._locks[key]
+                if lock_info["lock"].locked():
+                    lock_info["lock"].release()
+                if lock_info["expiration_timer"]:
+                    lock_info["expiration_timer"].cancel()
+                del self._locks[key]
+
+    def acquire_lock(
+        self,
+        key: str,
+        holder: str,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ) -> bool:
+        with self._locks_dict_lock:
+            if key not in self._locks:
+                lock = threading.Lock()
+                lock.acquire()
+                expiration_timer = None
+                if hold_timeout is not None:
+                    expiration_timer = threading.Timer(
+                        hold_timeout, self._expire_lock, args=(key,)
+                    )
+                    expiration_timer.start()
+                self._locks[key] = _LockInfo(
+                    holder=holder, lock=lock, expiration_timer=expiration_timer
+                )
+                return True
+            elif self._locks[key]["holder"] == holder:
+                return True
+            else:
+                existing_lock_info = self._locks[key]
+
+        if acquire_timeout is not None:
+            existing_lock_acquired = existing_lock_info["lock"].acquire(
+                timeout=acquire_timeout
+            )
+        else:
+            existing_lock_acquired = existing_lock_info["lock"].acquire()
+
+        if existing_lock_acquired:
+            with self._locks_dict_lock:
+                if (
+                    expiration_timer := existing_lock_info["expiration_timer"]
+                ) is not None:
+                    expiration_timer.cancel()
+                expiration_timer = None
+                if hold_timeout is not None:
+                    expiration_timer = threading.Timer(
+                        hold_timeout, self._expire_lock, args=(key,)
+                    )
+                    expiration_timer.start()
+                self._locks[key] = _LockInfo(
+                    holder=holder,
+                    lock=existing_lock_info["lock"],
+                    expiration_timer=expiration_timer,
+                )
+            return True
+        return False
+
+    async def aacquire_lock(
+        self,
+        key: str,
+        holder: str,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ) -> bool:
+        with self._locks_dict_lock:
+            if key not in self._locks:
+                lock = threading.Lock()
+                await asyncio.to_thread(lock.acquire)
+                expiration_timer = None
+                if hold_timeout is not None:
+                    expiration_timer = threading.Timer(
+                        hold_timeout, self._expire_lock, args=(key,)
+                    )
+                    expiration_timer.start()
+                self._locks[key] = _LockInfo(
+                    holder=holder, lock=lock, expiration_timer=expiration_timer
+                )
+                return True
+            elif self._locks[key]["holder"] == holder:
+                return True
+            else:
+                existing_lock_info = self._locks[key]
+
+        if acquire_timeout is not None:
+            existing_lock_acquired = await asyncio.to_thread(
+                existing_lock_info["lock"].acquire, timeout=acquire_timeout
+            )
+        else:
+            existing_lock_acquired = await asyncio.to_thread(
+                existing_lock_info["lock"].acquire
+            )
+
+        if existing_lock_acquired:
+            with self._locks_dict_lock:
+                if (
+                    expiration_timer := existing_lock_info["expiration_timer"]
+                ) is not None:
+                    expiration_timer.cancel()
+                expiration_timer = None
+                if hold_timeout is not None:
+                    expiration_timer = threading.Timer(
+                        hold_timeout, self._expire_lock, args=(key,)
+                    )
+                    expiration_timer.start()
+                self._locks[key] = _LockInfo(
+                    holder=holder,
+                    lock=existing_lock_info["lock"],
+                    expiration_timer=expiration_timer,
+                )
+            return True
+        return False
+
+    def release_lock(self, key: str, holder: str) -> None:
+        with self._locks_dict_lock:
+            if key in self._locks and self._locks[key]["holder"] == holder:
+                if (
+                    expiration_timer := self._locks[key]["expiration_timer"]
+                ) is not None:
+                    expiration_timer.cancel()
+                self._locks[key]["lock"].release()
+                del self._locks[key]
+            else:
+                raise ValueError(
+                    f"No lock held by {holder} for transaction with key {key}"
+                )
+
+    def is_locked(self, key: str) -> bool:
+        return key in self._locks and self._locks[key]["lock"].locked()
+
+    def is_lock_holder(self, key: str, holder: str) -> bool:
+        lock_info = self._locks.get(key)
+        return (
+            lock_info is not None
+            and lock_info["lock"].locked()
+            and lock_info["holder"] == holder
+        )
+
+    def wait_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+        if lock := self._locks.get(key, {}).get("lock"):
+            if timeout is not None:
+                lock_acquired = lock.acquire(timeout=timeout)
+            else:
+                lock_acquired = lock.acquire()
+            if lock_acquired:
+                lock.release()
+            return lock_acquired
+        return True
+
+    async def await_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+        if lock := self._locks.get(key, {}).get("lock"):
+            if timeout is not None:
+                lock_acquired = await asyncio.to_thread(lock.acquire, timeout=timeout)
+            else:
+                lock_acquired = await asyncio.to_thread(lock.acquire)
+            if lock_acquired:
+                lock.release()
+            return lock_acquired
+        return True

prefect/locking/protocol.py

@@ -0,0 +1,122 @@
+from typing import Optional, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class LockManager(Protocol):
+    def acquire_lock(
+        self,
+        key: str,
+        holder: str,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ) -> bool:
+        """
+        Acquire a lock for a transaction record with the given key. Will block other
+        actors from updating this transaction record until the lock is
+        released.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock.
+            acquire_timeout: Max number of seconds to wait for the record to become
+                available if it is locked while attempting to acquire a lock. Pass 0
+                to attempt to acquire a lock without waiting. Blocks indefinitely by
+                default.
+            hold_timeout: Max number of seconds to hold the lock for. Holds the lock
+                indefinitely by default.
+
+        Returns:
+            bool: True if the lock was successfully acquired; False otherwise.
+        """
+        ...
+
+    async def aacquire_lock(
+        self,
+        key: str,
+        holder: str,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ) -> bool:
+        """
+        Acquire a lock for a transaction record with the given key. Will block other
+        actors from updating this transaction record until the lock is
+        released.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock.
+            acquire_timeout: Max number of seconds to wait for the record to become
+                available if it is locked while attempting to acquire a lock. Pass 0
+                to attempt to acquire a lock without waiting. Blocks indefinitely by
+                default.
+            hold_timeout: Max number of seconds to hold the lock for. Holds the lock
+                indefinitely by default.
+
+        Returns:
+            bool: True if the lock was successfully acquired; False otherwise.
+        """
+        ...
+
+    def release_lock(self, key: str, holder: str):
+        """
+        Releases the lock on the corresponding transaction record.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock. Must match the
+                holder provided when acquiring the lock.
+        """
+        ...
+
+    def is_locked(self, key: str) -> bool:
+        """
+        Simple check to see if the corresponding record is currently locked.
+
+        Args:
+            key: Unique identifier for the transaction record.
+
+        Returns:
+            True is the record is locked; False otherwise.
+        """
+        ...
+
+    def is_lock_holder(self, key: str, holder: str) -> bool:
+        """
+        Check if the current holder is the lock holder for the transaction record.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock.
+
+        Returns:
+            bool: True if the current holder is the lock holder; False otherwise.
+        """
+        ...
+
+    def wait_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+        """
+        Wait for the corresponding transaction record to become free.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            timeout: Maximum time to wait. None means to wait indefinitely.
+
+        Returns:
+            bool: True if the lock becomes free within the timeout; False
+                otherwise.
+        """
+        ...
+
+    async def await_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+        """
+        Wait for the corresponding transaction record to become free.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            timeout: Maximum time to wait. None means to wait indefinitely.
+
+        Returns:
+            bool: True if the lock becomes free within the timeout; False
+                otherwise.
+        """
+        ...

prefect/logging/handlers.py

@@ -138,8 +138,6 @@ class APILogHandler(logging.Handler):
                 return  # Respect the global settings toggle
             if not getattr(record, "send_to_api", True):
                 return  # Do not send records that have opted out
-            if not getattr(record, "send_to_orion", True):
-                return  # Backwards compatibility
 
             log = self.prepare(record)
             APILogWorker.instance().send(log)

prefect/logging/loggers.py

@@ -1,7 +1,6 @@
 import io
 import logging
 import sys
-import warnings
 from builtins import print
 from contextlib import contextmanager
 from functools import lru_cache
@@ -34,23 +33,6 @@ class PrefectLogAdapter(logging.LoggerAdapter):
 
     def process(self, msg, kwargs):
         kwargs["extra"] = {**(self.extra or {}), **(kwargs.get("extra") or {})}
-
-        from prefect._internal.compatibility.deprecated import (
-            PrefectDeprecationWarning,
-            generate_deprecation_message,
-        )
-
-        if "send_to_orion" in kwargs["extra"]:
-            warnings.warn(
-                generate_deprecation_message(
-                    'The "send_to_orion" option',
-                    start_date="May 2023",
-                    help='Use "send_to_api" instead.',
-                ),
-                PrefectDeprecationWarning,
-                stacklevel=4,
-            )
-
         return (msg, kwargs)
 
     def getChild(

prefect/logging/logging.yml

@@ -77,6 +77,7 @@ loggers:
     prefect.extra:
         level: "${PREFECT_LOGGING_LEVEL}"
         handlers: [api]
+        propagate: false
 
     prefect.flow_runs:
         level: NOTSET

prefect/main.py

@@ -7,7 +7,7 @@ from prefect.transactions import Transaction
 from prefect.tasks import task, Task
 from prefect.context import tags
 from prefect.utilities.annotations import unmapped, allow_failure
-from prefect.results import BaseResult
+from prefect.results import BaseResult, ResultRecordMetadata
 from prefect.flow_runs import pause_flow_run, resume_flow_run, suspend_flow_run
 from prefect.client.orchestration import get_client, PrefectClient
 from prefect.client.cloud import get_cloud_client, CloudClient
@@ -26,12 +26,26 @@ import prefect.context
 import prefect.client.schemas
 
 prefect.context.FlowRunContext.model_rebuild(
-    _types_namespace={"Flow": Flow, "BaseResult": BaseResult}
+    _types_namespace={
+        "Flow": Flow,
+        "BaseResult": BaseResult,
+        "ResultRecordMetadata": ResultRecordMetadata,
+    }
+)
+prefect.context.TaskRunContext.model_rebuild(
+    _types_namespace={"Task": Task, "BaseResult": BaseResult}
+)
+prefect.client.schemas.State.model_rebuild(
+    _types_namespace={
+        "BaseResult": BaseResult,
+        "ResultRecordMetadata": ResultRecordMetadata,
+    }
 )
-prefect.context.TaskRunContext.model_rebuild(_types_namespace={"Task": Task})
-prefect.client.schemas.State.model_rebuild(_types_namespace={"BaseResult": BaseResult})
 prefect.client.schemas.StateCreate.model_rebuild(
-    _types_namespace={"BaseResult": BaseResult}
+    _types_namespace={
+        "BaseResult": BaseResult,
+        "ResultRecordMetadata": ResultRecordMetadata,
+    }
 )
 Transaction.model_rebuild()
 

prefect/records/base.py

@@ -6,11 +6,18 @@ from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Optional
 
+from prefect._internal.compatibility import deprecated
+
 if TYPE_CHECKING:
     from prefect.results import BaseResult
     from prefect.transactions import IsolationLevel
 
 
+@deprecated.deprecated_class(
+    start_date="Sep 2024",
+    end_date="Nov 2024",
+    help="Use `ResultRecord` instead to represent a result and its associated metadata.",
+)
 @dataclass
 class TransactionRecord:
     """
@@ -21,6 +28,11 @@ class TransactionRecord:
     result: "BaseResult"
 
 
+@deprecated.deprecated_class(
+    start_date="Sep 2024",
+    end_date="Nov 2024",
+    help="Use `ResultStore` and provide a `WritableFileSystem` for `metadata_storage` instead.",
+)
 class RecordStore(abc.ABC):
     @abc.abstractmethod
     def read(