prefect-client 3.1.12__py3-none-any.whl → 3.1.13__py3-none-any.whl

prefect/infrastructure/provisioners/coiled.py

@@ -0,0 +1,249 @@
+import importlib
+import shlex
+import sys
+from copy import deepcopy
+from types import ModuleType
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+from anyio import run_process
+from rich.console import Console
+from rich.progress import Progress, SpinnerColumn, TextColumn
+from rich.prompt import Confirm
+
+from prefect.client.schemas.actions import BlockDocumentCreate
+from prefect.client.schemas.objects import BlockDocument
+from prefect.client.utilities import inject_client
+from prefect.exceptions import ObjectNotFound
+from prefect.utilities.importtools import lazy_import
+
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
+
+coiled: ModuleType = lazy_import("coiled")
+
+
+class CoiledPushProvisioner:
+    """
+    A infrastructure provisioner for Coiled push work pools.
+    """
+
+    def __init__(self, client: Optional["PrefectClient"] = None):
+        self._console = Console()
+
+    @property
+    def console(self) -> Console:
+        return self._console
+
+    @console.setter
+    def console(self, value: Console) -> None:
+        self._console = value
+
+    @staticmethod
+    def _is_coiled_installed() -> bool:
+        """
+        Checks if the coiled package is installed.
+
+        Returns:
+            True if the coiled package is installed, False otherwise
+        """
+        try:
+            importlib.import_module("coiled")
+            return True
+        except ModuleNotFoundError:
+            return False
+
+    async def _install_coiled(self):
+        """
+        Installs the coiled package.
+        """
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[bold blue]Installing coiled..."),
+            transient=True,
+            console=self.console,
+        ) as progress:
+            task = progress.add_task("coiled install")
+            progress.start()
+            global coiled
+            await run_process(
+                [shlex.quote(sys.executable), "-m", "pip", "install", "coiled"]
+            )
+            coiled = importlib.import_module("coiled")
+            progress.advance(task)
+
+    async def _get_coiled_token(self) -> str:
+        """
+        Gets a Coiled API token from the current Coiled configuration.
+        """
+        import dask.config
+
+        return dask.config.get("coiled.token", "")
+
+    async def _create_new_coiled_token(self):
+        """
+        Triggers a Coiled login via the browser if no current token. Will create a new token.
+        """
+        await run_process(["coiled", "login"])
+
+    async def _create_coiled_credentials_block(
+        self,
+        block_document_name: str,
+        coiled_token: str,
+        client: "PrefectClient",
+    ) -> BlockDocument:
+        """
+        Creates a CoiledCredentials block containing the provided token.
+
+        Args:
+            block_document_name: The name of the block document to create
+            coiled_token: The Coiled API token
+
+        Returns:
+            The ID of the created block
+        """
+        assert client is not None, "client injection failed"
+        try:
+            credentials_block_type = await client.read_block_type_by_slug(
+                "coiled-credentials"
+            )
+        except ObjectNotFound:
+            # Shouldn't happen, but just in case
+            raise RuntimeError(
+                "Unable to find CoiledCredentials block type. Please ensure you are"
+                " using Prefect Cloud."
+            )
+        credentials_block_schema = (
+            await client.get_most_recent_block_schema_for_block_type(
+                block_type_id=credentials_block_type.id
+            )
+        )
+        assert (
+            credentials_block_schema is not None
+        ), f"Unable to find schema for block type {credentials_block_type.slug}"
+
+        block_doc = await client.create_block_document(
+            block_document=BlockDocumentCreate(
+                name=block_document_name,
+                data={
+                    "api_token": coiled_token,
+                },
+                block_type_id=credentials_block_type.id,
+                block_schema_id=credentials_block_schema.id,
+            )
+        )
+        return block_doc
+
+    @inject_client
+    async def provision(
+        self,
+        work_pool_name: str,
+        base_job_template: Dict[str, Any],
+        client: Optional["PrefectClient"] = None,
+    ) -> Dict[str, Any]:
+        """
+        Provisions resources necessary for a Coiled push work pool.
+
+        Provisioned resources:
+            - A CoiledCredentials block containing a Coiled API token
+
+        Args:
+            work_pool_name: The name of the work pool to provision resources for
+            base_job_template: The base job template to update
+
+        Returns:
+            A copy of the provided base job template with the provisioned resources
+        """
+        credentials_block_name = f"{work_pool_name}-coiled-credentials"
+        base_job_template_copy = deepcopy(base_job_template)
+        assert client is not None, "client injection failed"
+        try:
+            block_doc = await client.read_block_document_by_name(
+                credentials_block_name, "coiled-credentials"
+            )
+            self.console.print(
+                f"Work pool [blue]{work_pool_name!r}[/] will reuse the existing Coiled"
+                f" credentials block [blue]{credentials_block_name!r}[/blue]"
+            )
+        except ObjectNotFound:
+            if self._console.is_interactive and not Confirm.ask(
+                (
+                    "\n"
+                    "To configure your Coiled push work pool we'll need to store a Coiled"
+                    " API token with Prefect Cloud as a block. We'll pull the token from"
+                    " your local Coiled configuration or create a new token if we"
+                    " can't find one.\n"
+                    "\n"
+                    "Would you like to continue?"
+                ),
+                console=self.console,
+                default=True,
+            ):
+                self.console.print(
+                    "No problem! You can always configure your Coiled push work pool"
+                    " later via the Prefect UI."
+                )
+                return base_job_template
+
+            if not self._is_coiled_installed():
+                if self.console.is_interactive and Confirm.ask(
+                    (
+                        "The [blue]coiled[/] package is required to configure"
+                        " authentication for your work pool.\n"
+                        "\n"
+                        "Would you like to install it now?"
+                    ),
+                    console=self.console,
+                    default=True,
+                ):
+                    await self._install_coiled()
+
+            if not self._is_coiled_installed():
+                raise RuntimeError(
+                    "The coiled package is not installed.\n\nPlease try installing coiled,"
+                    " or you can use the Prefect UI to create your Coiled push work pool."
+                )
+
+            # Get the current Coiled API token
+            coiled_api_token = await self._get_coiled_token()
+            if not coiled_api_token:
+                # Create a new token one wasn't found
+                if self.console.is_interactive and Confirm.ask(
+                    "Coiled credentials not found. Would you like to create a new token?",
+                    console=self.console,
+                    default=True,
+                ):
+                    await self._create_new_coiled_token()
+                    coiled_api_token = await self._get_coiled_token()
+                else:
+                    raise RuntimeError(
+                        "Coiled credentials not found. Please create a new token by"
+                        " running [blue]coiled login[/] and try again."
+                    )
+
+            # Create the credentials block
+            with Progress(
+                SpinnerColumn(),
+                TextColumn("[bold blue]Saving Coiled credentials..."),
+                transient=True,
+                console=self.console,
+            ) as progress:
+                task = progress.add_task("create coiled credentials block")
+                progress.start()
+                block_doc = await self._create_coiled_credentials_block(
+                    credentials_block_name,
+                    coiled_api_token,
+                    client=client,
+                )
+                progress.advance(task)
+
+        base_job_template_copy["variables"]["properties"]["credentials"]["default"] = {
+            "$ref": {"block_document_id": str(block_doc.id)}
+        }
+        if "image" in base_job_template_copy["variables"]["properties"]:
+            base_job_template_copy["variables"]["properties"]["image"]["default"] = ""
+        self.console.print(
+            f"Successfully configured Coiled push work pool {work_pool_name!r}!",
+            style="green",
+        )
+        return base_job_template_copy

