prefect-client 3.0.0rc19__py3-none-any.whl → 3.0.1__py3-none-any.whl

prefect/concurrency/services.py

@@ -36,13 +36,18 @@ class ConcurrencySlotAcquisitionService(QueueService):
     async def _handle(
         self,
         item: Tuple[
-            int, str, Optional[float], concurrent.futures.Future, Optional[bool]
+            int,
+            str,
+            Optional[float],
+            concurrent.futures.Future,
+            Optional[bool],
+            Optional[int],
         ],
     ) -> None:
-        occupy, mode, timeout_seconds, future, create_if_missing = item
+        occupy, mode, timeout_seconds, future, create_if_missing, max_retries = item
         try:
             response = await self.acquire_slots(
-                occupy, mode, timeout_seconds, create_if_missing
+                occupy, mode, timeout_seconds, create_if_missing, max_retries
             )
         except Exception as exc:
             # If the request to the increment endpoint fails in a non-standard
@@ -59,6 +64,7 @@ class ConcurrencySlotAcquisitionService(QueueService):
         mode: str,
         timeout_seconds: Optional[float] = None,
         create_if_missing: Optional[bool] = False,
+        max_retries: Optional[int] = None,
     ) -> httpx.Response:
         with timeout_async(seconds=timeout_seconds):
             while True:
@@ -74,15 +80,19 @@ class ConcurrencySlotAcquisitionService(QueueService):
                         isinstance(exc, httpx.HTTPStatusError)
                         and exc.response.status_code == status.HTTP_423_LOCKED
                     ):
+                        if max_retries is not None and max_retries <= 0:
+                            raise exc
                         retry_after = float(exc.response.headers["Retry-After"])
                         await asyncio.sleep(retry_after)
+                        if max_retries is not None:
+                            max_retries -= 1
                     else:
                         raise exc
                 else:
                     return response
 
     def send(
-        self, item: Tuple[int, str, Optional[float], Optional[bool]]
+        self, item: Tuple[int, str, Optional[float], Optional[bool], Optional[int]]
     ) -> concurrent.futures.Future:
         with self._lock:
             if self._stopped:
@@ -91,9 +101,9 @@ class ConcurrencySlotAcquisitionService(QueueService):
             logger.debug("Service %r enqueuing item %r", self, item)
             future: concurrent.futures.Future = concurrent.futures.Future()
 
-            occupy, mode, timeout_seconds, create_if_missing = item
+            occupy, mode, timeout_seconds, create_if_missing, max_retries = item
             self._queue.put_nowait(
-                (occupy, mode, timeout_seconds, future, create_if_missing)
+                (occupy, mode, timeout_seconds, future, create_if_missing, max_retries)
             )
 
         return future

prefect/concurrency/sync.py

