prefect-client 2.14.18__tar.gz → 2.14.20__tar.gz

{prefect-client-2.14.18 → prefect-client-2.14.20}/PKG-INFO

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

{prefect-client-2.14.18 → prefect-client-2.14.20}/src/prefect/client/orchestration.py

@@ -1961,7 +1961,7 @@ class PrefectClient:
     async def create_task_run(
         self,
         task: "TaskObject",
-        flow_run_id: UUID,
+        flow_run_id: Optional[UUID],
         dynamic_key: str,
         name: str = None,
         extra_tags: Iterable[str] = None,

{prefect-client-2.14.18 → prefect-client-2.14.20}/src/prefect/client/schemas/filters.py

@@ -316,6 +316,19 @@ class FlowRunFilter(PrefectBaseModel, OperatorMixin):
     )
 
 
+class TaskRunFilterFlowRunId(PrefectBaseModel):
+    """Filter by `TaskRun.flow_run_id`."""
+
+    any_: Optional[List[UUID]] = Field(
+        default=None, description="A list of flow run ids to include"
+    )
+
+    is_null_: bool = Field(
+        default=False,
+        description="If true, only include task runs without a flow run id",
+    )
+
+
 class TaskRunFilterId(PrefectBaseModel):
     """Filter by `TaskRun.id`."""
 
@@ -428,6 +441,9 @@ class TaskRunFilter(PrefectBaseModel, OperatorMixin):
     subflow_runs: Optional[TaskRunFilterSubFlowRuns] = Field(
         default=None, description="Filter criteria for `TaskRun.subflow_run`"
     )
+    flow_run_id: Optional[TaskRunFilterFlowRunId] = Field(
+        default=None, description="Filter criteria for `TaskRun.flow_run_id`"
+    )
 
 
 class DeploymentFilterId(PrefectBaseModel):

{prefect-client-2.14.18 → prefect-client-2.14.20}/src/prefect/client/schemas/objects.py

@@ -113,6 +113,7 @@ class StateDetails(PrefectBaseModel):
     refresh_cache: bool = None
     retriable: bool = None
     transition_id: Optional[UUID] = None
+    task_parameters_id: Optional[UUID] = None
 
 
 class State(ObjectBaseModel, Generic[R]):

prefect-client-2.14.20/src/prefect/client/subscriptions.py

@@ -0,0 +1,82 @@
+import asyncio
+from typing import Generic, Type, TypeVar
+
+import orjson
+import websockets
+import websockets.exceptions
+from starlette.status import WS_1008_POLICY_VIOLATION
+from typing_extensions import Self
+
+from prefect._internal.schemas.bases import IDBaseModel
+from prefect.settings import PREFECT_API_KEY, PREFECT_API_URL
+
+S = TypeVar("S", bound=IDBaseModel)
+
+
+class Subscription(Generic[S]):
+    def __init__(self, model: Type[S], path: str):
+        self.model = model
+
+        base_url = PREFECT_API_URL.value().replace("http", "ws", 1)
+        self.subscription_url = f"{base_url}{path}"
+
+        self._connect = websockets.connect(
+            self.subscription_url,
+            subprotocols=["prefect"],
+        )
+        self._websocket = None
+
+    def __aiter__(self) -> Self:
+        return self
+
+    async def __anext__(self) -> S:
+        while True:
+            try:
+                await self._ensure_connected()
+                message = await self._websocket.recv()
+
+                message_data = orjson.loads(message)
+
+                if message_data.get("type") == "ping":
+                    await self._websocket.send(orjson.dumps({"type": "pong"}).decode())
+                    continue
+
+                return self.model.parse_raw(message)
+            except (
+                ConnectionRefusedError,
+                websockets.exceptions.ConnectionClosedError,
+            ):
+                self._websocket = None
+                if hasattr(self._connect, "protocol"):
+                    await self._connect.__aexit__(None, None, None)
+                await asyncio.sleep(0.5)
+
+    async def _ensure_connected(self):
+        if self._websocket:
+            return
+
+        websocket = await self._connect.__aenter__()
+
+        await websocket.send(
+            orjson.dumps({"type": "auth", "token": PREFECT_API_KEY.value()}).decode()
+        )
+
+        try:
+            auth = orjson.loads(await websocket.recv())
+            assert auth["type"] == "auth_success"
+        except (
+            AssertionError,
+            websockets.exceptions.ConnectionClosedError,
+        ) as e:
+            if isinstance(e, AssertionError) or e.code == WS_1008_POLICY_VIOLATION:
+                raise Exception(
+                    "Unable to authenticate to the subscription. Please "
+                    "ensure the provided `PREFECT_API_KEY` you are using is "
+                    "valid for this environment."
+                ) from e
+            raise
+        else:
+            self._websocket = websocket
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}[{self.model.__name__}]"

