prefect-client 3.3.4.dev3__tar.gz → 3.3.5__tar.gz

{prefect_client-3.3.4.dev3 → prefect_client-3.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prefect-client
-Version: 3.3.4.dev3
+Version: 3.3.5
 Summary: Workflow orchestration and management.
 Project-URL: Changelog, https://github.com/PrefectHQ/prefect/releases
 Project-URL: Documentation, https://docs.prefect.io

prefect_client-3.3.5/src/prefect/_build_info.py

@@ -0,0 +1,5 @@
+# Generated by versioningit
+__version__ = "3.3.5"
+__build_date__ = "2025-04-17 21:24:54.247813+00:00"
+__git_commit__ = "db4b7a338ee0030ad0905d803a5ea5167d2a0711"
+__dirty__ = False

{prefect_client-3.3.4.dev3 → prefect_client-3.3.5}/src/prefect/_experimental/bundles.py

@@ -8,6 +8,7 @@ import multiprocessing.context
 import os
 import subprocess
 import sys
+from pathlib import Path
 from typing import Any, TypedDict
 
 import cloudpickle
@@ -99,7 +100,9 @@ def extract_flow_from_bundle(bundle: SerializedBundle) -> Flow[Any, Any]:
 
 
 def _extract_and_run_flow(
-    bundle: SerializedBundle, env: dict[str, Any] | None = None
+    bundle: SerializedBundle,
+    cwd: Path | str | None = None,
+    env: dict[str, Any] | None = None,
 ) -> None:
     """
     Extracts a flow from a bundle and runs it.
@@ -108,6 +111,7 @@ def _extract_and_run_flow(
 
     Args:
         bundle: The bundle to extract and run.
+        cwd: The working directory to use when running the flow.
         env: The environment to use when running the flow.
     """
 
@@ -120,6 +124,9 @@ def _extract_and_run_flow(
     context = _deserialize_bundle_object(bundle["context"])
     flow_run = FlowRun.model_validate(bundle["flow_run"])
 
+    if cwd:
+        os.chdir(cwd)
+
     with SettingsContext(
         profile=settings_context.profile,
         settings=Settings(),
@@ -138,6 +145,8 @@ def _extract_and_run_flow(
 
 def execute_bundle_in_subprocess(
     bundle: SerializedBundle,
+    env: dict[str, Any] | None = None,
+    cwd: Path | str | None = None,
 ) -> multiprocessing.context.SpawnProcess:
     """
     Executes a bundle in a subprocess.
@@ -150,6 +159,7 @@ def execute_bundle_in_subprocess(
     """
 
     ctx = multiprocessing.get_context("spawn")
+    env = env or {}
 
     # Install dependencies if necessary
     if dependencies := bundle.get("dependencies"):
@@ -164,7 +174,9 @@ def execute_bundle_in_subprocess(
         kwargs={
             "bundle": bundle,
             "env": get_current_settings().to_environment_variables(exclude_unset=True)
-            | os.environ,
+            | os.environ
+            | env,
+            "cwd": cwd,
         },
     )
 

{prefect_client-3.3.4.dev3 → prefect_client-3.3.5}/src/prefect/cache_policies.py

@@ -33,9 +33,9 @@ def _register_stable_transforms() -> None:
     so that cache keys that utilize them are deterministic across invocations.
     """
     try:
-        import pandas as pd
+        import pandas as pd  # pyright: ignore
 
-        STABLE_TRANSFORMS[pd.DataFrame] = lambda df: [
+        STABLE_TRANSFORMS[pd.DataFrame] = lambda df: [  # pyright: ignore
             df[col] for col in sorted(df.columns)
         ]
     except (ImportError, ModuleNotFoundError):
@@ -183,7 +183,7 @@ class CompoundCachePolicy(CachePolicy):
     Any keys that return `None` will be ignored.
     """
 
-    policies: list[CachePolicy] = field(default_factory=list)
+    policies: list[CachePolicy] = field(default_factory=lambda: [])
 
     def __post_init__(self) -> None:
         # flatten any CompoundCachePolicies
@@ -349,7 +349,7 @@ class Inputs(CachePolicy):
     Policy that computes a cache key based on a hash of the runtime inputs provided to the task..
     """
 
-    exclude: list[str] = field(default_factory=list)
+    exclude: list[str] = field(default_factory=lambda: [])
 
     def compute_key(
         self,

{prefect_client-3.3.4.dev3 → prefect_client-3.3.5}/src/prefect/events/clients.py

@@ -39,6 +39,7 @@ import prefect.types._datetime
 from prefect.events import Event
 from prefect.logging import get_logger
 from prefect.settings import (
+    PREFECT_API_AUTH_STRING,
     PREFECT_API_KEY,
     PREFECT_API_SSL_CERT_FILE,
     PREFECT_API_TLS_INSECURE_SKIP_VERIFY,
@@ -367,7 +368,7 @@ class PrefectEventsClient(EventsClient):
         await self._connect.__aexit__(exc_type, exc_val, exc_tb)
         return await super().__aexit__(exc_type, exc_val, exc_tb)
 
-    def _log_debug(self, message: str, *args, **kwargs) -> None:
+    def _log_debug(self, message: str, *args: Any, **kwargs: Any) -> None:
         message = f"EventsClient(id={id(self)}): " + message
         logger.debug(message, *args, **kwargs)
 
@@ -578,6 +579,7 @@ class PrefectEventSubscriber:
     _seen_events: MutableMapping[UUID, bool]
 
     _api_key: Optional[str]
+    _auth_token: Optional[str]
 
     def __init__(
         self,
@@ -593,6 +595,8 @@ class PrefectEventSubscriber:
                 the client should attempt to reconnect
         """
         self._api_key = None
+        self._auth_token = PREFECT_API_AUTH_STRING.value()
+
         if not api_url:
             api_url = cast(str, PREFECT_API_URL.value())
 
@@ -641,8 +645,10 @@ class PrefectEventSubscriber:
         await pong
 
         logger.debug("  authenticating...")
+        # Use the API key (for Cloud) OR the auth token (for self-hosted with auth string)
+        token = self._api_key or self._auth_token
         await self._websocket.send(
-            orjson.dumps({"type": "auth", "token": self._api_key}).decode()
+            orjson.dumps({"type": "auth", "token": token}).decode()
         )
 
         try:
@@ -652,13 +658,13 @@ class PrefectEventSubscriber:
         except AssertionError as e:
             raise Exception(
                 "Unable to authenticate to the event stream. Please ensure the "
-                "provided api_key you are using is valid for this environment. "
+                "provided api_key or auth_token you are using is valid for this environment. "
                 f"Reason: {e.args[0]}"
             )
         except ConnectionClosedError as e:
             reason = getattr(e.rcvd, "reason", None)
             msg = "Unable to authenticate to the event stream. Please ensure the "
-            msg += "provided api_key you are using is valid for this environment. "
+            msg += "provided api_key or auth_token you are using is valid for this environment. "
             msg += f"Reason: {reason}" if reason else ""
             raise Exception(msg) from e
 

{prefect_client-3.3.4.dev3 → prefect_client-3.3.5}/src/prefect/flows.py

@@ -40,6 +40,7 @@ from typing import (
 from uuid import UUID
 
 import pydantic
+from exceptiongroup import BaseExceptionGroup, ExceptionGroup
 from pydantic.v1 import BaseModel as V1BaseModel
 from pydantic.v1.decorator import ValidatedFunction as V1ValidatedFunction
 from pydantic.v1.errors import ConfigError  # TODO
@@ -105,6 +106,9 @@ from ._internal.pydantic.v2_validated_func import (
     V2ValidatedFunction as ValidatedFunction,
 )
 
+if TYPE_CHECKING:
+    from prefect.workers.base import BaseWorker
+
 T = TypeVar("T")  # Generic type var for capturing the inner return type of async funcs
 R = TypeVar("R")  # The return type of the user's function
 P = ParamSpec("P")  # The parameters of the flow
@@ -1130,7 +1134,7 @@ class Flow(Generic[P, R]):
     @classmethod
     async def afrom_source(
         cls,
-        source: Union[str, "RunnerStorage", ReadableDeploymentStorage],
+        source: Union[str, Path, "RunnerStorage", ReadableDeploymentStorage],
         entrypoint: str,
     ) -> "Flow[..., Any]":
         """
@@ -1244,7 +1248,7 @@ class Flow(Generic[P, R]):
     @async_dispatch(afrom_source)
     def from_source(
         cls,
-        source: Union[str, "RunnerStorage", ReadableDeploymentStorage],
+        source: Union[str, Path, "RunnerStorage", ReadableDeploymentStorage],
         entrypoint: str,
     ) -> "Flow[..., Any]":
         """
@@ -1976,7 +1980,7 @@ class FlowDecorator:
         # manually here.
         @staticmethod
         def from_source(
-            source: Union[str, "RunnerStorage", ReadableDeploymentStorage],
+            source: Union[str, Path, "RunnerStorage", ReadableDeploymentStorage],
             entrypoint: str,
         ) -> Union["Flow[..., Any]", Coroutine[Any, Any, "Flow[..., Any]"]]: ...
 
@@ -1984,6 +1988,141 @@ class FlowDecorator:
 flow: FlowDecorator = FlowDecorator()
 
 
+class InfrastructureBoundFlow(Flow[P, R]):
+    """
+    EXPERIMENTAL: This class is experimental and may be removed or changed in future
+        releases.
+
+    A flow that is bound to running on a specific infrastructure.
+
+    Attributes:
+        work_pool: The name of the work pool to run the flow on. The base job
+            configuration of the work pool will determine the configuration of the
+            infrastructure the flow will run on.
+        job_variables: Infrastructure configuration that will override the base job
+            configuration of the work pool.
+        worker_cls: The class of the worker to use to spin up infrastructure and submit
+            the flow to it.
+    """
+
+    def __init__(
+        self,
+        *args: Any,
+        work_pool: str,
+        job_variables: dict[str, Any],
+        worker_cls: type["BaseWorker[Any, Any, Any]"],
+        **kwargs: Any,
+    ):
+        super().__init__(*args, **kwargs)
+        self.work_pool = work_pool
+        self.job_variables = job_variables
+        self.worker_cls = worker_cls
+
+    @overload
+    def __call__(self: "Flow[P, NoReturn]", *args: P.args, **kwargs: P.kwargs) -> None:
+        # `NoReturn` matches if a type can't be inferred for the function which stops a
+        # sync function from matching the `Coroutine` overload
+        ...
+
+    @overload
+    def __call__(
+        self: "Flow[P, Coroutine[Any, Any, T]]",
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> Coroutine[Any, Any, T]: ...
+
+    @overload
+    def __call__(
+        self: "Flow[P, T]",
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> T: ...
+
+    @overload
+    def __call__(
+        self: "Flow[P, Coroutine[Any, Any, T]]",
+        *args: P.args,
+        return_state: Literal[True],
+        **kwargs: P.kwargs,
+    ) -> Awaitable[State[T]]: ...
+
+    @overload
+    def __call__(
+        self: "Flow[P, T]",
+        *args: P.args,
+        return_state: Literal[True],
+        **kwargs: P.kwargs,
+    ) -> State[T]: ...
+
+    def __call__(
+        self,
+        *args: "P.args",
+        return_state: bool = False,
+        wait_for: Optional[Iterable[PrefectFuture[Any]]] = None,
+        **kwargs: "P.kwargs",
+    ):
+        async def modified_call(
+            *args: P.args,
+            return_state: bool = False,
+            # TODO: Handle wait_for once we have an asynchronous way to wait for futures
+            # We should wait locally for futures to resolve before spinning up
+            # infrastructure.
+            wait_for: Optional[Iterable[PrefectFuture[Any]]] = None,
+            **kwargs: P.kwargs,
+        ) -> R | State[R]:
+            try:
+                async with self.worker_cls(work_pool_name=self.work_pool) as worker:
+                    parameters = get_call_parameters(self, args, kwargs)
+                    future = await worker.submit(
+                        flow=self,
+                        parameters=parameters,
+                        job_variables=self.job_variables,
+                    )
+                    if return_state:
+                        await future.wait_async()
+                        return future.state
+                    return await future.aresult()
+            except (ExceptionGroup, BaseExceptionGroup) as exc:
+                # For less verbose tracebacks
+                exceptions = exc.exceptions
+                if len(exceptions) == 1:
+                    raise exceptions[0] from None
+                else:
+                    raise
+
+        if inspect.iscoroutinefunction(self.fn):
+            return modified_call(
+                *args, return_state=return_state, wait_for=wait_for, **kwargs
+            )
+        else:
+            return run_coro_as_sync(
+                modified_call(
+                    *args,
+                    return_state=return_state,
+                    wait_for=wait_for,
+                    **kwargs,
+                )
+            )
+
+
+def bind_flow_to_infrastructure(
+    flow: Flow[P, R],
+    work_pool: str,
+    worker_cls: type["BaseWorker[Any, Any, Any]"],
+    job_variables: dict[str, Any] | None = None,
+) -> InfrastructureBoundFlow[P, R]:
+    new = InfrastructureBoundFlow[P, R](
+        flow.fn,
+        work_pool=work_pool,
+        job_variables=job_variables or {},
+        worker_cls=worker_cls,
+    )
+    # Copy all attributes from the original flow
+    for attr, value in flow.__dict__.items():
+        setattr(new, attr, value)
+    return new
+
+
 def _raise_on_name_with_banned_characters(name: Optional[str]) -> Optional[str]:
     """
     Raise an InvalidNameError if the given name contains any invalid

{prefect_client-3.3.4.dev3 → prefect_client-3.3.5}/src/prefect/locking/memory.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import asyncio
 import threading
-from typing import Any, Optional, TypedDict
+from typing import Any, TypedDict
 
 from typing_extensions import Self
 
@@ -19,7 +21,7 @@ class _LockInfo(TypedDict):
 
     holder: str
     lock: threading.Lock
-    expiration_timer: Optional[threading.Timer]
+    expiration_timer: threading.Timer | None
 
 
 class MemoryLockManager(LockManager):
@@ -63,8 +65,8 @@ class MemoryLockManager(LockManager):
         self,
         key: str,
         holder: str,
-        acquire_timeout: Optional[float] = None,
-        hold_timeout: Optional[float] = None,
+        acquire_timeout: float | None = None,
+        hold_timeout: float | None = None,
     ) -> bool:
         with self._locks_dict_lock:
             if key not in self._locks:
@@ -116,13 +118,13 @@ class MemoryLockManager(LockManager):
         self,
         key: str,
         holder: str,
-        acquire_timeout: Optional[float] = None,
-        hold_timeout: Optional[float] = None,
+        acquire_timeout: float | None = None,
+        hold_timeout: float | None = None,
     ) -> bool:
         with self._locks_dict_lock:
             if key not in self._locks:
                 lock = threading.Lock()
-                await asyncio.to_thread(lock.acquire)
+                lock.acquire()
                 expiration_timer = None
                 if hold_timeout is not None:
                     expiration_timer = threading.Timer(
@@ -192,7 +194,7 @@ class MemoryLockManager(LockManager):
             and lock_info["holder"] == holder
         )
 
-    def wait_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+    def wait_for_lock(self, key: str, timeout: float | None = None) -> bool:
         lock_info: _LockInfo | None = self._locks.get(key)
         if lock_info is None:
             return True
@@ -206,7 +208,7 @@ class MemoryLockManager(LockManager):
             return lock_acquired
         return True
 
-    async def await_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+    async def await_for_lock(self, key: str, timeout: float | None = None) -> bool:
         lock_info: _LockInfo | None = self._locks.get(key, None)
         if lock_info is None:
             return True

{prefect_client-3.3.4.dev3 → prefect_client-3.3.5}/src/prefect/results.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import inspect
 import os
 import socket
@@ -33,6 +34,7 @@ from typing_extensions import ParamSpec, Self
 import prefect
 import prefect.types._datetime
 from prefect._internal.compatibility.async_dispatch import async_dispatch
+from prefect._internal.concurrency.event_loop import get_running_loop
 from prefect._result_records import R, ResultRecord, ResultRecordMetadata
 from prefect.blocks.core import Block
 from prefect.exceptions import (
@@ -475,10 +477,17 @@ class ResultStore(BaseModel):
         Returns:
             str: A unique identifier string.
         """
+        current_loop = get_running_loop()
         hostname = socket.gethostname()
         pid = os.getpid()
         thread_name = threading.current_thread().name
         thread_id = threading.get_ident()
+        if current_loop:
+            current_task = asyncio.current_task()
+            if current_task:
+                # include the task id to ensure uniqueness because there might be
+                # multiple tasks running in the same thread
+                return f"{hostname}:{pid}:{thread_id}:{thread_name}:{id(current_task)}"
         return f"{hostname}:{pid}:{thread_id}:{thread_name}"
 
     @sync_compatible

{prefect_client-3.3.4.dev3 → prefect_client-3.3.5}/src/prefect/runner/runner.py

@@ -636,7 +636,12 @@ class Runner:
 
                     return process
 
-    async def execute_bundle(self, bundle: SerializedBundle) -> None:
+    async def execute_bundle(
+        self,
+        bundle: SerializedBundle,
+        cwd: Path | str | None = None,
+        env: dict[str, str | None] | None = None,
+    ) -> None:
         """
         Executes a bundle in a subprocess.
         """
@@ -651,7 +656,7 @@ class Runner:
             if not self._acquire_limit_slot(flow_run.id):
                 return
 
-            process = execute_bundle_in_subprocess(bundle)
+            process = execute_bundle_in_subprocess(bundle, cwd=cwd, env=env)
 
             if process.pid is None:
                 # This shouldn't happen because `execute_bundle_in_subprocess` starts the process
@@ -776,7 +781,7 @@ class Runner:
         if command is None:
             runner_command = [get_sys_executable(), "-m", "prefect.engine"]
         else:
-            runner_command = shlex.split(command)
+            runner_command = shlex.split(command, posix=(os.name != "nt"))
 
         flow_run_logger = self._get_flow_run_logger(flow_run)
 

{prefect_client-3.3.4.dev3 → prefect_client-3.3.5}/src/prefect/server/api/deployments.py

@@ -911,6 +911,7 @@ async def update_deployment_schedule(
             session=session,
             deployment_id=deployment_id,
             auto_scheduled_only=True,
+            future_only=True,
         )
 
 

{prefect_client-3.3.4.dev3 → prefect_client-3.3.5}/src/prefect/states.py

@@ -776,7 +776,7 @@ def AwaitingRetry(
     """Convenience function for creating `AwaitingRetry` states.
 
     Returns:
-        State: a AwaitingRetry state
+        State: an AwaitingRetry state
     """
     return Scheduled(
         cls=cls, scheduled_time=scheduled_time, name="AwaitingRetry", **kwargs
@@ -791,7 +791,7 @@ def AwaitingConcurrencySlot(
     """Convenience function for creating `AwaitingConcurrencySlot` states.
 
     Returns:
-        State: a AwaitingConcurrencySlot state
+        State: an AwaitingConcurrencySlot state
     """
     return Scheduled(
         cls=cls, scheduled_time=scheduled_time, name="AwaitingConcurrencySlot", **kwargs

{prefect_client-3.3.4.dev3 → prefect_client-3.3.5}/src/prefect/task_engine.py

@@ -81,7 +81,13 @@ from prefect.states import (
     return_value_to_state,
 )
 from prefect.telemetry.run_telemetry import RunTelemetry
-from prefect.transactions import IsolationLevel, Transaction, transaction
+from prefect.transactions import (
+    AsyncTransaction,
+    IsolationLevel,
+    Transaction,
+    atransaction,
+    transaction,
+)
 from prefect.utilities._engine import get_hook_name
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import run_coro_as_sync
@@ -1020,7 +1026,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             # otherwise, return the exception
             return self._raised
 
-    async def handle_success(self, result: R, transaction: Transaction) -> R:
+    async def handle_success(self, result: R, transaction: AsyncTransaction) -> R:
         if self.task.cache_expiration is not None:
             expiration = prefect.types._datetime.now("UTC") + self.task.cache_expiration
         else:
@@ -1302,7 +1308,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                         await self.call_hooks()
 
     @asynccontextmanager
-    async def transaction_context(self) -> AsyncGenerator[Transaction, None]:
+    async def transaction_context(self) -> AsyncGenerator[AsyncTransaction, None]:
         # refresh cache setting is now repurposes as overwrite transaction record
         overwrite = (
             self.task.refresh_cache
@@ -1317,7 +1323,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             else None
         )
 
-        with transaction(
+        async with atransaction(
             key=self.compute_transaction_key(),
             store=get_result_store(),
             overwrite=overwrite,
@@ -1349,7 +1355,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 await self.handle_exception(exc)
 
     async def call_task_fn(
-        self, transaction: Transaction
+        self, transaction: AsyncTransaction
     ) -> Union[R, Coroutine[Any, Any, R]]:
         """
         Convenience method to call the task function. Returns a coroutine if the
@@ -1357,7 +1363,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         """
         parameters = self.parameters or {}
         if transaction.is_committed():
-            result = transaction.read()
+            result = await transaction.read()
         else:
             result = await call_with_parameters(self.task.fn, parameters)
         await self.handle_success(result, transaction=transaction)

{prefect_client-3.3.4.dev3 → prefect_client-3.3.5}/src/prefect/tasks.py

@@ -764,11 +764,6 @@ class Task(Generic[P, R]):
     def on_rollback(
         self, fn: Callable[["Transaction"], None]
     ) -> Callable[["Transaction"], None]:
-        if asyncio.iscoroutinefunction(fn):
-            raise ValueError(
-                "Asynchronous rollback hooks are not yet supported. Rollback hooks must be synchronous functions."
-            )
-
         self.on_rollback_hooks.append(fn)
         return fn
 

{prefect_client-3.3.4.dev3 → prefect_client-3.3.5}/src/prefect/telemetry/instrumentation.py

@@ -1,6 +1,6 @@
 import logging
-import os
 import re
+import socket
 from typing import TYPE_CHECKING
 from urllib.parse import urljoin
 from uuid import UUID
@@ -74,7 +74,7 @@ def setup_exporters(
     resource = Resource.create(
         {
             "service.name": "prefect",
-            "service.instance.id": os.uname().nodename,
+            "service.instance.id": socket.gethostname(),
             "prefect.account": str(account_id),
             "prefect.workspace": str(workspace_id),
         }