@@ -40,7 +40,8 @@ def concurrency(
     names: Union[str, List[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
-    create_if_missing: Optional[bool] = True,
+    create_if_missing: bool = True,
+    max_retries: Optional[int] = None,
 ) -> Generator[None, None, None]:
     """A context manager that acquires and releases concurrency slots from the
     given concurrency limits.
@@ -51,6 +52,7 @@ def concurrency(
         timeout_seconds: The number of seconds to wait for the slots to be acquired before
             raising a `TimeoutError`. A timeout of `None` will wait indefinitely.
         create_if_missing: Whether to create the concurrency limits if they do not exist.
+        max_retries: The maximum number of retries to acquire the concurrency slots.
 
     Raises:
         TimeoutError: If the slots are not acquired within the given timeout.
@@ -80,6 +82,7 @@ def concurrency(
         occupy,
         timeout_seconds=timeout_seconds,
         create_if_missing=create_if_missing,
+        max_retries=max_retries,
     )
     acquisition_time = pendulum.now("UTC")
     emitted_events = _emit_concurrency_acquisition_events(limits, occupy)

prefect/context.py

@@ -40,11 +40,10 @@ from prefect.client.orchestration import PrefectClient, SyncPrefectClient, get_c
 from prefect.client.schemas import FlowRun, TaskRun
 from prefect.events.worker import EventsWorker
 from prefect.exceptions import MissingContextError
-from prefect.results import ResultFactory
+from prefect.results import ResultStore
 from prefect.settings import PREFECT_HOME, Profile, Settings
 from prefect.states import State
 from prefect.task_runners import TaskRunner
-from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.services import start_client_metrics_server
 
 T = TypeVar("T")
@@ -95,20 +94,15 @@ def hydrated_context(
                 flow_run_context = FlowRunContext(
                     **flow_run_context,
                     client=client,
-                    result_factory=run_coro_as_sync(ResultFactory.from_flow(flow)),
                     task_runner=task_runner,
                     detached=True,
                 )
                 stack.enter_context(flow_run_context)
             # Set up parent task run context
             if parent_task_run_context := serialized_context.get("task_run_context"):
-                parent_task = parent_task_run_context["task"]
                 task_run_context = TaskRunContext(
                     **parent_task_run_context,
                     client=client,
-                    result_factory=run_coro_as_sync(
-                        ResultFactory.from_autonomous_task(parent_task)
-                    ),
                 )
                 stack.enter_context(task_run_context)
             # Set up tags context
@@ -216,6 +210,7 @@ class SyncClientContext(ContextModel):
         self._context_stack += 1
         if self._context_stack == 1:
             self.client.__enter__()
+            self.client.raise_for_api_version_mismatch()
             return super().__enter__()
         else:
             return self
@@ -273,6 +268,7 @@ class AsyncClientContext(ContextModel):
         self._context_stack += 1
         if self._context_stack == 1:
             await self.client.__aenter__()
+            await self.client.raise_for_api_version_mismatch()
             return super().__enter__()
         else:
             return self
@@ -346,7 +342,7 @@ class EngineContext(RunContext):
     detached: bool = False
 
     # Result handling
-    result_factory: ResultFactory
+    result_store: ResultStore
 
     # Counter for task calls allowing unique
     task_run_dynamic_keys: Dict[str, int] = Field(default_factory=dict)
@@ -375,6 +371,7 @@ class EngineContext(RunContext):
                 "log_prints",
                 "start_time",
                 "input_keyset",
+                "result_store",
             },
             exclude_unset=True,
         )
@@ -399,7 +396,7 @@ class TaskRunContext(RunContext):
     parameters: Dict[str, Any]
 
     # Result handling
-    result_factory: ResultFactory
+    result_store: ResultStore
 
     __var__ = ContextVar("task_run")
 
@@ -412,6 +409,7 @@ class TaskRunContext(RunContext):
                 "log_prints",
                 "start_time",
                 "input_keyset",
+                "result_store",
             },
             exclude_unset=True,
         )

prefect/deployments/runner.py

@@ -508,7 +508,7 @@ class RunnerDeployment(BaseModel):
             no_file_location_error = (
                 "Flows defined interactively cannot be deployed. Check out the"
                 " quickstart guide for help getting started:"
-                " https://docs.prefect.io/latest/getting-started/quickstart"
+                " https://docs.prefect.io/latest/get-started/quickstart"
             )
             ## first see if an entrypoint can be determined
             flow_file = getattr(flow, "__globals__", {}).get("__file__")
@@ -851,14 +851,14 @@ async def deploy(
                 " or specify a remote storage location for the flow with `.from_source`."
                 " If you are attempting to deploy a flow to a local process work pool,"
                 " consider using `flow.serve` instead. See the documentation for more"
-                " information: https://docs.prefect.io/latest/concepts/flows/#serving-a-flow"
+                " information: https://docs.prefect.io/latest/deploy/run-flows-in-local-processes"
             )
         elif work_pool.type == "process" and not ignore_warnings:
             console.print(
                 "Looks like you're deploying to a process work pool. If you're creating a"
                 " deployment for local development, calling `.serve` on your flow is a great"
                 " way to get started. See the documentation for more information:"
-                " https://docs.prefect.io/latest/concepts/flows/#serving-a-flow. "
+                " https://docs.prefect.io/latest/deploy/run-flows-in-local-processes "
                 " Set `ignore_warnings=True` to suppress this message.",
                 style="yellow",
             )

prefect/exceptions.py

@@ -412,3 +412,15 @@ class PrefectImportError(ImportError):
 
     def __init__(self, message: str) -> None:
         super().__init__(message)
+
+
+class SerializationError(PrefectException):
+    """
+    Raised when an object cannot be serialized.
+    """
+
+
+class ConfigurationError(PrefectException):
+    """
+    Raised when a configuration is invalid.
+    """

prefect/filesystems.py

@@ -84,7 +84,7 @@ class LocalFileSystem(WritableFileSystem, WritableDeploymentStorage):
     _block_type_name = "Local File System"
     _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/ad39089fa66d273b943394a68f003f7a19aa850e-48x48.png"
     _documentation_url = (
-        "https://docs.prefect.io/concepts/filesystems/#local-filesystem"
+        "https://docs.prefect.io/latest/develop/results#specifying-a-default-filesystem"
     )
 
     basepath: Optional[str] = Field(
@@ -260,7 +260,7 @@ class RemoteFileSystem(WritableFileSystem, WritableDeploymentStorage):
     _block_type_name = "Remote File System"
     _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/e86b41bc0f9c99ba9489abeee83433b43d5c9365-48x48.png"
     _documentation_url = (
-        "https://docs.prefect.io/concepts/filesystems/#remote-file-system"
+        "https://docs.prefect.io/latest/develop/results#specifying-a-default-filesystem"
     )
 
     basepath: str = Field(
@@ -433,7 +433,9 @@ class SMB(WritableFileSystem, WritableDeploymentStorage):
 
     _block_type_name = "SMB"
     _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/3f624663f7beb97d011d011bffd51ecf6c499efc-195x195.png"
-    _documentation_url = "https://docs.prefect.io/concepts/filesystems/#smb"
+    _documentation_url = (
+        "https://docs.prefect.io/latest/develop/results#specifying-a-default-filesystem"
+    )
 
     share_path: str = Field(
         default=...,

prefect/flow_engine.py

@@ -47,7 +47,7 @@ from prefect.logging.loggers import (
     get_run_logger,
     patch_print,
 )
-from prefect.results import BaseResult, ResultFactory
+from prefect.results import BaseResult, ResultStore, get_current_result_store
 from prefect.settings import PREFECT_DEBUG_MODE
 from prefect.states import (
     Failed,
@@ -202,9 +202,12 @@ class FlowRunEngine(Generic[P, R]):
                 self.handle_exception(
                     exc,
                     msg=message,
-                    result_factory=run_coro_as_sync(ResultFactory.from_flow(self.flow)),
+                    result_store=get_current_result_store().update_for_flow(
+                        self.flow, _sync=True
+                    ),
                 )
                 self.short_circuit = True
+                self.call_hooks()
 
         new_state = Running()
         state = self.set_state(new_state)
@@ -260,14 +263,15 @@ class FlowRunEngine(Generic[P, R]):
         return _result
 
     def handle_success(self, result: R) -> R:
-        result_factory = getattr(FlowRunContext.get(), "result_factory", None)
-        if result_factory is None:
-            raise ValueError("Result factory is not set")
+        result_store = getattr(FlowRunContext.get(), "result_store", None)
+        if result_store is None:
+            raise ValueError("Result store is not set")
         resolved_result = resolve_futures_to_states(result)
         terminal_state = run_coro_as_sync(
             return_value_to_state(
                 resolved_result,
-                result_factory=result_factory,
+                result_store=result_store,
+                write_result=True,
             )
         )
         self.set_state(terminal_state)
@@ -278,15 +282,15 @@ class FlowRunEngine(Generic[P, R]):
         self,
         exc: Exception,
         msg: Optional[str] = None,
-        result_factory: Optional[ResultFactory] = None,
+        result_store: Optional[ResultStore] = None,
     ) -> State:
         context = FlowRunContext.get()
         terminal_state = run_coro_as_sync(
             exception_to_failed_state(
                 exc,
                 message=msg or "Flow run encountered an exception:",
-                result_factory=result_factory
-                or getattr(context, "result_factory", None),
+                result_store=result_store or getattr(context, "result_store", None),
+                write_result=True,
             )
         )
         state = self.set_state(terminal_state)
@@ -503,7 +507,9 @@ class FlowRunEngine(Generic[P, R]):
                     flow_run=self.flow_run,
                     parameters=self.parameters,
                     client=client,
-                    result_factory=run_coro_as_sync(ResultFactory.from_flow(self.flow)),
+                    result_store=get_current_result_store().update_for_flow(
+                        self.flow, _sync=True
+                    ),
                     task_runner=task_runner,
                 )
             )

prefect/flows.py

@@ -642,7 +642,7 @@ class Flow(Generic[P, R]):
         cron: Optional[Union[Iterable[str], str]] = None,
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
-        schedules: Optional[List["FlexibleScheduleList"]] = None,
+        schedules: Optional["FlexibleScheduleList"] = None,
         concurrency_limit: Optional[int] = None,
         parameters: Optional[dict] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
@@ -730,7 +730,7 @@ class Flow(Generic[P, R]):
                 work_pool_name=work_pool_name,
                 work_queue_name=work_queue_name,
                 job_variables=job_variables,
-            )
+            )  # type: ignore # TODO: remove sync_compatible
         else:
             return RunnerDeployment.from_flow(
                 self,
@@ -1032,8 +1032,6 @@ class Flow(Generic[P, R]):
             if not isinstance(storage, LocalStorage):
                 storage.set_base_path(Path(tmpdir))
                 await storage.pull_code()
-            storage.set_base_path(Path(tmpdir))
-            await storage.pull_code()
 
             full_entrypoint = str(storage.destination / entrypoint)
             flow: Flow = await from_async.wait_for_call_in_new_thread(

prefect/futures.py

@@ -179,7 +179,8 @@ class PrefectConcurrentFuture(PrefectWrappedFuture[R, concurrent.futures.Future]
             local_logger = logger
         local_logger.warning(
             "A future was garbage collected before it resolved."
-            " Please call `.wait()` or `.result()` on futures to ensure they resolve.",
+            " Please call `.wait()` or `.result()` on futures to ensure they resolve."
+            "\nSee https://docs.prefect.io/latest/develop/task-runners for more details.",
         )
 
 

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