prefect/infrastructure/provisioners/container_instance.py

@@ -10,6 +10,7 @@ Classes:
     ContainerInstancePushProvisioner: A class for provisioning infrastructure using Azure Container Instances.
 
 """
+from __future__ import annotations
 
 import json
 import random
@@ -64,7 +65,7 @@ class AzureCLI:
         failure_message: Optional[str] = None,
         ignore_if_exists: bool = False,
         return_json: bool = False,
-    ):
+    ) -> str | dict[str, Any] | None:
         """
         Runs an Azure CLI command and processes the output.
 
@@ -156,7 +157,7 @@ class ContainerInstancePushProvisioner:
         self._subscription_name = None
         self._location = "eastus"
         self._identity_name = "prefect-acr-identity"
-        self.azure_cli = AzureCLI(self.console)
+        self.azure_cli: AzureCLI = AzureCLI(self.console)
         self._credentials_block_name = None
         self._resource_group_name = "prefect-aci-push-pool-rg"
         self._app_registration_name = "prefect-aci-push-pool-app"
@@ -170,7 +171,7 @@ class ContainerInstancePushProvisioner:
     def console(self, value: Console) -> None:
         self._console = value
 
-    async def set_location(self):
+    async def set_location(self) -> None:
         """
         Set the Azure resource deployment location to the default or 'eastus' on failure.
 

prefect/infrastructure/provisioners/ecs.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import base64
 import contextlib
 import contextvars
@@ -9,9 +11,11 @@ import sys
 from copy import deepcopy
 from functools import partial
 from textwrap import dedent
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
+from types import ModuleType
+from typing import TYPE_CHECKING, Any, Callable, Dict, Generator, List, Optional
 
 import anyio
+import anyio.to_thread
 from anyio import run_process
 from rich.console import Console
 from rich.panel import Panel
@@ -33,13 +37,15 @@ from prefect.utilities.importtools import lazy_import
 if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient
 
-boto3 = lazy_import("boto3")
+boto3: ModuleType = lazy_import("boto3")
 
-current_console = contextvars.ContextVar("console", default=Console())
+current_console: contextvars.ContextVar[Console] = contextvars.ContextVar(
+    "console", default=Console()
+)
 
 
 @contextlib.contextmanager
-def console_context(value: Console):
+def console_context(value: Console) -> Generator[None, None, None]:
     token = current_console.set(value)
     try:
         yield
@@ -73,7 +79,7 @@ class IamPolicyResource:
         """
         return 1 if await self.requires_provisioning() else 0
 
