prefect-client 3.2.13__py3-none-any.whl → 3.2.15.dev7__py3-none-any.whl

prefect/results.py

@@ -552,7 +552,11 @@ class ResultStore(BaseModel):
         if self.result_storage_block_id is None and (
             _resolve_path := getattr(self.result_storage, "_resolve_path", None)
         ):
-            return str(_resolve_path(key))
+            path_key = _resolve_path(key)
+            if path_key is not None:
+                return str(_resolve_path(key))
+            else:
+                return key
         return key
 
     @sync_compatible
@@ -684,7 +688,9 @@ class ResultStore(BaseModel):
 
         if self.result_storage_block_id is None:
             if _resolve_path := getattr(self.result_storage, "_resolve_path", None):
-                key = str(_resolve_path(key))
+                path_key = _resolve_path(key)
+                if path_key is not None:
+                    key = str(_resolve_path(key))
 
         return ResultRecord(
             result=obj,
@@ -773,7 +779,7 @@ class ResultStore(BaseModel):
                 )
                 else Path(".").resolve()
             )
-            base_key = str(Path(key).relative_to(basepath))
+            base_key = key if basepath is None else str(Path(key).relative_to(basepath))
         else:
             base_key = key
         if (

prefect/runner/__init__.py

@@ -1,2 +1,4 @@
 from .runner import Runner
 from .submit import submit_to_runner, wait_for_submitted_runs
+
+__all__ = ["Runner", "submit_to_runner", "wait_for_submitted_runs"]

prefect/runner/runner.py

@@ -1425,7 +1425,7 @@ class Runner:
             )
             status_code = process.returncode
         except Exception as exc:
