prefect-client 2.19.4__tar.gz → 2.19.6__tar.gz

{prefect-client-2.19.4/src/prefect_client.egg-info → prefect-client-2.19.6}/PKG-INFO

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

{prefect-client-2.19.4 → prefect-client-2.19.6}/requirements.txt

@@ -7,7 +7,7 @@ asyncpg >= 0.23
 click >= 8.0, < 8.2
 cryptography >= 36.0.1
 dateparser >= 1.1.1, < 2.0.0
-docker >= 4.0, < 7.0
+docker >= 4.0
 graphviz >= 0.20.1
 griffe >= 0.20.0
 jinja2 >= 3.0.0, < 4.0.0

{prefect-client-2.19.4 → prefect-client-2.19.6}/src/prefect/client/schemas/actions.py

@@ -778,7 +778,7 @@ class GlobalConcurrencyLimitUpdate(ActionBaseModel):
 
     name: Optional[str] = Field(None)
     limit: Optional[NonNegativeInteger] = Field(None)
-    active: Optional[NonNegativeInteger] = Field(None)
+    active: Optional[bool] = Field(None)
     active_slots: Optional[NonNegativeInteger] = Field(None)
     slot_decay_per_second: Optional[NonNegativeFloat] = Field(None)
 

{prefect-client-2.19.4 → prefect-client-2.19.6}/src/prefect/concurrency/asyncio.py