-    def _get_policy_by_name(self, name):
+    def _get_policy_by_name(self, name: str) -> dict[str, Any] | None:
         paginator = self._iam_client.get_paginator("list_policies")
         page_iterator = paginator.paginate(Scope="Local")
 
@@ -119,9 +125,9 @@ class IamPolicyResource:
 
     async def provision(
         self,
-        policy_document: Dict[str, Any],
+        policy_document: dict[str, Any],
         advance: Callable[[], None],
-    ):
+    ) -> str:
         """
         Provisions an IAM policy.
 
@@ -153,7 +159,7 @@ class IamPolicyResource:
             return policy["Arn"]
 
     @property
-    def next_steps(self):
+    def next_steps(self) -> list[str]:
         return []
 
 
@@ -215,7 +221,7 @@ class IamUserResource:
     async def provision(
         self,
         advance: Callable[[], None],
-    ):
+    ) -> None:
         """
         Provisions an IAM user.
 
@@ -231,7 +237,7 @@ class IamUserResource:
             advance()
 
     @property
-    def next_steps(self):
+    def next_steps(self) -> list[str]:
         return []
 
 
@@ -241,7 +247,7 @@ class CredentialsBlockResource:
         self._user_name = user_name
         self._requires_provisioning = None
 
-    async def get_task_count(self):
+    async def get_task_count(self) -> int:
         """
         Returns the number of tasks that will be executed to provision this resource.
 
@@ -357,7 +363,7 @@ class CredentialsBlockResource:
         }
 
     @property
-    def next_steps(self):
+    def next_steps(self) -> list[str]:
         return []
 
 
@@ -374,7 +380,7 @@ class AuthenticationResource:
             credentials_block_name or f"{work_pool_name}-aws-credentials"
         )
         self._policy_name = policy_name
-        self._policy_document = {
+        self._policy_document: dict[str, Any] = {
             "Version": "2012-10-17",
             "Statement": [
                 {
@@ -417,7 +423,11 @@ class AuthenticationResource:
         self._execution_role_resource = ExecutionRoleResource()
 
     @property
-    def resources(self):
+    def resources(
+        self,
+    ) -> list[
+        "ExecutionRoleResource | IamUserResource | IamPolicyResource | CredentialsBlockResource"
+    ]:
         return [
             self._execution_role_resource,
             self._iam_user_resource,
@@ -425,7 +435,7 @@ class AuthenticationResource:
             self._credentials_block_resource,
         ]
 
-    async def get_task_count(self):
+    async def get_task_count(self) -> int:
         """
         Returns the number of tasks that will be executed to provision this resource.
 