-            if not task_status._future.done():
+            if not task_status._future.done():  # type: ignore
                 # This flow run was being submitted and did not start successfully
                 run_logger.exception(
                     f"Failed to start process for flow run '{flow_run.id}'."

prefect/runner/server.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import uuid
-from typing import TYPE_CHECKING, Any, Callable, Coroutine, Hashable, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Coroutine, Hashable, Optional
 
 import uvicorn
 from fastapi import APIRouter, FastAPI, HTTPException, status
@@ -33,7 +35,7 @@ if TYPE_CHECKING:
 
 from pydantic import BaseModel
 
-logger: "logging.Logger" = get_logger("webserver")
+logger: "logging.Logger" = get_logger("runner.webserver")
 
 RunnableEndpoint = Literal["deployment", "flow", "task"]
 
@@ -45,7 +47,7 @@ class RunnerGenericFlowRunRequest(BaseModel):
 
 
 def perform_health_check(
-    runner: "Runner", delay_threshold: Optional[int] = None
+    runner: "Runner", delay_threshold: int | None = None
 ) -> Callable[..., JSONResponse]:
     if delay_threshold is None:
         delay_threshold = (
@@ -57,6 +59,9 @@ def perform_health_check(
         now = DateTime.now("utc")
         poll_delay = (now - runner.last_polled).total_seconds()
 
+        if TYPE_CHECKING:
+            assert delay_threshold is not None
+
         if poll_delay > delay_threshold:
             return JSONResponse(
                 status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
@@ -120,7 +125,7 @@ async def _build_endpoint_for_deployment(
 
 async def get_deployment_router(
     runner: "Runner",
-) -> Tuple[APIRouter, dict[Hashable, Any]]:
+) -> tuple[APIRouter, dict[Hashable, Any]]:
     router = APIRouter()
     schemas: dict[Hashable, Any] = {}
     async with get_client() as client:
@@ -216,14 +221,14 @@ def _build_generic_endpoint_for_flows(
         # Verify that the flow we're loading is a subflow this runner is
         # managing
         if not _flow_in_schemas(flow, schemas):
-            runner._logger.warning(
+            logger.warning(
                 f"Flow {flow.name} is not directly managed by the runner. Please "
                 "include it in the runner's served flows' import namespace."
             )
         # Verify that the flow we're loading hasn't changed since the webserver
         # was started
         if _flow_schema_changed(flow, schemas):
-            runner._logger.warning(
+            logger.warning(
                 "A change in flow parameters has been detected. Please "
                 "restart the runner."
             )
@@ -291,7 +296,7 @@ async def build_server(runner: "Runner") -> FastAPI:
     return webserver
 
 
-def start_webserver(runner: "Runner", log_level: Optional[str] = None) -> None:
+def start_webserver(runner: "Runner", log_level: str | None = None) -> None:
     """
     Run a FastAPI server for a runner.
 

prefect/runner/storage.py

@@ -6,7 +6,6 @@ from copy import deepcopy
 from pathlib import Path
 from typing import (
     Any,
-    Dict,
     Optional,
     Protocol,
     TypedDict,
@@ -16,7 +15,7 @@ from typing import (
 from urllib.parse import urlparse, urlsplit, urlunparse
 from uuid import uuid4
 
-import fsspec
+import fsspec  # pyright: ignore[reportMissingTypeStubs]
 from anyio import run_process
 from pydantic import SecretStr
 
@@ -79,7 +78,7 @@ class RunnerStorage(Protocol):
 
 class GitCredentials(TypedDict, total=False):
     username: str
-    access_token: Union[str, Secret[str]]
+    access_token: str | Secret[str]
 
 
 class GitRepository:
@@ -117,12 +116,12 @@ class GitRepository:
     def __init__(
         self,
         url: str,
-        credentials: Union[GitCredentials, Block, Dict[str, Any], None] = None,
-        name: Optional[str] = None,
-        branch: Optional[str] = None,
+        credentials: Union[GitCredentials, Block, dict[str, Any], None] = None,
+        name: str | None = None,
+        branch: str | None = None,
         include_submodules: bool = False,
-        pull_interval: Optional[int] = 60,
-        directories: Optional[str] = None,
+        pull_interval: int | None = 60,
+        directories: list[str] | None = None,
     ):
         if credentials is None:
             credentials = {}
@@ -198,7 +197,7 @@ class GitRepository:
     @property
     def _git_config(self) -> list[str]:
         """Build a git configuration to use when running git commands."""
-        config = {}
+        config: dict[str, str] = {}
 
         # Submodules can be private. The url in .gitmodules
         # will not include the credentials, we need to
@@ -220,7 +219,7 @@ class GitRepository:
             result = await run_process(
                 ["git", "config", "--get", "core.sparseCheckout"], cwd=self.destination
             )
-            return result.strip().lower() == "true"
+            return result.stdout.decode().strip().lower() == "true"
         except Exception:
             return False
 
@@ -243,8 +242,7 @@ class GitRepository:
                 cwd=str(self.destination),
             )
             existing_repo_url = None
-            if result.stdout is not None:
-                existing_repo_url = _strip_auth_from_url(result.stdout.decode().strip())
+            existing_repo_url = _strip_auth_from_url(result.stdout.decode().strip())
 
             if existing_repo_url != self._url:
                 raise ValueError(
@@ -255,7 +253,7 @@ class GitRepository:
             # Sparsely checkout the repository if directories are specified and the repo is not in sparse-checkout mode already
             if self._directories and not await self.is_sparsely_checked_out():
                 await run_process(
-                    ["git", "sparse-checkout", "set"] + self._directories,
+                    ["git", "sparse-checkout", "set", *self._directories],
                     cwd=self.destination,
                 )
 
@@ -323,7 +321,7 @@ class GitRepository:
         if self._directories:
             self._logger.debug("Will add %s", self._directories)
             await run_process(
-                ["git", "sparse-checkout", "set"] + self._directories,
+                ["git", "sparse-checkout", "set", *self._directories],
                 cwd=self.destination,
             )
 
@@ -343,7 +341,7 @@ class GitRepository:
         )
 
     def to_pull_step(self) -> dict[str, Any]:
-        pull_step = {
+        pull_step: dict[str, Any] = {
             "prefect.deployments.steps.git_clone": {
                 "repository": self._url,
                 "branch": self._branch,
@@ -357,13 +355,14 @@ class GitRepository:
             pull_step["prefect.deployments.steps.git_clone"]["credentials"] = (
                 f"{{{{ {self._credentials.get_block_placeholder()} }}}}"
             )
-        elif isinstance(self._credentials, dict):
-            if isinstance(self._credentials.get("access_token"), Secret):
+        elif isinstance(self._credentials, dict):  # pyright: ignore[reportUnnecessaryIsInstance]
+            if isinstance(
+                access_token := self._credentials.get("access_token"), Secret
+            ):
                 pull_step["prefect.deployments.steps.git_clone"]["credentials"] = {
                     **self._credentials,
                     "access_token": (
-                        "{{"
-                        f" {self._credentials['access_token'].get_block_placeholder()} }}}}"
+                        f"{{{{ {access_token.get_block_placeholder()} }}}}"
                     ),
                 }
             elif self._credentials.get("access_token") is not None:
@@ -455,10 +454,10 @@ class RemoteStorage:
 
         def replace_blocks_with_values(obj: Any) -> Any:
             if isinstance(obj, Block):
-                if hasattr(obj, "get"):
-                    return obj.get()
+                if get := getattr(obj, "get", None):
+                    return get()
                 if hasattr(obj, "value"):
-                    return obj.value
+                    return getattr(obj, "value")
                 else:
                     return obj.model_dump()
             return obj
@@ -467,7 +466,7 @@ class RemoteStorage:
             self._settings, replace_blocks_with_values, return_data=True
         )
 
-        return fsspec.filesystem(scheme, **settings_with_block_values)
+        return fsspec.filesystem(scheme, **settings_with_block_values)  # pyright: ignore[reportUnknownMemberType] missing type stubs
 
     def set_base_path(self, path: Path) -> None:
         self._storage_base_path = path
@@ -513,7 +512,7 @@ class RemoteStorage:
         try:
             await from_async.wait_for_call_in_new_thread(
                 create_call(
-                    self._filesystem.get,
+                    self._filesystem.get,  # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] missing type stubs
                     remote_path,
                     str(self.destination),
                     recursive=True,
@@ -580,18 +579,14 @@ class BlockStorageAdapter:
         self._block = block
         self._pull_interval = pull_interval
         self._storage_base_path = Path.cwd()
-        if not isinstance(block, Block):
+        if not isinstance(block, Block):  # pyright: ignore[reportUnnecessaryIsInstance]
             raise TypeError(
                 f"Expected a block object. Received a {type(block).__name__!r} object."
             )
         if not hasattr(block, "get_directory"):
             raise ValueError("Provided block must have a `get_directory` method.")
 
-        self._name = (
-            f"{block.get_block_type_slug()}-{block._block_document_name}"
-            if block._block_document_name
-            else str(uuid4())
-        )
+        self._name = f"{block.get_block_type_slug()}-{getattr(block, '_block_document_name', uuid4())}"
 
     def set_base_path(self, path: Path) -> None:
         self._storage_base_path = path
@@ -610,11 +605,11 @@ class BlockStorageAdapter:
         await self._block.get_directory(local_path=str(self.destination))
 
     def to_pull_step(self) -> dict[str, Any]:
-        # Give blocks the change to implement their own pull step
+        # Give blocks the chance to implement their own pull step
         if hasattr(self._block, "get_pull_step"):
-            return self._block.get_pull_step()
+            return getattr(self._block, "get_pull_step")()
         else:
-            if not self._block._block_document_name:
+            if getattr(self._block, "_block_document_name", None) is None:
                 raise BlockNotSavedError(
                     "Block must be saved with `.save()` before it can be converted to a"
                     " pull step."
@@ -622,7 +617,7 @@ class BlockStorageAdapter:
             return {
                 "prefect.deployments.steps.pull_with_block": {
                     "block_type_slug": self._block.get_block_type_slug(),
-                    "block_document_name": self._block._block_document_name,
+                    "block_document_name": getattr(self._block, "_block_document_name"),
                 }
             }
 
@@ -723,7 +718,9 @@ def create_storage_from_source(
         return LocalStorage(path=source, pull_interval=pull_interval)
 
 
-def _format_token_from_credentials(netloc: str, credentials: dict) -> str:
+def _format_token_from_credentials(
+    netloc: str, credentials: dict[str, Any] | GitCredentials
+) -> str:
     """
     Formats the credentials block for the git provider.
 
@@ -736,7 +733,10 @@ def _format_token_from_credentials(netloc: str, credentials: dict) -> str:
     token = credentials.get("token") if credentials else None
     access_token = credentials.get("access_token") if credentials else None
 
-    user_provided_token = access_token or token or password
+    user_provided_token: str | Secret[str] | None = access_token or token or password
+
+    if isinstance(user_provided_token, Secret):
+        user_provided_token = user_provided_token.get()
 
     if not user_provided_token:
         raise ValueError(
@@ -787,7 +787,7 @@ def _strip_auth_from_url(url: str) -> str:
 
     # Construct a new netloc without the auth info
     netloc = parsed.hostname
-    if parsed.port:
+    if parsed.port and netloc:
         netloc += f":{parsed.port}"
 
     # Build the sanitized URL

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,