@@ -13,7 +13,6 @@ except ImportError:
 
 from prefect import get_client
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
-from prefect.utilities.timeout import timeout_async
 
 from .events import (
     _emit_concurrency_acquisition_events,
@@ -26,6 +25,10 @@ class ConcurrencySlotAcquisitionError(Exception):
     """Raised when an unhandlable occurs while acquiring concurrency slots."""
 
 
+class AcquireConcurrencySlotTimeoutError(TimeoutError):
+    """Raised when acquiring a concurrency slot times out."""
+
+
 @asynccontextmanager
 async def concurrency(
     names: Union[str, List[str]],
@@ -58,8 +61,9 @@ async def concurrency(
     ```
     """
     names = names if isinstance(names, list) else [names]
-    with timeout_async(seconds=timeout_seconds):
-        limits = await _acquire_concurrency_slots(names, occupy)
+    limits = await _acquire_concurrency_slots(
+        names, occupy, timeout_seconds=timeout_seconds
+    )
     acquisition_time = pendulum.now("UTC")
     emitted_events = _emit_concurrency_acquisition_events(limits, occupy)
 
@@ -91,12 +95,18 @@ async def _acquire_concurrency_slots(
     names: List[str],
     slots: int,
     mode: Union[Literal["concurrency"], Literal["rate_limit"]] = "concurrency",
+    timeout_seconds: Optional[float] = None,
 ) -> List[MinimalConcurrencyLimitResponse]:
     service = ConcurrencySlotAcquisitionService.instance(frozenset(names))
-    future = service.send((slots, mode))
+    future = service.send((slots, mode, timeout_seconds))
     response_or_exception = await asyncio.wrap_future(future)
 
     if isinstance(response_or_exception, Exception):
+        if isinstance(response_or_exception, TimeoutError):
+            raise AcquireConcurrencySlotTimeoutError(
+                f"Attempt to acquire concurrency slots timed out after {timeout_seconds} second(s)"
+            ) from response_or_exception
+
         raise ConcurrencySlotAcquisitionError(
             f"Unable to acquire concurrency slots on {names!r}"
         ) from response_or_exception

{prefect-client-2.19.4 → prefect-client-2.19.6}/src/prefect/concurrency/services.py

@@ -3,6 +3,7 @@ import concurrent.futures
 from contextlib import asynccontextmanager
 from typing import (
     FrozenSet,
+    Optional,
     Tuple,
 )
 
@@ -13,6 +14,7 @@ from prefect import get_client
 from prefect._internal.concurrency import logger
 from prefect._internal.concurrency.services import QueueService
 from prefect.client.orchestration import PrefectClient
+from prefect.utilities.timeout import timeout_async
 
 
 class ConcurrencySlotAcquisitionService(QueueService):
@@ -27,10 +29,12 @@ class ConcurrencySlotAcquisitionService(QueueService):
             self._client = client
             yield
 
-    async def _handle(self, item: Tuple[int, str, concurrent.futures.Future]):
-        occupy, mode, future = item
+    async def _handle(
+        self, item: Tuple[int, str, Optional[float], concurrent.futures.Future]
+    ):
+        occupy, mode, timeout_seconds, future = item
         try:
-            response = await self.acquire_slots(occupy, mode)
+            response = await self.acquire_slots(occupy, mode, timeout_seconds)
         except Exception as exc:
             # If the request to the increment endpoint fails in a non-standard
             # way, we need to set the future's result so that the caller can
@@ -40,25 +44,28 @@ class ConcurrencySlotAcquisitionService(QueueService):
         else:
             future.set_result(response)
 
-    async def acquire_slots(self, slots: int, mode: str) -> httpx.Response:
-        while True:
-            try:
-                response = await self._client.increment_concurrency_slots(
-                    names=self.concurrency_limit_names, slots=slots, mode=mode
-                )
-            except Exception as exc:
-                if (
-                    isinstance(exc, httpx.HTTPStatusError)
-                    and exc.response.status_code == status.HTTP_423_LOCKED
-                ):
-                    retry_after = float(exc.response.headers["Retry-After"])
-                    await asyncio.sleep(retry_after)
+    async def acquire_slots(
+        self, slots: int, mode: str, timeout_seconds: Optional[float] = None
+    ):
+        with timeout_async(timeout_seconds):
+            while True:
+                try:
+                    response = await self._client.increment_concurrency_slots(
+                        names=self.concurrency_limit_names, slots=slots, mode=mode
+                    )
+                except Exception as exc:
+                    if (
+                        isinstance(exc, httpx.HTTPStatusError)
+                        and exc.response.status_code == status.HTTP_423_LOCKED
+                    ):
+                        retry_after = float(exc.response.headers["Retry-After"])
+                        await asyncio.sleep(retry_after)
+                    else:
+                        raise exc
                 else:
-                    raise exc
-            else:
-                return response
+                    return response
 
-    def send(self, item: Tuple[int, str]):
+    def send(self, item: Tuple[int, str, Optional[float]]) -> concurrent.futures.Future:
         with self._lock:
             if self._stopped:
                 raise RuntimeError("Cannot put items in a stopped service instance.")
@@ -66,7 +73,7 @@ class ConcurrencySlotAcquisitionService(QueueService):
             logger.debug("Service %r enqueuing item %r", self, item)
             future: concurrent.futures.Future = concurrent.futures.Future()
 
-            occupy, mode = item
-            self._queue.put_nowait((occupy, mode, future))
+            occupy, mode, timeout_seconds = item
+            self._queue.put_nowait((occupy, mode, timeout_seconds, future))
 
         return future

{prefect-client-2.19.4 → prefect-client-2.19.6}/src/prefect/concurrency/sync.py

@@ -12,7 +12,6 @@ except ImportError:
 from prefect._internal.concurrency.api import create_call, from_sync
 from prefect._internal.concurrency.event_loop import get_running_loop
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
-from prefect.utilities.timeout import timeout
 
 from .asyncio import (
     _acquire_concurrency_slots,
@@ -57,10 +56,9 @@ def concurrency(
     """
     names = names if isinstance(names, list) else [names]
 
-    with timeout(seconds=timeout_seconds):
-        limits: List[MinimalConcurrencyLimitResponse] = _call_async_function_from_sync(
-            _acquire_concurrency_slots, names, occupy
-        )
+    limits: List[MinimalConcurrencyLimitResponse] = _call_async_function_from_sync(
+        _acquire_concurrency_slots, names, occupy, timeout_seconds=timeout_seconds
+    )
     acquisition_time = pendulum.now("UTC")
     emitted_events = _emit_concurrency_acquisition_events(limits, occupy)
 

{prefect-client-2.19.4 → prefect-client-2.19.6}/src/prefect/flows.py

@@ -83,8 +83,9 @@ from prefect.logging import get_logger
 from prefect.results import ResultSerializer, ResultStorage
 from prefect.runner.storage import (
     BlockStorageAdapter,
+    LocalStorage,
     RunnerStorage,
-    create_storage_from_url,
+    create_storage_from_source,
 )
 from prefect.settings import (
     PREFECT_DEFAULT_WORK_POOL_NAME,
@@ -106,7 +107,7 @@ from prefect.utilities.callables import (
 )
 from prefect.utilities.collections import listrepr
 from prefect.utilities.hashing import file_hash
-from prefect.utilities.importtools import import_object
+from prefect.utilities.importtools import import_object, safe_load_namespace
 from prefect.utilities.visualization import (
     FlowVisualizationError,
     GraphvizExecutableNotFoundError,
@@ -910,7 +911,7 @@ class Flow(Generic[P, R]):
             ```
         """
         if isinstance(source, str):
-            storage = create_storage_from_url(source)
+            storage = create_storage_from_source(source)
         elif isinstance(source, RunnerStorage):
             storage = source
         elif hasattr(source, "get_directory"):
@@ -920,9 +921,11 @@ class Flow(Generic[P, R]):
                 f"Unsupported source type {type(source).__name__!r}. Please provide a"
                 " URL to remote storage or a storage object."
             )
+
         with tempfile.TemporaryDirectory() as tmpdir:
-            storage.set_base_path(Path(tmpdir))
-            await storage.pull_code()
+            if not isinstance(storage, LocalStorage):
+                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(
@@ -1850,11 +1853,33 @@ def load_flow_argument_from_entrypoint(
         ):
             for keyword in decorator.keywords:
                 if keyword.arg == arg:
-                    return (
-                        keyword.value.value
-                    )  # Return the string value of the argument
+                    if isinstance(keyword.value, ast.Constant):
+                        return (
+                            keyword.value.value
+                        )  # Return the string value of the argument
+
+                    # if the arg value is not a raw str (i.e. a variable or expression),
+                    # then attempt to evaluate it
+                    namespace = safe_load_namespace(source_code)
+                    literal_arg_value = ast.get_source_segment(
+                        source_code, keyword.value
+                    )
+                    try:
+                        evaluated_value = eval(literal_arg_value, namespace)  # type: ignore
+                    except Exception as e:
+                        logger.info(
+                            "Failed to parse @flow argument: `%s=%s` due to the following error. Ignoring and falling back to default behavior.",
+                            arg,
+                            literal_arg_value,
+                            exc_info=e,
+                        )
+                        # ignore the decorator arg and fallback to default behavior
+                        break
+                    return str(evaluated_value)
 
     if arg == "name":
         return func_name.replace(
             "_", "-"
         )  # If no matching decorator or keyword argument is found
+
+    return None

{prefect-client-2.19.4 → prefect-client-2.19.6}/src/prefect/infrastructure/container.py

@@ -124,6 +124,12 @@ class BaseDockerLogin(Block, ABC):
 )
 class DockerRegistry(BaseDockerLogin):
     """
+    DEPRECATION WARNING:
+
+    This class is deprecated as of March 2024 and will not be available after September 2024.
+    It has been replaced by `DockerRegistryCredentials` from the `prefect-docker` package, which
+    offers enhanced functionality and better a better user experience.
+
     Connects to a Docker registry.
 
     Requires a Docker Engine to be connectable.

{prefect-client-2.19.4 → prefect-client-2.19.6}/src/prefect/runner/storage.py

@@ -564,8 +564,74 @@ class BlockStorageAdapter:
         return False
 
 
-def create_storage_from_url(
-    url: str, pull_interval: Optional[int] = 60
+class LocalStorage:
+    """
+    Sets the working directory in the local filesystem.
+
+    Parameters:
+        Path: Local file path to set the working directory for the flow
+
+    Examples:
+        Sets the working directory for the local path to the flow:
+
+        ```python
+        from prefect.runner.storage import Localstorage
+
+        storage = LocalStorage(
+            path="/path/to/local/flow_directory",
+        )
+        ```
+    """
+
+    def __init__(
+        self,
+        path: str,
+        pull_interval: Optional[int] = None,
+    ):
+        self._path = Path(path).resolve()
+        self._logger = get_logger("runner.storage.local-storage")
+        self._storage_base_path = Path.cwd()
+        self._pull_interval = pull_interval
+
+    @property
+    def destination(self) -> Path:
+        return self._path
+
+    def set_base_path(self, path: Path):
+        self._storage_base_path = path
+
+    @property
+    def pull_interval(self) -> Optional[int]:
+        return self._pull_interval
+
+    async def pull_code(self):
+        # Local storage assumes the code already exists on the local filesystem
+        # and does not need to be pulled from a remote location
+        pass
+
+    def to_pull_step(self) -> dict:
+        """
+        Returns a dictionary representation of the storage object that can be
+        used as a deployment pull step.
+        """
+        step = {
+            "prefect.deployments.steps.set_working_directory": {
+                "directory": str(self.destination)
+            }
+        }
+        return step
+
+    def __eq__(self, __value) -> bool:
+        if isinstance(__value, LocalStorage):
+            return self._path == __value._path
+        return False
+
+    def __repr__(self) -> str:
+        return f"LocalStorage(path={self._path!r})"
+
+
+def create_storage_from_source(
+    source: str, pull_interval: Optional[int] = 60
 ) -> RunnerStorage:
     """
     Creates a storage object from a URL.
@@ -579,11 +645,18 @@ def create_storage_from_url(
     Returns:
         RunnerStorage: A runner storage compatible object
     """
-    parsed_url = urlparse(url)
-    if parsed_url.scheme == "git" or parsed_url.path.endswith(".git"):
-        return GitRepository(url=url, pull_interval=pull_interval)
+    logger = get_logger("runner.storage")
+    parsed_source = urlparse(source)
+    if parsed_source.scheme == "git" or parsed_source.path.endswith(".git"):
+        return GitRepository(url=source, pull_interval=pull_interval)
+    elif parsed_source.scheme in ("file", "local"):
+        source_path = source.split("://", 1)[-1]
+        return LocalStorage(path=source_path, pull_interval=pull_interval)
+    elif parsed_source.scheme in fsspec.available_protocols():
+        return RemoteStorage(url=source, pull_interval=pull_interval)
     else:
-        return RemoteStorage(url=url, pull_interval=pull_interval)
+        logger.debug("No valid fsspec protocol found for URL, assuming local storage.")
+        return LocalStorage(path=source, pull_interval=pull_interval)
 
 
 def _format_token_from_credentials(netloc: str, credentials: dict) -> str:

{prefect-client-2.19.4 → prefect-client-2.19.6}/src/prefect/utilities/importtools.py

@@ -381,6 +381,15 @@ def safe_load_namespace(source_code: str):
 
     namespace = {"__name__": "prefect_safe_namespace_loader"}
 
+    # Remove the body of the if __name__ == "__main__": block from the AST to prevent
+    # execution of guarded code
+    new_body = []
+    for node in parsed_code.body:
+        if _is_main_block(node):
+            continue
+        new_body.append(node)
+    parsed_code.body = new_body
+
     # Walk through the AST and find all import statements
     for node in ast.walk(parsed_code):
         if isinstance(node, ast.Import):
@@ -427,3 +436,23 @@ def safe_load_namespace(source_code: str):
             except Exception as e:
                 logger.debug("Failed to compile: %s", e)
     return namespace
+
+
+def _is_main_block(node: ast.AST):
+    """
+    Check if the node is an `if __name__ == "__main__":` block.
+    """
+    if isinstance(node, ast.If):
+        try:
+            # Check if the condition is `if __name__ == "__main__":`
+            if (
+                isinstance(node.test, ast.Compare)
+                and isinstance(node.test.left, ast.Name)
+                and node.test.left.id == "__name__"
+                and isinstance(node.test.comparators[0], ast.Constant)
+                and node.test.comparators[0].value == "__main__"
+            ):
+                return True
+        except AttributeError:
+            pass
+    return False

{prefect-client-2.19.4 → prefect-client-2.19.6/src/prefect_client.egg-info}/PKG-INFO

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