@@ -461,9 +471,9 @@ class AuthenticationResource:
 
     async def provision(
         self,
-        base_job_template: Dict[str, Any],
+        base_job_template: dict[str, Any],
         advance: Callable[[], None],
-    ):
+    ) -> None:
         """
         Provisions the authentication resources.
 
@@ -507,7 +517,7 @@ class AuthenticationResource:
         )
 
     @property
-    def next_steps(self):
+    def next_steps(self) -> list[str]:
         return [
             next_step
             for resource in self.resources
@@ -521,7 +531,7 @@ class ClusterResource:
         self._cluster_name = cluster_name
         self._requires_provisioning = None
 
-    async def get_task_count(self):
+    async def get_task_count(self) -> int:
         """
         Returns the number of tasks that will be executed to provision this resource.
 
@@ -566,9 +576,9 @@ class ClusterResource:
 
     async def provision(
         self,
-        base_job_template: Dict[str, Any],
+        base_job_template: dict[str, Any],
         advance: Callable[[], None],
-    ):
+    ) -> None:
         """
         Provisions an ECS cluster.
 
@@ -592,7 +602,7 @@ class ClusterResource:
         ] = self._cluster_name
 
     @property
-    def next_steps(self):
+    def next_steps(self) -> list[str]:
         return []
 
 
@@ -608,7 +618,7 @@ class VpcResource:
         self._requires_provisioning = None
         self._ecs_security_group_name = ecs_security_group_name
 
-    async def get_task_count(self):
+    async def get_task_count(self) -> int:
         """
         Returns the number of tasks that will be executed to provision this resource.
 
@@ -642,7 +652,9 @@ class VpcResource:
         response = await anyio.to_thread.run_sync(self._ec2_client.describe_vpcs)
         return [vpc["CidrBlock"] for vpc in response["Vpcs"]]
 
-    async def _find_non_overlapping_cidr(self, default_cidr="172.31.0.0/16"):
+    async def _find_non_overlapping_cidr(
+        self, default_cidr: str = "172.31.0.0/16"
+    ) -> str:
         """Find a non-overlapping CIDR block"""
         response = await anyio.to_thread.run_sync(self._ec2_client.describe_vpcs)
         existing_cidrs = [vpc["CidrBlock"] for vpc in response["Vpcs"]]
@@ -708,9 +720,9 @@ class VpcResource:
 
     async def provision(
         self,
-        base_job_template: Dict[str, Any],
+        base_job_template: dict[str, Any],
         advance: Callable[[], None],
-    ):
+    ) -> None:
         """
         Provisions a VPC.
 
@@ -768,7 +780,7 @@ class VpcResource:
                 )
             )["AvailabilityZones"]
             zones = [az["ZoneName"] for az in azs]
-            subnets = []
+            subnets: list[Any] = []
             for i, subnet_cidr in enumerate(subnet_cidrs[0:3]):
                 subnets.append(
                     await anyio.to_thread.run_sync(
@@ -828,7 +840,7 @@ class VpcResource:
             )
 
     @property
-    def next_steps(self):
+    def next_steps(self) -> list[str]:
         return []
 
 
@@ -838,9 +850,9 @@ class ContainerRepositoryResource:
         self._repository_name = repository_name
         self._requires_provisioning = None
         self._work_pool_name = work_pool_name
-        self._next_steps = []
+        self._next_steps: list[str | Panel] = []
 
-    async def get_task_count(self):
+    async def get_task_count(self) -> int:
         """
         Returns the number of tasks that will be executed to provision this resource.
 