{prefect-client-2.14.18 → prefect-client-2.14.20}/src/prefect/context.py

@@ -214,7 +214,7 @@ class RunContext(ContextModel):
     client: PrefectClient
 
 
-class FlowRunContext(RunContext):
+class EngineContext(RunContext):
     """
     The context for a flow run. Data in this context is only available from within a
     flow run function.
@@ -233,6 +233,7 @@ class FlowRunContext(RunContext):
 
     flow: Optional["Flow"] = None
     flow_run: Optional[FlowRun] = None
+    autonomous_task_run: Optional[TaskRun] = None
     task_runner: BaseTaskRunner
     log_prints: bool = False
     parameters: Dict[str, Any]
@@ -266,6 +267,9 @@ class FlowRunContext(RunContext):
     __var__ = ContextVar("flow_run")
 
 
+FlowRunContext = EngineContext  # for backwards compatibility
+
+
 class TaskRunContext(RunContext):
     """
     The context for a task run. Data in this context is only available from within a

{prefect-client-2.14.18 → prefect-client-2.14.20}/src/prefect/engine.py

@@ -155,7 +155,7 @@ from prefect.exceptions import (
 )
 from prefect.flows import Flow, load_flow_from_entrypoint
 from prefect.futures import PrefectFuture, call_repr, resolve_futures_to_states
-from prefect.input import RunInput, keyset_from_paused_state
+from prefect.input import keyset_from_paused_state
 from prefect.input.run_input import run_input_subclass_from_type
 from prefect.logging.configuration import setup_logging
 from prefect.logging.handlers import APILogHandler
