prefect-client 3.2.12__py3-none-any.whl → 3.2.14__py3-none-any.whl

prefect/runner/submit.py

@@ -3,15 +3,19 @@ from __future__ import annotations
 import asyncio
 import inspect
 import uuid
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, overload
+from typing import TYPE_CHECKING, Any, Union, overload
 
 import anyio
 import httpx
 from typing_extensions import Literal, TypeAlias
 
 from prefect.client.orchestration import get_client
-from prefect.client.schemas.filters import FlowRunFilter, TaskRunFilter
-from prefect.client.schemas.objects import FlowRun
+from prefect.client.schemas.filters import (
+    FlowRunFilter,
+    FlowRunFilterParentFlowRunId,
+    TaskRunFilter,
+)
+from prefect.client.schemas.objects import Constant, FlowRun, Parameter, TaskRunResult
 from prefect.context import FlowRunContext
 from prefect.flows import Flow
 from prefect.logging import get_logger
@@ -60,18 +64,20 @@ async def _submit_flow_to_runner(
 
         parent_flow_run_context = FlowRunContext.get()
 
-        task_inputs = {
-            k: await collect_task_run_inputs(v) for k, v in parameters.items()
+        task_inputs: dict[str, list[TaskRunResult | Parameter | Constant]] = {
+            k: list(await collect_task_run_inputs(v)) for k, v in parameters.items()
         }
         parameters = await resolve_inputs(parameters)
         dummy_task = Task(name=flow.name, fn=flow.fn, version=flow.version)
         parent_task_run = await client.create_task_run(
             task=dummy_task,
             flow_run_id=(
-                parent_flow_run_context.flow_run.id if parent_flow_run_context else None
+                parent_flow_run_context.flow_run.id
+                if parent_flow_run_context and parent_flow_run_context.flow_run
+                else None
             ),
             dynamic_key=(
-                dynamic_key_for_task_run(parent_flow_run_context, dummy_task)
+                str(dynamic_key_for_task_run(parent_flow_run_context, dummy_task))
                 if parent_flow_run_context
                 else str(uuid.uuid4())
             ),
@@ -79,14 +85,15 @@ async def _submit_flow_to_runner(
             state=Pending(),
         )
 
-        response = await client._client.post(
+        httpx_client = getattr(client, "_client")
+        response = await httpx_client.post(
             (
                 f"http://{PREFECT_RUNNER_SERVER_HOST.value()}"
                 f":{PREFECT_RUNNER_SERVER_PORT.value()}"
                 "/flow/run"
             ),
             json={
-                "entrypoint": flow._entrypoint,
+                "entrypoint": getattr(flow, "_entrypoint"),
                 "parameters": flow.serialize_parameters(parameters),
                 "parent_task_run_id": str(parent_task_run.id),
             },
@@ -98,15 +105,15 @@ async def _submit_flow_to_runner(
 
 @overload
 def submit_to_runner(
-    prefect_callable: Union[Flow[Any, Any], Task[Any, Any]],
-    parameters: Dict[str, Any],
+    prefect_callable: Flow[Any, Any] | Task[Any, Any],
+    parameters: dict[str, Any],
     retry_failed_submissions: bool = True,
 ) -> FlowRun: ...
 
 
 @overload
 def submit_to_runner(
-    prefect_callable: Union[Flow[Any, Any], Task[Any, Any]],
+    prefect_callable: Flow[Any, Any] | Task[Any, Any],
     parameters: list[dict[str, Any]],
     retry_failed_submissions: bool = True,
 ) -> list[FlowRun]: ...
@@ -114,10 +121,10 @@ def submit_to_runner(
 
 @sync_compatible
 async def submit_to_runner(
-    prefect_callable: FlowOrTask,
-    parameters: Optional[Union[dict[str, Any], list[dict[str, Any]]]] = None,
+    prefect_callable: Flow[Any, Any],
+    parameters: dict[str, Any] | list[dict[str, Any]] | None = None,
     retry_failed_submissions: bool = True,
-) -> Union[FlowRun, list[FlowRun]]:
+) -> FlowRun | list[FlowRun]:
     """
     Submit a callable in the background via the runner webserver one or more times.
 
@@ -127,22 +134,22 @@ async def submit_to_runner(
             each dictionary represents a discrete invocation of the callable
         retry_failed_submissions: Whether to retry failed submissions to the runner webserver.
     """
-    if not isinstance(prefect_callable, (Flow, Task)):
+    if not isinstance(prefect_callable, Flow):  # pyright: ignore[reportUnnecessaryIsInstance]
         raise TypeError(
             "The `submit_to_runner` utility only supports submitting flows and tasks."
         )
 
     parameters = parameters or {}
-    if isinstance(parameters, List):
+    if isinstance(parameters, list):
         return_single = False
-    elif isinstance(parameters, dict):
+    elif isinstance(parameters, dict):  # pyright: ignore[reportUnnecessaryIsInstance]
         parameters = [parameters]
         return_single = True
     else:
         raise TypeError("Parameters must be a dictionary or a list of dictionaries.")
 
-    submitted_runs = []
-    unsubmitted_parameters = []
+    submitted_runs: list[FlowRun] = []
+    unsubmitted_parameters: list[dict[str, Any]] = []
 
     for p in parameters:
         try:
@@ -181,9 +188,9 @@ async def submit_to_runner(
 
 @sync_compatible
 async def wait_for_submitted_runs(
-    flow_run_filter: Optional[FlowRunFilter] = None,
-    task_run_filter: Optional[TaskRunFilter] = None,
-    timeout: Optional[float] = None,
+    flow_run_filter: FlowRunFilter | None = None,
+    task_run_filter: TaskRunFilter | None = None,
+    timeout: float | None = None,
     poll_interval: float = 3.0,
 ):
     """
@@ -197,7 +204,9 @@ async def wait_for_submitted_runs(
         poll_interval: How long to wait between polling each run's state (seconds).
     """
 
-    parent_flow_run_id = ctx.flow_run.id if (ctx := FlowRunContext.get()) else None
+    parent_flow_run_id = (
+        ctx.flow_run.id if ((ctx := FlowRunContext.get()) and ctx.flow_run) else None
+    )
 
     if task_run_filter:
         raise NotImplementedError("Waiting for task runs is not yet supported.")
@@ -223,7 +232,9 @@ async def wait_for_submitted_runs(
             if parent_flow_run_id is not None:
                 subflow_runs = await client.read_flow_runs(
                     flow_run_filter=FlowRunFilter(
-                        parent_flow_run_id=dict(any_=[parent_flow_run_id])
+                        parent_flow_run_id=FlowRunFilterParentFlowRunId(
+                            any_=[parent_flow_run_id]
+                        )
                     )
                 )
 

prefect/runner/utils.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 from copy import deepcopy
-from typing import Any
+from typing import Any, Hashable
 
 from fastapi import FastAPI
 from fastapi.openapi.utils import get_openapi
@@ -8,7 +10,7 @@ from prefect import __version__ as PREFECT_VERSION
 
 
 def inject_schemas_into_openapi(
-    webserver: FastAPI, schemas_to_inject: dict[str, Any]
+    webserver: FastAPI, schemas_to_inject: dict[Hashable, Any]
 ) -> dict[str, Any]:
     """
     Augments the webserver's OpenAPI schema with additional schemas from deployments / flows / tasks.
@@ -29,7 +31,7 @@ def inject_schemas_into_openapi(
 
 
 def merge_definitions(
-    injected_schemas: dict[str, Any], openapi_schema: dict[str, Any]
+    injected_schemas: dict[Hashable, Any], openapi_schema: dict[str, Any]
 ) -> dict[str, Any]:
     """
     Integrates definitions from injected schemas into the OpenAPI components.
@@ -51,7 +53,9 @@ def merge_definitions(
     return openapi_schema_copy
 
 
-def update_refs_in_schema(schema_item: Any, new_ref: str) -> None:
+def update_refs_in_schema(
+    schema_item: dict[str, Any] | list[Any], new_ref: str
+) -> None:
     """
     Recursively replaces `$ref` with a new reference base in a schema item.
 
@@ -64,7 +68,7 @@ def update_refs_in_schema(schema_item: Any, new_ref: str) -> None:
             schema_item["$ref"] = schema_item["$ref"].replace("#/definitions/", new_ref)
         for value in schema_item.values():
             update_refs_in_schema(value, new_ref)
-    elif isinstance(schema_item, list):
+    elif isinstance(schema_item, list):  # pyright: ignore[reportUnnecessaryIsInstance]
         for item in schema_item:
             update_refs_in_schema(item, new_ref)
 

prefect/server/api/collections_data/views/aggregate-worker-metadata.json

@@ -559,7 +559,7 @@
               "description": "To use any private container registry with a username and password, choose DockerRegistry. To use a private Azure Container Registry with a managed identity, choose ACRManagedIdentity.",
               "anyOf": [
                 {
-                  "$ref": "#/definitions/DockerRegistry"
+                  "$ref": "#/definitions/DockerRegistryCredentials"
                 },
                 {
                   "$ref": "#/definitions/ACRManagedIdentity"
@@ -637,8 +637,8 @@
             "subscription_id"
           ],
           "definitions": {
-            "DockerRegistry": {
-              "title": "DockerRegistry",
+            "DockerRegistryCredentials": {
+              "title": "DockerRegistryCredentials",
               "description": "Connects to a Docker registry.\n\nRequires a Docker Engine to be connectable.",
               "type": "object",
               "properties": {
@@ -671,7 +671,7 @@
                 "password",
                 "registry_url"
               ],
-              "block_type_slug": "docker-registry",
+              "block_type_slug": "docker-registry-credentials",
               "secret_fields": [
                 "password"
               ],

prefect/server/api/flow_runs.py

@@ -29,6 +29,7 @@ import prefect.server.schemas as schemas
 from prefect.logging import get_logger
 from prefect.server.api.run_history import run_history
 from prefect.server.api.validation import validate_job_variables_for_deployment_flow_run
+from prefect.server.api.workers import WorkerLookups
 from prefect.server.database import PrefectDBInterface, provide_database_interface
 from prefect.server.exceptions import FlowRunGraphTooLarge
 from prefect.server.models.flow_runs import (
@@ -68,6 +69,7 @@ async def create_flow_run(
         orchestration_dependencies.provide_flow_orchestration_parameters
     ),
     api_version: str = Depends(dependencies.provide_request_api_version),
+    worker_lookups: WorkerLookups = Depends(WorkerLookups),
 ) -> schemas.responses.FlowRunResponse:
     """
     Create a flow run. If a flow run with the same flow_id and
@@ -91,6 +93,25 @@ async def create_flow_run(
     right_now = now("UTC")
 
     async with db.session_context(begin_transaction=True) as session:
+        if flow_run.work_pool_name:
+            if flow_run.work_queue_name:
+                work_queue_id = await worker_lookups._get_work_queue_id_from_name(
+                    session=session,
+                    work_pool_name=flow_run.work_pool_name,
+                    work_queue_name=flow_run.work_queue_name,
+                )
+            else:
+                work_queue_id = (
+                    await worker_lookups._get_default_work_queue_id_from_work_pool_name(
+                        session=session,
+                        work_pool_name=flow_run.work_pool_name,
+                    )
+                )
+        else:
+            work_queue_id = None
+
+        flow_run_object.work_queue_id = work_queue_id
+
         model = await models.flow_runs.create_flow_run(
             session=session,
             flow_run=flow_run_object,

prefect/server/api/server.py

@@ -294,12 +294,16 @@ async def prefect_object_not_found_exception_handler(
     )
 
 
+API_APP_CACHE: dict[tuple[str, str | None], FastAPI] = {}
+
+
 def create_api_app(
     dependencies: list[Any] | None = None,
     health_check_path: str = "/health",
     version_check_path: str = "/version",
     fast_api_app_kwargs: dict[str, Any] | None = None,
     final: bool = False,
+    ignore_cache: bool = False,
 ) -> FastAPI:
     """
     Create a FastAPI app that includes the Prefect REST API
@@ -310,10 +314,20 @@ def create_api_app(
         fast_api_app_kwargs: kwargs to pass to the FastAPI constructor
         final: whether this will be the last instance of the Prefect server to be
             created in this process, so that additional optimizations may be applied
+        ignore_cache: if set, a new app will be created even if the settings and fast_api_app_kwargs match
+            an existing app in the cache
 
     Returns:
         a FastAPI app that serves the Prefect REST API
     """
+    cache_key = (
+        prefect.settings.get_current_settings().hash_key(),
+        hash_objects(fast_api_app_kwargs) if fast_api_app_kwargs else None,
+    )
+
+    if cache_key in API_APP_CACHE and not ignore_cache:
+        return API_APP_CACHE[cache_key]
+
     fast_api_app_kwargs = fast_api_app_kwargs or {}
     api_app = FastAPI(title=API_TITLE, **fast_api_app_kwargs)
     api_app.add_middleware(GZipMiddleware)
@@ -389,6 +403,8 @@ def create_api_app(
                 )
             return await call_next(request)
 
+    API_APP_CACHE[cache_key] = api_app
+
     return api_app
 
 
@@ -655,6 +671,7 @@ def create_app(
             }
         },
         final=final,
+        ignore_cache=ignore_cache,
     )
     ui_app = create_ui_app(ephemeral)
 

prefect/server/api/task_runs.py

@@ -16,6 +16,7 @@ from fastapi import (
     WebSocket,
     status,
 )
+from fastapi.responses import ORJSONResponse
 from starlette.websockets import WebSocketDisconnect
 
 import prefect.server.api.dependencies as dependencies
@@ -27,7 +28,10 @@ from prefect.server.database import PrefectDBInterface, provide_database_interfa
 from prefect.server.orchestration import dependencies as orchestration_dependencies
 from prefect.server.orchestration.core_policy import CoreTaskPolicy
 from prefect.server.orchestration.policies import TaskRunOrchestrationPolicy
-from prefect.server.schemas.responses import OrchestrationResult
+from prefect.server.schemas.responses import (
+    OrchestrationResult,
+    TaskRunPaginationResponse,
+)
 from prefect.server.task_queue import MultiQueue, TaskQueue
 from prefect.server.utilities import subscriptions
 from prefect.server.utilities.server import PrefectRouter
@@ -214,6 +218,53 @@ async def read_task_runs(
         )
 
 
+@router.post("/paginate", response_class=ORJSONResponse)
+async def paginate_task_runs(
+    sort: schemas.sorting.TaskRunSort = Body(schemas.sorting.TaskRunSort.ID_DESC),
+    limit: int = dependencies.LimitBody(),
+    page: int = Body(1, ge=1),
+    flows: Optional[schemas.filters.FlowFilter] = None,
+    flow_runs: Optional[schemas.filters.FlowRunFilter] = None,
+    task_runs: Optional[schemas.filters.TaskRunFilter] = None,
+    deployments: Optional[schemas.filters.DeploymentFilter] = None,
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> TaskRunPaginationResponse:
+    """
+    Pagination query for task runs.
+    """
+    offset = (page - 1) * limit
+
+    async with db.session_context() as session:
+        runs = await models.task_runs.read_task_runs(
+            session=session,
+            flow_filter=flows,
+            flow_run_filter=flow_runs,
+            task_run_filter=task_runs,
+            deployment_filter=deployments,
+            offset=offset,
+            limit=limit,
+            sort=sort,
+        )
+
+        total_count = await models.task_runs.count_task_runs(
+            session=session,
+            flow_filter=flows,
+            flow_run_filter=flow_runs,
+            task_run_filter=task_runs,
+            deployment_filter=deployments,
+        )
+
+        return TaskRunPaginationResponse.model_validate(
+            dict(
+                results=runs,
+                count=total_count,
+                limit=limit,
+                pages=(total_count + limit - 1) // limit,
+                page=page,
+            )
+        )
+
+
 @router.delete("/{id}", status_code=status.HTTP_204_NO_CONTENT)
 async def delete_task_run(
     task_run_id: UUID = Path(..., description="The task run id", alias="id"),

prefect/settings/models/tasks.py

@@ -57,6 +57,11 @@ class TasksSettings(PrefectBaseSettings):
         description="If `True`, enables a refresh of cached results: re-executing the task will refresh the cached results.",
     )
 
+    default_no_cache: bool = Field(
+        default=False,
+        description="If `True`, sets the default cache policy on all tasks to `NO_CACHE`.",
+    )
+
     default_retries: int = Field(
         default=0,
         ge=0,

prefect/task_engine.py

@@ -24,6 +24,7 @@ from typing import (
     Type,
     TypeVar,
     Union,
+    overload,
 )
 from uuid import UUID
 
@@ -1527,6 +1528,32 @@ async def run_generator_task_async(
         await engine.result()
 
 
+@overload
+def run_task(
+    task: "Task[P, R]",
+    task_run_id: Optional[UUID] = None,
+    task_run: Optional[TaskRun] = None,
+    parameters: Optional[dict[str, Any]] = None,
+    wait_for: Optional["OneOrManyFutureOrResult[Any]"] = None,
+    return_type: Literal["state"] = "state",
+    dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+    context: Optional[dict[str, Any]] = None,
+) -> State[R]: ...
+
+
+@overload
+def run_task(
+    task: "Task[P, R]",
+    task_run_id: Optional[UUID] = None,
+    task_run: Optional[TaskRun] = None,
+    parameters: Optional[dict[str, Any]] = None,
+    wait_for: Optional["OneOrManyFutureOrResult[Any]"] = None,
+    return_type: Literal["result"] = "result",
+    dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+    context: Optional[dict[str, Any]] = None,
+) -> R: ...
+
+
 def run_task(
     task: "Task[P, Union[R, Coroutine[Any, Any, R]]]",
     task_run_id: Optional[UUID] = None,

prefect/tasks.py

@@ -422,6 +422,11 @@ class Task(Generic[P, R]):
 
         self.task_key: str = _generate_task_key(self.fn)
 
+        # determine cache and result configuration
+        settings = get_current_settings()
+        if settings.tasks.default_no_cache and cache_policy is NotSet:
+            cache_policy = NO_CACHE
+
         if cache_policy is not NotSet and cache_key_fn is not None:
             logger.warning(
                 f"Both `cache_policy` and `cache_key_fn` are set on task {self}. `cache_key_fn` will be used."
@@ -458,7 +463,7 @@ class Task(Generic[P, R]):
                 )
         elif cache_policy is NotSet and result_storage_key is None:
             self.cache_policy = DEFAULT
-        elif result_storage_key:
+        elif cache_policy != NO_CACHE and result_storage_key:
             # TODO: handle this situation with double storage
             self.cache_policy = None
         else:
@@ -468,7 +473,6 @@ class Task(Generic[P, R]):
         # TODO: We can instantiate a `TaskRunPolicy` and add Pydantic bound checks to
         #       validate that the user passes positive numbers here
 
-        settings = get_current_settings()
         self.retries: int = (
             retries if retries is not None else settings.tasks.default_retries
         )
@@ -1426,6 +1430,8 @@ class Task(Generic[P, R]):
         else:
             return futures
 
+    # Background task methods
+
     def apply_async(
         self,
         args: Optional[tuple[Any, ...]] = None,

prefect/telemetry/run_telemetry.py

@@ -16,6 +16,7 @@ from opentelemetry.trace import (
 from typing_extensions import TypeAlias
 
 import prefect
+import prefect.settings
 from prefect.client.orchestration import PrefectClient, SyncPrefectClient
 from prefect.client.schemas import FlowRun, TaskRun
 from prefect.client.schemas.objects import State
@@ -50,15 +51,19 @@ class RunTelemetry:
         default_factory=lambda: get_tracer("prefect", prefect.__version__)
     )
     span: Span | None = None
+    _enabled: bool = field(
+        default_factory=lambda: prefect.settings.get_current_settings().cloud.enable_orchestration_telemetry
+    )
 
     async def async_start_span(
         self,
         run: FlowOrTaskRun,
         client: PrefectClient,
         parameters: dict[str, Any] | None = None,
-    ) -> Span:
+    ) -> Span | None:
+        if not self._enabled:
+            return None
         traceparent, span = self._start_span(run, parameters)
-
         if self._run_type(run) == "flow" and traceparent:
             # Only explicitly update labels if the run is a flow as task runs
             # are updated via events.
@@ -73,7 +78,10 @@ class RunTelemetry:
         run: FlowOrTaskRun,
         client: SyncPrefectClient,
         parameters: dict[str, Any] | None = None,
-    ) -> Span:
+    ) -> Span | None:
+        if not self._enabled:
+            return None
+
         traceparent, span = self._start_span(run, parameters)
 
         if self._run_type(run) == "flow" and traceparent:

prefect/transactions.py

@@ -23,6 +23,7 @@ from prefect.exceptions import (
     MissingContextError,
     SerializationError,
 )
+from prefect.filesystems import NullFileSystem
 from prefect.logging.loggers import LoggingAdapter, get_logger, get_run_logger
 from prefect.results import (
     ResultRecord,
@@ -453,6 +454,10 @@ def transaction(
     if key and not store:
         store = get_result_store()
 
+    # Avoid inheriting a NullFileSystem for metadata_storage from a flow's result store
+    if store and isinstance(store.metadata_storage, NullFileSystem):
+        store = store.model_copy(update={"metadata_storage": None})
+
     try:
         _logger: Union[logging.Logger, LoggingAdapter] = logger or get_run_logger()
     except MissingContextError:

prefect/utilities/engine.py

@@ -63,7 +63,7 @@ engine_logger: Logger = get_logger("engine")
 T = TypeVar("T")
 
 
-async def collect_task_run_inputs(expr: Any, max_depth: int = -1) -> set[TaskRunInput]:
+async def collect_task_run_inputs(expr: Any, max_depth: int = -1) -> set[TaskRunResult]:
     """
     This function recurses through an expression to generate a set of any discernible
     task run inputs it finds in the data structure. It produces a set of all inputs
@@ -76,7 +76,7 @@ async def collect_task_run_inputs(expr: Any, max_depth: int = -1) -> set[TaskRun
     """
     # TODO: This function needs to be updated to detect parameters and constants
 
-    inputs: set[TaskRunInput] = set()
+    inputs: set[TaskRunResult] = set()
 
     def add_futures_and_states_to_inputs(obj: Any) -> None:
         if isinstance(obj, PrefectFuture):

prefect/utilities/importtools.py

@@ -93,8 +93,11 @@ def load_script_as_module(path: str) -> ModuleType:
     parent_path = str(Path(path).resolve().parent)
     working_directory = os.getcwd()
 
-    # Generate unique module name for thread safety
-    module_name = f"__prefect_loader_{id(path)}__"
+    module_name = os.path.splitext(Path(path).name)[0]
+
+    # fall back in case of filenames with the same names as modules
+    if module_name in sys.modules:
+        module_name = f"__prefect_loader_{id(path)}__"
 
     spec = importlib.util.spec_from_file_location(
         module_name,
@@ -112,15 +115,9 @@ def load_script_as_module(path: str) -> ModuleType:
         with _get_sys_path_lock():
             sys.path.insert(0, working_directory)
             sys.path.insert(0, parent_path)
-            try:
-                spec.loader.exec_module(module)
-            finally:
-                sys.path.remove(parent_path)
-                sys.path.remove(working_directory)
+            spec.loader.exec_module(module)
     except Exception as exc:
         raise ScriptError(user_exc=exc, path=path) from exc
-    finally:
-        sys.modules.pop(module_name)
 
     return module