@@ -895,9 +907,9 @@ class ContainerRepositoryResource:
 
     async def provision(
         self,
-        base_job_template: Dict[str, Any],
+        base_job_template: dict[str, Any],
         advance: Callable[[], None],
-    ):
+    ) -> None:
         """
         Provisions an ECR repository.
 
@@ -978,7 +990,7 @@ class ContainerRepositoryResource:
         )
 
     @property
-    def next_steps(self):
+    def next_steps(self) -> list[str | Panel]:
         return self._next_steps
 
 
@@ -1000,7 +1012,7 @@ class ExecutionRoleResource:
         )
         self._requires_provisioning = None
 
-    async def get_task_count(self):
+    async def get_task_count(self) -> int:
         """
         Returns the number of tasks that will be executed to provision this resource.
 
@@ -1046,9 +1058,9 @@ class ExecutionRoleResource:
 
     async def provision(
         self,
-        base_job_template: Dict[str, Any],
+        base_job_template: dict[str, Any],
         advance: Callable[[], None],
-    ):
+    ) -> str:
         """
         Provisions an IAM role.
 
@@ -1087,7 +1099,7 @@ class ExecutionRoleResource:
         return response["Role"]["Arn"]
 
     @property
-    def next_steps(self):
+    def next_steps(self) -> list[str]:
         return []
 
 
@@ -1100,11 +1112,11 @@ class ElasticContainerServicePushProvisioner:
         self._console = Console()
 
     @property
-    def console(self):
+    def console(self) -> Console:
         return self._console
 
     @console.setter
-    def console(self, value):
+    def console(self, value: Console) -> None:
         self._console = value
 
     async def _prompt_boto3_installation(self):
@@ -1115,7 +1127,7 @@ class ElasticContainerServicePushProvisioner:
         boto3 = importlib.import_module("boto3")
 
     @staticmethod
-    def is_boto3_installed():
+    def is_boto3_installed() -> bool:
         """
         Check if boto3 is installed.
         """
@@ -1157,8 +1169,8 @@ class ElasticContainerServicePushProvisioner:
     async def provision(
         self,
         work_pool_name: str,
-        base_job_template: dict,
-    ) -> Dict[str, Any]:
+        base_job_template: dict[str, Any],
+    ) -> dict[str, Any]:
         """
         Provisions the infrastructure for an ECS push work pool.
 
@@ -1310,7 +1322,7 @@ class ElasticContainerServicePushProvisioner:
                 # provision calls will be no-ops, but update the base job template
 
             base_job_template_copy = deepcopy(base_job_template)
-            next_steps = []
+            next_steps: list[str | Panel] = []
             with Progress(console=self._console, disable=num_tasks == 0) as progress:
                 task = progress.add_task(
                     "Provisioning Infrastructure",

prefect/infrastructure/provisioners/modal.py

@@ -2,6 +2,7 @@ import importlib
 import shlex
 import sys
 from copy import deepcopy
+from types import ModuleType
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
 
 from anyio import run_process
@@ -19,7 +20,7 @@ if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient
 
 
-modal = lazy_import("modal")
+modal: ModuleType = lazy_import("modal")
 
 
 class ModalPushProvisioner:
@@ -28,14 +29,14 @@ class ModalPushProvisioner:
     """
 
     def __init__(self, client: Optional["PrefectClient"] = None):
-        self._console = Console()
+        self._console: Console = Console()
 
     @property
-    def console(self):
+    def console(self) -> Console:
         return self._console
 
     @console.setter
-    def console(self, value):
+    def console(self, value: Console) -> None:
         self._console = value
 
     @staticmethod

prefect/input/actions.py

@@ -1,3 +1,4 @@
+import inspect
 from typing import TYPE_CHECKING, Any, Optional, Set
 from uuid import UUID
 
@@ -44,9 +45,12 @@ async def create_flow_run_input_from_model(
     else:
         json_safe = orjson.loads(model_instance.json())
 
-    await create_flow_run_input(
+    coro = create_flow_run_input(
         key=key, value=json_safe, flow_run_id=flow_run_id, sender=sender
     )
+    if TYPE_CHECKING:
+        assert inspect.iscoroutine(coro)
+    await coro
 
 
 @sync_compatible