@@ -169,6 +169,7 @@ from prefect.logging.loggers import (
 from prefect.results import BaseResult, ResultFactory, UnknownResult
 from prefect.settings import (
     PREFECT_DEBUG_MODE,
+    PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING,
     PREFECT_LOGGING_LOG_PRINTS,
     PREFECT_TASK_INTROSPECTION_WARN_THRESHOLD,
     PREFECT_TASKS_REFRESH_CACHE,
@@ -179,6 +180,7 @@ from prefect.states import (
     Paused,
     Pending,
     Running,
+    Scheduled,
     State,
     Suspended,
     exception_to_crashed_state,
@@ -210,9 +212,10 @@ from prefect.utilities.pydantic import PartialModel
 from prefect.utilities.text import truncated_to
 
 R = TypeVar("R")
-T = TypeVar("T", bound=RunInput)
+T = TypeVar("T")
 EngineReturnType = Literal["future", "state", "result"]
 
+NUM_CHARS_DYNAMIC_KEY = 8
 
 API_HEALTHCHECKS = {}
 UNTRACKABLE_TYPES = {bool, type(None), type(...), type(NotImplemented)}
@@ -984,18 +987,6 @@ async def pause_flow_run(
     ...
 
 
-@overload
-async def pause_flow_run(
-    wait_for_input: Type[Any],
-    flow_run_id: UUID = None,
-    timeout: int = 3600,
-    poll_interval: int = 10,
-    reschedule: bool = False,
-    key: str = None,
-) -> Any:
-    ...
-
-
 @sync_compatible
 @deprecated_parameter(
     "flow_run_id", start_date="Dec 2023", help="Use `suspend_flow_run` instead."
@@ -1010,13 +1001,13 @@ async def pause_flow_run(
     "wait_for_input", group="flow_run_input", when=lambda y: y is not None
 )
 async def pause_flow_run(
-    wait_for_input: Optional[Union[Type[T], Type[Any]]] = None,
+    wait_for_input: Optional[Type[T]] = None,
     flow_run_id: UUID = None,
     timeout: int = 3600,
     poll_interval: int = 10,
     reschedule: bool = False,
     key: str = None,
-):
+) -> Optional[T]:
     """
     Pauses the current flow run by blocking execution until resumed.
 
@@ -1099,8 +1090,8 @@ async def _in_process_pause(
     reschedule=False,
     key: str = None,
     client=None,
-    wait_for_input: Optional[Union[Type[RunInput], Type[Any]]] = None,
-) -> Optional[RunInput]:
+    wait_for_input: Optional[T] = None,
+) -> Optional[T]:
     if TaskRunContext.get():
         raise RuntimeError("Cannot pause task runs.")
 
@@ -1231,29 +1222,18 @@ async def suspend_flow_run(
     ...
 
 
-@overload
-async def suspend_flow_run(
-    wait_for_input: Type[Any],
-    flow_run_id: Optional[UUID] = None,
-    timeout: Optional[int] = 3600,
-    key: Optional[str] = None,
-    client: PrefectClient = None,
-) -> Any:
-    ...
-
-
 @sync_compatible
 @inject_client
 @experimental_parameter(
     "wait_for_input", group="flow_run_input", when=lambda y: y is not None
 )
 async def suspend_flow_run(
-    wait_for_input: Optional[Union[Type[T], Type[Any]]] = None,
+    wait_for_input: Optional[Type[T]] = None,
     flow_run_id: Optional[UUID] = None,
     timeout: Optional[int] = 3600,
     key: Optional[str] = None,
     client: PrefectClient = None,
-):
+) -> Optional[T]:
     """
     Suspends a flow run by stopping code execution until resumed.
 
@@ -1382,16 +1362,19 @@ def enter_task_run_engine(
     return_type: EngineReturnType,
     task_runner: Optional[BaseTaskRunner],
     mapped: bool,
-) -> Union[PrefectFuture, Awaitable[PrefectFuture]]:
-    """
-    Sync entrypoint for task calls
-    """
+) -> Union[PrefectFuture, Awaitable[PrefectFuture], TaskRun]:
+    """Sync entrypoint for task calls"""
 
     flow_run_context = FlowRunContext.get()
+
     if not flow_run_context:
+        if PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING.value():
+            return _create_autonomous_task_run(task=task, parameters=parameters)
+
         raise RuntimeError(
-            "Tasks cannot be run outside of a flow. To call the underlying task"
-            " function outside of a flow use `task.fn()`."
+            "Tasks cannot be run outside of a flow"
+            " - if you meant to submit an autonomous task, you need to set"
+            " `prefect config set PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING=true`"
         )
 
     if TaskRunContext.get():
@@ -1603,14 +1586,22 @@ async def create_task_run_future(
 
     # Generate a name for the future
     dynamic_key = _dynamic_key_for_task_run(flow_run_context, task)
-    task_run_name = f"{task.name}-{dynamic_key}"
+    task_run_name = (
+        f"{task.name}-{dynamic_key}"
+        if flow_run_context and flow_run_context.flow_run
+        else f"{task.name}-{dynamic_key[:NUM_CHARS_DYNAMIC_KEY]}"  # autonomous task run
+    )
 
     # Generate a future
     future = PrefectFuture(
         name=task_run_name,
         key=uuid4(),
         task_runner=task_runner,
-        asynchronous=task.isasync and flow_run_context.flow.isasync,
+        asynchronous=(
+            task.isasync and flow_run_context.flow.isasync
+            if flow_run_context and flow_run_context.flow
+            else task.isasync
+        ),
     )
 
     # Create and submit the task run in the background
@@ -1650,14 +1641,18 @@ async def create_task_run_then_submit(
     task_runner: BaseTaskRunner,
     extra_task_inputs: Dict[str, Set[TaskRunInput]],
 ) -> None:
-    task_run = await create_task_run(
-        task=task,
-        name=task_run_name,
-        flow_run_context=flow_run_context,
-        parameters=parameters,
-        dynamic_key=task_run_dynamic_key,
-        wait_for=wait_for,
-        extra_task_inputs=extra_task_inputs,
+    task_run = (
+        await create_task_run(
+            task=task,
+            name=task_run_name,
+            flow_run_context=flow_run_context,
+            parameters=parameters,
+            dynamic_key=task_run_dynamic_key,
+            wait_for=wait_for,
+            extra_task_inputs=extra_task_inputs,
+        )
+        if not flow_run_context.autonomous_task_run
+        else flow_run_context.autonomous_task_run
     )
 
     # Attach the task run to the future to support `get_state` operations
@@ -1698,7 +1693,7 @@ async def create_task_run(
     task_run = await flow_run_context.client.create_task_run(
         task=task,
         name=name,
-        flow_run_id=flow_run_context.flow_run.id,
+        flow_run_id=flow_run_context.flow_run.id if flow_run_context.flow_run else None,
         dynamic_key=dynamic_key,
         state=Pending(),
         extra_tags=TagsContext.get().current_tags,
@@ -1721,7 +1716,10 @@ async def submit_task_run(
 ) -> PrefectFuture:
     logger = get_run_logger(flow_run_context)
 
-    if task_runner.concurrency_type == TaskConcurrencyType.SEQUENTIAL:
+    if (
+        task_runner.concurrency_type == TaskConcurrencyType.SEQUENTIAL
+        and not flow_run_context.autonomous_task_run
+    ):
         logger.info(f"Executing {task_run.name!r} immediately...")
 
     future = await task_runner.submit(
@@ -1799,7 +1797,7 @@ async def begin_task_run(
             # worker, the flow run timeout will not be raised in the worker process.
             interruptible = maybe_flow_run_context.timeout_scope is not None
         else:
-            # Otherwise, retrieve a new client
+            # Otherwise, retrieve a new clien`t
             client = await stack.enter_async_context(get_client())
             interruptible = False
             await stack.enter_async_context(anyio.create_task_group())
@@ -2153,7 +2151,6 @@ async def orchestrate_task_run(
                     await _check_task_failure_retriable(task, task_run, terminal_state)
                 )
             state = await propose_state(client, terminal_state, task_run_id=task_run.id)
-
             last_event = _emit_task_run_state_change_event(
                 task_run=task_run,
                 initial_state=last_state,
@@ -2203,7 +2200,7 @@ async def orchestrate_task_run(
         level=logging.INFO if state.is_completed() else logging.ERROR,
         msg=f"Finished in state {display_state}",
     )
-
+    logger.warning(f"Task run {task_run.name!r} finished in state {display_state}")
     return state
 
 
@@ -2572,7 +2569,12 @@ async def propose_state(
 
 
 def _dynamic_key_for_task_run(context: FlowRunContext, task: Task) -> int:
-    if task.task_key not in context.task_run_dynamic_keys:
+    if context.flow_run is None:  # this is an autonomous task run
+        context.task_run_dynamic_keys[task.task_key] = getattr(
+            task, "dynamic_key", str(uuid4())
+        )
+
+    elif task.task_key not in context.task_run_dynamic_keys:
         context.task_run_dynamic_keys[task.task_key] = 0
     else:
         context.task_run_dynamic_keys[task.task_key] += 1
@@ -2912,6 +2914,34 @@ def _emit_task_run_state_change_event(
     )
 
 
+@sync_compatible
+async def _create_autonomous_task_run(
+    task: Task, parameters: Dict[str, Any]
+) -> TaskRun:
+    async with get_client() as client:
+        scheduled = Scheduled()
+        if parameters:
+            parameters_id = uuid4()
+            scheduled.state_details.task_parameters_id = parameters_id
+
+            # TODO: We want to use result storage for parameters, but we'll need
+            # a better way to use it than this.
+            task.persist_result = True
+            factory = await ResultFactory.from_task(task, client=client)
+            await factory.store_parameters(parameters_id, parameters)
+
+        task_run = await client.create_task_run(
+            task=task,
+            flow_run_id=None,
+            dynamic_key=f"{task.task_key}-{str(uuid4())[:NUM_CHARS_DYNAMIC_KEY]}",
+            state=scheduled,
+        )
+
+        engine_logger.debug(f"Submitted run of task {task.name!r} for execution")
+
+    return task_run
+
+
 if __name__ == "__main__":
     try:
         flow_run_id = UUID(

{prefect-client-2.14.18 → prefect-client-2.14.20}/src/prefect/filesystems.py

@@ -683,12 +683,27 @@ class Azure(WritableFileSystem, WritableDeploymentStorage):
             " require ADLFS to use DefaultAzureCredentials."
         ),
     )
-
+    azure_storage_container: Optional[SecretStr] = Field(
+        default=None,
+        title="Azure storage container",
+        description=(
+            "Blob Container in Azure Storage Account. If set the 'bucket_path' will"
+            " be interpreted using the following URL format:"
+            "'az://<container>@<storage_account>.dfs.core.windows.net/<bucket_path>'."
+        ),
+    )
     _remote_file_system: RemoteFileSystem = None
 
     @property
     def basepath(self) -> str:
-        return f"az://{self.bucket_path}"
+        if self.azure_storage_container:
+            return (
+                f"az://{self.azure_storage_container.get_secret_value()}"
+                f"@{self.azure_storage_account_name.get_secret_value()}"
+                f".dfs.core.windows.net/{self.bucket_path}"
+            )
+        else:
+            return f"az://{self.bucket_path}"
 
     @property
     def filesystem(self) -> RemoteFileSystem:
@@ -713,7 +728,7 @@ class Azure(WritableFileSystem, WritableDeploymentStorage):
             )
         settings["anon"] = self.azure_storage_anon
         self._remote_file_system = RemoteFileSystem(
-            basepath=f"az://{self.bucket_path}", settings=settings
+            basepath=self.basepath, settings=settings
         )
         return self._remote_file_system
 

{prefect-client-2.14.18 → prefect-client-2.14.20}/src/prefect/flows.py

@@ -65,7 +65,6 @@ else:
     V2ValidationError = None
 
 from rich.console import Console
-from rich.panel import Panel
 from typing_extensions import Literal, ParamSpec
 
 from prefect._internal.schemas.validators import raise_on_name_with_banned_characters
@@ -785,7 +784,7 @@ class Flow(Generic[P, R]):
                 )
 
             console = Console()
-            console.print(Panel(help_message))
+            console.print(help_message, soft_wrap=True)
         await runner.start(webserver=webserver)
 
     @classmethod